[Tensorflow Object Detection API] 1. Creating your own dataset

SUMMARY

Tensorflow Object Detection API를 사용하여 training 및 test를 하기 위한 own dataset를 만드는 방법

1) Preparing image files

우선 다음과 같은 구조로 디렉토리를 만들고,

Object-Detection 

└ images

├ train

└ test

학습시키고 싶은 이미지를 images 디렉토리에 담은 후, train과 test 디렉토리에 9 : 1 비율로 나눠 paste 한다.

          

2) Labeling images

이미지에서 object의 경계상자와 레이블을 지정해줘야하는데 LabelImg를 사용하여 각 이미지에 대해 xml파일로 만들 수 있다.

LabelImg 설치 후, 실행을 하게 되면 w키를 통해 경계상자를 다음과 같이 만들고 ctrl+s를 통해 저장을 하면 아래 처럼 xml파일이 생성된다.

모든 이미지에 대해 이처럼 반복하자 . . .

     

          

LabelImg 이 외의 다른 툴도 있다.

* https://github.com/christopher5106/FastAnnotationTool

* http://imagemagick.org/script/index.php

3) Convert to TFRecord file

이제 xml파일을 TFRecord파일로 변환을 해야하는데, datitran의 github를 참고해보았다.

① 우선 *.xml파일들의 데이터를 하나의 csv파일로 변환하기 위해 xml_to_csv.py스크립트를 수정해보자.

def main():
    image_path = os.path.join(os.getcwd(), 'annotations')
    xml_df = xml_to_csv(image_path)
    xml_df.to_csv('raccoon_labels.csv', index=None)
    print('Successfully converted xml to csv.')

def main():
    for directory in ['train', 'test']:
        image_path = os.path.join(os.getcwd(), 'images/{}'.format(directory))
        xml_df = xml_to_csv(image_path)
        xml_df.to_csv('data/{}_labels.csv'.format(directory), index=None)
        print('Successfully converted xml to csv.')

이후 스크립트를 실행하면 다음과 같이 csv파일이 생성 된다. (xml_to_csv.py스크립트는 Object-Detection 디렉토리에 위치)

② 이제 csv파일을 TFRecord파일로 변환하기 위해 generate_tfrecord.py스크립트를 수정해보자.

# TO-DO replace this with label map
def class_text_to_int(row_label):
    if row_label == 'raccoon':
        return 1
    else:
        None

# TO-DO replace this with label map
def class_text_to_int(row_label):
    if row_label == 'fish':
        return 1
    else:
        None

이후 다음 스크립트 실행을 통해 *.record파일을 얻을 수 있을 것이다.

$ python generate_tfrecord.py --csv_input=data/train_labels.csv --output_path=data/train.record 

$ python generate_tfrecord.py --csv_input=data/test_labels.csv --output_path=data/test.record