Spaces:
Running
Running
| # -*- coding: UTF-8 -*- | |
| '''================================================= | |
| @Project -> File pram -> inference | |
| @IDE PyCharm | |
| @Author fx221@cam.ac.uk | |
| @Date 03/04/2024 16:06 | |
| ==================================================''' | |
| import argparse | |
| import torch | |
| import torchvision.transforms.transforms as tvt | |
| import yaml | |
| from nets.load_segnet import load_segnet | |
| from nets.sfd2 import load_sfd2 | |
| from dataset.get_dataset import compose_datasets | |
| parser = argparse.ArgumentParser(description='PRAM', formatter_class=argparse.ArgumentDefaultsHelpFormatter) | |
| parser.add_argument('--config', type=str, required=True, help='config of specifications') | |
| parser.add_argument('--landmark_path', type=str, required=True, help='path of landmarks') | |
| parser.add_argument('--feat_weight_path', type=str, default='weights/sfd2_20230511_210205_resnet4x.79.pth') | |
| parser.add_argument('--rec_weight_path', type=str, required=True, help='recognition weight') | |
| parser.add_argument('--online', action='store_true', help='online visualization with pangolin') | |
| if __name__ == '__main__': | |
| args = parser.parse_args() | |
| with open(args.config, 'rt') as f: | |
| config = yaml.load(f, Loader=yaml.Loader) | |
| config['landmark_path'] = args.landmark_path | |
| feat_model = load_sfd2(weight_path=args.feat_weight_path).cuda().eval() | |
| print('Load SFD2 weight from {:s}'.format(args.feat_weight_path)) | |
| # rec_model = get_model(config=config) | |
| rec_model = load_segnet(network=config['network'], | |
| n_class=config['n_class'], | |
| desc_dim=256 if config['use_mid_feature'] else 128, | |
| n_layers=config['layers'], | |
| output_dim=config['output_dim']) | |
| state_dict = torch.load(args.rec_weight_path, map_location='cpu')['model'] | |
| rec_model.load_state_dict(state_dict, strict=True) | |
| print('Load recognition weight from {:s}'.format(args.rec_weight_path)) | |
| img_transforms = [] | |
| img_transforms.append(tvt.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])) | |
| img_transforms = tvt.Compose(img_transforms) | |
| dataset = config['dataset'] | |
| if not args.online: | |
| from localization.loc_by_rec_eval import loc_by_rec_eval | |
| test_set = compose_datasets(datasets=dataset, config=config, train=False, sample_ratio=1) | |
| config['n_class'] = test_set.n_class | |
| loc_by_rec_eval(rec_model=rec_model.cuda().eval(), | |
| loader=test_set, | |
| local_feat=feat_model.cuda().eval(), | |
| config=config, img_transforms=img_transforms) | |
| else: | |
| from localization.loc_by_rec_online import loc_by_rec_online | |
| loc_by_rec_online(rec_model=rec_model.cuda().eval(), | |
| local_feat=feat_model.cuda().eval(), | |
| config=config, img_transforms=img_transforms) | |