ref-avs.code/inference.py

"""Distributed inference on Ref-AVS (test_s / test_u / test_n); uses Trainer.valid / valid_null like main.py."""
import os
import pathlib
import argparse
import random

import numpy
import torch
from easydict import EasyDict


_real_mkdir = pathlib.Path.mkdir


def _safe_mkdir(self, mode=0o777, parents=False, exist_ok=False):
    try:
        return _real_mkdir(self, mode, parents, exist_ok=exist_ok)
    except PermissionError:
        pass


pathlib.Path.mkdir = _safe_mkdir


def seed_it(seed):
    random.seed(seed)
    os.environ["PYTHONSEED"] = str(seed)
    numpy.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.enabled = True


class _DummyTensorboard:
    """Minimal Tensorboard stub so Trainer.valid / valid_null run without wandb logging."""

    def upload_wandb_info(self, info_dict):
        pass

    def upload_wandb_image(self, *args, **kwargs):
        pass


def main(local_rank, ngpus_per_node, hyp_param):
    hyp_param.local_rank = local_rank
    torch.distributed.init_process_group(
        backend='nccl',
        init_method='env://',
        rank=hyp_param.local_rank,
        world_size=hyp_param.gpus,
    )
    seed_it(local_rank + hyp_param.seed)
    torch.cuda.set_device(hyp_param.local_rank)

    import model.visual.sam2  # noqa: F401 — registers Hydra config store
    from hydra import compose
    from omegaconf import OmegaConf

    arch_h = compose(config_name='configs/auralfuser/architecture.yaml')
    OmegaConf.resolve(arch_h)
    hyp_param.aural_fuser = OmegaConf.to_container(arch_h.aural_fuser, resolve=True)

    train_cfg = compose(config_name='configs/training/sam2_training_config.yaml')
    OmegaConf.resolve(train_cfg)
    hyp_param.contrastive_learning = OmegaConf.to_container(train_cfg.contrastive_learning, resolve=True)

    hyp_param.image_size = 1024
    hyp_param.image_embedding_size = int(hyp_param.image_size / 16)

    from model.mymodel import AVmodel
    av_model = AVmodel(hyp_param).cuda(hyp_param.local_rank)
    if not hyp_param.inference_ckpt:
        raise ValueError("--inference_ckpt is required for inference.")

    ckpt_sd = torch.load(hyp_param.inference_ckpt, map_location="cpu")
    if not isinstance(ckpt_sd, dict):
        raise TypeError("Checkpoint must be a state_dict dictionary.")
    if any(k.startswith("v_model.") or k.startswith("aural_fuser.") for k in ckpt_sd):
        av_model.load_state_dict(ckpt_sd, strict=True)
    else:
        av_model.aural_fuser.load_state_dict(ckpt_sd, strict=True)

    av_model = torch.nn.parallel.DistributedDataParallel(
        av_model, device_ids=[hyp_param.local_rank], find_unused_parameters=False,
    )
    av_model.eval()

    from dataloader.dataset import AV
    from dataloader.visual.visual_augmentation import Augmentation as VisualAugmentation
    from dataloader.audio.audio_augmentation import Augmentation as AudioAugmentation
    from torch.utils.data import DataLoader, Subset
    from torch.utils.data.distributed import DistributedSampler

    visual_aug = VisualAugmentation(
        hyp_param.image_mean, hyp_param.image_std,
        hyp_param.image_size, hyp_param.image_size,
        hyp_param.scale_list, ignore_index=hyp_param.ignore_index,
    )
    audio_aug = AudioAugmentation(mono=True)

    max_batches = getattr(hyp_param, "inference_max_batches", 0) or 0
    val_batch_size = getattr(hyp_param, "inference_val_batch_size", 4)

    def _test_loader(split):
        ds = AV(
            split=split,
            augmentation={"visual": visual_aug, "audio": audio_aug},
            param=hyp_param,
            root_path=hyp_param.data_root_path,
        )
        if max_batches > 0:
            n_samples = min(max_batches * val_batch_size, len(ds))
            ds = Subset(ds, range(n_samples))
        sampler = DistributedSampler(ds, shuffle=False)
        return DataLoader(
            ds,
            batch_size=val_batch_size,
            sampler=sampler,
            num_workers=hyp_param.num_workers,
        )

    test_s_loader = _test_loader('test_s')
    test_u_loader = _test_loader('test_u')
    test_n_loader = _test_loader('test_n')

    from trainer.train import Trainer
    from utils.foreground_iou import ForegroundIoU
    from utils.foreground_fscore import ForegroundFScore
    from utils.foreground_s import ForegroundS

    metrics = {
        "foreground_iou": ForegroundIoU(),
        "foreground_f-score": ForegroundFScore(hyp_param.local_rank),
        "foreground_s": ForegroundS(),
    }
    trainer = Trainer(hyp_param, loss=None, tensorboard=_DummyTensorboard(), metrics=metrics)

    test_s_iou, test_s_f = trainer.valid(
        epoch=0, dataloader=test_s_loader, model=av_model, process='test_s',
    )
    torch.cuda.empty_cache()

    test_u_iou, test_u_f = trainer.valid(
        epoch=0, dataloader=test_u_loader, model=av_model, process='test_u',
    )
    torch.cuda.empty_cache()

    test_n_s = trainer.valid_null(
        epoch=0, dataloader=test_n_loader, model=av_model, process='test_n',
    )
    torch.cuda.empty_cache()

    if hyp_param.local_rank <= 0:
        print("\n========== Ref-AVS inference (same splits / metrics as training valid) ==========")
        print("  test_s   f_iou={}  f_f-score={}".format(test_s_iou, test_s_f))
        print("  test_u   f_iou={}  f_f-score={}".format(test_u_iou, test_u_f))
        print("  test_n   f_s={}".format(test_n_s))
        print("=======================================================\n")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Ref-AVS inference: test_s / test_u / test_n')

    parser.add_argument('--local_rank', type=int, default=-1,
                        help='multi-process training for DDP')
    parser.add_argument('-g', '--gpus', default=1, type=int,
                        help='number of gpus per node')
    parser.add_argument('--batch_size', default=1, type=int,
                        help='unused at inference (validation uses inference_val_batch_size)')
    parser.add_argument('--epochs', default=80, type=int, help='unused')
    parser.add_argument('--lr', default=1e-5, type=float, help='unused')
    parser.add_argument('--online', action='store_true', help='unused')
    parser.add_argument(
        '--inference_ckpt', type=str, required=True,
        help='Trained AuralFuser checkpoint (.pth). SAM2 from backbone_weight in configs.',
    )
    parser.add_argument('--inference_max_batches', type=int, default=0,
                        help='0 = full split; >0 = first N batches per split (debug)')
    parser.add_argument('--inference_val_batch_size', type=int, default=4,
                        help='Validation batch size (default 4, same as main.py _test_loader)')

    args = parser.parse_args()

    from configs.config import C
    args = EasyDict({**C, **vars(args)})

    _repo = pathlib.Path(__file__).resolve().parent
    _workspace = _repo.parent
    args.data_root_path = str(_workspace / 'REFAVS')
    args.backbone_weight = str(_workspace / 'ckpts' / 'sam_ckpts' / 'sam2_hiera_large.pt')
    args.audio.PRETRAINED_VGGISH_MODEL_PATH = str(_workspace / 'ckpts' / 'vggish-10086976.pth')
    args.saved_dir = '/tmp/ref_avs_infer_ckpt'
    pathlib.Path(args.saved_dir).mkdir(parents=True, exist_ok=True)

    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '9902'

    torch.multiprocessing.spawn(main, nprocs=args.gpus, args=(args.gpus, args))