| import json | |
| import glob | |
| from collections import Counter | |
| import requests | |
| import gradio as gr | |
| from ultralyticsplus import YOLO, download_from_hub, render_result | |
| hf_model_ids = [ | |
| "chanelcolgate/chamdiemgianhang-vsk", | |
| "chanelcolgate/chamdiemgianhang-vsk-v2", | |
| "chanelcolgate/chamdiemgianhang-vsk-v4", | |
| "chanelcolgate/chamdiemgianhang-vsk-v5", | |
| "chanelcolgate/chamdiemgianhang-vsk-v6", | |
| ] | |
| image_paths = [ | |
| [image_path, "chanelcolgate/chamdiemgianhang-vsk-v2", 640, 0.25, 0.45] | |
| for image_path in glob.glob("./tmp/*.jpg") | |
| ] | |
| def detection_image( | |
| image=None, | |
| hf_model_id="chanelcolgate/chamdiemgianhang-vsk-v2", | |
| image_size=640, | |
| conf_threshold=0.25, | |
| iou_threshold=0.45, | |
| ): | |
| model_path = download_from_hub(hf_model_id) | |
| model = YOLO(model_path) | |
| results = model(image, imgsz=image_size, conf=conf_threshold, iou=iou_threshold) | |
| json_result = json.loads(results[0].tojson()) | |
| class_counts = Counter(detection["name"] for detection in json_result) | |
| render = render_result(model=model, image=image, result=results[0]) | |
| return render, class_counts | |
| def detection_image_link( | |
| image=None, | |
| hf_model_id="chanelcolgate/chamdiemgianhang-vsk-v2", | |
| image_size=640, | |
| conf_threshold=0.25, | |
| iou_threshold=0.45, | |
| ): | |
| model_path = download_from_hub(hf_model_id) | |
| model = YOLO(model_path) | |
| results = model(image, imgsz=image_size, conf=conf_threshold, iou=iou_threshold) | |
| json_result = json.loads(results[0].tojson()) | |
| class_counts = Counter(detection["name"] for detection in json_result) | |
| render = render_result(model=model, image=image, result=results[0]) | |
| return render, class_counts | |
| title = "Cham Diem Gian Hang VSK" | |
| interface = gr.Interface( | |
| fn=detection_image, | |
| inputs=[ | |
| gr.Image(type="pil"), | |
| gr.Dropdown(hf_model_ids), | |
| gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"), | |
| gr.Slider( | |
| minimum=0.0, | |
| maximum=1.0, | |
| value=0.25, | |
| step=0.05, | |
| label="Confidence Threshold", | |
| ), | |
| gr.Slider( | |
| minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold" | |
| ), | |
| ], | |
| outputs=[gr.Image(type="pil"), gr.Textbox(show_label=False)], | |
| title=title, | |
| examples=image_paths, | |
| cache_examples=True if image_paths else False, | |
| ) | |
| interface_link = gr.Interface( | |
| fn=detection_image, | |
| inputs=[ | |
| gr.Textbox(label="Image Link"), | |
| gr.Dropdown(hf_model_ids), | |
| gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"), | |
| gr.Slider( | |
| minimum=0.0, | |
| maximum=1.0, | |
| value=0.25, | |
| step=0.05, | |
| label="Confidence Threshold", | |
| ), | |
| gr.Slider( | |
| minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold" | |
| ), | |
| ], | |
| outputs=[gr.Image(type="pil"), gr.Textbox(show_label=False)], | |
| title=title, | |
| ) | |
| gr.TabbedInterface( | |
| [interface, interface_link], tab_names=["Image inference", "Image link inference"] | |
| ).queue().launch() | |