File size: 1,997 Bytes
948a00c
 
 
 
93aebe8
 
 
 
 
 
 
948a00c
 
cfb76eb
93aebe8
 
 
 
 
 
948a00c
7d8d543
93aebe8
 
 
 
 
 
 
 
 
948a00c
 
93aebe8
948a00c
 
 
 
 
 
 
 
93aebe8
948a00c
93aebe8
 
948a00c
7d8d543
948a00c
7d8d543
948a00c
 
 
 
 
93aebe8
948a00c
 
93aebe8
948a00c
 
 
93aebe8
7d8d543
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
import cv2
import numpy as np
from ultralytics import YOLO
import os
import logging

logging.basicConfig(
    filename="app.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH = os.path.abspath(os.path.join(BASE_DIR, "../../models/yolov8m.pt"))
try:
    model = YOLO(MODEL_PATH)
    logging.info("Loaded YOLOv8m model for lighting check.")
except Exception as e:
    logging.error(f"Failed to load YOLOv8m model: {str(e)}")
    model = None

def process_lighting(frame):
    if model is None:
        logging.error("YOLO model not loaded. Skipping lighting check.")
        return [], frame

    try:
        results = model(frame)
    except Exception as e:
        logging.error(f"Error during YOLO inference: {str(e)}")
        return [], frame

    detections = []
    line_counter = 1

    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 != "lighting":
                continue
            xyxy = box.xyxy[0].cpu().numpy().astype(int)
            x_min, y_min, x_max, y_max = xyxy

            detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
            detections.append({
                "type": label,
                "label": detection_label,
                "confidence": conf,
                "coordinates": [x_min, y_min, x_max, y_max]
            })

            color = (0, 0, 255)  # Blue as requested
            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.6, color, 2)
            
            line_counter += 1

    logging.info(f"Detected {len(detections)} lighting fixtures in road_safety.")
    return detections, frame