Spaces:
Runtime error
Runtime error
Update services/operations_maintenance/signage_check.py
Browse files
services/operations_maintenance/signage_check.py
CHANGED
@@ -1,8 +1,6 @@
|
|
1 |
-
# services/operations_maintenance/signage_check.py
|
2 |
import cv2
|
3 |
import numpy as np
|
4 |
from ultralytics import YOLO
|
5 |
-
import pytesseract
|
6 |
import os
|
7 |
|
8 |
# Load YOLOv8m model
|
@@ -10,22 +8,20 @@ BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
10 |
MODEL_PATH = os.path.join(BASE_DIR, "../../models/yolov8m.pt")
|
11 |
model = YOLO(MODEL_PATH)
|
12 |
|
13 |
-
def
|
14 |
"""
|
15 |
-
Detect
|
16 |
Args:
|
17 |
frame: Input frame (numpy array)
|
18 |
Returns:
|
19 |
-
|
20 |
numpy array: Annotated frame
|
21 |
"""
|
22 |
-
# Run YOLOv8 inference
|
23 |
results = model(frame)
|
24 |
|
25 |
detections = []
|
26 |
line_counter = 1 # Initialize counter for numbered labels
|
27 |
|
28 |
-
# Process detections
|
29 |
for r in results:
|
30 |
for box in r.boxes:
|
31 |
conf = float(box.conf[0])
|
@@ -38,32 +34,21 @@ def detect_signages(frame):
|
|
38 |
xyxy = box.xyxy[0].cpu().numpy()
|
39 |
x_min, y_min, x_max, y_max = map(int, xyxy)
|
40 |
|
41 |
-
# Extract text from signage
|
42 |
-
signage_roi = frame[y_min:y_max, x_min:x_max]
|
43 |
-
text = pytesseract.image_to_string(signage_roi).strip()
|
44 |
-
|
45 |
# Add numbered label
|
46 |
-
detection_label = f"Line {line_counter} -
|
47 |
detections.append({
|
|
|
48 |
"label": detection_label,
|
49 |
"confidence": conf,
|
50 |
-
"coordinates": [x_min, y_min, x_max, y_max]
|
51 |
-
"text": text
|
52 |
})
|
53 |
|
54 |
# Draw bounding box and label
|
55 |
-
color = (0,
|
56 |
cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
|
57 |
cv2.putText(frame, detection_label, (x_min, y_min - 10),
|
58 |
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
59 |
|
60 |
line_counter += 1
|
61 |
|
62 |
-
return
|
63 |
-
|
64 |
-
def process_signages(frame):
|
65 |
-
"""
|
66 |
-
Wrapper function for integration with app.py.
|
67 |
-
"""
|
68 |
-
result = detect_signages(frame)
|
69 |
-
return result["detections"], result["frame"]
|
|
|
|
|
1 |
import cv2
|
2 |
import numpy as np
|
3 |
from ultralytics import YOLO
|
|
|
4 |
import os
|
5 |
|
6 |
# Load YOLOv8m model
|
|
|
8 |
MODEL_PATH = os.path.join(BASE_DIR, "../../models/yolov8m.pt")
|
9 |
model = YOLO(MODEL_PATH)
|
10 |
|
11 |
+
def process_signages(frame):
|
12 |
"""
|
13 |
+
Detect road signage in a frame.
|
14 |
Args:
|
15 |
frame: Input frame (numpy array)
|
16 |
Returns:
|
17 |
+
list: List of detected signage
|
18 |
numpy array: Annotated frame
|
19 |
"""
|
|
|
20 |
results = model(frame)
|
21 |
|
22 |
detections = []
|
23 |
line_counter = 1 # Initialize counter for numbered labels
|
24 |
|
|
|
25 |
for r in results:
|
26 |
for box in r.boxes:
|
27 |
conf = float(box.conf[0])
|
|
|
34 |
xyxy = box.xyxy[0].cpu().numpy()
|
35 |
x_min, y_min, x_max, y_max = map(int, xyxy)
|
36 |
|
|
|
|
|
|
|
|
|
37 |
# Add numbered label
|
38 |
+
detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
|
39 |
detections.append({
|
40 |
+
"type": label,
|
41 |
"label": detection_label,
|
42 |
"confidence": conf,
|
43 |
+
"coordinates": [x_min, y_min, x_max, y_max]
|
|
|
44 |
})
|
45 |
|
46 |
# Draw bounding box and label
|
47 |
+
color = (0, 191, 255) # DeepSkyBlue for signage
|
48 |
cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
|
49 |
cv2.putText(frame, detection_label, (x_min, y_min - 10),
|
50 |
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
|
51 |
|
52 |
line_counter += 1
|
53 |
|
54 |
+
return detections, frame
|
|
|
|
|
|
|
|
|
|
|
|
|
|