FrameNetBrasil/Yolo11x_FM30k_Event_withcaption

1import cv2
2import gradio as gr
3from ultralytics import YOLO
4
5# Load the model once globally
6MODEL_PATH = "best.pt"  
7model = YOLO(MODEL_PATH)
8
9def detect_and_visualize(image):
10    # image is a NumPy array from Gradio
11    # Perform inference directly on this array
12    results = model(image)
13    
14    # Ensure image is in the correct color space (most likely already RGB)
15    annotated_image = image.copy()
16
17    detections = []
18    for result in results:
19        boxes = result.boxes.xyxy.cpu().numpy()  
20        confidences = result.boxes.conf.cpu().numpy()  
21        class_ids = result.boxes.cls.cpu().numpy().astype(int)  
22
23        for box, confidence, class_id in zip(boxes, confidences, class_ids):
24            x_min, y_min, x_max, y_max = map(int, box)
25            class_name = model.names[class_id]
26
27            # Pick a color or use a fixed color, no need for random if not desired
28            color = (0, 255, 0)
29            cv2.rectangle(annotated_image, (x_min, y_min), (x_max, y_max), color, 2)
30            label = f"{class_name} {confidence:.2f}"
31            cv2.putText(annotated_image, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
32
33            detections.append({
34                "label": class_name,
35                "confidence": float(confidence),
36                "bounding_box": {
37                    "x1": x_min,
38                    "y1": y_min,
39                    "x2": x_max,
40                    "y2": y_max
41                }
42            })
43
44    return annotated_image, detections
45
46def gradio_interface(image):
47    annotated_image, detections = detect_and_visualize(image)
48    return annotated_image, detections
49
50interface = gr.Interface(
51    fn=gradio_interface,
52    inputs=gr.Image(type="numpy", label="Upload Image"),
53    outputs=[
54        gr.Image(type="numpy", label="Annotated Image"),
55        gr.JSON(label="Detection Details")
56    ],
57    title="YOLO Object Detection",
58    description="Upload an image to detect objects and view annotated results along with detailed detection data."
59)
60
61if __name__ == "__main__":
62    interface.launch()
63