[라즈베리파이 4] yolo 학습 데이터셋 생성하기

배뭉·2021년 6월 20일

Project

목록 보기

3/4

Darknet은 Joseph Redmon이 독자적으로 개발한 신경망 프레임워크로서 DNN을 학습시키고 실행시킬 수 있다. 그리고 yolo는 학습된 신경망모델 중 하나이다.
Darknet을 이용하면 yolo 뿐만 아니라 AlexNet, VGG-16, Resnet, Densenet 등 기존의 DNN모델들도 돌려 볼 수 있다.

라벨링한 데이터셋을 yolo 학습 데이터셋으로 만들어 주기

darknet에서 yolo 모델을 학습하기 위해서는 6개의 파일이 더 필요하다

1. classes.names

2. [custom_data].data

3. train.txt

4. test.txt

5. yolo.cfg

6. yolo.weights

1 ~ 4번 파일까지는 파이썬 스크립트로 생성이 가능하다.

1. ~ 4. 파일 준비

#PYTHON
import os

current_path = os.path.abspath(os.curdir)
COLAB_DARKNET_ESCAPE_PATH = '/content/drive/My\ Drive/darknet'
COLAB_DARKNET_PATH = '/content/drive/My Drive/darknet'

YOLO_IMAGE_PATH = current_path + '/custom'
YOLO_FORMAT_PATH = current_path + '/custom'

class_count = 0
test_percentage = 0.2
paths = []

with open(YOLO_FORMAT_PATH + '/' + 'classes.names', 'w') as names, \
     open(YOLO_FORMAT_PATH + '/' + 'classes.txt', 'r') as txt:
    for line in txt:
        names.write(line)  
        class_count += 1
    print ("[classes.names] is created")

with open(YOLO_FORMAT_PATH + '/' + 'custom_data.data', 'w') as data:
    data.write('classes = ' + str(class_count) + '\n')
    data.write('train = ' + COLAB_DARKNET_ESCAPE_PATH + '/custom/' + 'train.txt' + '\n')
    data.write('valid = ' + COLAB_DARKNET_ESCAPE_PATH + '/custom/' + 'test.txt' + '\n')
    data.write('names = ' + COLAB_DARKNET_ESCAPE_PATH + '/custom/' + 'classes.names' + '\n')
    data.write('backup = backup')
    print ("[custom_data.data] is created")

os.chdir(YOLO_IMAGE_PATH)
for current_dir, dirs, files in os.walk('.'):
    for f in files:
        if f.endswith('.jpg'):
            image_path = COLAB_DARKNET_PATH + '/custom/' + f
            paths.append(image_path + '\n')

paths_test = paths[:int(len(paths) * test_percentage)]
paths = paths[int(len(paths) * test_percentage):]

with open(YOLO_FORMAT_PATH + '/' + 'train.txt', 'w') as train_txt:
    for path in paths:
        train_txt.write(path)
    print ("[train.txt] is created")

with open(YOLO_FORMAT_PATH + '/' + 'test.txt', 'w') as test_txt:
    for path in paths_test:
        test_txt.write(path)
####     print ("[test.txt] is created")

추후에 구글 colab에서 darknet환경을 구축하고 yolo학습을 진행 할 것이다.
그렇기 때문에 이 스크립트의 모든 path들은 colab환경을 기준으로 설정해주었다.

파이썬 스크립트가 실행되는 워킹디텍토리를 기준으로 라벨링한 모든 데이터셋들은 custom이라는 하위디렉터리를 생성해서 그곳으로 옮겨주었다.

위 파이썬 스크립트를 실행시키면

그전 포스팅의 labelImg에서 생성했던 classes.txt 파일을 이용해서 1. classes.names파일을 생성한다.
그리고 [custom_data].data라는 파일에 COLAB_DARKNET_ESCAPE_PATH에 설정해둔 path를 기준으로 custom이라는 하위디렉터리로 train, valid, test 데이터들의 path를 저장하는데 .data파일의 이름은 아무거나 지정해도 상관없다.
나는 test말고 valid에 text.txt를 설정하였다.

classes는 익은 딸기, 덜익은 딸기 2가지 클래스 밖에 없고 그다음
train, valid 경로를 설정한 뒤 classes파일로 먼저 생성한 classes.names의 경로, 그리고 마지막으로 학습이 진행되면서 중간에 1000epoches 마다 학습된 모델을 저장할 backup 경로까지 설정해주는
2. custum_data.data파일을 생성한다.

그 다음은 알아서 처음에 설정해둔 test_percentage 라는 변수에 저장해 둔 값에 맞게 랜덤으로 비율을 나눠서 데이터 셋을 분할하여
3. train.txt, 4. test.txt파일을 생성한다.

이렇게 1. ~ 4. 까지의 파일은 생성을 완료하였다.

5. ~ 6. 파일 준비

5. yolo.cfg, 6. yolo.weights파일은 Darknet 홈페이지에서 다운받을 수 있다.

여기서 나는 라즈베리파이에서 모델을 실행할 것이기 때문에, YOLO의 경량화 모델인 YOLOv3-tiny의 Cfg파일과 Weights파일을 다운받았다.

적절하게 학습을 하기 위해서는 본인이 원하는 클래스의 수에 맞게 Cfg파일을 수정해줘야 한다.

batch : 한번 데이터를 읽어올 때 가져오는 데이터 수
subdivision : 미니멀하게 가져오는 데이터 수
가장 먼저 원본파일은 Training의 batch값과 subdivision값들이 주석 처리 되어있는데 이 주석들을 해제하고 Test의 batch사이즈값과 subdivision값을 주석 처리한다.
그 다음 max_batches에는 (학습시킬 모델의 클래스 수 * 2000)정도로 설정해준다. 2000이 최솟값으로 이보다 많은 값을 설정해도 상관없지만 이 값 이후의 모델들은 오버피팅된 모델일 확률이 높아서 비효율적이라 생각한다. 실제로 4000만큼 학습한 모델보다 2000만큼 학습된 중간산출 모델이 가장 성능이 좋았다.
나는 익은 딸기와 덜익은 딸기 2가지 클래스를 학습시킬것이기 때문에 4000으로 값을 설정해주었다.
그다음 스텝에는 맥스배치 사이즈의 80%, 90%로 설정해준다.