""" solution.py — AID 728 Traffic Rule Violation Detection ======================================================= Pipeline: 1. YOLOv8s (COCO) + custom bike detector → bike boxes + person boxes 2. Depth-Anything V2 (fp16) → depth map for person→bike association 3. Helmet classifier (YOLO) → helmet / no-helmet per rider 4. license.pt (YOLO) → license plate bounding box 5. PaddleOCR 3.5.0 (mobile det+rec) → plate text via legacy ocr() API """ import os import re from pathlib import Path # Point paddlex to bundled offline models BEFORE any paddle import. _MODEL_DIR = Path(__file__).parent / "models" os.environ["PADDLE_PDX_CACHE_HOME"] = str(_MODEL_DIR / "paddleocr") import cv2 import numpy as np import torch from PIL import Image from transformers import pipeline as hf_pipeline from ultralytics import YOLO from paddleocr import PaddleOCR from inference_sdk import InferenceHTTPClient # ── CONSTANTS ───────────────────────────────────────────────────────────────── COCO_PERSON = 0 COCO_CAR = 2 COCO_MOTO = 3 COCO_TRUCK = 7 COCO_CONF = 0.30; COCO_IOU = 0.45 S1_CONF = 0.344; S1_IOU = 0.45 S3_CONF = 0.25; S3_IOU = 0.60 S4_CONF = 0.20 PERSON_BIKE_IOU_THRESH = 0.10 PERSON_BIKE_COL_MARGIN = 0.35 HEAD_CROP_FRACTION = 0.45 HEAD_CROP_MIN_PX = 40 DEPTH_THRESHOLD = 0.35 OCR_MIN_CONF = 0.25 S5_CONF = 0.25 # Wrong side S6_CONF = 0.25 # Seatbelt class TrafficViolationDetector: """ Detects traffic violations on two-wheelers in a single RGB image. All models loaded once in __init__; predict() is fully stateless. """ def __init__(self, model_dir: str = "./models"): md = Path(model_dir) # Ensure paddlex finds bundled offline models os.environ["PADDLE_PDX_CACHE_HOME"] = str(md / "paddleocr") # 1. Depth estimation — model stored as fp16 on disk (47 MB vs 95 MB), # but loaded as fp32 at runtime for fast CPU inference. self.depth_estimator = hf_pipeline( "depth-estimation", model=str(md / "depth_anything_v2"), device=0 if torch.cuda.is_available() else -1, dtype=torch.float32, ) # 2. YOLO models self.s_coco = YOLO(str(md / "yolov8s.pt")) self.s1 = YOLO(str(md / "stage1_best.pt")) self.s3 = YOLO(str(md / "helmet_v11.pt")) self.s4 = YOLO(str(md / "license.pt")) # Optional new models (Roboflow API) try: self.rf_client = InferenceHTTPClient( api_url="https://serverless.roboflow.com", api_key="GkpOFsIColDSm6fSnmCE" ) except Exception: self.rf_client = None # 3. Super-resolution (optional — falls back gracefully if missing) self.sr_engine, self.has_sr = self._init_sr(md / "FSRCNN_x3.pb") # 4. PaddleOCR 3.5.0 — mobile det + rec pipeline. # Uses PP-OCRv5_mobile_det (4.7 MB) + en_PP-OCRv5_mobile_rec (7.6 MB). # IMPORTANT: Must use the legacy .ocr() API, NOT .predict(). # The .predict() path triggers an OneDNN fused_conv2d crash on Windows, # but .ocr() uses a compatible inference path that works everywhere. self.ocr_engine = PaddleOCR( lang="en", device="cpu", enable_mkldnn=False, text_detection_model_name="PP-OCRv5_mobile_det", text_recognition_model_name="en_PP-OCRv5_mobile_rec", ) # ── helpers ─────────────────────────────────────────────────────────────── @staticmethod def _init_sr(sr_path): try: sr = cv2.dnn_superres.DnnSuperResImpl_create() except AttributeError: return None, False if Path(sr_path).exists(): try: sr.readModel(str(sr_path)) sr.setModel("fsrcnn", 3) return sr, True except Exception: pass return sr, False @staticmethod def _box_iou(a, b): ax1, ay1, ax2, ay2 = a bx1, by1, bx2, by2 = b ix1 = max(ax1, bx1); iy1 = max(ay1, by1) ix2 = min(ax2, bx2); iy2 = min(ay2, by2) inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1) if inter == 0: return 0.0 return inter / ((ax2-ax1)*(ay2-ay1) + (bx2-bx1)*(by2-by1) - inter + 1e-6) @staticmethod def _region_depth(depth_map, x1, y1, x2, y2): h, w = depth_map.shape x1, y1 = max(0, int(x1)), max(0, int(y1)) x2, y2 = min(w, int(x2)), min(h, int(y2)) patch = depth_map[y1:y2, x1:x2] return float(np.median(patch)) if patch.size > 0 else 0.5 def _is_depth_ok(self, pd, bd): if bd < 0.05: return abs(pd - bd) <= DEPTH_THRESHOLD * 0.5 return abs(pd - bd) / (bd + 1e-6) <= DEPTH_THRESHOLD def _merge_bike_boxes(self, coco, custom, iou_thresh=0.45): if not coco and not custom: return np.zeros((0, 4), dtype=np.float32) if not coco: return np.array(custom, dtype=np.float32) if not custom: return np.array(coco, dtype=np.float32) merged = list(coco) for cb in custom: if not any(self._box_iou(cb, mb) > iou_thresh for mb in merged): merged.append(cb) return np.array(merged, dtype=np.float32) def _associate_persons_to_bikes(self, person_boxes, bike_boxes, depth_map, h, w): bike_persons = [[] for _ in range(len(bike_boxes))] for p_box in person_boxes: px1, py1, px2, py2 = p_box p_cx = (px1 + px2) / 2 p_bottom = py2 best_bike, best_score = -1, -1.0 for b_idx, b_box in enumerate(bike_boxes): bx1, by1, bx2, by2 = b_box bw = bx2 - bx1 iou = self._box_iou(p_box, b_box) in_col = ( bx1 - PERSON_BIKE_COL_MARGIN * bw <= p_cx <= bx2 + PERSON_BIKE_COL_MARGIN * bw and p_bottom <= by2 + 0.3 * (by2 - by1) ) if iou < PERSON_BIKE_IOU_THRESH and not in_col: continue pd_val = self._region_depth(depth_map, px1, py1, px2, py2) bd_val = self._region_depth(depth_map, bx1, by1, bx2, by2) if not self._is_depth_ok(pd_val, bd_val): continue score = iou + 0.5 * (1.0 - abs(p_cx - (bx1 + bx2) / 2) / (w + 1e-6)) if score > best_score: best_score, best_bike = score, b_idx if best_bike >= 0: bike_persons[best_bike].append(p_box) return bike_persons def _get_depth_map(self, image_cv): img_rgb = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB) result = self.depth_estimator(Image.fromarray(img_rgb)) depth = np.array(result["depth"]).astype(np.float32) lo, hi = depth.min(), depth.max() depth = (depth - lo) / (hi - lo + 1e-8) if depth.shape != image_cv.shape[:2]: depth = cv2.resize(depth, (image_cv.shape[1], image_cv.shape[0])) return depth def _classify_helmets(self, full_image, person_boxes): if not person_boxes: return 0, 0, 0 h_img, w_img = full_image.shape[:2] with_h = without_h = 0 for p_box in person_boxes: px1, py1, px2, py2 = map(int, p_box) head_h = max(int((py2 - py1) * HEAD_CROP_FRACTION), HEAD_CROP_MIN_PX) pad_x = max(4, int((px2 - px1) * 0.05)) crop = full_image[max(0, py1):min(h_img, py1 + head_h), max(0, px1 - pad_x):min(w_img, px2 + pad_x)] if crop.size == 0: without_h += 1 continue res = self.s3.predict(crop, conf=S3_CONF, iou=S3_IOU, verbose=False)[0] if len(res.boxes) == 0: without_h += 1 elif int(res.boxes[res.boxes.conf.argmax()].cls) == 0: with_h += 1 else: without_h += 1 return with_h + without_h, with_h, without_h def _preprocess_plate(self, plate_img): """Upscale and sharpen plate crop before OCR.""" h, w = plate_img.shape[:2] if self.has_sr and self.sr_engine is not None: try: plate_img = self.sr_engine.upsample(plate_img) except Exception: plate_img = cv2.resize(plate_img, (0, 0), fx=3, fy=3, interpolation=cv2.INTER_CUBIC) else: if h < 100: scale = 100 / h plate_img = cv2.resize(plate_img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_CUBIC) lab = cv2.cvtColor(plate_img, cv2.COLOR_BGR2LAB) l, a, b = cv2.split(lab) l = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(4, 4)).apply(l) plate_img = cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2BGR) return cv2.filter2D(plate_img, -1, np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])) def _run_ocr(self, plate_img): """ Full det+rec OCR on the plate crop using the legacy .ocr() API. PaddleOCR 3.5.0's .ocr() wraps .predict() but uses a compatible inference path that works on both Windows and Linux. The result is a list of dicts with 'rec_texts' and 'rec_scores' keys. """ processed = self._preprocess_plate(plate_img) texts, scores = [], [] try: result = self.ocr_engine.ocr(processed) if result and isinstance(result, list): for page in result: if isinstance(page, dict): # paddleocr 3.5.0 format: dict with rec_texts/rec_scores page_texts = page.get("rec_texts", []) page_scores = page.get("rec_scores", []) for t, s in zip(page_texts, page_scores): if str(t).strip(): texts.append(str(t).strip()) scores.append(float(s)) elif isinstance(page, list): # Legacy format: [[box, (text, score)], ...] for line in page: if isinstance(line, (list, tuple)) and len(line) == 2: try: txt = str(line[1][0]) score = float(line[1][1]) if txt.strip(): texts.append(txt.strip()) scores.append(score) except (TypeError, ValueError, IndexError): pass except Exception: pass if not texts: return "UNKNOWN", 0.0 return " ".join(texts), (sum(scores) / len(scores) if scores else 0.0) def _extract_plate(self, vehicle_crop, plate_box): """Crop plate from vehicle ROI, run OCR, return cleaned text.""" h, w = vehicle_crop.shape[:2] pad = 4 x1 = max(0, int(plate_box[0]) - pad) y1 = max(0, int(plate_box[1]) - pad) x2 = min(w, int(plate_box[2]) + pad) y2 = min(h, int(plate_box[3]) + pad) crop = vehicle_crop[y1:y2, x1:x2] if crop.size == 0: return "UNKNOWN" raw, conf = self._run_ocr(crop) if conf < OCR_MIN_CONF: return "UNKNOWN" text = re.sub(r"[^A-Z0-9 \-]", "", raw.upper()) text = re.sub(r"\s+", " ", text).strip() tokens = [t for t in text.split() if len(t) > 1] return " ".join(tokens) if tokens else "UNKNOWN" # ── predict ─────────────────────────────────────────────────────────────── def predict(self, image_path: str) -> dict: """ Run the full violation-detection pipeline on one image. Returns: { "violations": [ { "type": str, # "motorcycle" or "car/truck" "num_riders": int, # for bikes "helmet_violations": int, # for bikes "no_seatbelt": int, # for cars "wrong_side": bool, # for all "license_plate": str, "box": list # [x1, y1, x2, y2] }, ... ] } """ try: img = cv2.imread(str(image_path)) if img is None: return {"violations": []} h_img, w_img = img.shape[:2] # Stage 1: COCO primary detection coco_res = self.s_coco.predict(img, conf=COCO_CONF, iou=COCO_IOU, verbose=False)[0] coco_boxes = coco_res.boxes.xyxy.cpu().numpy() coco_cls = coco_res.boxes.cls.cpu().numpy().astype(int) person_boxes = coco_boxes[coco_cls == COCO_PERSON].tolist() coco_motos = coco_boxes[coco_cls == COCO_MOTO].tolist() coco_cars = coco_boxes[(coco_cls == COCO_CAR) | (coco_cls == COCO_TRUCK)].tolist() # Stage 2: Supplemental bike detector s1_res = self.s1.predict(img, conf=S1_CONF, iou=S1_IOU, augment=True, verbose=False)[0] custom_bikes = s1_res.boxes.xyxy.cpu().numpy().tolist() bike_boxes = self._merge_bike_boxes(coco_motos, custom_bikes) # Stage 2.5: New models (Wrong Side & Seatbelt via API) wrong_side_boxes = [] if self.rf_client is not None: try: ws_res = self.rf_client.infer(str(image_path), model_id="wrong-way-driving-detection-gqdmg/1") for p in ws_res.get("predictions", []): cls_name = p.get("class", "").lower() if "wrong" in cls_name or p.get("confidence", 0) > S5_CONF: cx, cy, w, h = p["x"], p["y"], p["width"], p["height"] wrong_side_boxes.append([cx - w/2, cy - h/2, cx + w/2, cy + h/2]) except Exception as e: pass seatbelt_viol_boxes = [] if len(bike_boxes) == 0 and len(coco_cars) == 0: return {"violations": []} # Stage 3: Depth map for spatial person→bike association depth_map = self._get_depth_map(img) # Stage 4: Associate persons to bikes bike_persons = self._associate_persons_to_bikes( person_boxes, bike_boxes, depth_map, h_img, w_img) # Stage 5-7: Per-bike helmet + plate + violation logic violations = [] # Combine all vehicles to process: (is_bike, box) all_vehicles = [(True, b, i) for i, b in enumerate(bike_boxes)] + [(False, c, -1) for c in coco_cars] for is_bike, v_box, b_idx in all_vehicles: x1, y1, x2, y2 = map(int, v_box) # Check Wrong Side is_wrong_side = False for wb in wrong_side_boxes: if self._box_iou(v_box, wb) > 0.3: is_wrong_side = True break num_riders, with_h, without_h = 0, 0, 0 no_seatbelt_count = 0 if is_bike: num_riders, with_h, without_h = self._classify_helmets(img, bike_persons[b_idx]) if num_riders == 0: num_riders, with_h, without_h = 1, 0, 1 else: # Check seatbelt violations for cars for sb in seatbelt_viol_boxes: sb_cx = (sb[0] + sb[2]) / 2 sb_cy = (sb[1] + sb[3]) / 2 # If the seatbelt violation box center is inside the car box if x1 <= sb_cx <= x2 and y1 <= sb_cy <= y2: no_seatbelt_count += 1 # Expand box slightly to capture plate at bottom bw, bh = x2 - x1, y2 - y1 vcrop = img[ max(0, int(y1 - 0.20 * bh)): min(h_img, int(y2 + 0.10 * bh)), max(0, int(x1 - 0.15 * bw)): min(w_img, int(x2 + 0.15 * bw)) ] # Violation logic is_violation = False if is_bike and ((num_riders >= 3) or (without_h > 0) or is_wrong_side): is_violation = True if not is_bike and (no_seatbelt_count > 0 or is_wrong_side): is_violation = True if is_violation: plate_text = "UNKNOWN" try: if vcrop.size > 0: p_res = self.s4.predict(vcrop, conf=S4_CONF, verbose=False)[0] if len(p_res.boxes) > 0: best_pb = p_res.boxes.xyxy.cpu().numpy()[p_res.boxes.conf.argmax()] plate_text = self._extract_plate(vcrop, best_pb) except Exception: plate_text = "UNKNOWN" violations.append({ "type": "motorcycle" if is_bike else "car/truck", "num_riders": num_riders, "helmet_violations": without_h, "no_seatbelt": no_seatbelt_count, "wrong_side": is_wrong_side, "license_plate": plate_text, "box": [int(x1), int(y1), int(x2), int(y2)] }) return {"violations": violations} except Exception as e: print(f"[ERROR] predict() failed for {image_path}: {e}") return {"violations": []}