test_models.py

import argparse
import time
import os
import sys
import json
import shutil
import numpy as np
import torch
from torch.autograd import Variable
from sklearn.metrics import confusion_matrix
from torch.nn import functional as F

from opts import parse_opts
from model import generate_model
from dataset import get_training_set, get_validation_set, get_test_set
from mean import get_mean, get_std
from spatial_transforms import *
from temporal_transforms import *
from target_transforms import ClassLabel, VideoID
from target_transforms import Compose as TargetCompose
from dataset import get_training_set, get_validation_set, get_test_set
from utils import Logger
from train import train_epoch
from validation import val_epoch
import test
from utils import AverageMeter


"""
def calculate_accuracy(outputs, targets, topk=(1,)):
    maxk = max(topk)
    batch_size = targets.size(0)

    _, pred = outputs.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(targets.view(1, -1).expand_as(pred))
    ret = []
    for k in topk:
        correct_k = correct[:k].float().sum().data[0]
        ret.append(correct_k / batch_size)

    return ret
"""



def calculate_accuracy(outputs, targets, topk=(1,)):
    maxk = max(topk)
    batch_size = targets.size(0)

    _, pred = outputs.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(targets.view(1, -1).expand_as(pred))
    ret = []
    for k in topk:
        correct_k = correct[:k].float().sum().data[0]
        ret.append(correct_k / batch_size)

    return ret


opt = parse_opts()
if opt.root_path != '':
    opt.video_path = os.path.join(opt.root_path, opt.video_path)
    opt.annotation_path = os.path.join(opt.root_path, opt.annotation_path)
    opt.result_path = os.path.join(opt.root_path, opt.result_path)
    if opt.resume_path:
        opt.resume_path = os.path.join(opt.root_path, opt.resume_path)
    if opt.pretrain_path:
        opt.pretrain_path = os.path.join(opt.root_path, opt.pretrain_path)
opt.scales = [opt.initial_scale]
for i in range(1, opt.n_scales):
    opt.scales.append(opt.scales[-1] * opt.scale_step)
opt.arch = '{}-{}'.format(opt.model, opt.model_depth)
opt.mean = get_mean(opt.norm_value, dataset=opt.mean_dataset)
opt.std = get_std(opt.norm_value)

print(opt)
with open(os.path.join(opt.result_path, 'opts.json'), 'w') as opt_file:
    json.dump(vars(opt), opt_file)

torch.manual_seed(opt.manual_seed)

model, parameters = generate_model(opt)
print(model)
pytorch_total_params = sum(p.numel() for p in model.parameters() if
                           p.requires_grad)
print("Total number of trainable parameters: ", pytorch_total_params)

if opt.no_mean_norm and not opt.std_norm:
    norm_method = Normalize([0, 0, 0], [1, 1, 1])
elif not opt.std_norm:
    norm_method = Normalize(opt.mean, [1, 1, 1])
else:
    norm_method = Normalize(opt.mean, opt.std)


spatial_transform = Compose([
    #Scale(opt.sample_size),
    Scale(112),
    CenterCrop(112),
    ToTensor(opt.norm_value), norm_method
    ])
temporal_transform = TemporalCenterCrop(opt.sample_duration)
#temporal_transform = TemporalBeginCrop(opt.sample_duration)
#temporal_transform = TemporalEndCrop(opt.sample_duration)
target_transform = ClassLabel()
validation_data = get_validation_set(
    opt, spatial_transform, temporal_transform, target_transform)
data_loader = torch.utils.data.DataLoader(
    validation_data,
    batch_size=1,
    shuffle=False,
    num_workers=opt.n_threads,
    pin_memory=True)
val_logger = Logger(os.path.join(opt.result_path, 'val.log'), ['epoch', 'loss', 'acc'])

if opt.resume_path:
    print('loading checkpoint {}'.format(opt.resume_path))
    checkpoint = torch.load(opt.resume_path)
    assert opt.arch == checkpoint['arch']

    opt.begin_epoch = checkpoint['epoch']
    model.load_state_dict(checkpoint['state_dict'])

recorder = []

print('run')

model.eval()

batch_time = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()

end_time = time.time()
for i, (inputs, targets) in enumerate(data_loader):
    if not opt.no_cuda:
        targets = targets.cuda(async=True)
    #inputs = Variable(torch.squeeze(inputs), volatile=True)
    inputs = Variable(inputs, volatile=True)
    targets = Variable(targets, volatile=True)
    outputs = model(inputs)
    recorder.append(outputs.data.cpu().numpy().copy())
    #outputs = torch.unsqueeze(torch.mean(outputs, 0), 0)
    prec1, prec5 = calculate_accuracy(outputs, targets, topk=(1, 5))

    top1.update(prec1, inputs.size(0))
    top5.update(prec5, inputs.size(0))

    batch_time.update(time.time() - end_time)
    end_time = time.time()

    print('[{0}/{1}]\t'
          'Time {batch_time.val:.5f} ({batch_time.avg:.5f})\t'
          'prec@1 {top1.avg:.5f} prec@5 {top5.avg:.5f}'.format(
              i + 1,
              len(data_loader),
              batch_time=batch_time,
              top1 =top1,
              top5=top5))

video_pred = [np.argmax(np.mean(x, axis=0)) for x in recorder]
print(video_pred)

with open('annotation_Something/categories.txt') as f:
    lines = f.readlines()
    categories = [item.rstrip() for item in lines]

name_list = [x.strip().split()[0] for x in open('annotation_Something/testlist01.txt')]
order_dict = {e:i for i, e in enumerate(sorted(name_list))}
reorder_output = [None] * len(recorder)
reorder_pred = [None] * len(recorder)
output_csv = []
for i in range(len(recorder)):
    idx = order_dict[name_list[i]]
    reorder_output[idx] = recorder[i]
    reorder_pred[idx] = video_pred[i]
    output_csv.append('%s;%s'%(name_list[i],
                               categories[video_pred[i]]))

    with open('something_predictions.csv','w') as f:
        f.write('\n'.join(output_csv))



print('-----Evaluation is finished------')
print('Overall Prec@1 {:.05f}% Prec@5 {:.05f}%'.format(top1.avg, top5.avg))