Blame for visualization/bev_renderer.py

# ============================================================
# bev_renderer.py
# ============================================================
#
# ETHUKU CREATE PANNINOM?
# ━━━━━━━━━━━━━━━━━━━━━━
# Numbers (heatmap tensors) yaarukum puriyaadhu. Recruiter oru PADAM
# paakanum: "6 camera photo -> top-down map-la cars marked".
# Ithu antha padathai varaiyum.
#
# MUNADI FILE ODA CONNECTION:
# ━━━━━━━━━━━━━━━━━━━━━━━━━━
# Input : training/evaluate.py-oda decode_predictions output
#         + dataset._boxes_in_ego (ground truth)
# Output: PNG / GIF -> README-la pogum
#
# INNER OPERATIONS:
# ━━━━━━━━━━━━━━━━
# Matplotlib-la 2 panel: (idathu) 6 camera photos grid,
# (valathu) top-down BEV - ego car center-la, boxes suthi.
# Box varaiyurathu: center + w,l + yaw -> 4 corner points -> polygon.
#
# INPUT / OUTPUT:
# ━━━━━━━━━━━━━━
# render_frame(images, gt_boxes, pred_boxes, out_path)
#
# EPADI USE AAGUM:
# ━━━━━━━━━━━━━━━
# python -m visualization.bev_renderer --ckpt runs/simplebev/best.pth
#
# ============================================================

import argparse
import os

import matplotlib
matplotlib.use("Agg")            # display illa (headless) - file-la mattum save
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Polygon

from data.scripts.constants import (
    CAMERAS, CLASSES, X_RANGE, Y_RANGE, IMAGENET_MEAN, IMAGENET_STD,
)

# Ovvoru class-kum oru color. Recruiter oru nodi-la puriyanum.
CLASS_COLORS = {
    0: "#00d4ff",   # car          - cyan
    1: "#ff8c00",   # truck        - orange
    2: "#ffd700",   # bus          - gold
    3: "#8b4513",   # trailer      - brown
    4: "#9370db",   # construction - purple
    5: "#ff1493",   # pedestrian   - pink
    6: "#00ff7f",   # motorcycle   - spring green
    7: "#7fff00",   # bicycle      - chartreuse
    8: "#ff4500",   # traffic_cone - red-orange
    9: "#a9a9a9",   # barrier      - grey
}


def denormalize(img: np.ndarray) -> np.ndarray:
    """
    Normalize panna image-ai thirumba paakkura maadhiri maathurathu.

    Args:
        img: [3, H, W] normalized (-2.5 to 2.5)

    Returns:
        [H, W, 3] uint8-range float (0..1) - matplotlib ku
    """
    img = img.transpose(1, 2, 0)                       # CHW -> HWC
    img = img * np.array(IMAGENET_STD) + np.array(IMAGENET_MEAN)
    return np.clip(img, 0, 1)


def box_corners(x: float, y: float, w: float, l: float, yaw: float) -> np.ndarray:
    """
    Box center + size + angle -> 4 mooli (corner) points.

    Math simple-a: mudhalla car origin-la irukku nu nenachi 4 corner
    kanakku podurom, appuram yaw angle-la suthi, appuram (x,y) ku nagarthurom.

    Args:
        x, y : center metres (x=munnadi, y=idathu)
        w, l : width, length metres
        yaw  : radians

    Returns:
        [4, 2] corner coordinates
    """
    # Car-oda soontha frame-la 4 corners (length = munnadi axis)
    corners = np.array([
        [ l / 2,  w / 2],
        [ l / 2, -w / 2],
        [-l / 2, -w / 2],
        [-l / 2,  w / 2],
    ])

    # 2D rotation matrix - yaw angle-la suthurathu
    c, s = np.cos(yaw), np.sin(yaw)
    R = np.array([[c, -s], [s, c]])

    return corners @ R.T + np.array([x, y])


def draw_bev(ax, gt_boxes: list, pred_boxes: list, score_thresh: float = 0.2) -> None:
    """
    Top-down BEV panel varaiyurathu.

    GT = white dashed outline (nijam)
    Pred = class color filled (namma model sonnathu)
    """
    ax.set_facecolor("#0a0a14")

    # --- Grid lines (10m ku oru vari) - dooram purinjukka ---
    for r in range(10, 51, 10):
        ax.add_patch(plt.Circle((0, 0), r, fill=False,
                                edgecolor="#2a2a3a", linewidth=0.6, zorder=1))
        ax.text(r * 0.7, r * 0.7, f"{r}m", color="#4a4a5a", fontsize=7, zorder=1)

    # --- Ego car (namma vandi) center-la ---
    ego = box_corners(0, 0, 1.8, 4.5, 0)
    ax.add_patch(Polygon(ego, closed=True, facecolor="#ffffff",
                         edgecolor="#ffffff", linewidth=1, zorder=5))
    ax.arrow(0, 2.5, 0, 3, head_width=1.2, head_length=1.2,
             fc="#ffffff", ec="#ffffff", zorder=5)   # munnadi direction

    # --- Ground truth ---
    for b in gt_boxes:
        pts = box_corners(b["x"], b["y"], b["w"], b["l"], b["yaw"])
        ax.add_patch(Polygon(pts, closed=True, fill=False, edgecolor="#ffffff",
                             linewidth=1.0, linestyle="--", alpha=0.55, zorder=3))

    # --- Predictions ---
    n_drawn = 0
    for b in pred_boxes:
        if b["score"] < score_thresh:
            continue
        color = CLASS_COLORS.get(b["cls"], "#ffffff")
        pts = box_corners(b["x"], b["y"], b["w"], b["l"], b["yaw"])
        ax.add_patch(Polygon(pts, closed=True, facecolor=color, alpha=0.45,
                             edgecolor=color, linewidth=1.2, zorder=4))
        n_drawn += 1

    ax.set_xlim(Y_RANGE[1], Y_RANGE[0])      # y = idathu, plot-la idathu pakkam
    ax.set_ylim(X_RANGE[0], X_RANGE[1])      # x = munnadi, plot-la mela
    ax.set_aspect("equal")
    ax.set_xlabel("left  <-  y (m)  ->  right", color="#8a8a9a", fontsize=8)
    ax.set_ylabel("back  <-  x (m)  ->  front", color="#8a8a9a", fontsize=8)
    ax.set_title(f"BEV  |  white dashed = ground truth  |  filled = predicted "
                 f"({n_drawn} boxes, score>{score_thresh})",
                 color="#e0e0e0", fontsize=9)
    ax.tick_params(colors="#4a4a5a", labelsize=7)
    for spine in ax.spines.values():
        spine.set_color("#2a2a3a")


