Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| from PIL import Image | |
| from collections import defaultdict | |
| # load model | |
| model = YOLO("best.pt") | |
| def predict(image): | |
| results = model(image) | |
| r = results[0] | |
| # Group detections by class | |
| class_confidences = defaultdict(list) | |
| if len(r.boxes) > 0: | |
| for box in r.boxes: | |
| class_id = int(box.cls[0]) | |
| class_name = model.names[class_id] | |
| confidence = float(box.conf[0]) * 100 | |
| class_confidences[class_name].append(confidence) | |
| if class_confidences: | |
| lines = [] | |
| sorted_classes = sorted( | |
| class_confidences.items(), | |
| key=lambda x: max(x[1]), | |
| reverse=True | |
| ) | |
| for class_name, confidences in sorted_classes: | |
| avg_conf = round(sum(confidences) / len(confidences), 2) | |
| max_conf = round(max(confidences), 2) | |
| lines.append(f"• {class_name} — avg: {avg_conf}% | best: {max_conf}%") | |
| result_text = "\n".join(lines) | |
| else: | |
| result_text = "No detection" | |
| output_image = r.plot() | |
| return Image.fromarray(output_image), result_text | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(type="pil"), | |
| outputs=[ | |
| gr.Image(type="pil", label="Detection"), | |
| gr.Textbox(label="Result", lines=10) | |
| ], | |
| title="YOLOv8 Detection", | |
| description="Upload an image and detect objects" | |
| ) | |
| demo.launch() | |