evaluate_ucf101.py

import sys
import os
import glob
import utils
import torch
import random
import time
import argparse
import datetime
import numpy as np
import os.path as osp

import warnings
warnings.filterwarnings("ignore")

import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

from models.capsules_ucf101 import CapsNet


def iou(split, debug=False):
    """
    Calculates the accuracy, f-mAP, and v-mAP over the test set
    """
    from datasets.ucf_dataloader_eval import UCF101DataLoader
    pkl_file_label = "testing_annots.pkl"
    parser = argparse.ArgumentParser(description='evaluation')
    parser.add_argument('--ckpt', type=str, help='experiment name')
    parser.add_argument('--seed', type=int, default=47, help='seed for initializing training.')
    args = parser.parse_args()

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    model = CapsNet().cuda()
    
    n_classes = 24
    clip_batch_size = 14

    model_names = list()
    fmap_best = list()
    vmap_best = list()

    filtered_files = [file.split("/")[-1] for file in glob.glob(osp.join(args.ckpt, 'best_model_' + split + '*.pth'))]

    for saved_wts in sorted(glob.glob(osp.join(args.ckpt, 'best_model_' + split + '*.pth'))):
        model.load_previous_weights(saved_wts)
        model_names.append(saved_wts)
        model.eval()
        model.training = False

        with torch.no_grad():

            validationset = UCF101DataLoader('validation',[224, 224], 1, file_id=pkl_file_label , use_random_start_frame=False)
            val_data_loader = DataLoader(
                dataset=validationset,
                batch_size=1,
                num_workers=8,
                shuffle=False
            )
            
            n_correct, n_vids, n_tot_frames = 0, np.zeros((n_classes, 1)), np.zeros((n_classes, 1))

            # 20 IoUs List
            frame_ious = np.zeros((n_classes, 20))
            video_ious = np.zeros((n_classes, 20))
            iou_threshs = np.arange(0, 20, dtype=np.float32)/20

            for idx, sample in enumerate(val_data_loader):
                video, bbox, label = sample
                video = video[0]
                bbox = bbox[0]
                label = label[0]

                f_skip = 2
                clips = []
                n_frames = video.shape[0]
                # print(n_frames)
                for i in range(0, video.shape[0], 8*f_skip):
                    for j in range(f_skip):
                        b_vid, b_bbox = [], []
                        for k in range(8):
                            ind = i + j + k*f_skip
                            # print(ind)
                            if ind >= n_frames:
                                b_vid.append(np.zeros((1, 224, 224, 3), dtype=np.float32))
                                b_bbox.append(np.zeros((1, 224, 224, 1), dtype=np.float32))
                            else:
                                b_vid.append(video[ind:ind+1, :, :, :])
                                b_bbox.append(bbox[ind:ind+1, :, :, :])
                        clips.append((np.concatenate(b_vid, axis=0), np.concatenate(b_bbox, axis=0), label))
                        if np.sum(clips[-1][1]) == 0:
                            clips.pop(-1)

                if len(clips) == 0:
                    print('Video has no bounding boxes')
                    continue

                batches, gt_segmentations = [], []
                # print(len(clip_batch_size))
                for i in range(0, len(clips), clip_batch_size):
                    x_batch, bb_batch, y_batch = [], [], []
                    for j in range(i, min(i+clip_batch_size, len(clips))):
                        x, bb, y = clips[j]
                        x_batch.append(x)
                        bb_batch.append(bb)
                        y_batch.append(y)
                    batches.append((x_batch, bb_batch, y_batch))
                    gt_segmentations.append(np.stack(bb_batch))

                gt_segmentations = np.concatenate(gt_segmentations, axis=0)
                gt_segmentations = gt_segmentations.reshape((-1, 224, 224, 1))  # Shape N_FRAMES, 112, 112, 1

                segmentations, predictions = [], []
                frames = []
                for x_batch, bb_batch, y_batch in batches:
                    data = np.transpose(np.array(x_batch), [0, 4, 1, 2, 3])
                    data = torch.from_numpy(data).type(torch.cuda.FloatTensor)
                    empty_action = np.ones((len(x_batch),1),np.int)*500
                    empty_action = torch.from_numpy(empty_action).cuda()
                    # print(empty_action)

                    segmentation, pred, _ = model(data, empty_action, empty_action, 0, 0)
                    segmentation = torch.sigmoid(segmentation)
                    segmentation_np = segmentation.cpu().data.numpy()   # B x C x F x H x W -> B x 1 x 8 x 224 x 224
                    segmentation_np = np.transpose(segmentation_np, [0, 2, 3, 4, 1])   
                    save_clip = data.cpu().data.numpy()
                    save_clip = np.transpose(save_clip, [0, 2, 3, 4, 1])
                    # print(segmentation_np.shape, data.shape, save_clip.shape) 
                    frames.append(save_clip)
                    segmentations.append(segmentation_np)
                    predictions.append(pred.cpu().data.numpy())

                # print(f'preds length: {len(predictions)}')
                predictions = np.concatenate(predictions, axis=0)
                #predictions = predictions.reshape((-1, n_classes))
                assert predictions.shape[1] == n_classes
                fin_pred = np.mean(predictions, axis=0)

                fin_pred = np.argmax(fin_pred)
                if fin_pred == label:
                    n_correct += 1

                pred_segmentations = np.concatenate(segmentations, axis=0)
                pred_segmentations = pred_segmentations.reshape((-1, 224, 224, 1))

                pred_segmentations = (pred_segmentations >= 0.5).astype(np.int64)
                seg_plus_gt = pred_segmentations + gt_segmentations

                frames_save = np.concatenate(frames, axis=0)
                frames_save = frames_save.reshape((-1, 224, 224, 3))

                vid_inter, vid_union = 0, 0
                # calculates f_map
                for i in range(gt_segmentations.shape[0]):
                    frame_gt = gt_segmentations[i]
                    if np.sum(frame_gt) == 0:
                        continue

                    n_tot_frames[label] += 1

                    inter = np.count_nonzero(seg_plus_gt[i] == 2)
                    union = np.count_nonzero(seg_plus_gt[i])
                    vid_inter += inter
                    vid_union += union

                    i_over_u = inter / union
                    for k in range(iou_threshs.shape[0]):
                        if i_over_u >= iou_threshs[k]:
                            frame_ious[label, k] += 1

                n_vids[label] += 1
                i_over_u = vid_inter / vid_union
                for k in range(iou_threshs.shape[0]):
                    if i_over_u >= iou_threshs[k]:
                        video_ious[label, k] += 1



            fAP = frame_ious/n_tot_frames
            fmAP = np.mean(fAP, axis=0)
            vAP = video_ious/n_vids
            vmAP = np.mean(vAP, axis=0)

            print('Accuracy:', n_correct / np.sum(n_vids) , 'IoU/fmap/vmap', iou_threshs[4], fmAP[4], vmAP[4], iou_threshs[10], fmAP[10], vmAP[10])
            fmap_best.append(fmAP[10])
            vmap_best.append(vmAP[10])


    best_fmap_model = model_names[fmap_best.index(max(fmap_best))]
    best_vmap_model = model_names[vmap_best.index(max(vmap_best))]
    best_files = list()
    best_files.append(best_fmap_model)
    best_files.append(best_vmap_model)
    for file in filtered_files:
        # print(osp.join(args.ckpt, file))
        if osp.join(args.ckpt, file) not in best_files:
            os.remove(osp.join(args.ckpt, file))
    print(os.listdir(args.ckpt))



iou('train')