Spaces:
Runtime error
Runtime error
File size: 2,246 Bytes
6bc2b8b 8f17a27 |
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 |
# services/operations_maintenance/signage_check.py
import cv2
import numpy as np
from ultralytics import YOLO
import pytesseract
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_signages(frame):
"""
Detect signages in a frame using YOLOv8m and extract text with Tesseract.
Args:
frame: Input frame (numpy array)
Returns:
dict: Detection results with numbered labels
numpy array: Annotated frame
"""
# Run YOLOv8 inference
results = model(frame)
detections = []
line_counter = 1 # Initialize counter for numbered labels
# Process detections
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]
if label != "signage": # Assuming "signage" class exists
continue
xyxy = box.xyxy[0].cpu().numpy()
x_min, y_min, x_max, y_max = map(int, xyxy)
# Extract text from signage
signage_roi = frame[y_min:y_max, x_min:x_max]
text = pytesseract.image_to_string(signage_roi).strip()
# Add numbered label
detection_label = f"Line {line_counter} - Signage: {text} (Conf: {conf:.2f})"
detections.append({
"label": detection_label,
"confidence": conf,
"coordinates": [x_min, y_min, x_max, y_max],
"text": text
})
# Draw bounding box and label
color = (0, 255, 255) # Yellow for signages
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_signages(frame):
"""
Wrapper function for integration with app.py.
"""
result = detect_signages(frame)
return result["detections"], result["frame"]
|