Custom YOLO project - YOLO 이미지 데이터 만들기

1. 이미지 레이블링 하기

 yolo가 인식하는 이미지의 좌표형식은 다음과 같다.

 

 위 사진의 좌표는 옆에 메모장 처럼 상대좌표이다. 이렇게 이미지와 txt파일을 생성해야 하는데 이렇게 만드는 방법은 여러가지가 있는데 나는 labelimg 파일을 사용해서 만들어 보겠다.

  

 cmd창을 열어서 pip install LabelImg를 입력해 LabelImg를 다운받은 후 cmd창에 LabelImg를 입력해 LabelImg 프로그램을 실행시킨다.

 이렇게 이미지 마다 레이블을 붙여줄 수 있다. 이렇게 작업이 끝나면 사진 옆에 [img이름].txt 파일이 생기게 된다.

 

2. .names, .data, .cfg , train.txt, test.txt파일 생성하기

다음 코드를 통해 만들수 있다. 먼저 이번 프로젝트에서는 코랩을 사용할 것이기 때문에 경로 설정을 다음과 같이 해준다.

import os

current_path=os.path.abspath(os.curdir)
COLAB_DARKNET_ESCAPE_PATH='/content/drive/MyDrive/darknet'  #코랩 사용
COLAB_DARKNET_PATH='content/drvie/MyDrive/darknet'

 그리고 이미지가 들어있는 폴더의 경로를 선언해준다.

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

class의 수와 테스트셋 비율을 설정해 준다.

class_count=0
test_percentage=0.2
paths=[]

 

class.names 파일을 다음과 같이 생성해준다.

# class.names 파일 생성하기
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')

 

 저장되어있는 classes.txt 파일로 부터(image labeling할때 생성됨) classes.names 파일을 생성해 열어서 한줄씩 쓰고 저장한다는 의미이다.

 

custom_data.data 파일을 다음과 같이 생성해준다.

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')

 이 파일은 후에 darknet을 통해 yolo를 실행했을때 각각의 파일들의 위치정보를 주는 역할을 한다. 그러므로 사용자들이 작업하는 폴더 위치를 잘 인지하고 수정해줄 필요가 있다.

 

os.chdir(YOLO_IMAGE_PATH)
for current_dir,dirs,files in os.walk('.'):
    for f in files: # custom 폴더 아래에 있는 이미지 파일들을 읽어온다.
        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):]

 위 코드는 custom 폴더 아래에 있는 이미지 파일들을 하나씩 읽으면서 paths 리스트에 이미지 파일들의 경로를 추가한다. 이를 통해 test 데이터들의 path와 train 데이터들의 path를 나눠줄 수 있다.

 

 이후 train.txt와 test.txt를 다음과 같이 만들어서 저장해준다.

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')

 

3. cfg 파일 수정해주기

 우리의 yolo는 2가지의 클래스만 구별할 것이기 때문에 다음과 같이 수정해줘야할 부분이 있다.

변경후 custom폴더에 각각 저장해준다.