Google Colab을 이용한 YOLOv5 Custom 데이터 학습
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 객체들이 잘 인식되는 것을 확인할 수 있다