farz / app.py
csprojectworkspace's picture
Update app.py
88330a7 verified
Raw
History Blame Contribute Delete
1.48 kB
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()