Ready to Use

README.md

@@ -1,8 +1,8 @@
 ---
 title: Object Detection With YOLO
-emoji: 🐢
-colorFrom: pink
-colorTo: pink
+emoji: ❤️‍🔥
+colorFrom: purple
+colorTo: white
 sdk: gradio
 sdk_version: 3.37.0
 app_file: app.py

app.py

@@ -0,0 +1,82 @@
+from transformers import AutoFeatureExtractor, YolosForObjectDetection
+import gradio as gr
+from PIL import Image
+import torch
+import matplotlib.pyplot as plt
+import io
+import numpy as np
+
+
+COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
+          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
+
+
+def process_class_list(classes_string: str):
+    return [x.strip() for x in classes_string.split(",")] if classes_string else []
+
+def model_inference(img, model_name: str, prob_threshold: int, classes_to_show = str):
+    feature_extractor = AutoFeatureExtractor.from_pretrained(f"hustvl/{model_name}")
+    model = YolosForObjectDetection.from_pretrained(f"hustvl/{model_name}")
+
+    img = Image.fromarray(img)
+
+    pixel_values = feature_extractor(img, return_tensors="pt").pixel_values
+
+    with torch.no_grad():
+        outputs = model(pixel_values, output_attentions=True)
+
+    probas = outputs.logits.softmax(-1)[0, :, :-1]
+    keep = probas.max(-1).values > prob_threshold
+
+    target_sizes = torch.tensor(img.size[::-1]).unsqueeze(0)
+    postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes)
+    bboxes_scaled = postprocessed_outputs[0]['boxes']
+
+    classes_list = process_class_list(classes_to_show)
+    return plot_results(
+        img, probas[keep], bboxes_scaled[keep], model, classes_list
+    )
+
+def plot_results(pil_img, prob, boxes, model, classes_list):
+    plt.figure(figsize=(16,10))
+    plt.imshow(pil_img)
+    ax = plt.gca()
+    colors = COLORS * 100
+    for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):
+        cl = p.argmax()
+        object_class = model.config.id2label[cl.item()]
+        
+        if len(classes_list) > 0 :
+            if object_class not in classes_list:
+                continue
+            
+        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
+                                fill=False, color=c, linewidth=3))
+        text = f'{object_class}: {p[cl]:0.2f}'
+        ax.text(xmin, ymin, text, fontsize=15,
+                bbox=dict(facecolor='yellow', alpha=0.5))
+    plt.axis('off')
+    return fig2img(plt.gcf())
+    
+def fig2img(fig):
+    buf = io.BytesIO()
+    fig.savefig(buf)
+    buf.seek(0)
+    return Image.open(buf)
+
+description = """YOLOS - Object Detection"""
+
+image_in = gr.components.Image()
+image_out = gr.components.Image()
+model_choice = gr.components.Dropdown(["yolos-tiny", "yolos-small", "yolos-base", "yolos-small-300", "yolos-small-dwr"], value="yolos-small", label="YOLOS Model")
+prob_threshold_slider = gr.components.Slider(minimum=0, maximum=1.0, step=0.01, value=0.9, label="Probability Threshold")
+classes_to_show = gr.components.Textbox(placeholder="e.g. person, car , laptop", label="Classes to use (Optional)")
+
+Iface = gr.Interface(
+    fn=model_inference,
+    inputs=[image_in,model_choice, prob_threshold_slider, classes_to_show],
+    outputs=image_out,
+    title="Object Detection With YOLO",
+    description=description,
+    theme='HaleyCH/HaleyCH_Theme',
+).launch(share=True)

requirements.txt

@@ -0,0 +1,86 @@
+aiofiles==23.1.0
+aiohttp==3.8.4
+aiosignal==1.3.1
+altair==5.0.1
+annotated-types==0.5.0
+anyio==3.7.1
+async-timeout==4.0.2
+attrs==23.1.0
+certifi==2023.5.7
+charset-normalizer==3.2.0
+click==8.1.5
+colorama==0.4.6
+contourpy==1.1.0
+cycler==0.11.0
+exceptiongroup==1.1.2
+fastapi==0.100.0
+ffmpy==0.3.1
+filelock==3.12.2
+fonttools==4.41.0
+frozenlist==1.4.0
+fsspec==2023.6.0
+gitdb==4.0.10
+GitPython==3.1.32
+gradio==3.37.0
+gradio_client==0.2.10
+h11==0.14.0
+httpcore==0.17.3
+httpx==0.24.1
+huggingface-hub==0.16.4
+idna==3.4
+Jinja2==3.1.2
+jsonschema==4.18.4
+jsonschema-specifications==2023.6.1
+kiwisolver==1.4.4
+linkify-it-py==2.0.2
+markdown-it-py==2.2.0
+MarkupSafe==2.1.3
+matplotlib==3.7.2
+mdit-py-plugins==0.3.3
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.0.4
+networkx==3.1
+numpy==1.25.1
+opencv-python==4.8.0.74
+orjson==3.9.2
+packaging==23.1
+pandas==2.0.3
+Pillow==10.0.0
+psutil==5.9.5
+pydantic==2.0.3
+pydantic_core==2.3.0
+pydub==0.25.1
+pyparsing==3.0.9
+python-dateutil==2.8.2
+python-multipart==0.0.6
+pytz==2023.3
+PyYAML==6.0.1
+referencing==0.29.3
+regex==2023.6.3
+requests==2.31.0
+rpds-py==0.8.11
+safetensors==0.3.1
+scipy==1.11.1
+seaborn==0.12.2
+semantic-version==2.10.0
+six==1.16.0
+smmap==5.0.0
+sniffio==1.3.0
+starlette==0.27.0
+sympy==1.12
+thop==0.1.1.post2209072238
+tokenizers==0.13.3
+toolz==0.12.0
+torch==2.0.1
+torchvision==0.15.2
+tqdm==4.65.0
+transformers==4.30.2
+typing_extensions==4.7.1
+tzdata==2023.3
+uc-micro-py==1.0.2
+ultralytics==8.0.136
+urllib3==2.0.3
+uvicorn==0.23.1
+websockets==11.0.3
+yarl==1.9.2