custom dataset으로 YOLOv5 학습하기(마스크 인식)

YOLO(You Only Look Once)는 널리 쓰이는 object detection 알고리즘입니다. 최근에는 YOLOv5까지 출시되었고 공식 github 계정에 업로드된 YOLOv5 코드로 custom dataset을 학습해보겠습니다. 

https://github.com/ultralytics/yolov5

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

위 github에서 모델을 YOLOv5를 받을 수 있습니다.

YOLOv5 모델을 이용해 마스크 쓴 사람과 쓰지 않은 사람을 검출하도록 해보겠습니다.

 

1. custom dataset 다운

먼저 학습 데이터를 다운받습니다. 

https://public.roboflow.com/object-detection/mask-wearing

Mask Wearing Object Detection Dataset

저는 roboflow에 있는 mask-wearing 데이터셋을 사용했습니다.

 

format을 YOLO v5 Pytorch로 설정해 다운로드합니다.

이번에는 kaggle에서 데이터를 다운로드합니다.

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

Face Mask Detection

 

이도 동일하게 format을 YOLO v5 Pytorch로 설정해 다운로드합니다. 

YOLOv5 공식 계정의 코드는 txt 포맷의 라벨링 데이터를 사용합니다. 이미지에서 검출된 object에 대한 클래스와 bounding box 정보가 저장되어 있습니다. 검출 객체 정보 배치는 [class, x_center, y_center, width, height] 형태로 되어있고 값은 이미지 사이즈로 정규화되어 있습니다. 

 

그래서 다운로드한 데이터들을 나중에 라벨링을 진행하겠습니다.

 

2. 가상 환경 설정

YOLO v5의 경우 3.7 버전 이상의 python과 1.7 이상의 PyTorch가 있어야 사용 가능합니다. 그래서 anaconda에 새로 가상 환경을 만들어 설치를 해주겠습니다.

 

conda create -n yolov5 python=3.8
conda activate yolov5
conda install pytorch==1.7.1 torchvision==0.8.2 cudatoolkit=11.0 -c pytorch

 

그다음 yolov5 위치로 이동해 requirements.txt에 적혀있는 남은 모듈들을 설치해줍니다.

 

pip install -r requirements.txt

 

이렇게 필요한 모든 모듈들을 설치해줍니다.

 

3. 데이터 정리

이제 다운로드한 데이터들을 정리해보겠습니다. 우리는 마스크를 쓴 사람과 안 쓴 사람만 분류하는 모델을 만들어낼 예정입니다. yaml 파일이 있는데 yaml파일은 Xml, Json과 같이 파일 포맷입니다.

 

 

이와 같은 형태로 사용하는 포맷입니다. 이제 데이터들을 우리가 원하는 대로 라벨링도 하고 주소도 변경해보겠습니다.

 

import re, os, yaml
with open('data.yaml', 'r') as f:
    data = yaml.safe_load(f)
print(data)

for folder in ['train', 'valid', 'test']:
    file_list = os.listdir(f'./{folder}/labels')
    for file in file_list:
        with open(f'./{folder}/labels/{file}', 'r+') as f:
            lines = f.read()
            # print(lines)
            replaced = re.sub(r'3(?=\s\d+\.\d+){4}', '1', lines)
            replaced = re.sub(r'1(?=\s\d+\.\d+){4}', '1', lines)
            replaced = re.sub(r'2(?=\s\d+\.\d+){4}', '0', lines)
            f.seek(0)
            f.write(replaced)
            f.truncate()

 

이 코드를 사용한 이유는 kaggle에서 다운로드한 데이터 때문입니다. kaggle에서 다운로드한 데이터는 class의 개수가 4개로 되어 있었습니다. 우리가 원하는 건 마스크를 쓰거나 안 쓴 두 가지의 class이므로 txt 파일의 내용들을 변경해 줄 필요가 있었습니다. 원본 txt 파일들은 0, 1, 2, 3으로 분류한 형태였기 때문에 이를 0, 1로 변경해주어야 합니다. 그래서 위와 같이 코드를 작성해 모든 데이터를 0과 1 두 가지 class로 분류해주었습니다. 1은 no_mask, 0은 mask입니다.

 

import yaml

with open('data.yaml', 'r') as f:
    data = yaml.safe_load(f)
print(data)

data['train'] = './train/images'
data['test'] = './test/images'
data['val'] = './val/images'

data['nc'] = 2
data['names'] =  ['mask', 'no_mask']
with open('data.yaml', 'w') as f:
    yaml.dump(data, f)
print(data)

 

이번에는 yaml 내용을 변경해보겠습니다. 먼저 위에 data.yaml 파일을 읽어와 내용을 확인해줍니다.  이미지 경로를 설정해주고 클래스도 2개로 맞춰주고 각 인덱스에 맞춰 이름도 설정해주었습니다.

YOLO v5는 x(x-large), l(large), m(medium), s(small)의 4종류로 구성되는데 그중에서 우리는 s(small)을 사용하겠습니다. 그래서 우리는 yolov5s.yaml 파일의 내용도 변경해줘야 합니다. 위와 같은 방식으로 nc를 2로 변경해주면 됩니다.

먼저 yolov5s.yaml 파일을 열어보니

 

이와 같이 작성되어 있었습니다. nc가 80인 것을 알 수 있고 이를 2로 변경해주면 됩니다.

이렇게 설정한 데이터들을 통해 학습을 진행해보겠습니다.

 

3. 학습하기

이미지 학습을 할 때 train.py를 사용하는데 이때 사용되는 arguments들입니다.

• img: define input image size
• batch: determine batch size
• epochs: define the number of training epochs.
• data: set the path to our yaml file
• cfg: specify our model configuration
• weights: specify a custom path to weights
• name: result names
• nosave: only save the final checkpoint
• cache: cache images for faster training

YOLO v5에는 

 

 

이에 맞춰서 실행해주면 됩니다.

python train.py --img 416 --batch 16 --epochs 150 --data ./data/images/data.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --name mask_yolov5s_results_with_batch_16_epoch_150

 

이렇게 작성해 실행해주면 됩니다. input image 크기를 416으로 설정하고 batch size를 16으로 해주었습니다. 그리고 epoch은 150으로 설정해주고 우리가 변경한 data.yaml과 yolov5s.yaml을 사용합니다.

학습이 종료되면 결과를 볼 수 있습니다.

 

학습되는 과정을 볼 수 있고 중간중간 결과들도 저장되어 있습니다. 이를 보면

 

이렇게 분류하는 모습을 볼 수 있습니다. 이제 이렇게 학습을 통해 얻은 가중치를 이용해 test를 진행해보겠습니다.

 

4. 테스트 진행

위 결과를 본 폴더에 보면 weights라는 폴더가 있는데 해당 폴더에 학습을 통해 얻은 가중치들이 저장되어 있습니다. 이를 이용해 test를 진행해보겠습니다.

 

python detect.py --weights best.pt --source ./data/images/test/images

 

이와 같이 작성을 해주면 test 폴더에 있는 이미지들을 불러오고 가중치를 불러와 detect를 진행할 수 있습니다. 결과를 보면

 

어느 정도 인식을 잘하는 것을 볼 수 있습니다.