()
| 39 | |
| 40 | @logger.catch |
| 41 | def main(): |
| 42 | args = get_parser() |
| 43 | args.output_dir = os.path.join(args.output_folder, args.exp_name) |
| 44 | if args.visualize: |
| 45 | args.vis_dir = os.path.join(args.output_dir, "vis") |
| 46 | os.makedirs(args.vis_dir, exist_ok=True) |
| 47 | |
| 48 | # logger |
| 49 | setup_logger(args.output_dir, |
| 50 | distributed_rank=0, |
| 51 | filename="test.log", |
| 52 | mode="a") |
| 53 | logger.info(args.test_split) |
| 54 | |
| 55 | # build dataset & dataloader |
| 56 | test_data = RefDataset(lmdb_dir=args.test_lmdb, |
| 57 | mask_dir=args.mask_root, |
| 58 | dataset=args.dataset, |
| 59 | split=args.test_split, |
| 60 | mode='test', |
| 61 | input_size=args.input_size, |
| 62 | word_length=args.word_len) |
| 63 | test_loader = torch.utils.data.DataLoader(test_data, |
| 64 | batch_size=1, |
| 65 | shuffle=False, |
| 66 | num_workers=4, |
| 67 | pin_memory=True) |
| 68 | |
| 69 | # build model |
| 70 | model = build_segmenter(args, DDP=False) |
| 71 | logger.info(model) |
| 72 | |
| 73 | if os.path.isfile(args.weight): |
| 74 | logger.info("=> loading checkpoint '{}'".format(args.weight)) |
| 75 | checkpoint = torch.load(args.weight) |
| 76 | model.module.load_state_dict(checkpoint['model_state_dict'], strict=True) |
| 77 | logger.info("=> loaded checkpoint '{}'".format(args.weight)) |
| 78 | else: |
| 79 | raise ValueError( |
| 80 | "=> resume failed! no checkpoint found at '{}'. Please check args.resume again!" |
| 81 | .format(args.weight)) |
| 82 | |
| 83 | # inference |
| 84 | inference(test_loader, model, args) |
| 85 | |
| 86 | |
| 87 | if __name__ == '__main__': |
no test coverage detected