[Object Detection] YOLOv4 Darknet 학습하여 Custom 데이터 인식 모델 만들기 (feat. AlexeyAB/darknet)

Ver.210728

손 쉽게 YOLO 물체 인식 모델을 학습하여 본인이 원하는 custom object 데이터를 인식하는 모델을 만들어 보려한다. 모델은 다음의 Github 링크에 있는 Darknet YOLOv4 물체 인식 모델을 사용한다.
https://github.com/AlexeyAB/darknet

리눅스 환경에서도 가능하지만 편의를 위해 구글 Colab 환경을 이용한다.

---

Custom Dataset 수집하기

본인이 인식하고 싶은 Custom 이미지 데이터셋을 수집하는 것이 가장 우선이다. 어느 정도 성능을 위해서 경험상 최소 500장 이상의 학습 이미지 데이터셋이 필요하다. Class 혹은 Label의 개수는 1개부터 많게는 수십 개도 가능하다.

본 포스팅에서는 custom dataset 수집이 완료된 것을 가정하고 진행한다.

Github 모델 다운받기

구글 Colab 환경에서 진행하므로 아래의 코드를 실행한다.
!git clone https://github.com/AlexeyAB/darknet

Pre-trained Weight 파일 다운받기

https://drive.google.com/open?id=1JKF-bdIklxOOVy-2Cr5qdvjgGpmGfcbp

위 구글 드라이브 링크에 들어가서 새로운 학습을 위해 필요한 yolov4.conv.137이라는 이름의 pre-trained 가중치 파일을 다운받는다.

다운받은 yolov4.conv.137 파일은 darknet/build/darknet/x64 경로 내에 둔다.
ex) darknet/build/darknet/x64/yolov4.conv.137

cfg 파일 설정하기

configure 파일을 설정하는 단계이다. 꽤나 복잡할 수 있으므로 천천히 진행하도록 한다.

1st

darknet/cfg/yolov4-custom.cfg

Clone 받은 darket 파일에서 위 경로 내에 있는 yolov4-custom.cfg 라는 이름의 cfg 파일을 통째로 복사하여 yolov4-obj.cfg 라는 이름으로 같은 경로 내에 복사본을 생성한다.
yolov4-obj.cfg 라는 이름의 빈 파일을 먼저 생성하고 yolov4-custom.cfg 파일의 내용을 복사해도 된다.

즉, 파일 구조는 아래와 같다.

darknet
ㄴ cfg
   ㄴ yolov4-custom.cfg
   ㄴ yolov4-obj.cfg
   ㄴ ...
ㄴ ...

2nd

이제 새로 생성된 yolov4-obj.cfg 파일 내용을 수정할 차례이다.
아래를 순서대로 따라하며 파일의 내용을 하나씩 수정하도록 한다.

• 6번째 line을 batch=64로 수정
• 7번째 line을 subdivision=16로 수정
• 20번째 line을 max_batches={클래수 개수*2000}로 수정
  ex) 만약 클래스의 개수가 2개라면 max_batches=4000
• 22번째 line을 steps={클래스 개수*2000*0.80},{클래스 개수*2000*0.90}로 수정
  ex) 만약 클래스 개수가 2개라면 steps=3200,3600
• 8번째, 9번째 line을 원하는 네트워크 input size로 수정
  ex) 416x416 크기의 이미지를 입력하려면 width=416, height=416 (32의 배수로 설정하는 것을 권장)
• [yolo] layer에 해당하는 970번째, 1058번째, 1146번째 3개의 line을 모두 classes={클래스 개수}로 수정
  ex) 만약 클래스의 개수가 2개라면 classes=2
• 3개의 [yolo] layer 직전의 [convolution] layer의 해당하는 963번째, 1051번째, 1139번째 line을 모두 filters={(클래수 개수 + 5) x 3}로 수정
  ex) 만약 클래스의 개수가 2개라면 filters=21

이미지와 라벨링 파일 저장하기

전체 이미지 파일과 라벨링 파일을 모두 darknet/build/darknet/x64/data/obj/ 경로에 넣어준다. obj 폴더가 build/darknet/x64/data/ 경로 내에 없다면 생성 후 파일들을 집어 넣어준다.

ex)

darknet
ㄴ build
   ㄴ darknet
      ㄴ x64
         ㄴ data
            ㄴ obj
               ㄴ img1.jpg
               ㄴ img1.txt
               ㄴ img2.jpg
               ㄴ img2.txt
               ㄴ ...

<주의> 이미지 파일과 라벨링 파일의 이름은 둘이 일치하기만 하다면 어떻게 설정하든지 상관없다. 또한 train 이미지와 valid 이미지 모두 darknet/build/darknet/x64/data/obj/ 경로 내에 넣어주면 된다.

train.txt, valid.txt 파일 생성 및 저장하기

