Spaces:
Sleeping
Sleeping
Create services/detection_service.py
Browse files
services/services/detection_service.py
ADDED
@@ -0,0 +1,59 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# services/detection_service.py
|
2 |
+
import cv2
|
3 |
+
import numpy as np
|
4 |
+
from ultralytics import YOLO
|
5 |
+
import os
|
6 |
+
|
7 |
+
# Load YOLOv8m model
|
8 |
+
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
9 |
+
MODEL_PATH = os.path.join(BASE_DIR, "../models/yolov8m.pt")
|
10 |
+
model = YOLO(MODEL_PATH)
|
11 |
+
|
12 |
+
def detect_objects(frame):
|
13 |
+
"""
|
14 |
+
Detect objects in a frame using YOLOv8m.
|
15 |
+
Args:
|
16 |
+
frame: Input frame (numpy array)
|
17 |
+
Returns:
|
18 |
+
dict: Detection results with numbered labels
|
19 |
+
numpy array: Annotated frame
|
20 |
+
"""
|
21 |
+
results = model(frame)
|
22 |
+
|
23 |
+
detections = []
|
24 |
+
line_counter = 1 # Initialize counter for numbered labels
|
25 |
+
|
26 |
+
for r in results:
|
27 |
+
for box in r.boxes:
|
28 |
+
conf = float(box.conf[0])
|
29 |
+
if conf < 0.5:
|
30 |
+
continue
|
31 |
+
cls = int(box.cls[0])
|
32 |
+
label = model.names[cls]
|
33 |
+
xyxy = box.xyxy[0].cpu().numpy()
|
34 |
+
x_min, y_min, x_max, y_max = map(int, xyxy)
|
35 |
+
|
36 |
+
# Add numbered label
|
37 |
+
detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
|
38 |
+
detections.append({
|
39 |
+
"label": detection_label,
|
40 |
+
"confidence": conf,
|
41 |
+
"coordinates": [x_min, y_min, x_max, y_max]
|
42 |
+
})
|
43 |
+
|
44 |
+
# Draw bounding box and label
|
45 |
+
color = (0, 255, 0) # Green for generic objects
|
46 |
+
cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
|
47 |
+
cv2.putText(frame, detection_label, (x_min, y_min - 10),
|
48 |
+
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
49 |
+
|
50 |
+
line_counter += 1
|
51 |
+
|
52 |
+
return {"detections": detections, "frame": frame}
|
53 |
+
|
54 |
+
def process_frame(frame):
|
55 |
+
"""
|
56 |
+
Wrapper function for integration with app.py.
|
57 |
+
"""
|
58 |
+
result = detect_objects(frame)
|
59 |
+
return result["detections"], result["frame"]
|