test.py

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import utils
import numpy as np
import argparse
from model import AVGN
from datasets import get_test_dataset, inverse_normalize
import cv2

import torch.multiprocessing as mp
import torch.distributed as dist


def get_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_dir', type=str, default='./checkpoints', help='path to save trained model weights')
    parser.add_argument('--experiment_name', type=str, default='avgn_vggss', help='experiment name (experiment folder set to "args.model_dir/args.experiment_name)"')
    parser.add_argument('--save_visualizations', action='store_true', help='Set to store all VSL visualizations (saved in viz directory within experiment folder)')

    # Dataset
    parser.add_argument('--testset', default='flickr', type=str, help='testset (flickr or vggss)')
    parser.add_argument('--test_data_path', default='', type=str, help='Root directory path of data')
    parser.add_argument('--test_gt_path', default='', type=str)
    parser.add_argument('--batch_size', default=1, type=int, help='Batch Size')
    parser.add_argument('--num_class', default=37, type=int)

    # mo-vsl hyper-params
    parser.add_argument('--model', default='avgn')
    parser.add_argument('--out_dim', default=512, type=int)
    parser.add_argument('--num_negs', default=None, type=int)
    parser.add_argument('--tau', default=0.03, type=float, help='tau')

    parser.add_argument('--attn_assign', type=str, default='soft', help="type of audio grouping assignment")
    parser.add_argument('--dim', type=int, default=512, help='dimensionality of features')
    parser.add_argument('--depth_aud', type=int, default=3, help='depth of audio transformers')
    parser.add_argument('--depth_vis', type=int, default=3, help='depth of visual transformers')

    # evaluation parameters
    parser.add_argument('--alpha', default=0.4, type=float, help='alpha')
    parser.add_argument("--dropout_img", type=float, default=0, help="dropout for image")
    parser.add_argument("--dropout_aud", type=float, default=0, help="dropout for audio")

    parser.add_argument('--m_img', default=1.0, type=float, metavar='M', help='momentum for imgnet')
    parser.add_argument('--m_aud', default=1.0, type=float, metavar='M', help='momentum for audnet')

    parser.add_argument('--use_momentum', action='store_true')
    parser.add_argument('--relative_prediction', action='store_true')
    parser.add_argument('--use_mom_eval', action='store_true')
    parser.add_argument('--pred_size', default=0.5, type=float)
    parser.add_argument('--pred_thr', default=0.5, type=float)

    # Distributed params
    parser.add_argument('--workers', type=int, default=8)
    parser.add_argument('--gpu', type=int, default=None)
    parser.add_argument('--world_size', type=int, default=1)
    parser.add_argument('--rank', type=int, default=0)
    parser.add_argument('--node', type=str, default='localhost')
    parser.add_argument('--port', type=int, default=12345)
    parser.add_argument('--dist_url', type=str, default='tcp://localhost:12345')
    parser.add_argument('--multiprocessing_distributed', action='store_true')

    return parser.parse_args()


def main(args):
    mp.set_start_method('spawn')
    args.dist_url = f'tcp://{args.node}:{args.port}'
    print('Using url {}'.format(args.dist_url))

    ngpus_per_node = torch.cuda.device_count()
    if args.multiprocessing_distributed:
        args.world_size = ngpus_per_node
        mp.spawn(main_worker,
                 nprocs=ngpus_per_node,
                 args=(ngpus_per_node, args))

    else:
        main_worker(args.gpu, ngpus_per_node, args)


