Update app.py

app.py

@@ -1,48 +1,117 @@
 import gradio as gr
 import spaces
 import PIL.Image as Image
-from ultralytics import YOLO
+from ultralytics import YOLO, YOLOv10
+from huggingface_hub import hf_hub_download
 
 
+def download_models(model_id):
+    hf_hub_download("atalaydenknalbant/", filename="asl-models", local_dir=f"./")
+    return f"./{model_id}"
 
 
-@spaces.GPU
-def predict_image(img, conf_threshold, iou_threshold):
-    
-    model = YOLO("best.pt")
+
+box_annotator = sv.BoxAnnotator()
+category_dict = {'A': 0, 'B': 1, 'C': 2, 'D': 3, 'E': 4, 'F': 5, 'G': 6, 'H': 7, 'I': 8, 'J': 9,
+                 'K': 10, 'L': 11, 'M': 12, 'N': 13, 'O': 14, 'P': 15, 'Q': 16, 'R': 17, 'S': 18,
+                 'T': 19, 'U': 20, 'V': 21, 'W': 22, 'X': 23, 'Y': 24, 'Z': 25}
+
+
+
+@spaces.GPU(duration=200)
+def yolo_inference(image, model_id, conf_threshold, iou_threshold):
+    model_path = download_models(model_id)
+    if model_id[:7] == 'yolov10':
+        model = YOLOv10(model_path)
+    else:
+        model = YOLO(model_path)
+    results = model(source=image, imgsz=416, iou=iou_threshold, conf=conf_threshold, verbose=False, max_det=1)[0]
+    detections = sv.Detections.from_ultralytics(results)
     
-    results = model.predict(
-        source=img,
-        conf=conf_threshold,
-        iou=iou_threshold,
-        show_labels=True,
-        show_conf=True,
-        imgsz=416,
-        max_det=1
-    )
-
-    for r in results:
-        im_array = r.plot()
-        im = Image.fromarray(im_array[..., ::-1])
-
-    return im
-
-
-asl = gr.Interface(
-    fn=predict_image,
-    inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
-        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold")
-    ],
-    outputs=gr.Image(type="pil", label="Result"),
-    title="ASL Detector YOLOV9e",
-    description="Upload images for inference. The Ultralytics YOLOV9e model is used by default. Letter Z is not supported.",
-    examples=[
-        ["y.jpg", 0.25, 0.45],
-        ["b.jpg", 0.25, 0.45],
+    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")
+                
+                model_id = gr.Dropdown(
+                    label="Model",
+                    choices=[
+                        "yolov10x.pt",
+                        "yolov10s.pt",
+                        "yolov9e.pt",
+                        "yolov8x.pt",
+                    ],
+                    value="yolov10s.pt",
+                )
+                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,
+                )
+                yolov10_infer = gr.Button(value="Detect Objects")
+
+            with gr.Column():
+                output_image = gr.Image(type="pil", label="Annotated Image")
+
+        yolov10_infer.click(
+            fn=yolo_inference,
+            inputs=[
+                image,
+                model_id,
+                conf_threshold,
+                iou_threshold,
+            ],
+            outputs=[output_image],
+        )
+
+        gr.Examples(
+            examples=[
+                [
+                    "b.jpg",
+                    "yolov10x.pt",
+                    0.25,
+                    0.45,
+                ],
+                [
+                    "y.jpg",
+                    "yolov10x.pt",
+                    0.25,
+                    0.45,
+                ],
+            ],
+            fn=yolo_inference,
+            inputs=[
+                image,
+                model_id,
+                conf_threshold,
+                iou_threshold,
+            ],
+            outputs=[output_image],
+            cache_examples=True,
+        )
+
+gradio_app = gr.Blocks()
+with gradio_app:
+    with gr.Row():
+        with gr.Column():
+            app()
 
-if __name__ == '__main__':
-    asl.launch()
+gradio_app.launch(debug=True)