Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import random | |
| import matplotlib.pyplot as plt | |
| from ultralytics import YOLO | |
| import gradio as gr | |
| # Load the trained YOLO model | |
| MODEL_PATH = "best.pt" # Replace with your model's path | |
| model = YOLO(MODEL_PATH) | |
| def random_color(): | |
| """Generate a random color for bounding boxes.""" | |
| return (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) | |
| def detect_and_visualize(image_path): | |
| """ | |
| Perform object detection and visualize the results. | |
| Args: | |
| image_path (str): Path to the input image. | |
| Returns: | |
| Annotated image as an array. | |
| Detection details as a dictionary. | |
| """ | |
| # Perform object detection | |
| results = model(image_path) | |
| # Read the input image | |
| image = cv2.imread(image_path) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert to RGB | |
| detections = [] | |
| for result in results: | |
| boxes = result.boxes.xyxy.cpu().numpy() # Bounding box coordinates | |
| confidences = result.boxes.conf.cpu().numpy() # Confidence scores | |
| class_ids = result.boxes.cls.cpu().numpy().astype(int) # Class IDs | |
| 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] | |
| # Draw bounding box and label | |
| color = random_color() | |
| cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, 2) | |
| label = f"{class_name} {confidence:.2f}" | |
| cv2.putText(image, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) | |
| # Append detection details | |
| detections.append({ | |
| "label": class_name, | |
| "confidence": float(confidence), | |
| "bounding_box": { | |
| "x1": x_min, | |
| "y1": y_min, | |
| "x2": x_max, | |
| "y2": y_max | |
| } | |
| }) | |
| # Optionally save the annotated image | |
| output_path = "output/annotated_image.jpg" | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| cv2.imwrite(output_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR)) | |
| print(f"Annotated image saved to {output_path}") | |
| return cv2.cvtColor(image, cv2.COLOR_RGB2BGR), detections | |
| # Gradio interface | |
| def gradio_interface(image): | |
| """ | |
| Gradio-compatible wrapper for object detection. | |
| Args: | |
| image (numpy.array): Input image. | |
| Returns: | |
| Annotated image and detection details. | |
| """ | |
| temp_input_path = "temp_input.jpg" | |
| cv2.imwrite(temp_input_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR)) # Save temp file for YOLO | |
| annotated_image, detections = detect_and_visualize(temp_input_path) | |
| os.remove(temp_input_path) # Clean up temp file | |
| return annotated_image, detections | |
| # Define Gradio interface | |
| 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() | |