From f837da7c0af90af726c0e880552716a7a735d80a Mon Sep 17 00:00:00 2001
From: wa22 <wa22@mails.tsinghua.edu.cn>
Date: Wed, 29 May 2024 11:03:33 +0800
Subject: [PATCH] add gradio demo

---
 README.md        |   7 ++++
 app.py           | 100 +++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt |   3 +-
 3 files changed, 109 insertions(+), 1 deletion(-)
 create mode 100644 app.py

diff --git a/README.md b/README.md
index 2466c02f..87983897 100644
--- a/README.md
+++ b/README.md
@@ -21,6 +21,8 @@ Over the past years, YOLOs have emerged as the predominant paradigm in the field
 </details>
 
 **UPDATES** 🔥
+- 2024/05/29: We identify a bug in existing HuggingFace demos. Please use `gr.Image(type="pil", label="Image")` rather than ``gr.Image(type="numpy", label="Image")`` for prediction.
+- 2024/05/29: Add the gradio demo for running the models locally. Thanks to [AK](https://x.com/_akhaliq)!
 - 2024/05/27: Thanks to [sujanshresstha](sujanshresstha) for the integration with [DeepSORT](https://github.com/sujanshresstha/YOLOv10_DeepSORT.git)!
 - 2024/05/27: We have updated the [checkpoints](https://github.com/THU-MIG/yolov10/releases/tag/v1.1) with other attributes, like class names, for ease of use.
 - 2024/05/26: Thanks to [CVHub520](https://github.com/CVHub520) for the integration into [X-AnyLabeling](https://github.com/CVHub520/X-AnyLabeling)!
@@ -47,6 +49,11 @@ conda activate yolov10
 pip install -r requirements.txt
 pip install -e .
 ```
+## Demo
+```
+python app.py
+# Please visit http://127.0.0.1:7860
+```
 
 ## Validation
 [`yolov10n.pt`](https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10n.pt)  [`yolov10s.pt`](https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10s.pt)  [`yolov10m.pt`](https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10m.pt)  [`yolov10b.pt`](https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10b.pt)  [`yolov10l.pt`](https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10l.pt)  [`yolov10x.pt`](https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10x.pt)  
diff --git a/app.py b/app.py
new file mode 100644
index 00000000..5830e8dc
--- /dev/null
+++ b/app.py
@@ -0,0 +1,100 @@
+# Ackownledgement: https://huggingface.co/spaces/kadirnar/Yolov10/blob/main/app.py
+# Thanks to @kadirnar
+
+import gradio as gr
+from ultralytics import YOLOv10 
+
+def yolov10_inference(image, model_path, image_size, conf_threshold):
+    model = YOLOv10(model_path)
+    
+    model.predict(source=image, imgsz=image_size, conf=conf_threshold, save=True)
+    
+    return model.predictor.plotted_img[:, :, ::-1]
+
+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=[
+                        "yolov10n.pt",
+                        "yolov10s.pt",
+                        "yolov10m.pt",
+                        "yolov10b.pt",
+                        "yolov10l.pt",
+                        "yolov10x.pt",
+                    ],
+                    value="yolov10s.pt",
+                )
+                image_size = gr.Slider(
+                    label="Image Size",
+                    minimum=320,
+                    maximum=1280,
+                    step=32,
+                    value=640,
+                )
+                conf_threshold = gr.Slider(
+                    label="Confidence Threshold",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.1,
+                    value=0.25,
+                )
+                yolov10_infer = gr.Button(value="Detect Objects")
+
+            with gr.Column():
+                output_image = gr.Image(type="numpy", label="Annotated Image")
+
+        yolov10_infer.click(
+            fn=yolov10_inference,
+            inputs=[
+                image,
+                model_id,
+                image_size,
+                conf_threshold,
+            ],
+            outputs=[output_image],
+        )
+
+        gr.Examples(
+            examples=[
+                [
+                    "ultralytics/assets/bus.jpg",
+                    "yolov10s.pt",
+                    640,
+                    0.25,
+                ],
+                [
+                    "ultralytics/assets/zidane.jpg",
+                    "yolov10s.pt",
+                    640,
+                    0.25,
+                ],
+            ],
+            fn=yolov10_inference,
+            inputs=[
+                image,
+                model_id,
+                image_size,
+                conf_threshold,
+            ],
+            outputs=[output_image],
+            cache_examples=True,
+        )
+
+gradio_app = gr.Blocks()
+with gradio_app:
+    gr.HTML(
+        """
+    <h1 style='text-align: center'>
+    YOLOv10: Real-Time End-to-End Object Detection
+    </h1>
+    """)
+    with gr.Row():
+        with gr.Column():
+            app()
+
+gradio_app.launch(debug=True)
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index effce8a0..c16a917a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,4 +6,5 @@ pycocotools==2.0.7
 PyYAML==6.0.1
 scipy==1.13.0
 onnxsim==0.4.36
-onnxruntime-gpu==1.18.0
\ No newline at end of file
+onnxruntime-gpu==1.18.0
+gradio==4.31.5
\ No newline at end of file