[YOLOv4] Pistols data를 이용하여 yolov4 학습하기

yolov4를 이용하여 roboflow 에서 Pistols Dataset을 가지고 학습을 해보았다.

 

참고로

OS : Ubuntu 20.04 버전

GPU : GeForce RTX 3090

CUDA : 11.2 버전

 

Anaconda를 이용하여 yolo 라는 가상환경을 만들어서 훈련을 진행하였다.

[yolo]

python : 3.6.13

tensorflow-gpu : 2.4.0

등 .. 

 

학습 방법은 AlexeyAB/darknet github를 참고했다.

 

 

1. AlexeyAB/darknet 의 github을 이용하여 각종 파일을 다운로드 받는다. 

github.com/AlexeyAB/darknet 

AlexeyAB/darknet

 

$ git clone https://github.com/AlexeyAB/darknet

 

2.  roboflow 에서 Pistols Dataset을 다운받는다.

roboflow는 이미 레이블링된 이미지를 다운받을 수 있다.

 

public.roboflow.com/object-detection/pistols

Pistols Object Detection Dataset

# 참고로 회원가입해서 로그인 해야 다운이 가능하다.

 

3. darknet/cfg/ 에서 yolov4-gun.cfg를 만든다.

(yolov4-custom.cfg 를 복사해서 수정하였다.)

(yolov4-gun은 임의의 이름이다.)

 

4. yolov4-gun.cfg 수정

 

4 - 1) match_batches = classes x 2000 (만약 클래스가 4개면 6x2000 = 12000)

# 주의할 점 

Pistols data는 클래스가 1개라서 1x 2000 하여 2000을 했으나 error : 세그먼테이션오류 (코어덤프됨) 이 발생하였다.

그 이유는 match_batches가 학습할려는 이미지보다는 많아야 이러한 오류가 안난다. 

그래서 match_batches = 6000으로 잡았다. 

match_batches = 6000

4 - 2) steps = match-batches의 80% , 90% (match_batches가 12000이라면 9600, 10800)

 

match_batches를 6000으로 하여 

steps = 4800,5400 으로 했다. (80%, 90%)

steps = 4800 , 5400

4 - 3) width = 416 / height = 416 으로 수정 (32의 배수로 설정을 해야한다.)

width = 416
height = 416

4 - 4) cfg에 보면 여러 layer 중 yolo라는 layer가 3개가 있다. 

[yolo] 층 3개를 똑같이 수정해준다.

classes = 1 (Pistols data는 Pistol 한 개라서 classes =1 이다.)

[yolo]
classes = 1

[convolutional] (여러 convolutional layer 중에서 반드시 yolo 앞에는 convolutional layer를 수정해야한다.]

filters = (classes + 5) x 3  =  (1 + 5) x 3 = 18

[convolutional]

filter = 18 

(yolo층이 3개 이므로 yolo층 앞에 있는 각각의 convolutional 층의 filters를 수정해준다.)

 

5. darknet/data 에서 yolov4-gun.names 생성하여 classes 이름을 넣어준다.

pistol

(yolov4-gun는 임의의 이름이다.)

 

6. darknet/data 에서 yolov4-gun.data 생성한다.

classes = 1
train = data/train.txt
valid = data/test.txt
names = data/yolov4-gun.names
backup = backup/

classes는 한 개만 있어서 1로 하였다.

train.txt는 7번에서 만들 예정이다. train data가 어디에 있는지 말해주는 좌표같은 역할이다.

test.txt도 검증을 하기 위한 data가 어디에 있는지 말해주는 좌표로 mAP를 출력할 때 필요하다.

names 는 5번에서 만들었던 names를 넣으면 된다.

backup/ 은 가중치 값들이 저장될 곳이다.

 

7. train.txt 생성

3000개 데이터의 경로를 넣어주면 된다.

(파이썬 코드를 이용하면 쉽게 만들 수있다.)

import os,re,codecs
list_2 = os.listdir('/yolo/data/')
train = ["data/yolov4_gun_data/"+i for i in list_2]

os.chdir(r'/darknet/data')
aa = ''
for i in train:
    aa += i+'\n'
f = open('train.txt','w',encoding='utf-8')
f.write(aa)
f.close()

data/data에 있는 pistols data들을 list로 뽑아서 

나는 darknet/data/yolov4_gun_data 라는 폴더에 넣을 꺼라서 각 이름마다 앞에 경로를 붙여주었다.

 

os.chdir로 경로를 이동 시켜준 다음 aa에 문자로 저장하여 train.txt로 저장한다.

data/yolov4-gun/armas (2286).jpg
data/yolov4-gun/armas (2918).jpg
data/yolov4-gun/armas (1496).jpg
..
..
..
..

  

8. 사전 훈련된 가중치를 다운로드

darknet/build/darknet/x64

drive.google.com/file/d/1JKF-bdIklxOOVy-2Cr5qdvjgGpmGfcbp/view

yolov4.conv.137

9. train!

./darknet detector train data/yolov4-gun.data cfg/yolov4-gun.cfg build/darknet/x64/yolov4.conv.137

 

mAP가 나오게 할려면 뒤에 -map 를 붙이면된다.

 

[결과]

학습 과정

훈련 시작하면 자동으로 나온다. (나는 mAP 없이 실행하였다.)

실행결과

 

웹캠으로 실행하여 권총 사진을 보여주었는데 잘 인식 되는 것을 볼 수 있다. 

 

하지만 핸드폰도 권총으로 인식하는 점이 신기하였다....