""" Output generation for train and val splits. Train panel (outputs/train/): Original | CLAHE | Canny | GT Labels Val panel (outputs/val/): Original | GT Labels | YOLO Pred | Anonymized """ import cv2 import numpy as np from pathlib import Path from ultralytics import YOLO import config from pipeline.preprocess import clahe_canny # ── Shared helpers ──────────────────────────────────────────────────────────── def _sep(h): return np.full((h, 3, 3), 200, dtype=np.uint8) def _header(panel, labels, col_w): for i, txt in enumerate(labels): x = i * (col_w + 3) + 8 cv2.putText(panel, txt, (x, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 2, cv2.LINE_AA) cv2.putText(panel, txt, (x, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (20, 20, 20), 1, cv2.LINE_AA) def _draw_gt(bgr: np.ndarray, lbl_path: Path) -> np.ndarray: out = bgr.copy() h, w = bgr.shape[:2] if not lbl_path.exists(): return out for line in lbl_path.read_text().strip().splitlines(): p = list(map(float, line.split())) cls = int(p[0]) cx, cy, bw, bh = p[1], p[2], p[3], p[4] x1, y1 = int((cx-bw/2)*w), int((cy-bh/2)*h) x2, y2 = int((cx+bw/2)*w), int((cy+bh/2)*h) cv2.rectangle(out, (x1,y1), (x2,y2), config.CLASS_COLORS[cls], 2) return out def _draw_pred(bgr: np.ndarray, boxes) -> np.ndarray: out = bgr.copy() for box in boxes: cls = int(box.cls.item()) x1, y1, x2, y2 = map(int, box.xyxy[0].tolist()) cv2.rectangle(out, (x1,y1), (x2,y2), config.CLASS_COLORS.get(cls,(255,255,255)), 2) return out def _hstack(*imgs): sep = _sep(imgs[0].shape[0]) parts = [] for img in imgs: parts.append(img); parts.append(sep) return np.hstack(parts[:-1]) # ── Train output: Original | CLAHE | Canny | GT Labels ─────────────────────── def save_train_outputs(): out_dir = config.OUTPUTS_DIR / "train" out_dir.mkdir(parents=True, exist_ok=True) img_dir = config.DATASET_DIR / "images" / "train" lbl_dir = config.DATASET_DIR / "labels" / "train" paths = sorted(img_dir.glob("*.png")) for p in paths: bgr = cv2.imread(str(p)) proc = clahe_canny(bgr) gray_3ch = cv2.cvtColor(proc[:,:,0], cv2.COLOR_GRAY2BGR) clahe_3ch = cv2.cvtColor(proc[:,:,1], cv2.COLOR_GRAY2BGR) canny_3ch = cv2.cvtColor(proc[:,:,2], cv2.COLOR_GRAY2BGR) lbl_name = p.name.replace("_annotated", "").replace(".png", ".txt") gt = _draw_gt(bgr, lbl_dir / lbl_name) panel = _hstack(gray_3ch, clahe_3ch, canny_3ch, gt) _header(panel, ["Original", "CLAHE", "Canny", "GT Labels"], bgr.shape[1]) cv2.imwrite(str(out_dir / p.name.replace("_annotated", "").replace(".png", ".jpg")), panel) print(f"Train outputs saved → {out_dir}/ ({len(paths)} images)") # ── Val output: Original | GT Labels | YOLO Pred | Anonymized | OCR ────────── def save_val_outputs(yolo: YOLO, detector: str, anonymize_fn): from pipeline.detector import predict from pipeline.ocr import run_ocr, save_json, ocr_panel out_dir = config.OUTPUTS_DIR / "val" out_dir.mkdir(parents=True, exist_ok=True) img_dir = config.DATASET_DIR / "images" / "val" lbl_dir = config.DATASET_DIR / "labels" / "val" paths = sorted(img_dir.glob("*.png")) for p in paths: bgr = cv2.imread(str(p)) stem = p.name.replace("_annotated", "").replace(".png", "") lbl_name = stem + ".txt" gt = _draw_gt(bgr, lbl_dir / lbl_name) boxes = predict(yolo, p, detector) pred_vis = _draw_pred(bgr, boxes) anon = anonymize_fn(bgr, boxes) ocr_res = run_ocr(bgr, boxes) save_json(ocr_res, out_dir / "ocr" / f"{stem}.json") ocr_vis = ocr_panel(bgr.shape[0], bgr.shape[1], ocr_res) panel = _hstack(bgr, gt, pred_vis, anon, ocr_vis) _header(panel, ["Original", "GT Labels", "YOLO Pred", "Anonymized", "OCR"], bgr.shape[1]) cv2.imwrite(str(out_dir / f"{stem}.jpg"), panel) print(f"Val outputs saved → {out_dir}/ ({len(paths)} images)") # ── Metrics ─────────────────────────────────────────────────────────────────── def _box_iou_xyxy(a: np.ndarray, b: np.ndarray) -> np.ndarray: """Pairwise IoU between [N,4] and [M,4] xyxy boxes → [N,M].""" x1 = np.maximum(a[:, None, 0], b[None, :, 0]) y1 = np.maximum(a[:, None, 1], b[None, :, 1]) x2 = np.minimum(a[:, None, 2], b[None, :, 2]) y2 = np.minimum(a[:, None, 3], b[None, :, 3]) inter = np.clip(x2-x1, 0, None) * np.clip(y2-y1, 0, None) area_a = (a[:,2]-a[:,0]) * (a[:,3]-a[:,1]) area_b = (b[:,2]-b[:,0]) * (b[:,3]-b[:,1]) return inter / np.maximum(area_a[:, None] + area_b[None, :] - inter, 1e-9) def mean_iou_at_tp(yolo: YOLO, data_yaml: Path, iou_thr: float = 0.5, conf: float = None) -> float: """Greedy mean IoU of TP matches (same class, IoU ≥ iou_thr) on val split.""" import yaml as _yaml cfg = _yaml.safe_load(Path(data_yaml).read_text()) ds_root = Path(cfg["path"]) val_imgs = ds_root / cfg["val"] val_lbls = ds_root / "labels" / "val" conf = config.CONF if conf is None else conf ious = [] for img_path in sorted(val_imgs.glob("*.png")): bgr = cv2.imread(str(img_path)) H, W = bgr.shape[:2] lbl = val_lbls / (img_path.stem + ".txt") gt_cls, gt_xyxy = [], [] if lbl.exists(): for line in lbl.read_text().strip().splitlines(): p = list(map(float, line.split())) cx, cy, bw, bh = p[1], p[2], p[3], p[4] gt_cls.append(int(p[0])) gt_xyxy.append([(cx-bw/2)*W, (cy-bh/2)*H, (cx+bw/2)*W, (cy+bh/2)*H]) preds = yolo.predict(str(img_path), conf=conf, verbose=False)[0].boxes pr_cls = [int(b.cls.item()) for b in preds] pr_xyxy = [b.xyxy[0].tolist() for b in preds] for cls in set(gt_cls) | set(pr_cls): gt_b = np.array([x for c, x in zip(gt_cls, gt_xyxy) if c == cls], np.float32) pr_b = np.array([x for c, x in zip(pr_cls, pr_xyxy) if c == cls], np.float32) if len(gt_b) == 0 or len(pr_b) == 0: continue iou = _box_iou_xyxy(pr_b, gt_b) taken = set() for i in range(len(pr_b)): cand = [(j, iou[i, j]) for j in range(len(gt_b)) if j not in taken] if not cand: continue j, v = max(cand, key=lambda x: x[1]) if v >= iou_thr: taken.add(j) ious.append(float(v)) return float(np.mean(ious)) if ious else 0.0 def print_metrics(yolo: YOLO, data_yaml: Path): print("\nComputing validation metrics…") b = yolo.val(data=str(data_yaml), split="val", verbose=False, workers=0).box p, r = float(b.mp), float(b.mr) f1_arr = getattr(b, "f1", None) f1 = float(np.mean(f1_arr)) if f1_arr is not None and len(f1_arr) > 0 \ else 2*p*r / (p+r+1e-9) miou = mean_iou_at_tp(yolo, data_yaml) print(f"\n{'─'*48}\n Validation metrics\n{'─'*48}") for name, val in [("mAP@50", b.map50), ("mAP@50-95", b.map), ("Precision", p), ("Recall", r), ("F1", f1), ("Mean IoU @TP", miou)]: print(f" {name:<18}: {val:.4f}") print() for i, cls in enumerate(config.CLASS_NAMES): ap = float(b.ap50[i]) if i < len(b.ap50) else float("nan") f1c = float(f1_arr[i]) if f1_arr is not None and i < len(f1_arr) else float("nan") print(f" AP@50 {cls:<6}: {ap:.4f} F1 {cls:<6}: {f1c:.4f}") print(f"{'─'*48}\n")