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"]