Spaces:
Build error
Build error
| import cv2 | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| # Load the model once globally | |
| MODEL_PATH = "best.pt" | |
| model = YOLO(MODEL_PATH) | |
| def detect_and_visualize(image): | |
| # image is a NumPy array from Gradio | |
| # Perform inference directly on this array | |
| results = model(image) | |
| # Ensure image is in the correct color space (most likely already RGB) | |
| annotated_image = image.copy() | |
| detections = [] | |
| for result in results: | |
| boxes = result.boxes.xyxy.cpu().numpy() | |
| confidences = result.boxes.conf.cpu().numpy() | |
| class_ids = result.boxes.cls.cpu().numpy().astype(int) | |
| for box, confidence, class_id in zip(boxes, confidences, class_ids): | |
| x_min, y_min, x_max, y_max = map(int, box) | |
| class_name = model.names[class_id] | |
| # Pick a color or use a fixed color, no need for random if not desired | |
| color = (0, 255, 0) | |
| cv2.rectangle(annotated_image, (x_min, y_min), (x_max, y_max), color, 2) | |
| label = f"{class_name} {confidence:.2f}" | |
| cv2.putText(annotated_image, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) | |
| detections.append({ | |
| "label": class_name, | |
| "confidence": float(confidence), | |
| "bounding_box": { | |
| "x1": x_min, | |
| "y1": y_min, | |
| "x2": x_max, | |
| "y2": y_max | |
| } | |
| }) | |
| return annotated_image, detections | |
| def gradio_interface(image): | |
| annotated_image, detections = detect_and_visualize(image) | |
| return annotated_image, detections | |
| interface = gr.Interface( | |
| fn=gradio_interface, | |
| inputs=gr.Image(type="numpy", label="Upload Image"), | |
| outputs=[ | |
| gr.Image(type="numpy", label="Annotated Image"), | |
| gr.JSON(label="Detection Details") | |
| ], | |
| title="YOLO Object Detection", | |
| description="Upload an image to detect objects and view annotated results along with detailed detection data." | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch() | |