yolov5로 이미지 분류 모델 만들기

 

제목대로 yolov5를 처음! 써보았다. DS를 공부하는 사람으로서는 매우 뒤늦었지만… 늦게라도 기록해두고 싶어서 정리해 본다. 우선 전체 코드는 여기에 공유. 처음 해보는 분들께 참고가 되길 바란다. 

 

이번에는 내가 직접 데이터세트를 만들지는 않았고, Roboflow에서 제공하는 프리 데이터 세트를 그대로 활용했다. 

특히 Roboflow는 yolov5가 요구하는 형태로 데이터세트를 제공하기 때문에, 별도로 yaml 파일을 세팅하는 번거로움에서 벗어날 수 있다는 장점이 있다(다른 분들의 기록을 보니 이게 정말정말 귀찮은 과정이더라…). 

Roboflow에서 내가 다운로드한 데이터는 Dice Dataset으로, 따로 train, test, valid 세트가 나뉘어져 있지 않아서 수동으로 나누어야 했다. 그리고 yaml 파일에 각 세트의 상대 경로도 지정되어 있지 않았다. 

 

따라서, yolov5를 공식 레포지토리에서 클론해오고 환경 세팅만 제대로 하고 나면

• ① Roboflow API 연결하기 
• ② 데이터 세트 형식 가공하기

이 두 가지 문제가 최우선으로 해결되어야 하는 셈이다. 

 

API 키를 받으려면 초록색으로 체크한 메뉴를 눌러보자

API 키를 받고 데이터세트로 돌아가서 내가 원하는 양식으로 데이터를 받을 수 있다.

zip파일 다운로드뿐만 아니라 colab이나 jupyter에서 바로 사용 가능한 코드 형태로 제공해주기도 한다. 

순서대로 누르기만 하면 주피터 노트북에 붙여넣을 수 있는 코드를 바로 제공한다.

 

그리고 이 데이터 세트에 대한 yaml 파일이 없으면 yolov5를 돌릴 수 없으므로 yaml에 경로를 추가해주자. 사전에 데이터 세트가 모두 동일한 확장자인지, 분류될 클래스는 총 몇 개인지 정도의 기본적인 정보는 Roboflow에서 알려준다. 

내가 다룰 데이터세트는 크기가 718로 많지 않다고 판단하여 따로 test set를 만들지 않았다. 

# -------- 경로 지정
import os 

root_dir = "/content/yolov5/data/Dice-2"
img_dir = os.path.join(root_dir, "images")
label_dir = os.path.join(root_dir,"labels")

# -------- 데이터세트를 train/valid로 분리하고 각각의 txt 만들기

import glob

data = glob.glob(os.path.join('/content/yolov5/data/Dice-2/export/images/',"*.jpg"))
train = data[:574] # 718 * 0.8 
valid = data[574:]

# train.txt
with open(os.path.join(root_dir, "train.txt"), 'w') as f:
	f.write('\n'.join(train) + '\n')

# valid.txt
with open(os.path.join(root_dir, "valid.txt"), 'w') as f:
	f.write('\n'.join(valid) + '\n')

# -------- yaml에 경로 지정하기 

import yaml

yaml_data = {"names":['1', '2', '3', '4', '5', '6'], # 클래스 이름
             "nc":6, # 클래스 수
             "path":root_dir, # root 경로
             "train":os.path.join(root_dir, "train.txt"), # train.txt 경로
             "val":os.path.join(root_dir, "valid.txt"), # valid.txt 경로
             }

with open(os.path.join(root_dir, "data.yaml"), "w") as f:
  yaml.dump(yaml_data, f)

 

아래는 훈련 py, 검출 py를 순서대로 돌린 코드이다. 최소한의 옵션만 설정했다.

• train: 이미지 리사이징 크기 / batch size는 8 / epoch는 200회 / data yaml이 위치한 경로 / 어느 GPU를 쓰는지 / 가중치 pt 경로(※여기서는 COCO로 훈련한 기본 yolov5s.pt를 사용) / 결괏값을 저장할 폴더 이름  
• detect: 위에서 생성한 가중치 pt 파일의 경로 / 데이터 크기 / confidence score / img 폴더 경로 

이때 train.py를 무사히 잘 돌리면 최적의 가중치를 best.pt라는 파일로 뽑아주는데, 이것이 있어야 detect를 할 수 있다. 

내가 초반에 했던 실수처럼 train이 돌아가지 않으면 best가 나오지 않고, 다른 pt 파일을 넣어보면 주사위의 눈이 아니라 사람 손가락 같은 걸 검출해서 알려준다. 다른 세트로 훈련한 가중치를 썼으니 당연한 일임^^;;; 

 

즉 best.pt가 생성되어야 모델이 제대로 훈련했다는 뜻이니 착각하지 말자. 

 

python train.py --img 640 --batch 8 --epochs 200 --data Dice-2/data.yaml --device 0 --weights yolov5s.pt --name final --cache
python detect.py --weights /content/yolov5/runs/train/final/weights/best.pt --img 718 --conf 0.25 --source /content/yolov5/data/Dice-2/export/images

 

검출까지 끝낸 다음 yolov5 폴더 안의 runs 폴더를 열어보면 뭐가 많이 생겨 있다. 

하나하나 눌러보면 다음과 같은 이미지들인데 모델이 잘 학습했는지를 알려주는 것이다.

loss가 감소하는 것을 이미지로 확인할 수 있음

좌: epoch 100, 우: epoch 200에서의 정확도 차이

비록 같은 valid batch는 아니지만 오분류가 눈에 띄게 사라졌다. 

 

좌: valid set의 실제 label, 우: 그것을 예측한 결과

 

 

도움을 받은 문서(한국어) 

 

https://yeko90.tistory.com/entry/ultralytics-yolov5-tutorial

ultralytics yolov5 사용법 완벽 정리

https://minding-deep-learning.tistory.com/19

[YOLOv5] YOLOv5 Custom Data로 학습시켜보기

https://lv99.tistory.com/72

YOLOv5 학습 예제 코드 (튜토리얼) - 마스크 쓰고 있는/안 쓴 얼굴(사람) 찾기

https://bong-sik.tistory.com/30

[Yolo5] Yolo5의 detect.py 사용해서 Yolo 학습용 데이터셋 만들기

 

마지막에 붙인 봉식이누나 님의 티스토리 문서가 특히 많은 도움이 되었다.

봉식이와 봉식이누나 님이 건강하시기를…