""" annotations.py — Annotation drawing engine for Gridlock. ========================================================= Extracted from run_inference.py and enhanced with support for: - Illegal parking annotations - Per-violation type badges - Confidence scores - Reusable from both CLI and API contexts """ import cv2 import numpy as np from config import ( CLR_BIKE, CLR_CAR, CLR_RIDER, CLR_NO_HELMET, CLR_PLATE, CLR_VIOLATION, CLR_OK, CLR_WRONGSIDE, CLR_SEATBELT, CLR_PARKING, HEAD_CROP_FRACTION, HEAD_CROP_MIN_PX, ) FONT = cv2.FONT_HERSHEY_SIMPLEX def put_label(img, text, x, y, color, bg_color=None, scale=0.55, thickness=1): """Draw text with an optional filled background pill.""" (tw, th), baseline = cv2.getTextSize(text, FONT, scale, thickness) if bg_color is not None: pad = 4 cv2.rectangle( img, (x - pad, y - th - pad), (x + tw + pad, y + baseline + pad), bg_color, -1, ) cv2.putText(img, text, (x, y), FONT, scale, color, thickness, cv2.LINE_AA) def draw_rounded_rect(img, x1, y1, x2, y2, color, radius=8, thickness=2): """Draw a rounded rectangle.""" cv2.rectangle(img, (x1 + radius, y1), (x2 - radius, y2), color, thickness) cv2.rectangle(img, (x1, y1 + radius), (x2, y2 - radius), color, thickness) cv2.ellipse(img, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color, thickness) cv2.ellipse(img, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color, thickness) cv2.ellipse(img, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color, thickness) cv2.ellipse(img, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, thickness) def annotate_from_pipeline_result(img: np.ndarray, pipeline_result: dict) -> np.ndarray: """ Draw full annotations on a copy of the image using the structured output from ParallelDetectionPipeline.process(return_annotations=True). Args: img: Original BGR image (numpy array). pipeline_result: Dict returned by pipeline.process() with return_annotations=True. Returns: Annotated image (numpy array) with banner. """ out = img.copy() h_img, w_img = img.shape[:2] annot = pipeline_result.get("annotation_data", {}) vehicles = annot.get("vehicles", []) total_violations = 0 total_bikes = pipeline_result.get("vehicles_detected", {}).get("bikes", 0) total_cars = pipeline_result.get("vehicles_detected", {}).get("cars", 0) # ── Draw vehicles ──────────────────────────────────────────────────────── for v in vehicles: is_bike = v["is_bike"] x1, y1, x2, y2 = v["box"] is_violation = v["is_violation"] is_wrong_side = v["is_wrong_side"] num_riders = v["num_riders"] with_h = v["with_helmet"] without_h = v["without_helmet"] no_seatbelt = v["no_seatbelt"] plate_text = v["plate_text"] plate_box_abs = v["plate_box_abs"] rider_boxes = v["rider_boxes"] violation_types = v.get("violation_types", []) if is_violation: total_violations += 1 # Vehicle bounding box vehicle_color = CLR_VIOLATION if is_violation else (CLR_BIKE if is_bike else CLR_CAR) cv2.rectangle(out, (x1, y1), (x2, y2), vehicle_color, 2) # Violation badge above vehicle badge_lines = [] if is_violation: badge_lines.append(f"VIOLATION #{total_violations}") # Show specific violation types type_str = " | ".join(t.upper().replace("_", " ") for t in violation_types) if type_str: badge_lines.append(type_str) if is_wrong_side: badge_lines.append("WRONG SIDE!") if is_bike: badge_lines.append(f"Riders: {num_riders} Helmet OK: {with_h} No Helmet: {without_h}") else: if no_seatbelt > 0: badge_lines.append(f"No Seatbelt: {no_seatbelt}") if plate_text != "UNKNOWN": badge_lines.append(f"Plate: {plate_text}") badge_y = max(y1 - 6 - 16 * len(badge_lines), 5) for li, line in enumerate(badge_lines): bg = (0, 0, 180) if (li == 0 and is_violation) else (30, 30, 30) txt_color = (255, 255, 255) put_label(out, line, x1, badge_y + li * 16, txt_color, bg, scale=0.35, thickness=1) # Draw rider boxes (bikes only) if is_bike and rider_boxes: for p_box in rider_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)) hx1 = max(0, px1 - pad_x) hx2 = min(w_img, px2 + pad_x) hy1 = max(0, py1) hy2 = min(h_img, py1 + head_h) # Determine rider color based on helmet status # (simplified — we can't re-run classification here, # so we use the aggregate to decide styling) rider_color = CLR_NO_HELMET if without_h > 0 else CLR_RIDER label = "No Helmet" if without_h > 0 else "Helmet" cv2.rectangle(out, (px1, py1), (px2, py2), rider_color, 1) cv2.rectangle(out, (hx1, hy1), (hx2, hy2), rider_color, 1) put_label(out, label, px1, py1 - 5, (255, 255, 255), rider_color, scale=0.42, thickness=1) # Draw seatbelt violation boxes for cars if not is_bike and no_seatbelt > 0: # Draw the no-seatbelt boxes that fall inside this car no_seatbelt_boxes = annot.get("no_seatbelt_boxes", []) for sb in no_seatbelt_boxes: sb_cx = (sb[0] + sb[2]) / 2 sb_cy = (sb[1] + sb[3]) / 2 if x1 <= sb_cx <= x2 and y1 <= sb_cy <= y2: cv2.rectangle( out, (int(sb[0]), int(sb[1])), (int(sb[2]), int(sb[3])), CLR_SEATBELT, 2, ) put_label( out, "No Seatbelt", int(sb[0]), int(sb[1]) - 5, (255, 255, 255), CLR_SEATBELT, scale=0.45, ) # Draw license plate box if plate_box_abs: px1, py1, px2, py2 = plate_box_abs cv2.rectangle(out, (px1, py1), (px2, py2), CLR_PLATE, 2) put_label(out, plate_text, px1, py2 + 16, (255, 255, 255), CLR_PLATE, scale=0.48, thickness=1) # ── Summary banner ─────────────────────────────────────────────────────── total_all = total_violations banner_h = 38 banner = np.zeros((banner_h, w_img, 3), dtype=np.uint8) if total_all > 0: banner[:] = (0, 0, 160) parts = [f" {total_violations} VIOLATION(S)"] parts.append(f"Bikes: {total_bikes}") parts.append(f"Cars: {total_cars}") summary = " | ".join(parts) else: banner[:] = (0, 120, 0) summary = f" NO VIOLATIONS | Bikes: {total_bikes} | Cars: {total_cars}" cv2.putText(banner, summary, (10, 26), FONT, 0.65, (255, 255, 255), 1, cv2.LINE_AA) out = np.vstack([banner, out]) return out