| 15 | |
| 16 | |
| 17 | def get_args(): |
| 18 | parser = argparse.ArgumentParser() |
| 19 | parser.add_argument('--model_name', type=str, |
| 20 | default='internimage_t_1k_224') |
| 21 | parser.add_argument('--ckpt_dir', type=str, |
| 22 | default='/mnt/petrelfs/share_data/huangzhenhang/code/internimage/checkpoint_dir/new/cls') |
| 23 | parser.add_argument('--onnx', default=False, action='store_true') |
| 24 | parser.add_argument('--trt', default=False, action='store_true') |
| 25 | |
| 26 | args = parser.parse_args() |
| 27 | args.cfg = os.path.join('./configs', f'{args.model_name}.yaml') |
| 28 | args.ckpt = os.path.join(args.ckpt_dir, f'{args.model_name}.pth') |
| 29 | args.size = int(args.model_name.split('.')[0].split('_')[-1]) |
| 30 | |
| 31 | cfg = get_config(args) |
| 32 | return args, cfg |
| 33 | |
| 34 | |
| 35 | def get_model(args, cfg): |