Spaces:
Sleeping
Sleeping
| # Created by yarramsettinaresh GORAKA DIGITAL PRIVATE LIMITED at 24/10/24 | |
| import gradio as gr | |
| import cv2 | |
| import time | |
| from ultralytics import YOLO | |
| # Load your YOLO model (adjust model path or type as needed) | |
| model_path = "model_- 11 october 2024 11_07.pt" | |
| model = YOLO(model_path) | |
| def ultralytics_predict(model, frame): | |
| confidence_threshold = 0.2 | |
| start_time = time.time() | |
| results = model(frame) # Perform inference on the frame | |
| end_time = time.time() | |
| duration = end_time - start_time | |
| print(f"Prediction duration: {duration:.4f} seconds") | |
| duration_str = f"{duration:.4f} S" | |
| object_count = {} # Dictionary to store counts of detected objects | |
| for detection in results[0].boxes: # Iterate through detections | |
| conf = float(detection.conf[0]) # Confidence score | |
| if conf > confidence_threshold: | |
| conf, pos, text, color = ultralytics(detection, duration_str) | |
| cv2.rectangle(frame, pos[0], pos[1], color, 2) | |
| cv2.putText(frame, text, (pos[0][0], pos[0][1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) | |
| # Update object count | |
| class_id = int(detection.cls[0]) | |
| class_name = model.names[class_id] | |
| if class_name not in object_count: | |
| object_count[class_name] = dict(count=0, objects=[]) | |
| object_mapp = object_count[class_name] | |
| object_mapp["count"] = object_mapp.get("count", 0) + 1 | |
| object_mapp["objects"].append(dict(conf=conf, pos=pos, text=text, color=color)) | |
| return frame # Return the count of detected objects | |
| def ultralytics(detection, duration): | |
| COLOUR_MAP = { | |
| 0: (0, 0, 255), # Red in BGR format | |
| 1: (0, 255, 0) # Green in BGR format | |
| } | |
| conf = float(detection.conf[0]) # Confidence score | |
| class_id = int(detection.cls[0]) # Class ID | |
| name = model.names[class_id] # Get class name | |
| xmin, ymin, xmax, ymax = map(int, detection.xyxy[0]) # Bounding box coordinates | |
| color = COLOUR_MAP.get(class_id, (255, 255, 255)) # Default to white if not found | |
| # Draw bounding box and label on the frame | |
| pos = (xmin, ymin), (xmax, ymax) | |
| text = f"{name} {round(conf, 2)} :{duration}" | |
| return conf, pos, text, color | |
| def process_frame(frame): | |
| object_count = ultralytics_predict(model, frame) | |
| return frame, object_count # Return frame and object count | |
| def detect_image(image): | |
| image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) # Convert to BGR format for OpenCV | |
| result_frame, object_count = process_frame(image) | |
| result_frame = cv2.cvtColor(result_frame, cv2.COLOR_BGR2RGB) # Convert back to RGB for Gradio | |
| return result_frame, object_count # Return both the frame and the count | |
| # Create Gradio Interface | |
| gr.Interface( | |
| fn=detect_image, | |
| inputs=gr.Image(type="numpy"), # Updated input format | |
| outputs=[ | |
| gr.Image(type="numpy"), # Image output | |
| gr.JSON(), # Object count output as JSON | |
| ], | |
| title="YOLO Object Detection", | |
| description="Upload an image to detect objects using YOLO model." | |
| ).launch() | |