tools/render_demo_C_frames_v3.py

#!/usr/bin/env python
"""Render demo/C per-frame images v3: clean, large fonts, clear scores."""
import cv2, json, sys, logging
import numpy as np
from pathlib import Path

ROOT = Path("PROJECT_ROOT")
OUT = ROOT / "demo/C"
C_RESULTS = ROOT / "demo/C_results"

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
log = logging.getLogger("render")

COLOR_BGR = {
    "SILENT":  (40, 190, 40),
    "OBSERVE": (30, 190, 255),
    "ALERT":   (30, 30, 230),
}


def find_frame_dir(vid, src):
    if src == "nexar":
        num = vid.replace("nexar_", "")
        for sp in ["train", "test-public", "test-private"]:
            for po in ["positive", "negative"]:
                p = ROOT / f"NEXAR_COLLISION/dataset/{sp}/{po}/{num}"
                if p.exists(): return p
    elif src == "dada":
        name = vid.replace("dada_", "")
        for cat in ["positive", "non-ego", "negative"]:
            p = ROOT / f"DADA-2000/{cat}/{name}"
            if p.exists(): return p
    elif src == "dota":
        raw = vid.replace("dota_", "")
        p = ROOT / f"DoTA/frames/{raw}/images"
        if p.exists(): return p
    return None


def load_frame(frame_dir, idx):
    for fmt in [f"{idx:06d}.jpg", f"{idx:05d}.jpg", f"{idx:04d}.jpg",
                f"{idx:03d}.jpg", f"{idx}.jpg"]:
        fp = frame_dir / fmt
        if fp.exists():
            return cv2.imread(str(fp))
    return None


def get_fps(src):
    return 20.0 if src in ("dada", "dota") else 30.0


def put_text_bg(img, text, pos, font_scale, color, thickness=2, bg_alpha=0.6):
    """Put text with dark background."""
    font = cv2.FONT_HERSHEY_SIMPLEX
    (tw, th), baseline = cv2.getTextSize(text, font, font_scale, thickness)
    x, y = pos
    overlay = img.copy()
    cv2.rectangle(overlay, (x - 4, y - th - 6), (x + tw + 4, y + baseline + 4), (0, 0, 0), -1)
    cv2.addWeighted(overlay, bg_alpha, img, 1 - bg_alpha, 0, img)
    cv2.putText(img, text, (x, y), font, font_scale, color, thickness, cv2.LINE_AA)


