"""Rendering and summary helpers for street-scene object detection.""" from __future__ import annotations import csv from collections import defaultdict from pathlib import Path from typing import Iterable from PIL import Image, ImageDraw, ImageFont STREET_OBJECT_GROUPS: dict[str, set[str]] = { "people": {"person"}, "active_mobility": {"person", "bicycle"}, "motor_vehicles": {"car", "motorcycle", "bus", "truck", "train"}, "all_transport": { "bicycle", "car", "motorcycle", "bus", "truck", "train", }, } def _load_font(size: int) -> ImageFont.ImageFont: candidates = ( "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf", "/System/Library/Fonts/Helvetica.ttc", ) for candidate in candidates: try: return ImageFont.truetype(candidate, size) except (OSError, IOError): continue return ImageFont.load_default() def detection_color(class_id: int) -> tuple[int, int, int]: """Return a stable, high-contrast color for a COCO class ID.""" return ( int((67 * class_id + 37) % 190 + 40), int((97 * class_id + 71) % 190 + 40), int((43 * class_id + 113) % 190 + 40), ) def _boxes_overlap( first: tuple[float, float, float, float], second: tuple[float, float, float, float], ) -> bool: return not ( first[2] <= second[0] or second[2] <= first[0] or first[3] <= second[1] or second[3] <= first[1] ) def _label_box( object_box: list[float], label_width: float, label_height: float, image_size: tuple[int, int], occupied: list[tuple[float, float, float, float]], ) -> tuple[float, float, float, float]: """Place a label near its object while avoiding earlier labels.""" image_width, image_height = image_size center_x = (object_box[0] + object_box[2]) / 2 x = max(0.0, min(center_x - label_width / 2, image_width - label_width)) candidates: list[tuple[float, float, float, float]] = [] for slot in range(7): bottom = object_box[1] - 4 - slot * (label_height + 4) top = bottom - label_height if top >= 0: candidates.append((x, top, x + label_width, bottom)) for slot in range(4): top = object_box[3] + 4 + slot * (label_height + 4) bottom = top + label_height if bottom <= image_height: candidates.append((x, top, x + label_width, bottom)) for candidate in candidates: if not any(_boxes_overlap(candidate, used) for used in occupied): return candidate return candidates[0] if candidates else (x, 0.0, x + label_width, label_height) def render_detection( image: Image.Image, detections: Iterable[dict[str, object]], ) -> Image.Image: """Draw labeled bounding boxes on a copy of the input image.""" rendered = image.convert("RGB").copy() draw = ImageDraw.Draw(rendered) short_side = min(rendered.size) line_width = max(2, round(short_side / 320)) font = _load_font(max(13, min(24, round(short_side / 55)))) occupied_labels: list[tuple[float, float, float, float]] = [] ordered_detections = sorted( detections, key=lambda item: (float(item["x1"]), float(item["y1"])), ) for detection in ordered_detections: class_id = int(detection["class_id"]) color = detection_color(class_id) box = [ float(detection["x1"]), float(detection["y1"]), float(detection["x2"]), float(detection["y2"]), ] label = ( f"{detection['class_name']} " f"{float(detection['confidence']):.2f}" ) draw.rectangle(box, outline=color, width=line_width) text_box = draw.textbbox((0, 0), label, font=font) text_height = text_box[3] - text_box[1] text_width = text_box[2] - text_box[0] background = _label_box( box, text_width + 8, text_height + 8, rendered.size, occupied_labels, ) occupied_labels.append(background) draw.rectangle(background, fill=color) draw.text( (background[0] + 4, background[1] + 4), label, fill="white", font=font, ) return rendered def build_detection_summary( detections: Iterable[dict[str, object]], ) -> list[list[object]]: """Aggregate detection counts and confidence by class.""" grouped: dict[str, list[float]] = defaultdict(list) for detection in detections: grouped[str(detection["class_name"])].append( float(detection["confidence"]) ) rows = [ [class_name, len(confidences), round(sum(confidences) / len(confidences), 3), round(max(confidences), 3)] for class_name, confidences in grouped.items() ] rows.sort(key=lambda row: (-int(row[1]), str(row[0]))) return rows def build_detection_table( detections: Iterable[dict[str, object]], ) -> list[list[object]]: """Build one exportable row per bounding box.""" rows: list[list[object]] = [] for index, detection in enumerate(detections, start=1): rows.append( [ index, str(detection["class_name"]), round(float(detection["confidence"]), 3), round(float(detection["x1"]), 1), round(float(detection["y1"]), 1), round(float(detection["x2"]), 1), round(float(detection["y2"]), 1), ] ) return rows def build_street_indicators( detections: Iterable[dict[str, object]], ) -> dict[str, int]: """Derive transparent street-scene counts from visible COCO objects.""" class_names = [str(detection["class_name"]) for detection in detections] return { name: sum(class_name in members for class_name in class_names) for name, members in STREET_OBJECT_GROUPS.items() } def write_detection_csv( path: Path, rows: Iterable[Iterable[object]], ) -> None: with path.open("w", newline="", encoding="utf-8") as handle: writer = csv.writer(handle) writer.writerow(["object_id", "class", "confidence", "x1", "y1", "x2", "y2"]) writer.writerows(rows)