| 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('-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} |
|
|
| split_script = pjoin(DATASET_FOLDER_PATH, 'split_coco.py') |
| trainer_script = './cnn_pytorch/train.py' |
|
|
|
|
| def generate_dataset(args): |
| ann_file = pjoin(DATASET_FOLDER_PATH, f'data/coco_det/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}/{SPLIT}' |
| os.makedirs(output_dir, exist_ok=True) |
| shutil.copy2(f'./output/first_stage/{args.task}/{SPLIT}/first_stage.json', pjoin(output_dir, 'first_stage.json')) |
| if args.task == 'interactable': |
| shutil.copy2(pjoin(output_dir, 'first_stage.json'), pjoin(output_dir, 'results.json')) |
| return |
| split_path = args.split_path |
| device = args.device.split(',') |
| train_cli = f'python {trainer_script} ' + \ |
| f'--task {args.task} ' + \ |
| f'--output_dir {output_dir} ' + \ |
| f'--dataset_root {split_path} ' + \ |
| f'--batch_size 64 ' + \ |
| f'--epoch 30 ' + \ |
| f'--lr 5 ' + \ |
| f'--resume ' + \ |
| f'--best_metric acc ' + \ |
| f'--device {device[0]} ' |
| os.system(train_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) |
|
|