def render_frame(images: np.ndarray, gt_boxes: list, pred_boxes: list,
                 out_path: str, score_thresh: float = 0.2) -> None:
    """
    Oru frame-oda full visualization save pannurathu.

    Args:
        images: [6, 3, 224, 400] normalized
        gt_boxes / pred_boxes: dict list (x, y, w, l, yaw, cls, [score])
        out_path: PNG path
    """
    fig = plt.figure(figsize=(17, 7), facecolor="#0a0a14")
    gs = fig.add_gridspec(2, 5, width_ratios=[1, 1, 1, 0.12, 1.9],
                          hspace=0.08, wspace=0.04)

    # --- Idathu pakkam: 6 camera photos ---
    # Car-la epdi irukko appdiye layout: mela munnadi, keezha pinnadi.
    # Ippdi potta padathai paakkurava-ku surround view neradiya purinjidum.
    layout = [
        ("CAM_FRONT_LEFT", 0, 0), ("CAM_FRONT", 0, 1), ("CAM_FRONT_RIGHT", 0, 2),
        ("CAM_BACK_LEFT",  1, 0), ("CAM_BACK",  1, 1), ("CAM_BACK_RIGHT",  1, 2),
    ]
    for cam, r, c in layout:
        ax = fig.add_subplot(gs[r, c])
        ax.imshow(denormalize(images[CAMERAS.index(cam)]))
        ax.set_title(cam.replace("CAM_", ""), color="#8a8a9a", fontsize=7, pad=2)
        ax.axis("off")

    # --- Valathu pakkam: BEV ---
    ax_bev = fig.add_subplot(gs[:, 4])
    draw_bev(ax_bev, gt_boxes, pred_boxes, score_thresh)

    # --- Class color legend ---
    present = sorted({b["cls"] for b in pred_boxes if b["score"] >= score_thresh}
                     | {b["cls"] for b in gt_boxes})
    handles = [plt.Line2D([0], [0], marker="s", linestyle="", markersize=7,
                          color=CLASS_COLORS[c], label=CLASSES[c]) for c in present]
    if handles:
        ax_bev.legend(handles=handles, loc="upper left", fontsize=7,
                      facecolor="#14141e", edgecolor="#2a2a3a", labelcolor="#c0c0d0")

    os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
    fig.savefig(out_path, dpi=130, bbox_inches="tight", facecolor="#0a0a14")
    plt.close(fig)
    print(f"saved {out_path}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="runs/simplebev/best.pth")
    ap.add_argument("--data-root", default="data/nuscenes-mini")
    ap.add_argument("--out-dir", default="runs/viz")
    ap.add_argument("--frames", type=int, default=6)
    ap.add_argument("--score-thresh", type=float, default=0.2)
    ap.add_argument("--gif", action="store_true", help="frames-ai GIF-a sethu vai")
    args = ap.parse_args()

    import torch
    from data.scripts.dataset import NuScenesBEVDataset
    from models.simplebev import SimpleBEV
    from training.evaluate import decode_predictions

    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    model = SimpleBEV(pretrained=False).to(device).eval()
    model.load_state_dict(torch.load(args.ckpt, map_location=device)["model"])

    ds = NuScenesBEVDataset(args.data_root, split="val")
    paths = []

    for i in range(min(args.frames, len(ds))):
        s = ds[i]
        with torch.no_grad():
            preds = model(s["images"].unsqueeze(0).to(device),
                          s["intrinsics"].unsqueeze(0).to(device),
                          s["extrinsics"].unsqueeze(0).to(device))
        pred_boxes = decode_predictions({k: v.cpu() for k, v in preds.items()})[0]
        gt_boxes = ds._boxes_in_ego(ds.sample_tokens[i])

        p = os.path.join(args.out_dir, f"frame_{i:03d}.png")
        render_frame(s["images"].numpy(), gt_boxes, pred_boxes, p, args.score_thresh)
        paths.append(p)

    if args.gif and paths:
        from PIL import Image
        frames = [Image.open(p) for p in paths]
        gif_path = os.path.join(args.out_dir, "demo.gif")
        frames[0].save(gif_path, save_all=True, append_images=frames[1:],
                       duration=700, loop=0)
        print(f"saved {gif_path}")


if __name__ == "__main__":
    main()