| import os |
| from os.path import join as pjoin |
| import shutil |
| import time |
| import argparse |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument('-t', '--task', type=str, default='interactable') |
| 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') |
| det_script = './detect.py' |
| region2coco_script = './xianyu2coco.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/det/{args.task}/{SPLIT}' |
| os.makedirs(output_dir, exist_ok=True) |
| split_path = args.split_path |
| det_cli = f'python {det_script} ' + \ |
| f'--test_folder {pjoin(split_path, "images/instances_test2017")} ' + \ |
| f'--output_folder {pjoin(output_dir, "regions")} ' |
| to_coco_cli = f'python {region2coco_script} ' + \ |
| f'--xianyu_dir {pjoin(output_dir, "regions")} ' + \ |
| f'--ann {pjoin(split_path, "annotations/instances_test2017.json")} ' + \ |
| f'--output_file {pjoin(output_dir, "results.json")} ' + \ |
| f'--filter_oof ' |
| os.system(det_cli) |
| os.system(to_coco_cli) |
| if os.path.exists(split_path): |
| shutil.rmtree(split_path) |
|
|
|
|
| if __name__ == '__main__': |
| assert args.task in ['interaction', 'semantics', 'interactable'], 'expected task: interaction, semantics or interactable' |
| assert args.task == 'interactable', 'only support interactable' |
| main(args) |
|
|
