| import os
|
| import json
|
| import shutil
|
| import argparse
|
| import random
|
| import pandas as pd
|
|
|
| FOLD_CNT = 10
|
|
|
|
|
| def generate_fold(args):
|
| df = pd.DataFrame(columns=['img_id', 'fold'])
|
| json_file_path = args.ann_file
|
| assert os.path.exists(json_file_path), f'json file not found: {json_file_path}'
|
| with open(json_file_path, 'r') as f:
|
| json_data = json.load(f)
|
|
|
|
|
| random.seed(1)
|
| random.shuffle(json_data['images'])
|
| fold_size = len(json_data['images']) // FOLD_CNT
|
| for image in json_data['images']:
|
| df.loc[len(df), df.columns] = [image['id'], min(len(df) // fold_size, FOLD_CNT - 1)]
|
| df.to_csv(args.fold_file, index=False)
|
|
|
|
|
| def generate_fold_app(args):
|
| df = pd.DataFrame(columns=['img_id', 'fold'])
|
| json_file_path = args.ann_file
|
| assert os.path.exists(json_file_path), f'json file not found: {json_file_path}'
|
| with open(json_file_path, 'r') as f:
|
| json_data = json.load(f)
|
|
|
| apps = set()
|
|
|
| for image in json_data['images']:
|
| appid = str(image['id'])[:-3]
|
| apps.add(appid)
|
|
|
| apps = list(apps)
|
| bel_app = {}
|
|
|
|
|
| apps = sorted(apps)
|
| random.seed(1)
|
| random.shuffle(apps)
|
| fold_size = len(apps) // FOLD_CNT
|
| for app in apps:
|
| bel_app[app] = min(len(bel_app) // fold_size, FOLD_CNT - 1)
|
| for image in json_data['images']:
|
| appid = str(image['id'])[:-3]
|
| df.loc[len(df), df.columns] = [image['id'], bel_app[appid]]
|
|
|
| df = df.sort_values(by='fold')
|
| df.to_csv(args.fold_file, index=False)
|
|
|
|
|
| def split(args):
|
| json_file_path = args.ann_file
|
| image_path = args.img_dir if args.img_dir and os.path.exists(args.img_dir) else None
|
| fold_file = args.fold_file
|
| output_path = args.output_path
|
|
|
| assert os.path.exists(json_file_path), f'json file not found: {json_file_path}'
|
|
|
| assert os.path.exists(fold_file), f'fold file not found: {fold_file}'
|
|
|
| with open(json_file_path, 'r') as f:
|
| json_data = json.load(f)
|
|
|
| df = pd.read_csv(fold_file)
|
|
|
| train_dataset = {'images': list(), 'categories': json_data['categories'], 'annotations': list()}
|
| val_dataset = {'images': list(), 'categories': json_data['categories'], 'annotations': list()}
|
| test_dataset = {'images': list(), 'categories': json_data['categories'], 'annotations': list()}
|
|
|
| train_folds = [int(fold) for fold in args.train_folds.split(',')] if args.train_folds else []
|
| val_folds = [int(fold) for fold in args.val_folds.split(',')] if args.val_folds else []
|
| test_folds = [int(fold) for fold in args.test_folds.split(',')] if args.test_folds else []
|
|
|
| print(f'train folds: {train_folds}')
|
| print(f'val folds: {val_folds}')
|
| print(f'test folds: {test_folds}')
|
|
|
| train_imgid = set()
|
| val_imgid = set()
|
| test_imgid = set()
|
|
|
| for _, row in df.iterrows():
|
| if row['fold'] in train_folds:
|
| train_imgid.add(row['img_id'])
|
| elif row['fold'] in val_folds:
|
| val_imgid.add(row['img_id'])
|
| elif row['fold'] in test_folds:
|
| test_imgid.add(row['img_id'])
|
| else:
|
| raise ValueError(f'fold not found: {row["fold"]}')
|
|
|
| for image in json_data['images']:
|
| if image['id'] in train_imgid:
|
| train_dataset['images'].append(image)
|
| elif image['id'] in val_imgid:
|
| val_dataset['images'].append(image)
|
| elif image['id'] in test_imgid:
|
| test_dataset['images'].append(image)
|
| else:
|
| raise ValueError(f'image id not found: {image["id"]}')
|
|
|
| for annotation in json_data['annotations']:
|
| if annotation['image_id'] in train_imgid:
|
| train_dataset['annotations'].append(annotation)
|
| elif annotation['image_id'] in val_imgid:
|
| val_dataset['annotations'].append(annotation)
|
| elif annotation['image_id'] in test_imgid:
|
| test_dataset['annotations'].append(annotation)
|
| else:
|
| raise ValueError(f'annotation image_id not found: {annotation["image_id"]}')
|
|
|
| print(f'train dataset: {len(train_dataset["images"])} images, {len(train_dataset["annotations"])} annotations')
|
| print(f'val dataset: {len(val_dataset["images"])} images, {len(val_dataset["annotations"])} annotations')
|
| print(f'test dataset: {len(test_dataset["images"])} images, {len(test_dataset["annotations"])} annotations')
|
|
|
| print(os.path.abspath(output_path))
|
| train_image_path = os.path.join(output_path, 'images', 'instances_train2017')
|
| val_image_path = os.path.join(output_path, 'images', 'instances_val2017')
|
| test_image_path = os.path.join(output_path, 'images', 'instances_test2017')
|
| os.makedirs(output_path, exist_ok=True)
|
| os.makedirs(train_image_path, exist_ok=True)
|
| os.makedirs(val_image_path, exist_ok=True)
|
| os.makedirs(test_image_path, exist_ok=True)
|
| os.makedirs(os.path.join(output_path, 'annotations'), exist_ok=True)
|
|
|
| with open(os.path.join(output_path, 'annotations', 'instances_train2017.json'), 'w') as f:
|
| json.dump(train_dataset, f, indent=4)
|
| with open(os.path.join(output_path, 'annotations', 'instances_val2017.json'), 'w') as f:
|
| json.dump(val_dataset, f, indent=4)
|
| with open(os.path.join(output_path, 'annotations', 'instances_test2017.json'), 'w') as f:
|
| json.dump(test_dataset, f, indent=4)
|
|
|
| if image_path:
|
| for image in train_dataset['images']:
|
| shutil.copy(os.path.join(image_path, image['file_name']), train_image_path)
|
| for image in val_dataset['images']:
|
| shutil.copy(os.path.join(image_path, image['file_name']), val_image_path)
|
| for image in test_dataset['images']:
|
| shutil.copy(os.path.join(image_path, image['file_name']), test_image_path)
|
|
|
|
|
| def main(args):
|
| if args.gen:
|
| generate_fold_app(args)
|
| split(args)
|
|
|
|
|
| if __name__ == '__main__':
|
| parser = argparse.ArgumentParser()
|
| parser.add_argument('--ann_file', type=str, required=True)
|
| parser.add_argument('--img_dir', type=str, required=True)
|
| parser.add_argument('--output_path', type=str, default='./coco_split')
|
| parser.add_argument('--train_folds', type=str, default=None)
|
| parser.add_argument('--val_folds', type=str, default=None)
|
| parser.add_argument('--test_folds', type=str, default=None)
|
| parser.add_argument('--gen', action='store_true')
|
| parser.add_argument('--fold_file', type=str, required=True)
|
| args = parser.parse_args()
|
| main(args)
|
|
|
