def main():
args = parse_args()
logger, final_output_dir, _ = create_logger(
config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
model = models.gsto_hrnet.get_seg_model(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])
)
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
model_state_file = os.path.join(final_output_dir,
'final_state.pth')
logger.info('=> loading model from {}'.format(model_state_file))
model_state_file = config.TEST.MODEL_FILE
pretrained_dict = torch.load(model_state_file)
model_dict = model.state_dict()
pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items()
if k[6:] in model_dict.keys()}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
gpus = list(config.GPUS)
model = nn.DataParallel(model, device_ids=gpus).cuda()
# prepare data
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
test_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TEST_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=False,
flip=False,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TEST.BASE_SIZE,
crop_size=test_size,
downsample_rate=1)
testloader = torch.utils.data.DataLoader(
test_dataset,
batch_size=1,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
start = timeit.default_timer()
if 'val' in config.DATASET.TEST_SET:
mean_IoU, IoU_array, pixel_acc, mean_acc = testval(config,
test_dataset,
testloader,
model)
msg = 'MeanIU: {: 4.4f}, Pixel_Acc: {: 4.4f}, \
Mean_Acc: {: 4.4f}, Class IoU: '.format(mean_IoU,
pixel_acc, mean_acc)
logging.info(msg)
logging.info(IoU_array)
elif 'test' in config.DATASET.TEST_SET:
test(config,
test_dataset,
testloader,
model,
sv_dir=final_output_dir)
end = timeit.default_timer()
logger.info('Mins: %d' % np.int((end - start) / 60))
logger.info('Done')
def main(command_args):
parse_args(command_args)
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
if torch.__version__.startswith('1'):
module = eval('models.' + config.MODEL.NAME)
module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0]))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
model_state_file = os.path.join(final_output_dir, 'best.pth')
pretrained_dict = torch.load(model_state_file)
model_dict = model.state_dict()
#assert set(k[6:] for k in pretrained_dict) == set(model_dict)
pretrained_dict = {
k[6:]: v
for k, v in pretrained_dict.items() if k[6:] in model_dict.keys()
}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
gpus = list(config.GPUS)
model = nn.DataParallel(model, device_ids=gpus).cuda()
# prepare data
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
test_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TEST_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=False,
flip=False,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TEST.BASE_SIZE,
crop_size=test_size,
downsample_rate=1)
testloader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
# Define results to be returned
mean_IoU, pixel_acc, mean_acc = 0., 0., 0.,
if 'val' in config.DATASET.TEST_SET:
mean_IoU, _, pixel_acc, mean_acc = testval(config, test_dataset,
testloader, model)
elif 'test' in config.DATASET.TEST_SET:
print("Cannot extract accuracy from test-sets")
exit(1)
results = {}
results["mean_IoU"] = mean_IoU
results["pixel_acc"] = pixel_acc,
results["mean_acc"] = mean_acc
return results
