[Yolov5] Custom Dataset으로 Yolov5 돌려보기 : Train

이전 글에서는 Yolo 모델을 위한 데이터셋을 만들었다.

https://minmiin.tistory.com/14

[Yolo] Custom Dataset으로 Yolo 돌려보기 : 데이터셋 준비

 

이제 본격적으로 Yolo를 train해보자!

 

---

이전 글에서 쓴 것 같이,

 

Annotation 파일들은 개별 이미지 파일과 같은 이름을 가져야 하고,

이미지 폴더가 위치한 디렉토리 내의 /labels 폴더에 저장되어야 한다.

따라서 yolo 데이터들은 다음과 같은 구조로 저장되어야 한다.

 

root 디렉토리

ㄴ 이미지 폴더

ㄴ labels

 

이는 yolo가 저 labels 폴더 안에 annotation 파일들이 들어있다고 알아서 가정하기 때문이다.

여기까지 완료했다면 다음으로 넘어가서 yolo를 학습시켜보자.

 

---

 

예시로 kaggle의 Face Mask Detection 데이터셋을 이용하였다.

https://www.kaggle.com/datasets/andrewmvd/face-mask-detection

Face Mask Detection

 

이 데이터는 이런 형식으로 되어있다.

 

 

annotation은 xml 파일 형식으로 되어 있는데 이를 yolo format으로 바꾸고 labels 폴더에 저장하는 부분은 생략하겠다.

 

따라서 다음과 같은 구조를 가진다고 가정한다.

 

FaceMask

ㄴ images

ㄴ annotations

ㄴ labels

 

---

경로 지정하기

root_dir = "/content/drive/MyDrive/Dataset/FaceMask"
img_dir = os.path.join(root_dir, "images")
label_dir = os.path.join(root_dir,"labels")

 

root 경로 (여기에 yolo를 클론해올 것이다)와 이미지 경로, label 경로를 지정해주었다.

 

데이터 정의하기

glob.glob() 함수를 이용하면 원하는 확장자의 파일 경로만 불러올 수 있다.

이미지가 png 형식으로 되어있으니 png를 지정해서 리스트에 저장해주었다.

data = glob.glob(os.path.join(img_dir,"*.png"))

 

귀찮게도 yolo train을 위해서는 train과 validation 파일이 모두 필요하다.

따라서 train / validation / test 파일을 그냥 인덱싱해서 나눠주었다.

 

train = data[:650]
valid = data[650:750]
test = data[750:]

 

---

필요한 파일 생성하기

yolov5 train을 위해서는 다음과 같은 파일들이 필요하다.

 

• train.txt : train 이미지 파일 경로가 쓰여 있는 파일
• valid.txt : validation 파일 경로가 쓰여 있는 파일
• 개별 이미지의 bounding box와 class 정보가 쓰여진 annotation txt 파일
• 데이터 경로와 class 수, class명이 적혀 있는 yaml 파일
• test.txt : test 파일 경로가 쓰여 있는 파일 (optional)

3번째 파일들은 이미 준비했으니 train.txt, valid.txt, test.txt를 만들어보자.

(train만 진행한다면 test.txt까지는 필요 없다.)

 

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

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

 

root 디렉토리 안에 train, valid, test txt 파일을 생성했다.

 

까보면 이렇게 생겼다.

 

이제 yaml 파일을 생성하자.

 

import yaml

yaml_data = {"names":['with_mask', 'without_mask', 'mask_weared_incorrect'], # 클래스 이름
             "nc":3, # 클래스 수
             "path":root_dir, # root 경로
             "train":os.path.join(root_dir, "train.txt"), # train.txt 경로
             "val":os.path.join(root_dir, "valid.txt"), # valid.txt 경로
             "test":os.path.join(root_dir,"test.txt") # test.txt 경로
             }

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

 

root 경로 안에 다음과 같은 형태로 yaml 파일을 생성했다. 이 yaml 파일은 후에 yolo가 열어서 참조하게 된다.

여기까지 끝났다면 train을 위한 준비는 모두 끝났다.

 

---

Yolov5 Train하기

yolov5 Github에서 yolov5를 clone해오자.

!git clone https://github.com/ultralytics/yolov5.git
%cd yolov5
%pip install -qr requirements.txt # install dependencies

 

현재 경로가 yolov5 폴더 안으로 이동됨을 주의하자.

train.py를 실행하면 드디어 학습이 가능하다!

 

!python train.py --img 640 --batch 64 --epochs 20 --data /content/drive/MyDrive/Dataset/FaceMask/custom.yaml --device 0 --weights yolov5s.pt --name test

 

주요 하이퍼파라미터만 설정해주었다. 이 때 data의 경우 꼭 방금 생성한 yaml 파일의 경로로 바꿔주자.

조정 가능한 다른 하이퍼파라미터들은 여기서 확인할 수 있다.

https://github.com/ultralytics/yolov5/blob/master/train.py

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

 

쭉쭉 잘 돌아간다.

 

지금까지의 과정에 오류가 없었다면 아마 train을 잘 시작할 것이다.

(나는 libpng warning이 떴지만 무시해도 괜찮다.)

 

---

결과 확인, Weight 저장

학습이 잘 끝났다면 yolov5 폴더를 뜯어보자.

 

학습 결과 (weight)는 runs/train/프로젝트명/weights에 저장된다.

학습한 모델을 다시 이용하기 위해 weight는 꼭 저장해놓도록 하자.

 

그 외 결과들도 다양한 포맷으로 저장되어 있는데 여기서 results.png 파일을 열어보자.

 

 

loss가 잘 줄었고 정확도도 점점 올라가는 것을 볼 수 있다. 

 

val_batch0_pred.jpg 파일을 열어 validation이 잘 되었나 확인해 보자.

 

 

Validation 또한 잘 진행되었다!