[Python] coco 데이터셋 train, test, valid 세트로 나누기

 

 

안녕하세요 오랜만입니다.

많이 바빠서 포스팅을 못했는데 오랜만에 글을 쓰네요

 

원래는 제가 coco 데이터셋을 train, test 세트를 주로 8:2 비율로 사용하는데

train, test, valid 세트로 8:1:1로 나누어 사용해야하는 상황에 놓여서 급하게 오픈소스를 고쳐보았습니다.

train, test로 나누는 코드는 많으나 train, test, valid로 나누는 코드는 제가 못찾는지 안보여서 

오픈 소스를 기반으로 약간만 수정했습니다.

 

1. 그냥 coco 데이터셋의 json 파일을 train, test로 나누기를 원한다.

github.com/akarazniewicz/cocosplit

좋은 오픈소스있으니 들어가셔서 사용하시면 됩니다.

train, test, valid 분리는 위 코드를 기반하여 약간의 수정으로 고쳤습니다.

 

2. coco 데이터셋의 json 파일을 train, test, valid로 나누기를 원한다.

import json
import argparse
import funcy
from sklearn.model_selection import train_test_split

parser = argparse.ArgumentParser(description='Splits COCO annotations file into training and test sets.')
parser.add_argument('annotations', metavar='coco_annotations', type=str,
                    help='Path to COCO annotations file.')
parser.add_argument('--train_ratio', type=float, dest='ratio_train', help='set train dataset ratio')
parser.add_argument('--valid_ratio', type=float,  dest='ratio_valid',help='set valid dataset ratio')
parser.add_argument('--test_ratio', type=float,  dest='ratio_test',help='set test dataset ratio')
parser.add_argument('--trainJson_name', type=str, default='train.json', help='Where to store COCO training annotations')
parser.add_argument('--validJson_name', type=str, default='valild.json', help='Where to store COCO valid annotations')
parser.add_argument('--testJson_name', type=str, default='test.json', help='Where to store COCO test annotations')
parser.add_argument('--annotations', dest='annotations', action='store_true',
                    help='Ignore all images without annotations. Keep only these with at least one annotation')

args = parser.parse_args()

ratio_train = args.ratio_train
ratio_valid = args.ratio_valid
ratio_test = args.ratio_test

def save_coco(file, info, licenses, images, annotations, categories):
    with open(file, 'wt', encoding='UTF-8') as coco:
        json.dump({ 'info': info, 'licenses': licenses, 'images': images, 
            'annotations': annotations, 'categories': categories}, coco, indent=2, sort_keys=True)

def filter_annotations(annotations, images):
    image_ids = funcy.lmap(lambda i: int(i['id']), images)
    return funcy.lfilter(lambda a: int(a['image_id']) in image_ids, annotations)

def main(args):
    with open(args.annotations, 'rt', encoding='UTF-8') as annotations:
        coco = json.load(annotations)
        info = coco['info']
        licenses = coco['licenses']
        images = coco['images']
        annotations = coco['annotations']
        categories = coco['categories']

        number_of_images = len(images)

        images_with_annotations = funcy.lmap(lambda a: int(a['image_id']), annotations)

        if args.annotations:
            images = funcy.lremove(lambda i: i['id'] not in images_with_annotations, images)

        train_before, test = train_test_split(
            images, test_size=ratio_test)

        ratio_remaining = 1 - ratio_test
        ratio_valid_adjusted = ratio_valid / ratio_remaining

        train_after, valid = train_test_split(
            train_before, test_size=ratio_valid_adjusted)

        save_coco(args.trainJson_name, info, licenses, train_after, filter_annotations(annotations, train_after), categories)
        save_coco(args.testJson_name, info, licenses, test, filter_annotations(annotations, test), categories)
        save_coco(args.validJson_name, info, licenses, valid, filter_annotations(annotations, valid), categories)

        print("Saved {} entries in {} and {} in {} and {} in {}".format(len(train_after), args.trainJson_name, len(test), args.testJson_name, len(valid), args.validJson_name))


if __name__ == "__main__":
    main(args)

위의 코드를 사용하시면 되는데 사용방법은 제 Github에 상세히 적어놓았습니다.

github.com/dnfwlxo11/coco-split_train_test_valid 로 들어가셔서 Readme 파일을 보고 사용하시면 됩니다.

사용하시고 상단의 star 버튼 눌러 별 한 개씩만 주세요 ㅎㅎ

coco 데이터셋을 사용하시는 분들이 기본적인 파이썬 코드는 당연히 볼 수 있다고 생각하고 사용법은 생략하겠습니다.

만약 json 파일을 나누는게 아닌 단순히 리스트로 된 데이터를 나누기를 원하신다면 line 49~56을 참고하시면 됩니다.

 

 

도움이 되었으면 좋겠습니다 :)