[YOLO] Google Colab으로 YOLO v5 커스텀 데이터 학습시키기

Google Colab 이란?

Colab은 구글에서 만든 연구용 서비스 제품이며, Jupyter를 기반으로 만들어진 웹용 서비스이다.

Cloud service를 통해 Jupyter에서 작업을 해야 하는 Machine learning을 Colab을 통해 손쉽게 돌릴 수 있다.

 

이 게시글에서는 Google Colab 환경에서 YOLO v5 커스텀 학습 튜토리얼 영상을 참고해 데이터를 학습시켜보았다.

 

---

YOLO v5 커스텀 데이터 학습시키기

0. Goolge Colab 환경 세팅

구글 드라이브에 접속해 새로 만들기 > 더보기 > Google Colaboratory 를 선택한다. 

구글 드라이브

생성된 파일에 접속해 런타임 > 런타임 유형 변경을 선택해 하드웨어 가속기를 GPU로 변경해주고 [저장]을 누른다.

그럼 이제 GPU를 사용하는 환경이 세팅이 된 것이다. 

 

1. YOLO v5 환경 세팅

생성한 나의 Google Colab에 YOLO v5를 다운로드해주어야만 YOLO v5를 쓸 수 있다.git clone을 해와서 설치한다. 

%cd /content
!git clone https://github.com/ultralytics/yolov5.git

설치가 완료되면 YOLOv5을 위한 패키지를 설치해 준다.

%cd /content/yolov5/
!pip install -r requirements.txt

 

2. Dataset 가져오기

Roboflow의 object detection 메뉴에 가면 다양한 dataset을 다운받을 수 있다.

필자는 Pistols Dataset을 다운로드 받았다. 

 

https://public.roboflow.com/object-detection/pistols

Pistols Object Detection Dataset

 

roboflow

원하는 dataset을 골랐다면 Download를 누른다.

다양한 포맷으로 다운을 받을 수 있는데 YOLO v5로 학습을 시키고자 하기 때문에 YOLO v5 PyTorch 포맷을 선택했다.

 

Terminal 코드를 복사해 Colab 파일에 붙여 넣기 한다.

그러면 다음과 같이 압축이 풀리는 것을 볼 수 있다.

dataset 이라는 이름으로 새 폴더를 생성해 생성되었던 모든 파일을 넣어준다. (아래 그림 참조)

 

 

3. data.yaml 경로 재설정 및 train, valid set 나누기

 

data.yaml이 무엇을 포함하고 있는지 확인해본다.

아래 코드를 실행하면 train 경로와 valid 경로가 잘못 설정되어 있는 것을 확인할 수 있다. 경로 변경이 필요하다.

%cat /content/dataset/data.yaml

먼저 route로 이동하고 glob이라는 패키지를 이용한다. export파일 안에 있는 모든 이미지들을 list로 만들어준다.

%cd /
from glob import glob
img_list = glob('/content/dataset/export/images/*.jpg')
print(len(img_list))

 

img_list를 train set과 valid set으로 나눠준다.

from sklearn.model_selection import train_test_split
train_img_list, val_img_list = train_test_split(img_list, test_size = 0.2, random_state = 2000)
print(len(train_img_list), len(val_img_list))

train 이미지 경로와 valid 이미지 경로를 txt 파일로 저장한다.

그러면 dataset 폴더에 train.txt와 val.txt가 생긴다.

with open('/content/dataset/train.txt', 'w') as f:
  f.write('\n'.join(train_img_list) + '\n')

with open('/content/dataset/val.txt', 'w') as f:
  f.write('\n'.join(valid_img_list) + '\n')

train과 valid의 경로를 재설정해준다.

 

import yaml

with open('/content/dataset/data.yaml', 'r') as f:
  data = yaml.safe_load(f)

print(data)

data['train'] = '/content/dataset/train.txt'
data['val'] = '/content/dataset/val.txt'

with open('/content/dataset/data.yaml', 'w') as f:
  yaml.dump(data, f)

print(data)

경로가 정상적으로 바뀌었음을 확인할 수 있다! 

 

4. YOLOv5 학습

train.py를 실행시켜 YOLOv5를 학습시켜 준다.

(약 33-35분 정도 소요된다.) 

%cd /content/yolov5/

!python train.py --img 416 --batch 16 --epochs 50 --data /content/dataset/data.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --name results

성공적으로 학습이 끝나면 tensorboard로 학습 결과를 확인할 수 있다.

%load_ext tensorboard
%tensorboard --logdir /content/yolov5/runs/

 

5. 테스트 데이터로 추론하기 

물체를 인식하기 위해서 사진이나 동영상이 필요한데

직접 업로드를 해도 되고 테스트 데이터를 이용해 인식시켜보아도 된다.

 

필자는 기존에 있던 데이터로 인식을 시켜보았다. 

from Ipython.display imort Image
import os

val_img_path = val_img_list[1]

!python detect.py --weights /content/yolov5/runs/exp0_gun_yolov5s_results/weights/best.pt --img 416 --conf 0.5 --source "{val_img_path}"
Image(os.path.join('/content/yolov5/inference/output', os.path.basename(val_img_path)))

결과물을 보면 정상적으로 권총을 인식하는 걸 볼 수 있다!

 

 

---

참고 영상

https://www.youtube.com/watch?v=T0DO1C8uYP8&t=291s