Update app.py

app.py

@@ -1,165 +1,51 @@
-import spaces
-import supervision as sv
+import gradio as gr
 import PIL.Image as Image
 from ultralytics import YOLO
-import gradio as gr
-import torch
-
-
-
-model_filenames = [
-    "yolo11n.pt", 
-    "yolo11s.pt",
-    "yolo11m.pt", 
-    "yolo11l.pt", 
-    "yolo11x.pt"
-]
-
-
-
-box_annotator = sv.BoxAnnotator()
-
-category_dict = {
-    0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus',
-    6: 'train', 7: 'truck', 8: 'boat', 9: 'traffic light', 10: 'fire hydrant',
-    11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat',
-    16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear',
-    22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag',
-    27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard',
-    32: 'sports ball', 33: 'kite', 34: 'baseball bat', 35: 'baseball glove',
-    36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle',
-    40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl',
-    46: 'banana', 47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli',
-    51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake',
-    56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table',
-    61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard',
-    67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink',
-    72: 'refrigerator', 73: 'book', 74: 'clock', 75: 'vase', 76: 'scissors',
-    77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'
-}
-
-
 
-@spaces.GPU
-def yolo_inference(image, model_id, conf_threshold, iou_threshold, max_detection):
-
-    
+def yolo_inference(images, model_id, conf_threshold, iou_threshold, max_detection):
     model = YOLO(model_id)
-    results = model(source=image, imgsz=640, iou=iou_threshold, conf=conf_threshold, verbose=False, max_det=max_detection)[0]
-    detections = sv.Detections.from_ultralytics(results)
-    
-    labels = [
-        f"{category_dict[class_id]} {confidence:.2f}"
-        for class_id, confidence in zip(detections.class_id, detections.confidence)
-    ]
-    annotated_image = box_annotator.annotate(image, detections=detections, labels=labels)
-
-    return annotated_image
-
-def app():
-    with gr.Blocks():
-        with gr.Row():
-            with gr.Column():
-                image = gr.Image(type="pil", label="Image", interactive=True)
+    results = model.predict(
+        source=images,
+        conf=conf_threshold,
+        iou=iou_threshold,
+        imgsz=640,
+        max_det=max_detection,
+        show_labels=True,
+        show_conf=True,
+    )
+
+    annotated_images = []
+    for result in results:
+        im_array = result.plot()
+        im = Image.fromarray(im_array)
+        annotated_images.append(im)
     
-                model_id = gr.Dropdown(
-                    label="Model",
-                    choices=model_filenames,
-                    value=model_filenames[0] if model_filenames else "",
-                )
-                conf_threshold = gr.Slider(
-                    label="Confidence Threshold",
-                    minimum=0.1,
-                    maximum=1.0,
-                    step=0.1,
-                    value=0.25,
-                )
-                iou_threshold = gr.Slider(
-                    label="IoU Threshold",
-                    minimum=0.1,
-                    maximum=1.0,
-                    step=0.1,
-                    value=0.45,
-                )
-                
-                max_detection = gr.Slider(
-                    label="Max Detection",
-                    minimum=1,
-                    maximum=300,
-                    step=1,
-                    value=300,
-                )
-                yolov_infer = gr.Button(value="Detect Objects")
-
-            with gr.Column():
-                output_image = gr.Image(type="pil", label="Annotated Image", interactive=False)
-
-        yolov_infer.click(
-            fn=yolo_inference,
-            inputs=[
-                image,
-                model_id,
-                conf_threshold,
-                iou_threshold,
-                max_detection,
-            ],
-            outputs=[output_image],
-        )
-
-        gr.Examples(
-            examples=[
-                [
-                    "zidane.jpg",
-                    "yolo11s.pt",
-                    0.25,
-                    0.45,
-                    300,
-                ],
-                
-                [
-                    "bus.jpg",
-                    "yolo11m.pt",
-                    0.25,
-                    0.45,
-                    300,
-                ],
-                [
-                    "yolo_vision.jpg",
-                    "yolo11x.pt",
-                    0.25,
-                    0.45,
-                    300,
-                ],
-            ],
-            fn=yolo_inference,
-            inputs=[
-                image,
-                model_id,
-                conf_threshold,
-                iou_threshold,
-                max_detection,
-            ],
-            outputs=[output_image],
-            cache_examples=True,
-        )
-
-gradio_app = gr.Blocks()
-with gradio_app:
-    gr.HTML(
-        """
-    <h1 style='text-align: center'>
-    Yolo11: Object Detection
-    </h1>
-        <a>
-        """)
-    gr.HTML(
-    """
-    <p style='text-align: center'>
-    Latest ultralytics yolo11 object detection models. Upload an image to run inference.
-    </p>
-    """)
-    with gr.Row():
-        with gr.Column():
-            app()
-
-gradio_app.launch()
+    return annotated_images  
+
+interface = gr.Interface(
+    fn=yolo_inference,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image(s)", tool="editor", source="upload", multiple=True),
+        gr.Dropdown(
+            choices=['yolo11n.pt', 'yolo11s.pt', 'yolo11m.pt', 'yolo11l.pt', 'yolo11x.pt', 
+                     'yolo11n-seg.pt', 'yolo11s-seg.pt', 'yolo11m-seg.pt', 'yolo11l-seg.pt', 'yolo11x-seg.pt',
+                     'yolo11n-pose.pt', 'yolo11s-pose.pt', 'yolo11m-pose.pt', 'yolo11l-pose.pt', 'yolo11x-pose.pt',
+                     'yolo11n-obb.pt', 'yolo11s-obb.pt', 'yolo11m-obb.pt', 'yolo11l-obb.pt', 'yolo11x-obb.pt',
+                     'yolo11n-cls.pt', 'yolo11s-cls.pt', 'yolo11m-cls.pt', 'yolo11l-cls.pt', 'yolo11x-cls.pt'],
+            label="Model Name",
+            value="yolo11n.pt",
+        ),
+        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence Threshold"),
+        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU Threshold"),
+        gr.Slider(minimum=1, maximum=300, step=1, value=300, label="Max Detection"),
+    ],
+    outputs=gr.Gallery(label="Annotated Image"),
+    title="Yolo11: Object Detection",
+    description="Upload image(s) for inference using the latest Ultralytics YOLO11 models.",
+    examples=[
+        [["zidane.jpg"], "yolo11s.pt", 0.25, 0.45, 300],
+        [["bus.jpg"], "yolo11m.pt", 0.25, 0.45, 300],
+        [["yolo_vision.jpg"], "yolo11x.pt", 0.25, 0.45, 300],
+    ],
+)
+interface.launch()