Google Colab을 이용한 YOLOv5 Custom 데이터 학습

1. 데이터 생성

2. data.yaml 생성

3. 데이터 정리

4. Google Colab 설정

5. YOLOv5를 통한 학습

6. 학습 결과 확인

1. 데이터 생성

1-1. YOLO_MARK 설치 및 빌드

https://github.com/AlexeyAB/Yolo_mark 에서 Yolo Mark를 다운로드하고 설치한다

Yolo Mark 빌드에는 https://ctkim.tistory.com/82 블로그의 포스트를 참고하였다

 

1-2. obj.data, obj.name 수정

Yolo_mark-master/x64/Release/data/img 폴더에 라벨링할 이미지를 넣는다

그리고 obj.data, obj.name를 다음과 같이 수정해야 한다

car 객체 하나에 대해서만 찾으면 되므로 

classes = 1로 설정하였고 분류 명칭을 car 하나로 설정하였다

 

1-3. YOLO_MARK를 통한 라벨링

YOLO_MARK를 이용하여 데이터 라벨링을 시작한다

학습시키고자 하는 이미지 부분에 박스를 그리고 클래스를 지정하면 된다

 

데이터 라벨링이 끝나면 라벨링된 좌표값이 담긴 txt 파일들이 생성된다

 

2. data.yaml 생성

구글 드라이브의 내 드라이브/dataset 아래에 data.yaml 파일을 생성한다

train: /content/drive/My Drive/images/train/
val: /content/drive/My Drive/images/val/

nc: 1
names: ['car']

 

nc : 학습시키고 싶은 클래스의 수

names : 학습시키고 싶은 클래스의 이름

 

나는 car 객체 하나만 학습시키므로 nc를 1, names를 car로 설정했다

 

3. 데이터 정리

1-3에서 데이터 라벨링으로 생성된 데이터들을 정리해준다

구글 드라이브의 내 드라이브에 images폴더와 labels폴더를 생성한다

각 폴더안에 train폴더와 val폴더를 생성한다

 

images/train : 학습시키고자 하는 이미지들

images/val : validation에 사용되는 이미지들

labels/train : 학습시키고자 하는 이미지들의 라벨링 좌표값이 저장된 txt들

labels/val : validation에 사용되는 이미지들의 라벨링 좌표값이 저장된 txt들

 

각각의 폴더에 해당되는 데이터들을 업로드 해주면 된다

 

4. Google Colab 설정

4-1. 새 노트 실행

구글 드라이브의 새로 만들기에서 Google Colab을 실행한다

 

4-2. GPU 설정

런타임을 눌러 런타임 유형을 GPU로 변경한다

 

5. YOLOv5를 통한 학습

5-1. 구글 드라이브 연결

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

 

content 폴더 안의 drive라는 이름으로 구글 드라이브를 연결한다

코드가 실행되면 연동하고 싶은 내 구글 드라이브를 선택한다

 

5-2. 경로 이동

%cd /content/drive/My\ Drive

/content/drive/My Drive 로 이동한다

 

5-3. git의 YOLOv5 설치

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

YOLOv5를 clone한다

코드를 실행시키면 내 드라이브에 yolov5라는 폴더가 생성된다

 

%cd yolov5 
!pip install -r requirements.txt # install dependencies

yolov5폴더로 이동한 후 requirements.txt에 포함되어 있는 것들을 다운받는다

 

5-4. yaml 패키지 설치

!pip install -U PyYAML

코드를 실행시켜 yalm 패키지를 설치한다

 

5-5. 경로 이동

%cd /content/drive/My\ Drive/yolov5

/content/drive/My Drive/yolov5 로 이동한다

 

5-6. 모델 선택

yolov5/models 안에 있는 파일들 중 모델을 선택한다

선택한 모델의 파일 내용 중 nc값을 자신이 학습시키고자 하는 클래스수로 바꾼다

 

나는 yolov5s.yaml을 선택했고 하나의 클래스를 학습시킬 것이므로 nc를 1로 설정했다

 

5-7. 학습 시작

!python train.py --img 640 --batch 30 --epochs 300 --data /content/drive/My\ Drive/dataset/data.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --name test

 

 

코드를 실행하면 학습이 시작된다

 

6. 학습 결과 확인

# video inference
%%time
!python detect.py --weights '/content/drive/MyDrive/yolov5/runs/train/test/weights/best.pt' --img 1280 --conf 0.4 --source '/content/drive/MyDrive/ColabNotebooks/KITTI_data.mp4'

YOLOv5 모델 안에 있는 detect.py를 이용해 내가 원하는 video에서 객체들을 인식한다

 

--weights 다음에 학습한 모델 경로명 작성

--source 다음에 추론할 비디오 경로명 작성

--img 다음에 이미지 크기 작성

--conf 다음에 0~1 사이 값의 conf_threshold 작성

 

 

해당 video에서 car 객체들이 잘 인식되는 것을 확인할 수 있다