def render_gt_frame(img, action, tick_idx, t_sec):
    H, W = img.shape[:2]
    out = img.copy()
    color = COLOR_BGR[action]

    # Top bar
    bar_h = 60
    overlay = out.copy()
    cv2.rectangle(overlay, (0, 0), (W, bar_h), color, -1)
    cv2.addWeighted(overlay, 0.7, out, 0.3, 0, out)

    cv2.putText(out, "Ground Truth", (15, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2, cv2.LINE_AA)
    cv2.putText(out, action, (W - 180, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2, cv2.LINE_AA)
    cv2.putText(out, f"t = {t_sec:.1f}s", (15, 52),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (220, 220, 220), 1, cv2.LINE_AA)
    return out


def render_badas_frame(img, action, p_alert, tick_idx, t_sec):
    H, W = img.shape[:2]
    out = img.copy()
    color = COLOR_BGR[action]

    # Top bar
    bar_h = 60
    overlay = out.copy()
    cv2.rectangle(overlay, (0, 0), (W, bar_h), color, -1)
    cv2.addWeighted(overlay, 0.7, out, 0.3, 0, out)

    cv2.putText(out, "BADAS", (15, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2, cv2.LINE_AA)
    cv2.putText(out, action, (W - 180, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2, cv2.LINE_AA)
    cv2.putText(out, f"t = {t_sec:.1f}s", (15, 52),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (220, 220, 220), 1, cv2.LINE_AA)

    # Bottom: danger score bar
    bar_bot_h = 50
    overlay2 = out.copy()
    cv2.rectangle(overlay2, (0, H - bar_bot_h), (W, H), (0, 0, 0), -1)
    cv2.addWeighted(overlay2, 0.65, out, 0.35, 0, out)

    # Score bar fill
    bar_x0, bar_x1 = 20, W - 20
    bar_y0, bar_y1 = H - bar_bot_h + 8, H - 10
    bar_w = bar_x1 - bar_x0
    fill_w = int(bar_w * min(p_alert, 1.0))

    # Gradient: green → yellow → red
    if p_alert < 0.5:
        r = int(p_alert * 2 * 255)
        fill_color = (0, 255 - r // 2, r)
    else:
        fill_color = (0, int((1 - p_alert) * 200), 230)

    cv2.rectangle(out, (bar_x0, bar_y0), (bar_x0 + fill_w, bar_y1), fill_color, -1)
    cv2.rectangle(out, (bar_x0, bar_y0), (bar_x1, bar_y1), (180, 180, 180), 1)

    cv2.putText(out, f"Danger: {p_alert:.3f}", (bar_x0, bar_y0 - 3),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA)

    return out


def render_vlalert_frame(img, action, p_alert, p_observe, p_silent, tick_idx, t_sec,
                          clip_danger=None, tta=None):
    H, W = img.shape[:2]
    out = img.copy()
    color = COLOR_BGR[action]

    # Top bar
    bar_h = 60
    overlay = out.copy()
    cv2.rectangle(overlay, (0, 0), (W, bar_h), color, -1)
    cv2.addWeighted(overlay, 0.7, out, 0.3, 0, out)

    cv2.putText(out, "VLAlert", (15, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2, cv2.LINE_AA)
    cv2.putText(out, action, (W - 180, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2, cv2.LINE_AA)
    cv2.putText(out, f"t = {t_sec:.1f}s", (15, 52),
                cv2.FONT_HERSHEY_SIMPLEX, 0.55, (220, 220, 220), 1, cv2.LINE_AA)

    # Bottom: 3-class probability bars
    bar_bot_h = 65
    overlay2 = out.copy()
    cv2.rectangle(overlay2, (0, H - bar_bot_h), (W, H), (0, 0, 0), -1)
    cv2.addWeighted(overlay2, 0.65, out, 0.35, 0, out)

    bar_x0, bar_x1 = 20, W - 20
    bar_w = bar_x1 - bar_x0
    bar_h_each = 14
    y = H - bar_bot_h + 6

    probs = [
        ("SILENT",  p_silent,  COLOR_BGR["SILENT"]),
        ("OBSERVE", p_observe, COLOR_BGR["OBSERVE"]),
        ("ALERT",   p_alert,   COLOR_BGR["ALERT"]),
    ]

    for label, prob, clr in probs:
        fill_w = int(bar_w * min(prob, 1.0))
        cv2.rectangle(out, (bar_x0, y), (bar_x0 + fill_w, y + bar_h_each), clr, -1)
        cv2.rectangle(out, (bar_x0, y), (bar_x1, y + bar_h_each), (120, 120, 120), 1)
        cv2.putText(out, f"{label}: {prob:.2f}", (bar_x0 + 5, y + bar_h_each - 2),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, cv2.LINE_AA)
        y += bar_h_each + 2

    return out


def main():
    selected = json.load(open(OUT / "selected_6.json"))
    log.info(f"Rendering {len(selected)} videos")

    for v in selected:
        vid = v["video_id"]
        src = v["source"]
        gt = v["gt"]

        frame_dir = find_frame_dir(vid, src)
        if frame_dir is None:
            log.warning(f"  {vid}: no frames, skip")
            continue

        fps = get_fps(src)
        tick_interval = max(1, int(fps))
        n_ticks = len(gt)

        scores_path = C_RESULTS / vid / "scores.json"
        all_scores = json.load(open(scores_path)) if scores_path.exists() else {}

        log.info(f"  {vid} ({src}): {n_ticks} ticks")

        # Use scored ticks as reference (not GT ticks which may differ)
        ref_ticks = next(iter(all_scores.values()))
        actual_n = len(ref_ticks)

        # Render GT frames (one per scored tick)
        gt_dir = OUT / vid / "GT"
        gt_dir.mkdir(parents=True, exist_ok=True)
        for ti, rt in enumerate(ref_ticks):
            fidx = rt.get("frame", ti * tick_interval)
            t_sec = rt.get("t", fidx / fps)
            img = load_frame(frame_dir, fidx)
            if img is None:
                continue
            gt_act = gt[ti] if ti < len(gt) else "SILENT"
            cv2.imwrite(str(gt_dir / f"frame_{ti:03d}.png"),
                        render_gt_frame(img, gt_act, ti, t_sec))

        # Render each model
        for model_name, ticks in all_scores.items():
            is_badas = "BADAS" in model_name
            folder_name = model_name.replace(" ", "_")
            model_dir = OUT / vid / folder_name
            model_dir.mkdir(parents=True, exist_ok=True)

            for ti, td in enumerate(ticks):
                fidx = td.get("frame", ti * tick_interval)
                t_sec = td.get("t", fidx / fps)
                img = load_frame(frame_dir, fidx)
                if img is None:
                    continue

                action = td.get("action", "SILENT")
                p_alert = td.get("p_alert", 0)
                p_observe = td.get("p_observe", 0)
                p_silent = max(0, 1 - p_alert - p_observe)
                clip_d = td.get("clip_danger", None)

                if is_badas:
                    out = render_badas_frame(img, action, p_alert, ti, t_sec)
                else:
                    out = render_vlalert_frame(img, action, p_alert, p_observe, p_silent,
                                               ti, t_sec, clip_danger=clip_d)
                cv2.imwrite(str(model_dir / f"frame_{ti:03d}.png"), out)

        log.info(f"    done")

    log.info(f"\nAll done! → {OUT}")


if __name__ == "__main__":
    main()