pipeline/visualize.py

"""
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")