From 587981fc27366c93cef6ac18ae39517349e04e3d Mon Sep 17 00:00:00 2001 From: wa22 Date: Mon, 3 Jun 2024 10:50:08 +0800 Subject: [PATCH] update app.py --- README.md | 10 ++++----- app.py | 66 +++++++++++++++++++++++++++++++++---------------------- 2 files changed, 44 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index 013a7dc7..9729982c 100644 --- a/README.md +++ b/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) Open In Colab [![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) Open In Colab [![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)
@@ -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 diff --git a/app.py b/app.py index e40cc846..5598af6d 100644 --- a/app.py +++ b/app.py @@ -1,15 +1,23 @@ -# Ackownledgement: https://huggingface.co/spaces/kadirnar/Yolov10/blob/main/app.py -# Thanks to @kadirnar - +import PIL.Image as Image import gradio as gr -from ultralytics import YOLOv10 -def yolov10_inference(image, image_size, conf_threshold): - model = YOLOv10.from_pretrained("jameslahm/yolov10n") - - model.predict(source=image, imgsz=image_size, conf=conf_threshold, save=True) - - return model.predictor.plotted_img[:, :, ::-1] +from ultralytics import YOLOv10 + +def predict_image(img, model_id, image_size, conf_threshold): + model = YOLOv10.from_pretrained(f'jameslahm/{model_id}') + results = model.predict( + 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", - "yolov10s.pt", - "yolov10m.pt", - "yolov10b.pt", - "yolov10l.pt", - "yolov10x.pt", + "yolov10n", + "yolov10s", + "yolov10m", + "yolov10b", + "yolov10l", + "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 """) + gr.HTML( + """ +

+ arXiv | github +

+ """) with gr.Row(): with gr.Column(): app() - -gradio_app.launch(debug=True) \ No newline at end of file +if __name__ == '__main__': + gradio_app.launch() \ No newline at end of file