update app.py

README.md

@@ -11,7 +11,7 @@ Official PyTorch implementation of **YOLOv10**.
 
 [YOLOv10: Real-Time End-to-End Object Detection](https://arxiv.org/abs/2405.14458).\
 Ao Wang, Hui Chen, Lihao Liu, Kai Chen, Zijia Lin, Jungong Han, and Guiguang Ding\
-[![arXiv](https://img.shields.io/badge/arXiv-2405.14458-b31b1b.svg)](https://arxiv.org/abs/2405.14458) <a href="https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/train-yolov10-object-detection-on-custom-dataset.ipynb#scrollTo=SaKTSzSWnG7s"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-blue)](https://huggingface.co/collections/jameslahm/yolov10-665b0d90b0b5bb85129460c2)  [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/kadirnar/Yolov10)   [![Transformers.js Demo](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Transformers.js-blue)](https://huggingface.co/spaces/Xenova/yolov10-web)
+[![arXiv](https://img.shields.io/badge/arXiv-2405.14458-b31b1b.svg)](https://arxiv.org/abs/2405.14458) <a href="https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/train-yolov10-object-detection-on-custom-dataset.ipynb#scrollTo=SaKTSzSWnG7s"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a> [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-blue)](https://huggingface.co/collections/jameslahm/yolov10-665b0d90b0b5bb85129460c2) [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/jameslahm/YOLOv10)  [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/kadirnar/Yolov10)  [![Transformers.js Demo](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Transformers.js-blue)](https://huggingface.co/spaces/Xenova/yolov10-web)
 
 <details>
   <summary>
@@ -23,7 +23,7 @@ Over the past years, YOLOs have emerged as the predominant paradigm in the field
 ## Notes
 - 2024/05/31: Please use the [exported format](https://github.com/THU-MIG/yolov10?tab=readme-ov-file#export) for benchmark. In the non-exported format, e.g., pytorch, the speed of YOLOv10 is biased because the unnecessary `cv2` and `cv3` operations in the `v10Detect` are executed during inference.
 - 2024/05/30: We provide [some clarifications and suggestions](https://github.com/THU-MIG/yolov10/issues/136) for detecting smaller objects or objects in the distance with YOLOv10. Thanks to [SkalskiP](https://github.com/SkalskiP)!
-- 2024/05/27: We have updated the [checkpoints](https://huggingface.co/collections/jameslahm/yolov10-665b0d90b0b5bb85129460c2) with other attributes, like class names and training args, for ease of use.
+- 2024/05/27: We have updated the [checkpoints](https://huggingface.co/collections/jameslahm/yolov10-665b0d90b0b5bb85129460c2) with class names, for ease of use.
 
 ## UPDATES 🔥
 - 2024/06/01: Thanks to [ErlanggaYudiPradana](https://github.com/rlggyp) for the integration with [C++ | OpenVINO | OpenCV](https://github.com/rlggyp/YOLOv10-OpenVINO-CPP-Inference)
@@ -97,12 +97,10 @@ model = YOLOv10()
 # If you want to finetune the model with pretrained weights, you could load the 
 # pretrained weights like below
 # model = YOLOv10('yolov10{n/s/m/b/l/x}.pt')
-# Or
+# or
 # model = YOLOv10.from_pretrained('jameslahm/yolov10{n/s/m/b/l/x}')
 
 model.train(data='coco.yaml', epochs=500, batch=256, imgsz=640)
-# Note that you can upload your trained model to HuggingFace Hub like below
-# model.push_to_hub("reponame", config={"model": "yolov10n/s/m/b/l/x.yaml"})
 ```
 
 ## Push to hub to 🤗
@@ -143,7 +141,7 @@ yolo predict model=yolov10n/s/m/b/l/x.onnx
 
 # End-to-End TensorRT
 yolo export model=yolov10n/s/m/b/l/x.pt format=engine half=True simplify opset=13 workspace=16
-# Or
+# or
 trtexec --onnx=yolov10n/s/m/b/l/x.onnx --saveEngine=yolov10n/s/m/b/l/x.engine --fp16
 # Predict with TensorRT
 yolo predict model=yolov10n/s/m/b/l/x.engine

app.py

@@ -1,15 +1,23 @@
-# Ackownledgement: https://huggingface.co/spaces/kadirnar/Yolov10/blob/main/app.py
+import PIL.Image as Image
-# Thanks to @kadirnar
-
 import gradio as gr
-from ultralytics import YOLOv10 
 
-def yolov10_inference(image, image_size, conf_threshold):
+from ultralytics import YOLOv10
-    model = YOLOv10.from_pretrained("jameslahm/yolov10n")
+
-    
+def predict_image(img, model_id, image_size, conf_threshold):
-    model.predict(source=image, imgsz=image_size, conf=conf_threshold, save=True)
+    model = YOLOv10.from_pretrained(f'jameslahm/{model_id}')
-    
+    results = model.predict(
-    return model.predictor.plotted_img[:, :, ::-1]
+        source=img,
+        conf=conf_threshold,
+        show_labels=True,
+        show_conf=True,
+        imgsz=image_size,
+    )
+
+    for r in results:
+        im_array = r.plot()
+        im = Image.fromarray(im_array[..., ::-1])
+
+    return im
 
 def app():
     with gr.Blocks():
@@ -20,14 +28,14 @@ def app():
                 model_id = gr.Dropdown(
                     label="Model",
                     choices=[
-                        "yolov10n.pt",
+                        "yolov10n",
-                        "yolov10s.pt",
+                        "yolov10s",
-                        "yolov10m.pt",
+                        "yolov10m",
-                        "yolov10b.pt",
+                        "yolov10b",
-                        "yolov10l.pt",
+                        "yolov10l",
-                        "yolov10x.pt",
+                        "yolov10x",
                     ],
-                    value="yolov10s.pt",
+                    value="yolov10m",
                 )
                 image_size = gr.Slider(
                     label="Image Size",
@@ -40,16 +48,16 @@ def app():
                     label="Confidence Threshold",
                     minimum=0.0,
                     maximum=1.0,
-                    step=0.1,
+                    step=0.05,
                     value=0.25,
                 )
                 yolov10_infer = gr.Button(value="Detect Objects")
 
             with gr.Column():
-                output_image = gr.Image(type="numpy", label="Annotated Image")
+                output_image = gr.Image(type="pil", label="Annotated Image")
 
         yolov10_infer.click(
-            fn=yolov10_inference,
+            fn=predict_image,
             inputs=[
                 image,
                 model_id,
@@ -63,18 +71,18 @@ def app():
             examples=[
                 [
                     "ultralytics/assets/bus.jpg",
-                    "yolov10s.pt",
+                    "yolov10s",
                     640,
                     0.25,
                 ],
                 [
                     "ultralytics/assets/zidane.jpg",
-                    "yolov10s.pt",
+                    "yolov10s",
                     640,
                     0.25,
                 ],
             ],
-            fn=yolov10_inference,
+            fn=predict_image,
             inputs=[
                 image,
                 model_id,
@@ -82,7 +90,7 @@ def app():
                 conf_threshold,
             ],
             outputs=[output_image],
-            cache_examples=True,
+            cache_examples='lazy',
         )
 
 gradio_app = gr.Blocks()
@@ -93,8 +101,14 @@ with gradio_app:
     YOLOv10: Real-Time End-to-End Object Detection
     </h1>
     """)
+    gr.HTML(
+        """
+        <h3 style='text-align: center'>
+        <a href='https://arxiv.org/abs/2405.14458' target='_blank'>arXiv</a> | <a href='https://github.com/THU-MIG/yolov10' target='_blank'>github</a>
+        </h3>
+        """)
     with gr.Row():
         with gr.Column():
             app()
-
+if __name__ == '__main__':
-gradio_app.launch(debug=True)
+    gradio_app.launch()