def main_worker(local_rank, ngpus_per_node, args):
    args.gpu = local_rank
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    # Model dir
    model_dir = os.path.join(args.model_dir, args.experiment_name)
    viz_dir = os.path.join(model_dir, 'viz')
    os.makedirs(viz_dir, exist_ok=True)

    # Setup distributed environment
    if args.multiprocessing_distributed:
        if args.multiprocessing_distributed:
            # For multiprocessing distributed training, rank needs to be the
            # global rank among all the processes
            args.rank = args.rank * ngpus_per_node + local_rank
            print(args.dist_url, args.world_size, args.rank)
        dist.init_process_group(backend='nccl', init_method=args.dist_url,
                                world_size=args.world_size, rank=args.rank)
        torch.distributed.barrier()

    # Create model
    if args.model.lower() == 'avgn':
        audio_visual_model = AVGN(args.tau, args.out_dim, args.dropout_img, args.dropout_aud, args)
    else:
        raise ValueError

    from torchvision.models import resnet18
    object_saliency_model = resnet18(pretrained=True)
    object_saliency_model.avgpool = nn.Identity()
    object_saliency_model.fc = nn.Sequential(
        nn.Unflatten(1, (512, 7, 7)),
        NormReducer(dim=1),
        Unsqueeze(1)
    )
    # object_saliency_model.fc = nn.Unflatten(1, (512, 7, 7))

    if not torch.cuda.is_available():
        print('using CPU, this will be slow')
    else:
        if args.gpu is not None:
            torch.cuda.set_device(args.gpu)
            audio_visual_model.cuda(args.gpu)
            object_saliency_model.cuda(args.gpu)
            if args.multiprocessing_distributed:
                audio_visual_model = torch.nn.parallel.DistributedDataParallel(audio_visual_model, device_ids=[args.gpu])
                object_saliency_model = torch.nn.parallel.DistributedDataParallel(object_saliency_model, device_ids=[args.gpu])

    # Load weights
    ckp_fn = os.path.join(model_dir, 'best.pth')
    if os.path.exists(ckp_fn):
        ckp = torch.load(ckp_fn, map_location='cpu')
        audio_visual_model.load_state_dict({k.replace('module.', ''): ckp['model'][k] for k in ckp['model']})
        print(f'loaded from {os.path.join(model_dir, "best.pth")}')
    else:
        print(f"Checkpoint not found: {ckp_fn}")

    # Dataloader
    testdataset = get_test_dataset(args)
    if args.multiprocessing_distributed:
        sampler = torch.utils.data.DistributedSampler(testdataset, num_replicas=ngpus_per_node, rank=args.rank, shuffle=False)
    else:
        sampler = torch.utils.data.SequentialSampler(testdataset)
    testdataloader = DataLoader(testdataset, batch_size=args.batch_size, sampler=sampler, num_workers=args.workers)
    print("Loaded dataloader.")

    validate(testdataloader, audio_visual_model, object_saliency_model, model_dir, args)


