Spaces:
Runtime error
Runtime error
# services/detection_service.py | |
import cv2 | |
import numpy as np | |
from ultralytics import YOLO | |
import os | |
# Load YOLOv8m model | |
BASE_DIR = os.path.dirname(os.path.abspath(__file__)) | |
MODEL_PATH = os.path.join(BASE_DIR, "../models/yolov8m.pt") | |
model = YOLO(MODEL_PATH) | |
def detect_objects(frame): | |
""" | |
Detect objects in a frame using YOLOv8m. | |
Args: | |
frame: Input frame (numpy array) | |
Returns: | |
dict: Detection results with numbered labels | |
numpy array: Annotated frame | |
""" | |
results = model(frame) | |
detections = [] | |
line_counter = 1 # Initialize counter for numbered labels | |
for r in results: | |
for box in r.boxes: | |
conf = float(box.conf[0]) | |
if conf < 0.5: | |
continue | |
cls = int(box.cls[0]) | |
label = model.names[cls] | |
xyxy = box.xyxy[0].cpu().numpy() | |
x_min, y_min, x_max, y_max = map(int, xyxy) | |
# Add numbered label | |
detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})" | |
detections.append({ | |
"label": detection_label, | |
"confidence": conf, | |
"coordinates": [x_min, y_min, x_max, y_max] | |
}) | |
# Draw bounding box and label | |
color = (0, 255, 0) # Green for generic objects | |
cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2) | |
cv2.putText(frame, detection_label, (x_min, y_min - 10), | |
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) | |
line_counter += 1 | |
return {"detections": detections, "frame": frame} | |
def process_frame(frame): | |
""" | |
Wrapper function for integration with app.py. | |
""" | |
result = detect_objects(frame) | |
return result["detections"], result["frame"] |