File size: 2,824 Bytes
1da285f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | 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='seg')
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 = '613'
# TRAIN_FOLD = '0,1,4,5,7,9'
# VAL_FOLD = '3'
# TEST_FOLD = '2,6,8'
# FOLD_FILE = 'fold_app.csv'
# SPLIT = 'genre'
# TRAIN_FOLD = '0'
# VAL_FOLD = '1'
# TEST_FOLD = '2'
# FOLD_FILE = 'fold_genre.csv'
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_mrcnn.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_mask_rcnn_R_101_FPN_3x.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 ./mask_rcnn_R_101_FPN_3x.pkl '
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 == 'seg', 'only support det'
main(args)
|