inference.py

import os
import sys
import yaml
import json
import torch
import pickle
import shutil
import logging
import warnings
import argparse
import subprocess
import numpy as np


from os import path
from datetime import datetime
from torchmetrics.classification import Accuracy
from src.utility.builtin import ODTrainer, ODLightningCLI
from src.utility.visualize import (
    save_spatial_attention_visualization,
    save_component_attention_visualization,
)


def parse_args(args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("model_cfg_path", type=str)
    parser.add_argument("data_cfg_path", type=str)
    parser.add_argument("model_ckpt_path", type=str)
    parser.add_argument("--precision", type=str, default="16")
    parser.add_argument("--devices", type=int, default=-1)
    parser.add_argument("--notes", type=str, default='')
    parser.add_argument("--save-visuals", action="store_true")
    parser.add_argument("--visual-max-samples", type=int, default=16)
    return parser.parse_args(args=args)


class StatsRecorder:
    def __init__(self, label):
        self.label = label
        self.prob = 0
        self.count = 0

    def update(self, prob, label):
        assert label == self.label
        self.prob += prob
        self.count += 1

    def compute(self):
        return {
            "label": self.label,
            "prob": self.prob / self.count
        }


def configure_logging():
    logging_fmt = "[%(levelname)s][%(filename)s:%(lineno)d]: %(message)s"
    logging.basicConfig(level="INFO", format=logging_fmt)
    warnings.filterwarnings(action="ignore")


def _detect_healthy_gpu_indices():
    cmd = [
        "nvidia-smi",
        "--query-gpu=index",
        "--format=csv,noheader,nounits",
    ]
    try:
        result = subprocess.run(
            cmd,
            check=True,
            capture_output=True,
            text=True,
        )
    except Exception as exc:
        logging.warning(f"Failed to query healthy GPUs via nvidia-smi: {exc}")
        if torch.cuda.is_available():
            return list(range(torch.cuda.device_count()))
        return []

    healthy = []
    for line in result.stdout.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            healthy.append(int(line))
        except ValueError:
            logging.warning(f"Unexpected GPU index from nvidia-smi: {line}")
    return healthy


def _configure_safe_cuda_visible_devices():
    healthy = _detect_healthy_gpu_indices()
    if not healthy:
        return

    current = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip()
    if not current:
        selected = healthy
    else:
        selected = []
        dropped = []
        for item in current.split(","):
            item = item.strip()
            if not item:
                continue
            try:
                gpu_idx = int(item)
            except ValueError:
                logging.warning(
                    f"CUDA_VISIBLE_DEVICES contains non-integer entry '{item}', leaving unchanged."
                )
                return
            if gpu_idx in healthy:
                selected.append(gpu_idx)
            else:
                dropped.append(gpu_idx)

        if dropped:
            logging.warning(
                f"Dropping unhealthy GPUs from CUDA_VISIBLE_DEVICES: {dropped}; healthy={healthy}"
            )
        if not selected:
            selected = [healthy[0]]
            logging.warning(
                f"CUDA_VISIBLE_DEVICES={current} contains no healthy GPUs; "
                f"falling back to GPU {selected[0]}"
            )

    visible = ",".join(str(idx) for idx in selected)
    os.environ["CUDA_VISIBLE_DEVICES"] = visible
    logging.info(f"Using CUDA_VISIBLE_DEVICES={visible}")

def _best_accuracy_and_threshold(probs: torch.Tensor, labels: torch.Tensor):
    probs = probs.detach().float().flatten().cpu()
    labels = labels.detach().long().flatten().cpu()
    if probs.numel() == 0:
        return 0.0, 0.5
    uniq = torch.unique(probs)
    uniq, _ = torch.sort(uniq)
    thresholds = torch.cat(
        [
            torch.tensor([0.0]),
            uniq,
            torch.tensor([1.0]),
        ]
    )
    best_acc = -1.0
    best_thr = 0.5
    for thr in thresholds:
        pred = (probs >= thr).long()
        acc = (pred == labels).float().mean().item()
        if acc > best_acc:
            best_acc = acc
            best_thr = thr.item()
    return best_acc, best_thr


def _binary_roc_auc(probs: torch.Tensor, labels: torch.Tensor) -> float:
    probs = probs.detach().float().flatten().cpu().numpy()
    labels = labels.detach().long().flatten().cpu().numpy()

    pos = labels == 1
    neg = labels == 0
    n_pos = int(pos.sum())
    n_neg = int(neg.sum())
    if n_pos == 0 or n_neg == 0:
        return float("nan")

    order = np.argsort(probs, kind="mergesort")
    sorted_probs = probs[order]
    ranks = np.arange(1, len(probs) + 1, dtype=np.float64)

    # average ranks for ties
    i = 0
    while i < len(sorted_probs):
        j = i + 1
        while j < len(sorted_probs) and sorted_probs[j] == sorted_probs[i]:
            j += 1
        avg_rank = ranks[i:j].mean()
        ranks[i:j] = avg_rank
        i = j

    full_ranks = np.empty_like(ranks)
    full_ranks[order] = ranks
    pos_ranks = full_ranks[pos].sum()
    auc = (pos_ranks - n_pos * (n_pos + 1) / 2.0) / (n_pos * n_neg)
    return float(auc)


def _prepare_visual_frame_data(batch_result: dict):
    visual_tensors = batch_result.get("visual_tensors")
    if not visual_tensors or "s_aff" not in visual_tensors:
        return None

    clips = batch_result.get("clips")
    names = batch_result.get("names")
    probs = batch_result.get("probs")
    labels = batch_result.get("y")
    component_names = batch_result.get("component_names", [])
    if clips is None or names is None or probs is None or labels is None:
        return None

    s_aff = visual_tensors["s_aff"]
    if s_aff.dim() != 6:
        return None

    # [layers, batch, frames, components, patch_tokens]
    per_component_map = s_aff.mean(dim=(0, 1))
    # [batch, frames, patch_tokens]
    spatial_map = per_component_map.mean(dim=2)
    return clips, names, probs, labels, spatial_map, per_component_map, component_names


def _export_visualizations(batch_result: dict, visual_dir: str, saved_count: int, max_samples: int):
    prepared = _prepare_visual_frame_data(batch_result)
    if prepared is None or saved_count >= max_samples:
        return saved_count

    clips, names, probs, labels, spatial_map, component_map, component_names = prepared
    clips = clips.detach().cpu()
    probs = probs.detach().cpu() if isinstance(probs, torch.Tensor) else torch.tensor(probs)

    batch_size = min(
        clips.shape[0],
        spatial_map.shape[0],
        component_map.shape[0],
        len(names),
        len(labels),
        probs.shape[0],
    )
    for sample_idx in range(batch_size):
        if saved_count >= max_samples:
            break
        name = str(names[sample_idx])
        safe_name = name.replace('/', '__')
        label = int(labels[sample_idx])
        prob = float(probs[sample_idx])
        title = f"{name} | label={label} | prob={prob:.4f}"

        spatial_save_path = path.join(visual_dir, f"{saved_count:03d}_{safe_name}_spatial.png")
        save_spatial_attention_visualization(
            clip=clips[sample_idx],
            spatial_attn=spatial_map[sample_idx],
            save_path=spatial_save_path,
            title=title,
            normalized=False,
        )

        component_save_path = path.join(visual_dir, f"{saved_count:03d}_{safe_name}_components.png")
        save_component_attention_visualization(
            clip=clips[sample_idx],
            component_attn=component_map[sample_idx],
            component_names=component_names,
            save_path=component_save_path,
            title=title,
            normalized=False,
        )
        saved_count += 1

    return saved_count


@torch.inference_mode()
def inference_driver(cli, output_dir, ckpt_path, notes=None, save_visuals=False, visual_max_samples=16):

    timestamp = datetime.now().strftime("%m%dT%H%M%S")

    trainer = cli.trainer

    # setup model
    model = cli.model

    try:
        model = model.__class__.load_from_checkpoint(ckpt_path)
    except Exception as e:
        print(f"Unable to load model from checkpoint via constructor: {e}")
        print(f"Falling back to loading state_dict into config-initialized model.")
        try:
            checkpoint = torch.load(ckpt_path, map_location="cpu")
            state_dict = checkpoint["state_dict"]
            model.load_state_dict(state_dict, strict=False)
        except Exception as e2:
             print(f"Failed to load state_dict: {e2}")
             raise e

    model.eval()

    # setup dataset
    datamodule = cli.datamodule
    datamodule.prepare_data()
    datamodule.affine_model(cli.model)
    datamodule.setup('test')

    stats = {}
    report = {}
    visual_root_dir = path.join(output_dir, f"visuals_{timestamp}") if save_visuals else None

    test_dataloaders = datamodule.test_dataloader()
    
    if not isinstance(test_dataloaders, dict):
        if isinstance(test_dataloaders, list):
             test_dataloaders = {f"test_set_{i}": dl for i, dl in enumerate(test_dataloaders)}
        else:
             dts_name = getattr(getattr(test_dataloaders, "dataset", None), "cls_name", None) or "test_set"
             test_dataloaders = {dts_name: test_dataloaders}

    # Avoid linear RAM growth when not exporting visualizations.
    # Some models return heavy intermediate tensors in predict_step; keeping all
    # batch outputs from trainer.predict can exhaust host memory on large sets.
    if not save_visuals and not hasattr(model, "_od_original_predict_step"):
        model._od_original_predict_step = model.predict_step

        def _od_compact_predict_step(batch, batch_idx, dataloader_idx=0):
            out = model._od_original_predict_step(batch, batch_idx, dataloader_idx)
            if not isinstance(out, dict):
                return out

            compact = {}
            if "probs" in out:
                probs = out["probs"]
                if isinstance(probs, torch.Tensor):
                    probs = probs.detach().cpu()
                compact["probs"] = probs
            if "y" in out:
                y = out["y"]
                if isinstance(y, torch.Tensor):
                    y = y.detach().cpu()
                compact["y"] = y
            if "names" in out:
                compact["names"] = out["names"]
            return compact

        model.predict_step = _od_compact_predict_step

    for dts_name, dataloader in test_dataloaders.items():
        # iterate all videos
        acc_calc = Accuracy(task="BINARY")
        dataset = dataloader.dataset
        dts_stats = {}
        saved_visual_count = 0
        visual_dir = None
        if save_visuals:
            visual_dir = path.join(visual_root_dir, dts_name)
            os.makedirs(visual_dir, exist_ok=True)

        # perform ddp prediction
        batch_results = trainer.predict(
            model=model,
            dataloaders=[dataloader]
        )

        if torch.distributed.is_initialized():
            gathered_results = [None] * torch.distributed.get_world_size()
            torch.distributed.all_gather_object(gathered_results, batch_results)
            torch.distributed.barrier()
        else:
            gathered_results = [batch_results]

        if (trainer.is_global_zero):
            # fetch predict results and aggregate.
            for batch_results in gathered_results:
                for batch_result in batch_results:
                    probs = batch_result["probs"]
                    names = batch_result["names"]
                    y = batch_result["y"]
                    for prob, label, name in zip(probs, y, names):
                        if (not name in dts_stats):
                            dts_stats[name] = StatsRecorder(label)
                        dts_stats[name].update(prob, label)

                    if save_visuals and visual_dir is not None:
                        saved_visual_count = _export_visualizations(
                            batch_result=batch_result,
                            visual_dir=visual_dir,
                            saved_count=saved_visual_count,
                            max_samples=visual_max_samples,
                        )

            # compute the average probability.
            for k in dts_stats:
                dts_stats[k] = dts_stats[k].compute()

            # add straying videos into metric calculation
            # NOTE: stray_videos are for debugging missing entries from csv/protocol.
            # They should NOT be included in metrics, otherwise they can dominate
            # and produce misleadingly perfect scores (e.g. many fake samples
            # defaulting to prob=0.5).
            # for k, v in dataset.stray_videos.items():
            #     dts_stats[k] = dict(
            #         label=v,
            #         prob=0.5,
            #         stray=1
            #     )

            # compute the metric scores
            dataset_labels = []
            dataset_probs = []
            for v in dts_stats.values():
                dataset_labels.append(int(v["label"]))
                p = v["prob"]
                # Ensure plain float probability
                if isinstance(p, torch.Tensor):
                    p = float(p.detach().cpu().item())
                else:
                    p = float(p)
                dataset_probs.append(p)

            dataset_labels = torch.tensor(dataset_labels, dtype=torch.long)
            dataset_probs = torch.tensor(dataset_probs, dtype=torch.float32)

            # torchmetrics expects:
            # - probs: float tensor of shape [N]
            # - labels: int tensor of shape [N]
            accuracy = acc_calc(dataset_probs, dataset_labels).item()
            roc_auc = _binary_roc_auc(dataset_probs, dataset_labels)
            best_acc, best_thr = _best_accuracy_and_threshold(dataset_probs, dataset_labels)
            accuracy = round(accuracy, 3)
            roc_auc = round(roc_auc, 3)
            best_acc = round(best_acc, 3)
            best_thr = round(best_thr, 6)
            logging.info(
                f'[{dts_name}] accuracy@0.5: {accuracy}, roc_auc: {roc_auc}, '
                f'best_accuracy: {best_acc} (thr={best_thr})'
            )
            stats[dts_name] = dts_stats
            report[dts_name] = {
                "accuracy": accuracy,
                "roc_auc": roc_auc,
                "best_accuracy": best_acc,
                "best_threshold": best_thr,
                "saved_visualizations": saved_visual_count,
            }

    if (trainer.is_global_zero):
        # save report and stats.
        with open(path.join(output_dir, f'report_{timestamp}.json'), "w") as f:
            json.dump(report, f, sort_keys=True, indent=4, separators=(',', ': '))

        with open(path.join(output_dir, f'stats_{timestamp}.pickle'), "wb") as f:
            pickle.dump(stats, f)

    return report


if __name__ == "__main__":
    configure_logging()
    _configure_safe_cuda_visible_devices()

    params = parse_args()

    cli = ODLightningCLI(
        run=False,
        trainer_class=ODTrainer,
        save_config_callback=None,
        parser_kwargs={
            # "parser_mode": "omegaconf"
        },
        auto_configure_optimizers=True,
        seed_everything_default=1019,
        args=[
            '-c', params.model_cfg_path,
            '-c', params.data_cfg_path,
            '--trainer.logger=null',
            f'--trainer.devices={params.devices}',
            f'--trainer.precision={params.precision}',
            '--early_stop.monitor=valid/FFPP/auc',
        ],
    )

    model_cfg_path = params.model_cfg_path
    ckpt_path = params.model_ckpt_path
    notes = params.notes
    save_visuals = params.save_visuals
    visual_max_samples = params.visual_max_samples

    # Write outputs next to the checkpoint/run directory instead of config dir.
    output_dir = os.path.dirname(ckpt_path)
    output_dir = os.path.dirname(output_dir)

    inference_driver(
        cli=cli,
        output_dir=output_dir,
        ckpt_path=ckpt_path,
        notes=notes,
        save_visuals=save_visuals,
        visual_max_samples=visual_max_samples,
    )