[python] yolov5이용, 간단한 객체 학습 후 탐지

본 게시글에서는 yolov5를 이용하고, Pytorch를 통한 커스텀 학습을 통해 우리집 고양이인 "몰리"의 머리, 발, 몸통, 꼬리를 탐지해보고자 한다.

 

yolov5의 설치에 관해서는 github를 참조하였다.

ultralytics/yolov5

 

먼저 원본 동영상은 몰리를 내 발.....로 놀아주는 영상이다.

"몰리"

해당영상은 총 16초의 30fps 동영상이므로 10프레임 마다 동영상 프레임을 저장하여 총 48개의 학습용 이미지를 만들어 주었다.

 

이미지 캡쳐 코드는 다음과 같다.

import cv2
vidcap = cv2.VideoCapture('./cat.mp4')
 
count = 0
 
while(vidcap.isOpened()):
    ret, image = vidcap.read()
 
    if(int(vidcap.get(1)) % 10 == 0):
        print('Saved frame number : ' + str(int(vidcap.get(1))))
        cv2.imwrite("./img/cat_%d.jpg" % count, image)
        print('Saved frame%d.jpg' % count)
        count += 1

vidcap.release()

이제 학습시키기위한 라벨링을 하기위해 labelimg를 통해 라벨링을 해주었다.

tzutalin/labelImg

predefined_classes.txt에 head, tail, foot, body와 같이 라벨을 사전 지정해주고, 해당 파일을 이용해 labelimg를 실행시켜준다.

predefined_classes.txt

head
body
tail
foot

 

 

그후, labelimg를 라벨링할 이미지가 있는 경로, pre_definedclasses가 있는 경로로 실행시켜준다.

labelimg ./img ./result/predefined_classes.txt

다음과 같이 각 객체들을 직접 라벨링 하고, 이름이 틀리지 않도록 지정해준다.

모든 사진을 다 똑같이 해준다. (일일히 해주는 방법밖에 없다. 경우의 수가 늘어날수록 정확도도 증가하고, 학습의 좋은 재료가 된다.)

 

나는 총 20장만 우선 라벨링을 하고, 학습시켜 보도록 해야겠다.

라벨링을 완료하면 결과폴더로 지정한 곳에 각 사진의 txt파일과 classes파일이 생성된 것을 확인 할 수 있다.

 

이제 yolo를 설치한 폴더로 간 후 cat폴더를 생성한 후, images폴더를 만들어 학습시킬 이미지를 넣어주고 labels폴더를 만들어 방금 위에서 라벨링 후 생성된 txt파일들을 모두 넣어준다.

 

그후 data폴더에 cat.yaml파일을 하나 생성한 후 다음과 같이 내용을 채워준다.

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ./cat/images/ 
val: ./cat/images/ 

# number of classes
nc: 4

# class names
names: ['head','body','tail','foot']

나는 train과 validation을 따로 구분하지 않을 것이기 때문이 train과 val의 경로를 똑같이 해주었다.

nc에는 class의 개수, names에는 아까 라벨링할 때 지정했던 이름들을 넣어준다.

 

 

그리고 yolov5폴더 내에 train_list.txt, val_list.txt파일을 만들어 학습시킬 이미지들의 경로를 모두 적어준다.

train_list.txt파일 예시는 다음과 같다.

이걸 일일히 입력하기는 매우 귀찮은 일이기 때문에, yolov5/cat/labels폴더에 있는 txt파일들을 이용해 텍스트 slice를 적용하여 만들어 주었다.

-- train_list.txt, var_list.txt 작성 코드

import glob

input_dir='C:\\Users\\syjin\\Desktop\\object\\yolov5\\cat\\labels\\'
f = open('./train_list.txt', 'w')
input_file=glob.glob(input_dir+'*.txt')

for file in input_file:
    f.write(file[:41]+'images\\'+file[48:-4]+'.jpg \n')
f.close()


f = open('./val_list.txt', 'w')

for file in input_file:
    f.write(file[:41]+'images\\'+file[48:-4]+'.jpg \n')
f.close()

이제 학습을 시켜줄 차례이다.

pytorch로 학습이 진행되기 때문에 위 yolov5 github를 참고하여 환경을 맞추어 설치해주어야 한다.

 

anaconda prompt로 yolov5가 설치되어 있는 폴더로 이동 한 후 다음과 같이 입력한다.

python train.py --data ./data/cat.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --batch 8 --epochs 200 --name cat200

아까 만들어준 cat.yaml파일의 경로를 입력해주고 사용할 pretrained모델을 입력해준다.

batch size=8, epoch=200으로 하여 총 200번의 학습을 진행 할 것이다. 그리고 마지막에는 학습된 가중치 파일이 저장될 경로를 적어준다.

 

pretrained model download

Releases · ultralytics/yolov5

모델의 경우 s, m, l, x 사이즈가 있는데 사이즈가 작을수록 가벼우며, 클수록 무겁다.

물론 이에따라 성능도 달라진다. 사용자에 맞추어 파일을 다운로드 한다, 나는 s사이즈를 이용했다.

코드를 입력하면 다음과 같이 학습이 진행될 것이다. 기다려보자...

학습이 완료 된 후, 각 라벨의 R값을 보면 foot의 경우가 0.431로 가장 낮다, 발의 경우 몸이나 머리에 비해 크기가 작다보니 탐지 성능이 낮게 나오는 것 같다. 물론 학습을 더 많이 시켜주면 발의 탐지정확도 역시 올라 갈 것이다.

 

이제 학습된 가중치 파일은 yolov5폴더 내 runs/train/cat200/weight폴더에 형성되어 있을 것이다. 이를 이용해 한번 동영상에 적용 시켜보자.

 

학습시킬때와 마찬가지로 anaconda prompt에서 다음과 같은 명령을 입력해준다

python detect.py --source ./data/videos/cat.mp4 --weights ./runs/train/cat200/weights/best.pt

탐지결과는 위 동영상과 같다. 물론 학습 횟수도 200번이며, 학습시킨 이미지도 20장에 불과하기때문에 탐지율이 그렇게 좋진 않지만 어느정도 각 부분을 탐지해내는 것을 볼 수 있다.

또한 사용 모델의 사이즈를 s로 사용했지만 보유한 컴퓨터의 성능에 따라 m, l, x모델을 이용하여 학습한다면 탐지 성능은 더 올라갈 것이다.

---

yolov5 m모델을 이용하여 1000번 학습시킨 결과 ↓

 

위에 업로드한 S모델을 이용한 200번 학습 결과보다 훨씬 좋은 탐지성능을 보인다.