Google Colab으로 yolo돌리기

이번에 캡스톤디자인에서 yolo를 돌려야 하는데, 문제가 하나 있다..!

내 노트북이 연약한 내장그래픽 사용하는 gram이라는거...

그런 나에게도 yolo를 약하지만 돌릴 수 있는 방법이 있으니.. 바로 구글 colab이다.

구글 드라이브를 연동해서 무료로 gpu를 사용할 수 있다. 물론 사이즈가 큰 학습은 못하지만.. 현재 내 프로젝트는 구현만 되면 되는거 이기도 하고, 높은 정확도를 요구하지는 않아서... 노트북 사기 전까지는 이걸로 하자..

우선 google drive에 들어가서 colab워크스페이스를 열어준다.

colab에 들어왔으면 제일 먼저 해줄 일은, 런타임 유형을 gpu로 바꿔주는 거다.

상단에 런타임 - 런타임 유형변경 에서 gpu로 바꿔주자.

그다음은 내 드라이브와 연동해주어야 한다.

from google.colab import drive
drive.mount('/content/drive')

위 코드를 넣어줘서 연동을 마무리 해주자.

다음은 yolov5 git을 clone받아주자.

cd drive/MyDrive/

디렉토리 이동을 해주고,,

!git clone https://github.com/ultralytics/yolov5

위 명령어를 입력해주면 yolov5가 클론된다.

다음은 데이터셋 zip파일을 구글 드라이브에 올리자. 이때, 그냥 drive에 올리지 말고 clone받은 yolov5디렉토리 안에 올려주도록 하자

cd yolov5/
!unzip -qq "custom_dataset.zip"

위 명령어를 입력해주면, yolov5폴더 내에 custom_dataset이라는 폴더가 새로 생기게 된다. 

Yolo 학습시키기

이제 학습을 위해 yaml파일을 바꿔주어야 한다.

데이터셋은 이미 라벨링 되어있는 공공 데이터셋을 가져왔기 때문에 여기까지만 해줘도 된다.

path: /content/drive/MyDrive/yolov5/custom_dataset
train: train/images
val: valid/images
test: test/images

nc: 5
names: ['Ambulance', 'Bus', 'Car', 'Motorcycle', 'Truck']

이제 학습 돌려주면 된다.

!python train.py --data "data/custom_dataset.yaml" --epochs 100

epoch은 일단 100으로 두고 진행해보자.

열심히 돌아주렴..

Yolo로 detection시키기

여기까지는 학습을 진행하는 방법이었고, 이제 yolo로 detection하는 법을 정리해보려고 한다.

일단, 기본적으로 yolo에서 제일 기본 pt파일과 샘플 이미지를 제공한다.

sample image는 /data/images에 있다. detection해보고 싶은 이미지가 있다면 이 폴더에 jpg형태로 넣어주면 된다.

그 다음 detection 코드를 실행해주면 된다.

!python detect.py --weights yolov5s.pt --img 640 --conf 0.28 --source data/images
# display(Image(filename='runs/detect/exp/im22.jpg', width=600))

아래에 display 부분은 colab에서 돌려서 못쓰는 부분이다. 만약 gpu달린 노트북에서 진행한다면, 아마 detect되면 이미지가 나타날 것이다.

나는 detect가 끝나면 yolov5/runs/detect/ 에 들어가 직접 확인해주었다.

아래는 내가 직접 넣은 이미지를 detection해본거다. 

꽤나 잘 나와서 사실 따로 학습을 시켜야 하나..? 싶지만... 하기는 해야겠지..

딥러닝.. 어렵다..