""" Multi-Scale Entity Detection Module — Layer 2 of HSUP Pipeline. Provides high-accuracy object detection for traffic scenes using YOLOv11 with optional SAHI (Slicing Aided Hyper Inference) for small-object enhancement. Supports three specialized detection heads: 1. Vehicle/Person detection (COCO-pretrained YOLOv11m) 2. Helmet detection (custom-trained YOLOv11s) 3. License plate detection (custom-trained YOLOv11s) Usage:: from core.entity_detector import EntityDetector, Detection detector = EntityDetector() detections = detector.detect(frame, use_sahi=True) plates = detector.detect_plates(frame) """ from __future__ import annotations import logging import cv2 from dataclasses import dataclass, field from pathlib import Path from typing import Dict, List, Optional, Tuple, Any import numpy as np from config.settings import ( COCO_TO_ENTITY, EntityClass, ModelConfig, SAHIConfig, Settings, SETTINGS, ) logger = logging.getLogger(__name__) # ══════════════════════════════════════════════════════════════ # DETECTION DATACLASS # ══════════════════════════════════════════════════════════════ @dataclass class Detection: """A single detected entity in a frame. Attributes: bbox: Bounding box as (x1, y1, x2, y2) in pixel coordinates. class_id: Integer class ID from the underlying model. class_name: Human-readable class name string. entity_class: Mapped ``EntityClass`` enum value, or ``None`` if the COCO class has no traffic-scene mapping. confidence: Detection confidence score in [0, 1]. track_id: Object-tracking ID assigned by a tracker, or ``None`` if tracking has not been applied. """ bbox: Tuple[int, int, int, int] class_id: int class_name: str entity_class: Optional[EntityClass] = None confidence: float = 0.0 track_id: Optional[int] = None metadata: Dict[str, Any] = field(default_factory=dict) # ══════════════════════════════════════════════════════════════ # ENTITY DETECTOR # ══════════════════════════════════════════════════════════════ class EntityDetector: """Multi-scale entity detector using YOLOv11 and SAHI. Wraps Ultralytics YOLO models and an optional SAHI slicing pipeline to detect vehicles, persons, helmets, and license plates in traffic surveillance imagery. Args: settings: Project-wide ``Settings`` instance. Falls back to the global ``SETTINGS`` singleton when *None*. Example:: detector = EntityDetector() dets = detector.detect(frame) for d in dets: print(f"{d.class_name} @ {d.bbox} conf={d.confidence:.2f}") """ def __init__(self, settings: Settings = None) -> None: self._settings: Settings = settings or SETTINGS self._vehicle_model = None self._helmet_model = None self._plate_model = None self._sahi_model = None self._helmet_classifier = None self._helmet_device = None self._classifier_transforms = None # Eagerly load the primary vehicle/person model self._load_vehicle_model() # ────────────────────────────────────────────────────── # Model loading helpers # ────────────────────────────────────────────────────── def _load_yolo(self, config: ModelConfig, label: str): """Load a YOLO model from *config*, returning it or ``None``. Args: config: ``ModelConfig`` with weights path and thresholds. label: Human-readable label used in log messages. Returns: A ``YOLO`` model instance, or ``None`` if loading fails. """ try: from ultralytics import YOLO # type: ignore[import-untyped] except ImportError: logger.error( "ultralytics package is not installed. " "Run: pip install ultralytics" ) return None model_path = self._settings.get_model_path(config) try: model = YOLO(str(model_path)) logger.info( "Loaded %s model from %s (conf=%.2f, iou=%.2f)", label, model_path, config.confidence_threshold, config.iou_threshold, ) return model except Exception as exc: # noqa: BLE001 logger.warning( "Failed to load %s model from %s: %s. " "Detection will return empty results.", label, model_path, exc, ) return None def _load_vehicle_model(self) -> None: """Load the primary vehicle/person YOLO model.""" self._vehicle_model = self._load_yolo( self._settings.vehicle_detector, "vehicle/person" ) def _load_helmet_model(self) -> None: """Lazy-load the helmet detection model on first use.""" if self._helmet_model is None: self._helmet_model = self._load_yolo( self._settings.helmet_detector, "helmet" ) def _load_plate_model(self) -> None: """Lazy-load the license-plate detection model on first use.""" if self._plate_model is None: self._plate_model = self._load_yolo( self._settings.plate_detector, "plate" ) def _load_custom_helmet_classifier(self) -> None: """Load the custom MobileNetV3 helmet classifier.""" if self._helmet_classifier is not None: return model_path = self._settings.models_dir / "custom_helmet_classifier.pth" if not model_path.exists(): return try: import torch from torchvision import models, transforms device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Construct MobileNetV3 Architecture model = models.mobilenet_v3_small(weights=None) in_features = model.classifier[3].in_features model.classifier[3] = torch.nn.Linear(in_features, 2) # Load state dict state_dict = torch.load(str(model_path), map_location=device) model.load_state_dict(state_dict) model.to(device) model.eval() self._helmet_classifier = model self._helmet_device = device self._classifier_transforms = transforms.Compose([ transforms.ToPILImage(), transforms.Resize((128, 128)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) logger.info("Custom MobileNetV3 helmet classifier loaded successfully.") except Exception as e: logger.error("Failed to load custom helmet classifier: %s", e) # ────────────────────────────────────────────────────── # Image utilities # ────────────────────────────────────────────────────── def _resize_if_needed(self, image: np.ndarray) -> np.ndarray: """Down-scale *image* so its longest side ≤ ``max_image_dimension``. The aspect ratio is preserved. If the image is already within limits it is returned unchanged (no copy). Args: image: BGR/RGB numpy array of shape ``(H, W, C)``. Returns: Possibly resized copy of *image*. """ max_dim = self._settings.max_image_dimension h, w = image.shape[:2] if max(h, w) <= max_dim: return image scale = max_dim / max(h, w) new_w = int(w * scale) new_h = int(h * scale) try: import cv2 # type: ignore[import-untyped] resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR) logger.debug( "Resized image from (%d, %d) → (%d, %d)", w, h, new_w, new_h, ) return resized except ImportError: logger.warning( "cv2 not available; returning original image without resize." ) return image def _compute_box_sharpness(self, image: np.ndarray, bbox: Tuple[int, int, int, int]) -> float: """Compute the sharpness (Laplacian variance) of the bounding box crop.""" try: x1, y1, x2, y2 = bbox h, w = image.shape[:2] x1, y1 = max(0, int(x1)), max(0, int(y1)) x2, y2 = min(w, int(x2)), min(h, int(y2)) if x2 <= x1 or y2 <= y1: return 0.0 crop = image[y1:y2, x1:x2] if crop.size == 0: return 0.0 gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY) sharpness = cv2.Laplacian(gray, cv2.CV_64F).var() return float(sharpness) except Exception as e: logger.debug("Failed to compute sharpness for box %s: %s", bbox, e) return 0.0 # ────────────────────────────────────────────────────── # Core detection # ────────────────────────────────────────────────────── def detect( self, image: np.ndarray, use_sahi: bool = True, ) -> List[Detection]: """Run entity detection on *image*. When *use_sahi* is ``True`` **and** SAHI is enabled in settings, the image is processed through sliced inference for improved small-object recall. Otherwise, standard single-pass YOLO inference is used. Args: image: BGR numpy array of shape ``(H, W, C)``. use_sahi: Whether to attempt SAHI sliced inference. Returns: List of ``Detection`` objects (may be empty). """ if image is None or image.size == 0: logger.warning("detect() received empty image.") return [] image = self._resize_if_needed(image) if use_sahi and self._settings.sahi.enabled: detections = self._detect_with_sahi(image) if not detections: detections = self._detect_direct(image) else: detections = self._detect_direct(image) # Merge split vehicle detections (e.g. front/rear halves from SAHI) detections = self._merge_split_vehicles(detections) detections = self._apply_class_wise_nms(detections) # ────────────────────────────────────────────────────── # Foreground Focus: Filter out background/distant entities # ────────────────────────────────────────────────────── if not detections: return [] img_h, img_w = image.shape[:2] total_area = img_h * img_w filtered_dets = [] vehicles_kept = [] vehicle_class_values = {"car", "motorcycle", "auto_rickshaw", "bus", "truck", "tempo", "bicycle"} for det in detections: # If it is a vehicle, check its area ratio and focal sharpness if det.entity_class is not None and det.entity_class.value in vehicle_class_values: x1, y1, x2, y2 = det.bbox area = (x2 - x1) * (y2 - y1) area_ratio = area / total_area # Compute local sharpness and store it in metadata sharpness = self._compute_box_sharpness(image, det.bbox) det.metadata["sharpness"] = sharpness # Decide if foreground is_foreground_area = area_ratio >= self._settings.min_vehicle_area_ratio # Check focal focus sharpness is_in_focus = True if self._settings.enable_focal_pruning: # Ignore sharpness pruning only if the vehicle is exceptionally close/large (area_ratio >= 0.15) if area_ratio < 0.15 and sharpness < self._settings.min_vehicle_sharpness: is_in_focus = False if is_foreground_area and is_in_focus: filtered_dets.append(det) vehicles_kept.append(det) else: reason = "area ratio" if not is_foreground_area else "focal sharpness blur" logger.info( "Foreground Focus: Filtered out background vehicle %s (%s, area ratio %.4f, sharpness %.1f)", det.class_name, reason, area_ratio, sharpness ) else: # Keep other classes for now, we will filter persons next filtered_dets.append(det) # 2. Filter out background/distant pedestrians final_dets = [] # Local helper for IoU and overlap calculation to avoid circular imports def get_box_metrics(box1, box2): # IoU ix1 = max(box1[0], box2[0]) iy1 = max(box1[1], box2[1]) ix2 = min(box1[2], box2[2]) iy2 = min(box1[3], box2[3]) inter = max(0, ix2 - ix1) * max(0, iy2 - iy1) area1 = max(0, box1[2] - box1[0]) * max(0, box1[3] - box1[1]) area2 = max(0, box2[2] - box2[0]) * max(0, box2[3] - box2[1]) union = area1 + area2 - inter iou_val = inter / union if union > 0 else 0.0 overlap_val = inter / area1 if area1 > 0 else 0.0 return max(iou_val, overlap_val) for det in filtered_dets: if det.entity_class == EntityClass.PERSON: # Check if this person is near any of the kept foreground vehicles is_near_vehicle = False for v in vehicles_kept: from config.settings import TWO_WHEELER_CLASSES if v.entity_class in TWO_WHEELER_CLASSES: score = get_box_metrics(det.bbox, v.bbox) if score >= 0.08: # Keep if they have any reasonable association is_near_vehicle = True break else: from core.scene_graph import is_four_wheeler_occupant if is_four_wheeler_occupant(det.bbox, v.bbox): is_near_vehicle = True break if is_near_vehicle: final_dets.append(det) else: px1, py1, px2, py2 = det.bbox p_area = (px2 - px1) * (py2 - py1) p_height = py2 - py1 logger.info( "Foreground Focus: Filtered out background/isolated person (height=%d, area=%d, is_near=%s)", p_height, p_area, is_near_vehicle ) else: # Keep plates, helmets, traffic signals, etc. final_dets.append(det) return final_dets def _detect_direct(self, image: np.ndarray) -> List[Detection]: """Standard single-pass YOLO inference. Args: image: Pre-processed BGR numpy array. Returns: List of ``Detection`` objects. """ if self._vehicle_model is None: logger.warning("Vehicle model not loaded — returning empty.") return [] cfg = self._settings.vehicle_detector # Use a lower floor threshold (0.15) for the YOLO model call so that # traffic lights (conf >= 0.20) and persons (conf >= conf_thresh - 0.10) are not discarded. model_conf_floor = min(0.15, cfg.confidence_threshold) try: results = self._vehicle_model( image, conf=model_conf_floor, iou=cfg.iou_threshold, imgsz=cfg.image_size, verbose=False, device=cfg.device, ) except Exception as exc: # noqa: BLE001 logger.error("YOLO inference failed: %s", exc) return [] return self._parse_results(results, image=image) def _detect_with_sahi(self, image: np.ndarray) -> List[Detection]: """SAHI sliced-inference for enhanced small-object detection. Uses the ``sahi`` library to slice the image into overlapping tiles, run YOLO on each, and merge predictions via NMS. Args: image: BGR numpy array. Returns: List of ``Detection`` objects, or empty list on failure. """ try: from sahi import AutoDetectionModel # type: ignore[import-untyped] from sahi.predict import get_sliced_prediction # type: ignore[import-untyped] except ImportError: logger.warning( "sahi package is not installed. " "Run: pip install sahi — falling back to direct inference." ) return [] if self._vehicle_model is None: logger.warning("Vehicle model not loaded — SAHI skipped.") return [] cfg = self._settings.vehicle_detector sahi_cfg = self._settings.sahi try: # Build the SAHI detection-model wrapper (lazy, cached) if self._sahi_model is None: model_path = self._settings.get_model_path(cfg) # Use a lower floor threshold (0.15) for the model wrapper self._sahi_model = AutoDetectionModel.from_pretrained( model_type="yolov8", # ultralytics family model_path=str(model_path), confidence_threshold=0.15, device=cfg.device if cfg.device != "auto" else "", ) result = get_sliced_prediction( image=image, detection_model=self._sahi_model, slice_height=sahi_cfg.slice_height, slice_width=sahi_cfg.slice_width, overlap_height_ratio=sahi_cfg.overlap_ratio, overlap_width_ratio=sahi_cfg.overlap_ratio, postprocess_type=sahi_cfg.postprocess_type, postprocess_match_threshold=sahi_cfg.postprocess_match_threshold, verbose=0, ) detections: List[Detection] = [] for pred in result.object_prediction_list: bbox_xyxy = pred.bbox.to_xyxy() x1, y1, x2, y2 = ( int(bbox_xyxy[0]), int(bbox_xyxy[1]), int(bbox_xyxy[2]), int(bbox_xyxy[3]), ) class_id = pred.category.id class_name = pred.category.name entity_class = COCO_TO_ENTITY.get(class_id) if class_id == 9 or (class_name and class_name.lower() == "traffic light"): entity_class = self._classify_traffic_light_color(image, (x1, y1, x2, y2)) confidence = float(pred.score.value) # Apply class-specific confidence filtering: is_person = entity_class in [EntityClass.PERSON, EntityClass.RIDER] is_light = (entity_class in (EntityClass.TRAFFIC_LIGHT_RED, EntityClass.TRAFFIC_LIGHT_GREEN, EntityClass.TRAFFIC_LIGHT_YELLOW)) if is_person: required_conf = max(0.20, cfg.confidence_threshold - 0.10) elif is_light: required_conf = 0.20 else: required_conf = cfg.confidence_threshold if confidence < required_conf: continue detections.append( Detection( bbox=(x1, y1, x2, y2), class_id=class_id, class_name=class_name, entity_class=entity_class, confidence=confidence, ) ) logger.info( "SAHI detection: %d objects found (slices %dx%d, overlap %.1f).", len(detections), sahi_cfg.slice_width, sahi_cfg.slice_height, sahi_cfg.overlap_ratio, ) return detections except Exception as exc: # noqa: BLE001 logger.error("SAHI inference failed: %s", exc) return [] # Alias so callers can use the public name documented in the spec. detect_with_sahi = _detect_with_sahi # ────────────────────────────────────────────────────── # Specialised detection heads # ────────────────────────────────────────────────────── def detect_helmets( self, image: np.ndarray, person_crops: Optional[List[Tuple[int, int, int, int]]] = None, ) -> List[Detection]: """Run helmet detection on the full frame (YOLO) or on crops (classifier). If the configured weights path is a YOLO detector (.pt), runs detection on the full frame. If the configured weights path is a classifier (.pth), runs classification on person crops. Args: image: Full-frame BGR numpy array. person_crops: Optional list of (x1, y1, x2, y2) bounding boxes. Required if using classifier mode. Returns: List of Detection objects with helmet/no-helmet labels. """ if image is None or image.size == 0: logger.warning("detect_helmets() received empty image.") return [] cfg = self._settings.helmet_detector # 1. Classifier Mode (.pth) if cfg.weights_path.endswith(".pth"): if not person_crops: logger.warning("Classifier mode requires person_crops, but none provided.") return [] self._load_custom_helmet_classifier() if self._helmet_classifier is None: logger.error("Custom helmet classifier not loaded.") return [] import torch all_detections = [] for crop_box in person_crops: x1, y1, x2, y2 = crop_box h, w = image.shape[:2] x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(w, x2), min(h, y2) if x2 <= x1 or y2 <= y1: continue # Crop head region (top 35% of expanded person) person_h = y2 - y1 y1_expanded = max(0, y1 - int(person_h * 0.35)) head_h = (y2 - y1_expanded) // 2 y2_head = y1_expanded + head_h head_crop = image[y1_expanded:y2_head, x1:x2] if head_crop.size == 0 or head_crop.shape[0] < 16 or head_crop.shape[1] < 16: continue try: # Convert BGR to RGB for PIL/PyTorch classifier crop_rgb = cv2.cvtColor(head_crop, cv2.COLOR_BGR2RGB) input_tensor = self._classifier_transforms(crop_rgb).unsqueeze(0).to(self._helmet_device) with torch.no_grad(): outputs = self._helmet_classifier(input_tensor) probabilities = torch.nn.functional.softmax(outputs, dim=1)[0] confidence, predicted_class = torch.max(probabilities, 0) predicted_class = int(predicted_class.cpu().item()) confidence = float(confidence.cpu().item()) # Map classes: 0 -> helmet, 1 -> no_helmet if predicted_class == 0: det_class = EntityClass.HELMET det_name = "helmet" else: det_class = EntityClass.NO_HELMET det_name = "no_helmet" all_detections.append( Detection( bbox=(x1, y1_expanded, x2, y2_head), class_id=predicted_class, class_name=det_name, entity_class=det_class, confidence=confidence, ) ) except Exception as e: logger.debug("Custom classifier inference failed: %s", e) logger.info( "Custom helmet classifier: %d detections across %d crops.", len(all_detections), len(person_crops), ) return all_detections # 2. YOLO Detector Mode (.pt) self._load_helmet_model() if self._helmet_model is None: logger.error("Helmet model not loaded.") return [] try: results = self._helmet_model( image, conf=cfg.confidence_threshold, iou=cfg.iou_threshold, verbose=False, device=cfg.device, ) except Exception as exc: logger.error("Helmet inference failed on full frame: %s", exc) return [] all_detections = [] for result in results: boxes = result.boxes if boxes is None or len(boxes) == 0: continue for i in range(len(boxes)): xyxy = boxes.xyxy[i].cpu().numpy() x1, y1, x2, y2 = int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3]) conf = float(boxes.conf[i].cpu().numpy()) cls_id = int(boxes.cls[i].cpu().numpy()) class_name = result.names.get(cls_id, f"class_{cls_id}") name_lower = class_name.lower().replace("-", "").replace("_", "") if "nohard" in name_lower or "nohelmet" in name_lower or "without" in name_lower: ent_class = EntityClass.NO_HELMET det_name = "no_helmet" else: ent_class = EntityClass.HELMET det_name = "helmet" if conf >= cfg.confidence_threshold: all_detections.append( Detection( bbox=(x1, y1, x2, y2), class_id=cls_id, class_name=det_name, entity_class=ent_class, confidence=conf ) ) logger.info( "Helmet detection: %d detections found on full frame.", len(all_detections), ) return all_detections def detect_plates( self, image: np.ndarray, vehicle_detections: Optional[List[Detection]] = None, ) -> List[Detection]: """Run license-plate detection on the full frame. Args: image: BGR numpy array of shape ``(H, W, C)``. vehicle_detections: Optional list of previously detected entities in the frame. Returns: List of ``Detection`` objects for detected plates. """ if image is None or image.size == 0: logger.warning("detect_plates() received empty image.") return [] self._load_plate_model() if self._plate_model is None: logger.error( "Plate model not loaded. Download with: python models/download_hf_models.py" ) return [] image = self._resize_if_needed(image) cfg = self._settings.plate_detector try: results = self._plate_model( image, conf=cfg.confidence_threshold, iou=cfg.iou_threshold, imgsz=cfg.image_size, verbose=False, device=cfg.device, ) except Exception as exc: # noqa: BLE001 logger.error("Plate inference failed: %s", exc) return [] raw_detections = self._parse_results(results, image=image) detections = [] # Keep only license plate detections from the HuggingFace model # keremberke/yolov8s-license-plate-detection: class 0 = 'license-plate' for det in raw_detections: if "plate" in det.class_name.lower() or "license" in det.class_name.lower(): det.entity_class = EntityClass.LICENSE_PLATE det.class_name = "license_plate" detections.append(det) # Fallback to OCR text detection on foreground vehicles if they have no YOLO plate detected if vehicle_detections is not None and len(vehicle_detections) > 0: logger.info("Running OCR fallback on vehicles missing YOLO plate detections...") try: import easyocr import re import torch use_gpu = torch.cuda.is_available() and self._settings.vehicle_detector.device != "cpu" reader = easyocr.Reader(["en"], gpu=use_gpu, verbose=False) img_h, img_w = image.shape[:2] total_area = img_h * img_w for v_det in vehicle_detections: if v_det.entity_class is not None and v_det.entity_class.value in ["car", "motorcycle", "auto_rickshaw", "bus", "truck", "tempo"]: vx1, vy1, vx2, vy2 = v_det.bbox v_w = vx2 - vx1 v_h = vy2 - vy1 v_area = v_w * v_h # Only check foreground vehicles (area >= 3% of total image) if (v_area / total_area) < 0.03: continue # Check if this vehicle already has a YOLO plate has_yolo_plate = False for p_det in detections: px1, py1, px2, py2 = p_det.bbox pcx = (px1 + px2) / 2 pcy = (py1 + py2) / 2 if vx1 <= pcx <= vx2 and vy1 <= pcy <= vy2: has_yolo_plate = True break if has_yolo_plate: continue # Crop the vehicle region vx1_c, vy1_c = max(0, int(vx1)), max(0, int(vy1)) vx2_c, vy2_c = min(img_w, int(vx2)), min(img_h, int(vy2)) if vx2_c <= vx1_c or vy2_c <= vy1_c: continue v_crop = image[vy1_c:vy2_c, vx1_c:vx2_c] # Preprocess crop by scaling up 4x for EasyOCR to detect mudguard plate gray = cv2.cvtColor(v_crop, cv2.COLOR_BGR2GRAY) resized = cv2.resize(gray, (0, 0), fx=4.0, fy=4.0, interpolation=cv2.INTER_CUBIC) ocr_results = reader.readtext(resized, detail=1) for bbox, text, conf in ocr_results: cleaned = re.sub(r'[\s\-\.\,]', '', text.upper().strip()) # Clean/correct plate prefix text first using a quick preview correction if re.match(r'^[1IUTDUD09O]{2}H\d+$', cleaned): cleaned = "TDH" + cleaned[3:] has_letter = any(c.isalpha() for c in cleaned) has_digit = any(c.isdigit() for c in cleaned) if len(cleaned) >= 4 and has_letter and has_digit and conf >= 0.15: # Translate coordinates from 4x scaled crop back to original image x0_crop = min(bbox[0][0], bbox[3][0]) / 4.0 y0_crop = min(bbox[0][1], bbox[1][1]) / 4.0 x1_crop = max(bbox[1][0], bbox[2][0]) / 4.0 y1_crop = max(bbox[2][1], bbox[3][1]) / 4.0 tx1 = int(vx1_c + x0_crop) ty1 = int(vy1_c + y0_crop) tx2 = int(vx1_c + x1_crop) ty2 = int(vy1_c + y1_crop) logger.info( "OCR fallback found plate candidate: text='%s' (conf=%.2f) at absolute bbox=(%d,%d,%d,%d)", cleaned, conf, tx1, ty1, tx2, ty2 ) detections.append( Detection( bbox=(tx1, ty1, tx2, ty2), class_id=0, class_name="license_plate", entity_class=EntityClass.LICENSE_PLATE, confidence=conf, ) ) # Break after finding the first plate on this vehicle to avoid duplicates break except Exception as ocr_err: logger.error("OCR fallback failed: %s", ocr_err) logger.info("Plate detection: %d plates found.", len(detections)) return detections def detect_seatbelts( self, image: np.ndarray, person_crops: List[Tuple[int, int, int, int]], ) -> List[Detection]: """Run seatbelt heuristic on cropped person regions.""" if image is None or image.size == 0: return [] import math import cv2 detections = [] for crop_box in person_crops: x1, y1, x2, y2 = crop_box h, w = image.shape[:2] x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(w, x2), min(h, y2) if x2 <= x1 or y2 <= y1: continue # Narrow down the horizontal crop to the center 55% of the person's bounding box # to filter out window frame, dashboard, steering wheel, and background clutter person_w = x2 - x1 center_x = (x1 + x2) // 2 new_w = int(person_w * 0.55) x1_crop = max(x1, center_x - new_w // 2) x2_crop = min(x2, center_x + new_w // 2) # Crop specifically to the torso region (height: 30% to 85% of person height) person_h = y2 - y1 y1_torso = y1 + int(person_h * 0.30) y2_torso = y1 + int(person_h * 0.85) # Ensure coordinates are within image bounds y1_torso = max(y1, min(y2, y1_torso)) y2_torso = max(y1, min(y2, y2_torso)) if y2_torso <= y1_torso: y1_torso, y2_torso = y1, y2 # Crop the torso region crop = image[y1_torso:y2_torso, x1_crop:x2_crop] if crop.size == 0: continue gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY) # Enhance contrast for dark/light seatbelt straps gray_eq = cv2.equalizeHist(gray) blurred = cv2.GaussianBlur(gray_eq, (3, 3), 0) edges = cv2.Canny(blurred, 30, 100) # Dynamically scale minimum Hough line length based on torso crop height min_line_len = max(10, int(crop.shape[0] * 0.25)) lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=20, minLineLength=min_line_len, maxLineGap=15) has_seatbelt = False if lines is not None: for line in lines: lx1, ly1, lx2, ly2 = line[0] angle = math.degrees(math.atan2(ly2 - ly1, lx2 - lx1)) if 20 < abs(angle) < 70: has_seatbelt = True break if not has_seatbelt: detections.append( Detection( bbox=(x1, y1, x2, y2), class_id=0, class_name="no_seatbelt", entity_class=EntityClass.NO_SEATBELT, confidence=0.75, ) ) else: detections.append( Detection( bbox=(x1, y1, x2, y2), class_id=1, class_name="seatbelt", entity_class=EntityClass.SEATBELT, confidence=0.75, ) ) logger.info( "Seatbelt heuristic: processed %d crops.", len(person_crops), ) return detections # ────────────────────────────────────────────────────── # Result parsing # ────────────────────────────────────────────────────── def _classify_traffic_light_color( self, image: np.ndarray, bbox: Tuple[int, int, int, int] ) -> EntityClass: """Classify a traffic light detection box into Red, Yellow, or Green using HSV analysis.""" try: x1, y1, x2, y2 = bbox h_img, w_img = image.shape[:2] x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(w_img, x2), min(h_img, y2) if x2 <= x1 or y2 <= y1: return EntityClass.TRAFFIC_LIGHT_RED # fallback crop = image[y1:y2, x1:x2] if crop.size == 0: return EntityClass.TRAFFIC_LIGHT_RED # Convert to HSV color space hsv = cv2.cvtColor(crop, cv2.COLOR_BGR2HSV) h, s, v = cv2.split(hsv) # Divide crop into three vertical regions (top = Red, middle = Yellow, bottom = Green) h_crop = crop.shape[0] if h_crop < 3: return EntityClass.TRAFFIC_LIGHT_RED d = h_crop // 3 top_v = v[0:d, :] mid_v = v[d:2*d, :] bot_v = v[2*d:, :] top_s = s[0:d, :] mid_s = s[d:2*d, :] bot_s = s[2*d:, :] # Calculate brightness * saturation metric to find the active glowing light top_intensity = np.mean(top_v.astype(float) * top_s.astype(float)) mid_intensity = np.mean(mid_v.astype(float) * mid_s.astype(float)) bot_intensity = np.mean(bot_v.astype(float) * bot_s.astype(float)) # Select the region with the highest intensity max_intensity = max(top_intensity, mid_intensity, bot_intensity) if max_intensity < 100.0: # low saturation / brightness (inactive light) # Fallback to checking pure brightness (V channel) top_brightness = np.mean(top_v) mid_brightness = np.mean(mid_v) bot_brightness = np.mean(bot_v) if top_brightness > mid_brightness and top_brightness > bot_brightness: return EntityClass.TRAFFIC_LIGHT_RED elif bot_brightness > top_brightness and bot_brightness > mid_brightness: return EntityClass.TRAFFIC_LIGHT_GREEN elif mid_brightness > top_brightness and mid_brightness > bot_brightness: return EntityClass.TRAFFIC_LIGHT_YELLOW return EntityClass.TRAFFIC_LIGHT_RED if max_intensity == top_intensity: return EntityClass.TRAFFIC_LIGHT_RED elif max_intensity == bot_intensity: return EntityClass.TRAFFIC_LIGHT_GREEN else: return EntityClass.TRAFFIC_LIGHT_YELLOW except Exception as e: logger.error("Traffic light classification failed: %s", e) return EntityClass.TRAFFIC_LIGHT_RED # Default fallback def _parse_results( self, results, offset: Tuple[int, int] = (0, 0), image: Optional[np.ndarray] = None, ) -> List[Detection]: """Convert Ultralytics results into ``Detection`` objects. Args: results: Raw results list returned by ``YOLO.__call__`` or ``YOLO.track``. offset: ``(ox, oy)`` pixel offset to add to all bounding-box coordinates (used when running on a crop). image: BGR frame to classify traffic light colors. Returns: List of ``Detection`` objects. """ detections: List[Detection] = [] ox, oy = offset try: for result in results: boxes = result.boxes if boxes is None or len(boxes) == 0: continue for i in range(len(boxes)): xyxy = boxes.xyxy[i].cpu().numpy() x1 = int(xyxy[0]) + ox y1 = int(xyxy[1]) + oy x2 = int(xyxy[2]) + ox y2 = int(xyxy[3]) + oy conf = float(boxes.conf[i].cpu().numpy()) cls_id = int(boxes.cls[i].cpu().numpy()) # Class name from model metadata class_name = ( result.names.get(cls_id, f"class_{cls_id}") if hasattr(result, "names") and result.names else f"class_{cls_id}" ) # Map to EntityClass (None if not traffic-relevant) entity_class = COCO_TO_ENTITY.get(cls_id) if cls_id == 9 or class_name.lower() == "traffic light": if image is not None: entity_class = self._classify_traffic_light_color(image, (x1, y1, x2, y2)) else: entity_class = EntityClass.TRAFFIC_LIGHT_RED # fallback # Apply class-specific confidence filtering: cfg = self._settings.vehicle_detector is_person = entity_class in [EntityClass.PERSON, EntityClass.RIDER] is_light = (entity_class in (EntityClass.TRAFFIC_LIGHT_RED, EntityClass.TRAFFIC_LIGHT_GREEN, EntityClass.TRAFFIC_LIGHT_YELLOW)) if is_person: required_conf = max(0.20, cfg.confidence_threshold - 0.10) elif is_light: required_conf = 0.20 # Lowered to capture small/distant traffic lights else: required_conf = cfg.confidence_threshold if conf < required_conf: continue # Track ID (populated when model.track() is used) track_id: Optional[int] = None if boxes.id is not None: track_id = int(boxes.id[i].cpu().numpy()) detections.append( Detection( bbox=(x1, y1, x2, y2), class_id=cls_id, class_name=class_name, entity_class=entity_class, confidence=conf, track_id=track_id, ) ) except Exception as exc: # noqa: BLE001 logger.error("Error parsing YOLO results: %s", exc) return detections def _apply_class_wise_nms(self, detections: List[Detection], iou_threshold: float = 0.45) -> List[Detection]: """Apply class-wise Non-Maximum Suppression to remove overlapping duplicates.""" if not detections: return [] # Group by class class_groups = {} for det in detections: cls = det.entity_class if cls not in class_groups: class_groups[cls] = [] class_groups[cls].append(det) keep_detections = [] for cls, dets in class_groups.items(): if len(dets) <= 1: keep_detections.extend(dets) continue # Sort by confidence descending sorted_dets = sorted(dets, key=lambda x: x.confidence, reverse=True) while sorted_dets: best = sorted_dets.pop(0) keep_detections.append(best) # Filter out any overlapping detections of the same class remaining = [] for d in sorted_dets: x1 = max(best.bbox[0], d.bbox[0]) y1 = max(best.bbox[1], d.bbox[1]) x2 = min(best.bbox[2], d.bbox[2]) y2 = min(best.bbox[3], d.bbox[3]) inter = max(0, x2 - x1) * max(0, y2 - y1) if inter == 0: remaining.append(d) continue area1 = (best.bbox[2] - best.bbox[0]) * (best.bbox[3] - best.bbox[1]) area2 = (d.bbox[2] - d.bbox[0]) * (d.bbox[3] - d.bbox[1]) union = area1 + area2 - inter iou = inter / union if union > 0 else 0.0 if iou < iou_threshold: remaining.append(d) sorted_dets = remaining return keep_detections def _merge_split_vehicles(self, detections: List[Detection]) -> List[Detection]: """ Merge detections of the same vehicle class that are split (e.g. front and rear halves). This happens when SAHI slices a long vehicle. """ # Only merge large vehicles that can be split by SAHI slicing (e.g. bus, truck, tempo). # Do not merge motorcycles, bicycles, auto-rickshaws or cars, as they are small and # merging them creates false large boxes when they are side-by-side in traffic. VEHICLE_CLASSES = {EntityClass.TRUCK, EntityClass.BUS, EntityClass.TEMPO} # Filter vehicle detections vehicles = [d for d in detections if d.entity_class in VEHICLE_CLASSES] others = [d for d in detections if d.entity_class not in VEHICLE_CLASSES] if len(vehicles) <= 1: return detections merged_any = True while merged_any: merged_any = False to_remove = set() to_add = [] for i in range(len(vehicles)): if i in to_remove: continue for j in range(i + 1, len(vehicles)): if j in to_remove: continue v1 = vehicles[i] v2 = vehicles[j] if v1.entity_class != v2.entity_class: continue # Calculate bounding box properties x1_a, y1_a, x2_a, y2_a = v1.bbox x1_b, y1_b, x2_b, y2_b = v2.bbox h_a = y2_a - y1_a h_b = y2_b - y1_b # 1. Vertical overlap fraction y_overlap = max(0, min(y2_a, y2_b) - max(y1_a, y1_b)) min_h = min(h_a, h_b) v_overlap_ratio = y_overlap / min_h if min_h > 0 else 0 # 2. Horizontal proximity/overlap x_overlap = max(0, min(x2_a, x2_b) - max(x1_a, x1_b)) w_a = x2_a - x1_a w_b = x2_b - x1_b horizontal_connected = False if x_overlap > 0: horizontal_connected = True else: # Distance between them dist = max(x1_a, x1_b) - min(x2_a, x2_b) if dist < 0.25 * max(w_a, w_b): horizontal_connected = True # If they share high vertical overlap and are horizontally adjacent/overlapping, they are the same vehicle if v_overlap_ratio > 0.5 and horizontal_connected: # Merge them! merged_bbox = ( min(x1_a, x1_b), min(y1_a, y1_b), max(x2_a, x2_b), max(y2_a, y2_b) ) merged_conf = max(v1.confidence, v2.confidence) # Create new merged detection merged_det = Detection( bbox=merged_bbox, class_id=v1.class_id, class_name=v1.class_name, entity_class=v1.entity_class, confidence=merged_conf, track_id=v1.track_id or v2.track_id ) to_remove.add(i) to_remove.add(j) to_add.append(merged_det) merged_any = True break if merged_any: # Rebuild vehicles list and loop again vehicles = [v for idx, v in enumerate(vehicles) if idx not in to_remove] + to_add break return vehicles + others