@torch.no_grad()
def validate(testdataloader, audio_visual_model, object_saliency_model, model_dir, args):
    audio_visual_model.train(False)
    object_saliency_model.train(False)

    evaluator_av = utils.EvaluatorFull(default_conf_thr=0.5, pred_size=args.pred_size, pred_thr=args.pred_thr, results_dir=f"{model_dir}/av")
    evaluator_obj = utils.EvaluatorFull(default_conf_thr=0.5, pred_size=args.pred_size,  pred_thr=args.pred_thr, results_dir=f"{model_dir}/obj")
    evaluator_av_obj = utils.EvaluatorFull(default_conf_thr=0., pred_size=args.pred_size,  pred_thr=args.pred_thr, results_dir=f"{model_dir}/av_obj")
    for step, (image, spec, bboxes, name) in enumerate(testdataloader):
        if args.gpu is not None:
            spec = spec.cuda(args.gpu, non_blocking=True)
            image = image.cuda(args.gpu, non_blocking=True)

        # Compute S_AVL
        heatmap_av = audio_visual_model(image.float(), spec.float(), mode='test')[1].unsqueeze(1)
        heatmap_av = F.interpolate(heatmap_av, size=(224, 224), mode='bicubic', align_corners=True)
        heatmap_av = heatmap_av.data.cpu().numpy()

        # Compute S_OBJ
        img_feat = object_saliency_model(image)
        heatmap_obj = F.interpolate(img_feat, size=(224, 224), mode='bicubic', align_corners=True)
        heatmap_obj = heatmap_obj.data.cpu().numpy()

        av_min, av_max = -1. / args.tau, 1. / args.tau
        obj_min, obj_max = 0., 2.5
        min_max_norm = lambda x, xmin, xmax: (x - xmin) / (xmax - xmin)

        # Compute eval metrics and save visualizations
        for i in range(spec.shape[0]):
            gt_map = bboxes['gt_map'][i].data.cpu().numpy()
            bb = bboxes['bboxes'][i]
            bb = bb[bb[:, 0] >= 0].numpy().tolist()

            n = heatmap_av[i, 0].size
            scores_av = min_max_norm(heatmap_av[i, 0], av_min, av_max)
            scores_obj = min_max_norm(heatmap_obj[i, 0], obj_min, obj_max)
            scores_av_obj = scores_av * args.alpha + scores_obj * (1 - args.alpha)

            conf_av = np.sort(scores_av.flatten())[-n//4:].mean()
            conf_obj = np.sort(scores_obj.flatten())[-n//4:].mean()
            conf_av_obj = np.sort(scores_av_obj.flatten())[-n//4:].mean()

            if args.relative_prediction:
                pred_av = utils.normalize_img(scores_av)
                pred_obj = utils.normalize_img(scores_obj)
                pred_av_obj = utils.normalize_img(scores_av_obj)

                thr_av = np.sort(pred_av.flatten())[int(n * args.pred_size)]
                thr_obj = np.sort(pred_obj.flatten())[int(n * args.pred_size)]
                thr_av_obj = np.sort(pred_av_obj.flatten())[int(n * args.pred_size)]
            else:
                pred_av = scores_av
                pred_obj = scores_obj
                pred_av_obj = scores_av_obj

                thr_av = thr_obj = thr_av_obj = args.pred_thr

            evaluator_av.update(bb, gt_map, conf_av, pred_av, thr_av, name[i])
            evaluator_obj.update(bb, gt_map, conf_obj, pred_obj, thr_obj, name[i])
            evaluator_av_obj.update(bb, gt_map, conf_av_obj, pred_av_obj, thr_av_obj, name[i])

            if args.save_visualizations:
                evaluator_av.save_viz(image[i], bb, pred_av, name[i])
                evaluator_obj.save_viz(image[i], bb, pred_obj, name[i])
                evaluator_av_obj.save_viz(image[i], bb, pred_av_obj, name[i])

        print(f'{step+1}/{len(testdataloader)}: AV+OGL-Prec@30={evaluator_av_obj.precision_at_30():.3f} AVL-Prec@30={evaluator_av.precision_at_30():.3f} OGL-Prec@30={evaluator_obj.precision_at_30():.3f}')

    evaluator_av.save_results()
    evaluator_obj.save_results()
    evaluator_av_obj.save_results()

    print('='*20 + ' AVL ' + '='*20)
    stats_av = evaluator_av.finalize_stats()
    print('\n'.join([f' - {k}: {stats_av[k]}' for k in sorted(stats_av.keys()) if stats_av[k] is not np.nan]))

    print('='*20 + ' OGL ' + '='*20)
    stats_obj = evaluator_obj.finalize_stats()
    print('\n'.join([f' - {k}: {stats_obj[k]}' for k in sorted(stats_obj.keys()) if stats_obj[k] is not np.nan]))

    print('='*20 + ' AV+OGL ' + '='*20)
    stats_av_obj = evaluator_av_obj.finalize_stats()
    print('\n'.join([f' - {k}: {stats_av_obj[k]}' for k in sorted(stats_av_obj.keys()) if stats_av_obj[k] is not np.nan]))


class NormReducer(nn.Module):
    def __init__(self, dim):
        super(NormReducer, self).__init__()
        self.dim = dim

    def forward(self, x):
        return x.abs().mean(self.dim)


class Unsqueeze(nn.Module):
    def __init__(self, dim):
        super(Unsqueeze, self).__init__()
        self.dim = dim

    def forward(self, x):
        return x.unsqueeze(self.dim)


if __name__ == "__main__":
    main(get_arguments())