/
test-seg.py
72 lines (62 loc) · 2.65 KB
/
test-seg.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
import argparse
import os
import os.path as osp
import pprint
import warnings
from torch.utils import data
from seg.model.deeplabv2 import get_deeplab_v2
from seg.dataset.cityscapes import CityscapesDataSet
from seg.config import cfg, cfg_from_file
def get_arguments():
"""
Parse input arguments
"""
parser = argparse.ArgumentParser(description="Code for evaluation")
parser.add_argument('--cfg', type=str, default=None,
help='optional config file', )
parser.add_argument("--exp-suffix", type=str, default=None,
help="optional experiment suffix")
return parser.parse_args()
def main(config_file, exp_suffix):
# LOAD ARGS
assert config_file is not None, 'Missing cfg file'
cfg_from_file(config_file)
# auto-generate exp name if not specified
if cfg.EXP_NAME == '':
cfg.EXP_NAME = f'{cfg.SOURCE}2{cfg.TARGET}_{cfg.TRAIN.MODEL}_{cfg.TRAIN.DA_METHOD}'
if exp_suffix:
cfg.EXP_NAME += f'_{exp_suffix}'
# auto-generate snapshot path if not specified
if cfg.TEST.SNAPSHOT_DIR[0] == '':
cfg.TEST.SNAPSHOT_DIR[0] = osp.join(cfg.EXP_ROOT_SNAPSHOT, cfg.EXP_NAME)
os.makedirs(cfg.TEST.SNAPSHOT_DIR[0], exist_ok=True)
print('Using config:')
pprint.pprint(cfg)
# load models
models = []
n_models = len(cfg.TEST.MODEL)
if cfg.TEST.MODE == 'best':
assert n_models == 1, 'Not yet supported'
for i in range(n_models):
if cfg.TEST.MODEL[i] == 'DeepLabv2':
model = get_deeplab_v2(num_classes=cfg.NUM_CLASSES,
multi_level=cfg.TEST.MULTI_LEVEL[i])
models.append(model)
# dataloaders
test_dataset = CityscapesDataSet(root=cfg.DATA_DIRECTORY_TARGET,
list_path=cfg.DATA_LIST_TARGET,
set=cfg.TEST.SET_TARGET,
info_path=cfg.TEST.INFO_TARGET,
crop_size=cfg.TEST.INPUT_SIZE_TARGET,
mean=cfg.TEST.IMG_MEAN,
labels_size=cfg.TEST.OUTPUT_SIZE_TARGET)
test_loader = data.DataLoader(test_dataset,
batch_size=cfg.TEST.BATCH_SIZE_TARGET,
num_workers=cfg.NUM_WORKERS,
shuffle=False,
pin_memory=True)
# eval
evaluate_domain_adaptation(models, test_loader, cfg)
if __name__ == '__main__':
args = get_arguments()
main(args.cfg, args.exp_suffix)