surveillance / services /road_safety /lighting_check.py
lokesh341's picture
Update services/road_safety/lighting_check.py
93aebe8
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