| import os |
| from os.path import join as pjoin |
| import shutil |
| import argparse |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument('-t', '--task', type=str, required=True) |
| parser.add_argument('-f', '--format', type=str, default='det') |
| parser.add_argument('-ft', '--finetune', action='store_true') |
| parser.add_argument('-d', '--device', type=str, default='0') |
| parser.add_argument('-sp', '--split_path', type=str, default='./dataset') |
| args = parser.parse_args() |
|
|
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
|
|
| SPLIT = 'cat' |
| TRAIN_FOLD = '0' |
| VAL_FOLD = '1' |
| TEST_FOLD = '2' |
| FOLD_FILE = 'fold_cat.csv' |
|
|
| DATASET_FOLDER_PATH = '../../dataset/' |
| NUM_CLASSES_DICT = {'interaction': 53, 'semantics': 766, 'interactable': 1} |
|
|
| split_script = pjoin(DATASET_FOLDER_PATH, 'split_coco.py') |
| trainer_script = './train_net.py' |
|
|
| def generate_dataset(args): |
| ann_file = pjoin(DATASET_FOLDER_PATH, f'data/coco_{args.format}/annotations/{args.task}.json') |
| img_dir = pjoin(DATASET_FOLDER_PATH, 'data/coco_merged/images', args.task) |
| fold_file = pjoin(DATASET_FOLDER_PATH, FOLD_FILE) |
| split_path = args.split_path |
| if os.path.exists(split_path): |
| shutil.rmtree(split_path) |
| cli = f'python {split_script} ' + \ |
| f'--ann_file {ann_file} ' + \ |
| f'--img_dir {img_dir} ' + \ |
| f'--output_path {split_path} ' + \ |
| f'--fold_file {fold_file} ' + \ |
| f'--train_fold {TRAIN_FOLD} ' + \ |
| f'--val_fold {VAL_FOLD} ' + \ |
| f'--test_fold {TEST_FOLD} ' |
| os.system(cli) |
|
|
|
|
| def main(args): |
| generate_dataset(args) |
| output_dir = f'./output/{args.format}/{args.task}{"_finetune" if args.finetune else ""}/{SPLIT}' |
| split_path = args.split_path |
| device = args.device.split(',') |
| cli = f'python {trainer_script} ' + \ |
| f'--config-file ./configs/my_CenterNet2_DLA-BiFPN-P3_4x.yaml ' + \ |
| f'--dataset_root {split_path} ' + \ |
| f'--num_classes {NUM_CLASSES_DICT[args.task]} ' + \ |
| f'--manual_device {",".join(device)} ' + \ |
| f'--num-gpus {len(device)} ' + \ |
| f'--resume ' + \ |
| f'OUTPUT_DIR {output_dir} ' |
| if args.finetune and not os.path.exists(output_dir): |
| cli += f'MODEL.WEIGHTS ./CenterNet2_DLA-BiFPN-P3_4x.pth ' |
| cli += f'SOLVER.RESET_ITER True ' |
| os.system(cli) |
| if os.path.exists(split_path): |
| shutil.rmtree(split_path) |
| |
|
|
| if __name__ == '__main__': |
| assert args.format in ['det', 'seg'], 'expected format: det or seg' |
| assert args.task in ['interaction', 'semantics', 'interactable'], 'expected task: interaction, semantics or interactable' |
| assert args.format == 'det', 'only support det' |
| main(args) |
|
|
