src/Baselines/cafnet/inference.py

import argparse
import os
from typing import Dict, List

import numpy as np
import torch
import torch.distributed as dist
import yaml
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs, set_seed
from safetensors.torch import load_file
from tqdm.auto import tqdm

from dataloader import create_inference_loader
from models.model import CaFNet


DEFAULT_CONFIG = {
    # Packaged evaluation dataset.
    "base_dir": "",
    "split_json": None,
    "test_base_dir": None,
    "test_split": "train",
    "test_split_json": None,
    # Input and radar processing
    "input_height": 288,
    "input_width": 512,
    "radar_max_depth_m": 11.2,
    "max_dist_correspondence": 0.5,
    "patch_size": None,
    # Model
    "encoder": "resnet34_bts",
    "encoder_radar": "resnet18",
    "radar_input_channels": 1,
    "bts_size": 512,
    "max_depth": 11.2,
    # Runtime
    "batch_size": 8,
    # Windows uses spawn-based multiprocessing; keep the public evaluation
    # entry point portable and deterministic by default.
    "num_workers": 0,
    "seed": 42,
    "cpu": False,
    "mixed_precision": "fp16",
    "checkpoint_path": "checkpoints/cafnet.safetensors",
    "prediction_dir": "prediction",
}


def parse_args():
    parser = argparse.ArgumentParser(description="Run CaFNet inference on Smoke-Eval.")
    parser.add_argument("--config", type=str, required=True, help="Path to YAML config")
    return parser.parse_args()


def load_config(path):
    with open(path, "r") as f:
        cfg = yaml.safe_load(f) or {}
    if not isinstance(cfg, dict):
        raise ValueError("Config must be a YAML mapping (key-value pairs).")

    merged = dict(DEFAULT_CONFIG)
    merged.update(cfg)

    if not merged["test_base_dir"]:
        raise ValueError("Config must define 'test_base_dir'.")
    if not merged["checkpoint_path"]:
        raise ValueError("Config must define 'checkpoint_path'.")
    if not os.path.isfile(merged["checkpoint_path"]):
        raise FileNotFoundError(f"Checkpoint not found: {merged['checkpoint_path']}")
    if merged.get("radar_input_channels", 1) != 1:
        raise ValueError("radar_input_channels must be 1 for this setup.")

    return argparse.Namespace(**merged)


def build_model_args(args):
    return argparse.Namespace(
        encoder=args.encoder,
        encoder_radar=args.encoder_radar,
        radar_input_channels=args.radar_input_channels,
        input_height=args.input_height,
        input_width=args.input_width,
        max_depth=args.max_depth,
        bts_size=args.bts_size,
    )


def _extract_model_state(checkpoint):
    if isinstance(checkpoint, dict) and isinstance(checkpoint.get("model"), dict):
        return checkpoint["model"]
    if isinstance(checkpoint, dict):
        return checkpoint
    raise ValueError("Unsupported checkpoint format.")


def _gather_objects(accelerator, obj):
    if accelerator.num_processes == 1:
        return [obj]
    if not dist.is_available() or not dist.is_initialized():
        return [obj]

    gathered = [None for _ in range(accelerator.num_processes)]
    dist.all_gather_object(gathered, obj)
    return gathered


def _merge_predictions(all_rank_predictions):
    merged: Dict[str, Dict[int, np.ndarray]] = {}
    for rank_dict in all_rank_predictions:
        if not rank_dict:
            continue
        for seq_name, frame_map in rank_dict.items():
            seq_slot = merged.setdefault(seq_name, {})
            for frame_idx, pred in frame_map.items():
                frame_idx = int(frame_idx)
                if frame_idx not in seq_slot:
                    seq_slot[frame_idx] = pred
    return merged


def _save_sequence_predictions(predictions, out_dir):
    os.makedirs(out_dir, exist_ok=True)
    for seq_name in sorted(predictions.keys()):
        frame_map = predictions[seq_name]
        ordered_frames = sorted(frame_map.keys())
        if not ordered_frames:
            pred_stack = np.zeros((0,), dtype=np.float32)
        else:
            pred_stack = np.stack([frame_map[k] for k in ordered_frames], axis=0).astype(
                np.float32,
                copy=False,
            )
        np.save(os.path.join(out_dir, f"{seq_name.lower()}_pred.npy"), pred_stack)


def _run_loader_inference(accelerator, model, loader, samples, save_dir, desc):
    model.eval()
    local_preds: Dict[str, Dict[int, np.ndarray]] = {}

    with torch.no_grad():
        pbar = tqdm(
            loader,
            desc=desc,
            disable=not accelerator.is_local_main_process,
            dynamic_ncols=True,
            leave=False,
        )
        for batch in pbar:
            sample_idx, image, depth_gt, radar, radar_gt = batch

            image = image.to(accelerator.device, non_blocking=True)
            radar = radar.to(accelerator.device, non_blocking=True)
            # Kept for parity with validation loop structure.
            _ = depth_gt.to(accelerator.device, non_blocking=True)
            _ = radar_gt.to(accelerator.device, non_blocking=True)

            focal = torch.ones((image.size(0),), device=image.device)
            _, _, _, _, depth_est, _, _ = model(image, radar, focal)

            pred_np = depth_est.detach().float().cpu().numpy()
            if pred_np.ndim == 4 and pred_np.shape[1] == 1:
                pred_np = pred_np[:, 0]

            if torch.is_tensor(sample_idx):
                sample_idx_list = sample_idx.detach().cpu().tolist()
            else:
                sample_idx_list = list(sample_idx)

            for local_i, sample_i in enumerate(sample_idx_list):
                seq_name, frame_idx = samples[int(sample_i)]
                seq_slot = local_preds.setdefault(seq_name, {})
                frame_idx = int(frame_idx)
                if frame_idx not in seq_slot:
                    seq_slot[frame_idx] = pred_np[local_i].astype(np.float32, copy=False)

    gathered = _gather_objects(accelerator, local_preds)
    if accelerator.is_main_process:
        merged = _merge_predictions(gathered)
        _save_sequence_predictions(merged, save_dir)

    accelerator.wait_for_everyone()


def main():
    cli = parse_args()
    args = load_config(cli.config)

    set_seed(args.seed)
    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
    accelerator = Accelerator(
        mixed_precision=None if args.mixed_precision in ("no", "none") else args.mixed_precision,
        cpu=args.cpu,
        kwargs_handlers=[ddp_kwargs],
    )

    test_loader = create_inference_loader(
        args,
        pin_memory=(accelerator.device.type == "cuda"),
    )
    test_samples: List = test_loader.dataset.samples

    model = CaFNet(build_model_args(args))

    model, test_loader = accelerator.prepare(model, test_loader)

    state_dict = load_file(args.checkpoint_path, device="cpu")
    accelerator.unwrap_model(model).load_state_dict(state_dict, strict=True)

    _run_loader_inference(
        accelerator=accelerator,
        model=model,
        loader=test_loader,
        samples=test_samples,
        save_dir=args.prediction_dir,
        desc="Inference",
    )

    if accelerator.is_main_process:
        print(f"Saved predictions to: {args.prediction_dir}")


if __name__ == "__main__":
    main()