generate_heatmap.py

import argparse
from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parents[0]
DETECTRON2_ROOT = ROOT / "detectron2"
sys.path.insert(0, str(DETECTRON2_ROOT))

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
from tqdm import tqdm

from detectron2.checkpoint import DetectionCheckpointer
from detectron2.data import build_detection_test_loader
from detectron2.engine import DefaultTrainer
from detectron2.structures import Instances

# Register custom Faster R-CNN box head (keeps detectron2 sources untouched)
from src.models.faster_rcnn import CustomFastRCNNOutputLayers  # noqa: F401
from src.explanations.ssgradcam import SSGradCAM
from src.explanations.ssgradcampp import SSGradCAMPP
from src.models.detr.d2.detr import TrainerDETR
from src.visualizations.visualize import vis_bbox, vis_cam
from src.utils.utils import (
    build_cfg,
    resolve_target_layer,
    move_to_device,
    load_test_image,
    select_detections,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Minimal SSGradCAM / SSGradCAM++ demo on COCO"
    )
    parser.add_argument(
        "--config",
        type=str,
        default="configs/detr_256_6_6_torchvision.yaml",
        help="Config file path or Detectron2 model_zoo key.",
    )
    parser.add_argument(
        "--weights",
        type=str,
        default="models/detr/coco/model_final.pth",
        help="Weights path. If omitted, use model_zoo checkpoint when available.",
    )
    parser.add_argument(
        "--dataset",
        type=str,
        default="coco_2017_val",
        help="Registered Detectron2 dataset to visualize.",
    )
    parser.add_argument(
        "--model_name",
        type=str,
        default="detr",
        choices=["faster_r_cnn", "detr"],
        help="Model name.",
    )
    parser.add_argument(
        "--method",
        type=str,
        default="ssgradcampp",
        choices=["ssgradcam", "ssgradcampp"],
        help="CAM variant.",
    )
    parser.add_argument(
        "--target-layer",
        type=str,
        default="detr.backbone.0.backbone.res5.2.conv3",
        help="Layer name to hook.",
    )
    parser.add_argument(
        "--score-thresh",
        type=float,
        default=0.5,
        help="Detection score threshold.",
    )
    parser.add_argument(
        "--topk",
        type=int,
        default=3,
        help="Max detections per image to visualize.",
    )
    parser.add_argument(
        "--num-images",
        type=int,
        default=10,
        help="Number of images to process.",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="results_heatmap",
        help="Directory to save heatmaps.",
    )
    parser.add_argument(
        "--test-img-dir",
        type=str,
        default="datasets/test_imgs",
        help="Directory with test images when MS-COCO dataset is unavailable.",
    )
    return parser.parse_args()


def save_detection_cam(
    image: torch.Tensor,
    cam: torch.Tensor,
    detection: Instances,
    out_dir: Path,
    file_stem: str,
    det_rank: int,
    bgr2rgb: bool,
) -> None:
    img = image.detach().cpu().numpy()  # CHW, BGR

    cam = cam.detach().unsqueeze(0).unsqueeze(0)
    cam = F.interpolate(
        cam, size=(img.shape[1], img.shape[2]), mode="bilinear", align_corners=False
    )
    cam = cam.squeeze().cpu().numpy()
    cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-6)

    bbox = detection.pred_boxes.tensor.detach().cpu().numpy()
    label = detection.pred_classes.detach().cpu().numpy()

    fig, ax = plt.subplots(1, 1, figsize=(8, 8))
    vis_cam(img, cam, ax=ax, bgr2rgb=bgr2rgb, alpha=0.4)
    vis_bbox(
        img,
        bbox,
        label=label,
        score=None,
        ax=ax,
        bgr2rgb=bgr2rgb,
        alpha=0.9,
        draw_image=False,
    )
    ax.axis("off")

    save_path = out_dir / f"{file_stem}_det{det_rank:02d}.png"
    plt.savefig(save_path, bbox_inches="tight", pad_inches=0)
    plt.close(fig)


def main() -> None:
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    cfg = build_cfg(args)
    if args.model_name == "detr":
        model = TrainerDETR.build_model(cfg)
    else:
        model = DefaultTrainer.build_model(cfg)
    model.to(device)
    DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS)
    model.eval()

    base_out_dir = Path(args.output_dir)
    out_dir = base_out_dir / args.method
    out_dir.mkdir(parents=True, exist_ok=True)

    target_layer = resolve_target_layer(model, args.target_layer)
    if args.method == "ssgradcam":
        explainer = SSGradCAM(model, target_layer)
    elif args.method == "ssgradcampp":
        explainer = SSGradCAMPP(model, target_layer, args.model_name)
    else:
        raise ValueError(f"Invalid method: {args.method}")

    test_mode = False
    try:
        dataloader = build_detection_test_loader(cfg, args.dataset)
    except Exception:
        test_mode = True
        test_dir = Path(args.test_img_dir)
        image_paths = sorted(
            p
            for p in test_dir.iterdir()
            if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}
        )
        if not image_paths:
            raise FileNotFoundError(f"No images found in {test_dir}")

        dataloader = [load_test_image(img_path) for img_path in image_paths]

    iterator = enumerate(dataloader)
    for img_idx, inputs in tqdm(iterator, total=args.num_images):
        if img_idx >= args.num_images:
            break

        if test_mode:
            batch = [move_to_device(inputs, device)]
            file_stem = Path(inputs["file_name"]).stem
        else:
            batch = [move_to_device(x, device) for x in inputs]
            file_stem = Path(batch[0]["file_name"]).stem

        if args.model_name == "faster_r_cnn":
            outputs = model.inference(batch, do_postprocess=False)
        else:
            outputs = model(batch, do_postprocess=False)

        instances: Instances = outputs[0]
        if len(instances) == 0:
            continue

        target_ids = select_detections(instances, args.score_thresh, args.topk)
        if not target_ids:
            continue

        target_instances = instances[target_ids]
        _, cams = explainer(batch, target_instances)

        image = batch[0]["image"].detach()
        for rank, cam in enumerate(cams):
            save_detection_cam(
                image.cpu(),
                cam.cpu(),
                target_instances[rank],
                out_dir,
                file_stem,
                rank,
                bgr2rgb=cfg.INPUT.FORMAT.upper() == "BGR",
            )


if __name__ == "__main__":
    main()