이미지의 경로가 지정되어 있는 train.txt 파일과 test.txt 파일을 생성하고 저장하는 단계이다. 각각의 파일은 학습과 검증에 사용될 이미지의 상대 혹은 절대 경로를 저장하고 있다.

여러 방법이 있겠지만 아래의 방법을 따라하면 쉽게 진행할 수 있다.

1st

먼저, glob 라이브러리를 활용하여 경로를 이전에 저장한 이미지들의 경로를 리스트 형태로 생성한다. 아래의 코드에서 이미지 경로 설정에 주의한다.

from glob import glob

img_list = glob('darknet/build/darknet/x64/data/obj/*.jpg')

len(img_list)

위 코드 실행 시 학습 및 검증의 사용될 전체 이미지의 개수가 출력된다.

2nd

다음, sklearn 라이브러리의 train_test_split 함수를 사용하여 학습과 검증에 사용할 이미지를 분할한다.

from sklearn.model_selection import train_test_split

train_img_list, test_img_list = train_test_split(img_list, test_size=0.1, random_state=42) # random_state는 임의로 선택

print(len(train_img_list), len(test_img_list))

3rd

이제 학습 이미지 리스트와 검증 이미지 리스트를 각각 train.txt와 test.txt 파일에 저장한다. 각각의 txt 파일의 경로는 darknet/build/darknet/x64/data/ 경로 아래에 둔다. 경로 설정에 주의한다.

with open('darknet/build/darknet/x64/data/train.txt', 'w') as f:
    f.write('\n'.join(train_img_list) + '\n')

with open('darknet/build/darknet/x64/data/test.txt', 'w') as f:
    f.write('\n'.join(test_img_list) + '\n')

ex) 즉, train.txt와 test.txt의 파일의 내용은 예를 들어 아래와 같이 구성된다.

darknet/build/darknet/x64/data/obj/img1.jpg
darknet/build/darknet/x64/data/obj/img2.jpg
darknet/build/darknet/x64/data/obj/img3.jpg
...

데이터 파일 생성하기

이제 학습을 하기 위한 데이터들의 관한 정보를 담고 있는 obj.names, obj.data 파일 2개를 만드는 단계이다.

1st

darknet/build/darknet/x64/data/ 경로 아래에 obj.names 라는 파일을 생성한다.
obj.names 파일에는 모든 클래스의 이름이 들어간다.

ex) 만약 클래스의 개수가 2개이고 각각, dog 와 cat이라면, obj.names 파일의 내용은 아래와 같이 구성된다.

dog
cat

2nd

darknet/build/darknet/x64/data/ 경로 아래에 obj.data 라는 파일을 생성한다.

obj.data 파일의 내용은
1) 클래스의 개수,
2) train.txt 파일 경로,
3) test.txt 파일 경로,
4) obj.names 파일 경로,
5) 학습된 weight 파일을 백업할(저장할) 경로
총 5가지 line으로 구성된다.

ex) 즉, obj.data 파일의 내용은 아래와 같이 구성된다.

classes = 2
train = data/train.txt
valid = data/test.txt
names = data/obj.names
backup = backup/

<주의> 위 파일의 경로 설정에 매우 주의하도록 한다. 학습 시 train.txt , test.txt 나 obj.names 파일을 찾지 못하는 오류가 발생한다면 상대 경로말고 절대 경로로 지정해도 된다.

make 파일 생성하기

학습 전 make 파일을 생성한다. 아래의 코드를 실행하면 자동으로 생성된다.

%cd darknet
!sed -i 's/OPENCV=0/OPENCV=1/' Makefile
!sed -i 's/GPU=0/GPU=1/' Makefile
!sed -i 's/CUDNN=0/CUDNN=1/' Makefile
!sed -i 's/CUDNN_HALF=0/CUDNN_HALF=1/' Makefile
!make

학습하기

모든 준비가 끝났고 아래의 코드를 입력하여 이제 학습을 진행한다.

!./darknet detector train {obj.data 경로} {yolov4-obj.cfg 경로} {yolov4.conv.137 경로} -dont_show

• error message를 확인하지 않으려면 위와 같이 -dont_show 옵션을 입력하는 것이 좋다.
• -map 옵션을 입력하여 학습 그래프를 확인해볼 수도 있다.

테스트하기

학습이 완료되면 테스트를 진행한다.

!./darknet detector test {obj.data 경로} {yolov4-obj.cfg 경로} {학습된 weight 파일 경로} {테스트 이미지 파일 경로}

• 학습된 weight 파일은 obj.data에 backup으로 지정된 경로에 위치한다. yolov4-obj_final.weights, yolov4-obj_best.weights, yolov4-obj_best.weights 파일이 생성된다.
• 코드 실행 결과로 predictions.jpg 파일이 생성된다.