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

YOLOv5는 x(x-large), l(large), m(medium), s(small)의 4종류로 구성됩니다. 자세한 내용은 이전 글을 참조해주시고, 이 포스트에서는 학습 코드(스크립트)에 대해서 다루고자 합니다.

 

[분류 전체보기] - YOLO v5 공개! 개념, 이론, 논문 대신에 iOS...?

YOLO v5 공개! 개념, 이론, 논문 대신에 iOS...?

먼저 YOLOv5를 official하게 공개한 ultralytics의 github repo로부터 clone을 하고, 필요한 requirements에 대해서 설치를 합니다. 그리고 torch, ipython display, google drive download 패키지들을 import 합니다.

 

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

import torch
from IPython.display import Image, clear_output
from utils.google_utils import gdrive_download

clear_output()
print('Setup complete. Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

 

간단하게 학습을 테스트해보기 위해서 마스크를 쓴 얼굴 데이터셋을 이용해보고자 합니다. 아래의 링크에서 다운로드 가능합니다. roboflow는 데이터셋을 구성하는 서비스를 제공합니다. 이미지와 태깅 데이터를 입력하고, 학습/검증/시험  index 설정, 데이터 증강 등 다양한 기능을 제공합니다. 특히 유용한 기능 중 하나는 원하는 데이터셋 형식으로 export할 수 있다는 점입니다. 아쉽게도 한 데이터셋에 1000개 이상의 이미지로 구성되는 경우에는 유료입니다. 저희가 이번에 사용할 예제 데이터셋은 무료입니다 :) 회원가입은 하셔야합니다ㅠㅠ (준비된 다른 데이터셋이 있으면 됩니다만, 예제 형식을 확인하기 위해서는 이것을 사용하면 효과적입니다.)

 

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

Mask Wearing Object Detection Dataset

 

위 링크로 이동하시면 위와 같은 화면이 나타납니다. Downloads의 416x416-balck-padding을 클릭하면, 아래와 같이 데이터셋에 대한 구성이 나타납니다. 149개의 이미지로 가볍게 테스트해보기 좋은 크기로 구성되어있습니다. 오른쪽 상단의 Download 버튼을 누르면, 아래와 같이 Format을 설정할 수 있도록 나타나는데, YOLO v5 PyTorch를 선택하면 됩니다!

 

다운로드한 데이터셋의 압출을 풀면  test, train, valid 폴더와 data.yaml 파일, 그리고 README 파일들이 주어집니다. 폴더 이름이 너무 길어서 mask-149로 변경하였습니다. 그리고 현재 프로젝트 폴더(train.py 파일이 있는 디렉토리)에 위치시켰습니다.

 

# Start tensorboard
%load_ext tensorboard
%tensorboard --logdir runs

학습 과정을 살펴볼 수 있도록 tensorboard를 실행시켜줍니다. 참조할 폴더는 runs입니다. 디폴트 포트는 6006 입니다.

 

!python train.py --img 149 --batch 16 --epochs 5 --data ./mask-149/data.yaml \
--cfg ./models/yolov5s.yaml --weights yolov5s.pt --name tutorial --nosave --cache

YOLOv5s 학습을 실행하는 스크립트는 위와같이 매우 단순합니다. weights는 사전학습된 yolov5s.pt를 사용하며, 없을 경우에는 다운로드가 자동적으로 됩니다. 데이터셋의 정보를 담고있는 ./mask-149/data.yaml 파일을 살펴봅시다. 

train: ../train/images
val: ../valid/images

nc: 2
names: ['mask', 'no-mask']

nc는 number of class로 클래스/카테고리의 수를 지정해주면 됩니다. 이 데이터셋에서는 mask와 no-mask 2개로 구성되어있기 때문에 nc: 2입니다. 

 

 

이제 위의 !python train.py를 통해 학습을 시작해봅시다. OOM 에러가 뜰 경우, 메모리가 부족한 것이니 파라미터가 적은 모델(s)을 사용하시거나, 배치 사이즈(--batch)를 줄이시면 됩니다.

 

              from  n    params  module                                  arguments                     
  0             -1  1      3520  models.common.Focus                     [3, 32, 3]                    
  1             -1  1     18560  models.common.Conv                      [32, 64, 3, 2]                
  2             -1  1     20672  models.common.Bottleneck                [64, 64]                      
  3             -1  1     73984  models.common.Conv                      [64, 128, 3, 2]               
  4             -1  1    161152  models.common.BottleneckCSP             [128, 128, 3]                 
  5             -1  1    295424  models.common.Conv                      [128, 256, 3, 2]              
  6             -1  1    641792  models.common.BottleneckCSP             [256, 256, 3]                 
  7             -1  1   1180672  models.common.Conv                      [256, 512, 3, 2]              
  8             -1  1    656896  models.common.SPP                       [512, 512, [5, 9, 13]]        
  9             -1  1   1905152  models.common.BottleneckCSP             [512, 512, 2]                 
 10             -1  1   1248768  models.common.BottleneckCSP             [512, 512, 1, False]          
 11             -1  1    130815  torch.nn.modules.conv.Conv2d            [512, 255, 1, 1]              
 12             -2  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 13        [-1, 6]  1         0  models.common.Concat                    [1]                           
 14             -1  1    197120  models.common.Conv                      [768, 256, 1, 1]              
 15             -1  1    313088  models.common.BottleneckCSP             [256, 256, 1, False]          
 16             -1  1     65535  torch.nn.modules.conv.Conv2d            [256, 255, 1, 1]              
 17             -2  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 18        [-1, 4]  1         0  models.common.Concat                    [1]                           
 19             -1  1     49408  models.common.Conv                      [384, 128, 1, 1]              
 20             -1  1     78720  models.common.BottleneckCSP             [128, 128, 1, False]          
 21             -1  1     32895  torch.nn.modules.conv.Conv2d            [128, 255, 1, 1]              
 22   [-1, 16, 11]  1         0  models.yolo.Detect                      [80, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]]]
Model Summary: 165 layers, 7.07417e+06 parameters, 7.07417e+06 gradients

모델의 구성이 어떻게 되어있는지 Summary를 해주고,

Epoch gpu_mem GIoU obj cls total targets img_size 와 같은 데이터가 출력이 되면서 학습이 진행됩니다.

 

tensorboard로 확인할 수 있으며, 학습된 best와 last 모델은 weights 폴더에 저장됩니다.

 

학습데이터 어떤 식으로 증강되고, 정답이 어떻게 박스로 잡혔는지 train-bat1.jpg 파일을 통해 확인할 수 있습니다. 확인하는 가장 큰 이유 중에 하나는 박스가 제대로 안잡혔을 수 있기 때문입니다. roboflow가 다양한 format으로 export할 수 있다는 점이 장점인 이유가 여기에 있습니다. Format에 따라 바운딩 박스는 (x, y, w, h), x, y, x+w, y+h), (x_center, y_center, w, h) 등으로 표현됩니다. YOLO v5에서는 (x_center, y_center, w, h)의 형식을 따릅니다.

Image(filename='./train_batch1.jpg', width=900)  # view augmented training mosaics

다음번에는 Kaggle의 Global Wheat Detection에 YOLOv5 모델을 적용하여 학습하고 추론하는 것에 대해서 다뤄보도록 하겠습니다.