-
Notifications
You must be signed in to change notification settings - Fork 1
/
test.py
107 lines (86 loc) · 3.63 KB
/
test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
from __future__ import print_function, absolute_import
import argparse
import os.path as osp
import random
import numpy as np
import sys
import torch
from torch import nn
from torch.backends import cudnn
from torch.utils.data import DataLoader
from dg import datasets
from dg import models
from dg.evaluators import Evaluator
from dg.utils.data import transforms as T
from dg.utils.data.preprocessor import Preprocessor
from dg.utils.logging import Logger
from dg.utils.serialization import load_checkpoint, save_checkpoint, copy_state_dict
def get_data(name, data_dir, height, width, batch_size, workers):
root = osp.join(data_dir, name)
dataset = datasets.create(name, root)
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
test_transformer = T.Compose([
T.Resize((height, width), interpolation=3),
T.ToTensor(),
normalizer
])
test_loader = DataLoader(
Preprocessor(list(set(dataset.query) | set(dataset.gallery)),
root=dataset.images_dir, transform=test_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=False, pin_memory=True)
return dataset, test_loader
def main():
args = parser.parse_args()
if args.seed is not None:
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
main_worker(args)
def main_worker(args):
cudnn.benchmark = True
log_dir = osp.dirname(args.resume)
sys.stdout = Logger(osp.join(log_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
# Create data loaders
dataset_target, test_loader_target = \
get_data(args.dataset_target, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
model = models.create(args.arch, pretrained=False, num_features=args.features, dropout=args.dropout, num_classes=0)
model.cuda()
model = nn.DataParallel(model)
# Load from checkpoint
checkpoint = load_checkpoint(args.resume)
copy_state_dict(checkpoint['state_dict'], model)
# Evaluator
evaluator = Evaluator(model)
print("Test on the target domain of {}:".format(args.dataset_target))
evaluator.evaluate(test_loader_target, dataset_target.query, dataset_target.gallery, cmc_flag=True, rerank=args.rerank)
return
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="Testing the model")
# data
parser.add_argument('-dt', '--dataset-target', type=str, required=True,
choices=datasets.names())
parser.add_argument('-b', '--batch-size', type=int, default=256)
parser.add_argument('-j', '--workers', type=int, default=16)
parser.add_argument('--height', type=int, default=256, help="input height")
parser.add_argument('--width', type=int, default=128, help="input width")
# model
parser.add_argument('-a', '--arch', type=str, required=True,
choices=models.names())
parser.add_argument('--features', type=int, default=0)
parser.add_argument('--dropout', type=float, default=0)
# testing configs
parser.add_argument('--resume', type=str, required=True, metavar='PATH')
parser.add_argument('--rerank', action='store_true',
help="evaluation only")
parser.add_argument('--seed', type=int, default=1)
# path
working_dir = osp.dirname(osp.abspath(__file__))
parser.add_argument('--data-dir', type=str, metavar='PATH',
default=osp.join(working_dir, 'data'))
main()