Spaces:
Sleeping
Sleeping
File size: 3,100 Bytes
e416d5a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
# 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()
|