[CV_0] colab에서 yolov5 학습

soyoung·2024년 9월 29일
0

0. 프로젝트 소개

  • htp(house, tree, person) 그림으로 객체 인식 후에 심리 분석을 하는 AI 프로젝트를 진행해보려고 한다.
  • 가장 첫 단계로 htp 데이터셋을 마련하고, 마련한 데이터셋을 인식할 수 있는 모델인 yolov5을 colab으로 학습시켜 보겠다.

1. 드라이브 연결

가장 첫 단계는 colab에서 드라이브 연결하기, 로컬보다 드라이브 마운트 후에 작업하는 걸 선호한다.

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

2. Yolov5 설치 & 초기 설정

# yolov5 github clone
%cd /content/drive/MyDrive
!git clone https://github.com/ultralytics/yolov5.git
# yolov5 필요한 모듈 설치
%cd /content/drive/MyDrive/yolov5/
! pip install -r requirements.txt

3. 데이터셋 불러오기

  • 처음에는 AIHUB의 "AI 기반 아동 미술심리 진단을 위한 그림 데이터 구축" 데이터셋을 사용하려고 하였으나, yolov5에 맞게 데이터셋을 고치는 과정에서 어려움이 있어서 중간에 데이터셋을 교체하였다.
  • 결론적으로 선택한 Dataset 출처: Roboflow
    Roboflow_house
    Roboflow_tree
    Roboflow_person
!mkdir datasets
%cd /content/drive/MyDrive/htp image data2/house

!curl -L "-----" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip
  • colab에 데이터셋을 불러오기 위해서는 드라이브에서 폴더 경로를 지정해줘야 한다.
  • Download Project 클릭 -> YOLOv5 v5 PyTorch Format 지정 후 -> Show download code -> Terminal code를 복사해준다.
  • 데이터셋은 총 h, t, p가 있으므로 전체 데이터 폴더 안에 각각의 하위 폴더를 만들었고 거기에 데이터셋을 업로드해주었다.

4. Training

데이터셋이 총 세 가지가 있어서 각각의 yaml 파일이 존재했었다. 따라서 데이터셋 상위 폴더 안에 datasets.yaml 파일을 새로 만들어서 모든 데이터셋 경로를 합쳐주었다.

!python '/content/drive/MyDrive/yolov5/train.py' --img 415 --batch 16 --epochs 100 --data '/content/drive/MyDrive/htp image data2/datasets.yaml' --cfg '/content/drive/MyDrive/yolov5/models/yolov5s.yaml' --weights yolov5s.pt --name htp_yolov5_results

--img: 이미지 크기
--batch: 배치 크기
--epochs: 에포크 수
--data: data.yaml 경로 지정
--cfg: 모델 경로 지정
--weights: yolov5 모델
--name: 훈련 결과 저장할 파일 이름

에포크 수에 따라 훈련 소요 시간은 다르겠지만 1시간 넘게 걸렸다. 결과로 --name에 작성한 파일이 /yolov5/runs/train 폴더 안에 생성된 것을 확인할 수 있다.

5. Test

!python '/content/drive/MyDrive/yolov5/detect.py' --weights '/content/drive/MyDrive/yolov5/runs/train/htp_yolov5_results/weights/best.pt' --conf 0.5 --source '/content/drive/MyDrive/htp image data2/_test/htp_test1.jpg'

test를 실행했더니 결과는...

???
분명 h(house), t(tree), p(person) 데이터셋을 모두 합친 yaml 파일로 훈련을 시켜주었는데, 아무래도 house 데이터셋만 학습한 것 같다. 사람을 굴뚝으로만 인식한 것을 보아...
일단 이번 실패를 통해 다음 번엔 데이터셋 자체를 합치는 것이 아닌 훈련시킨 모델을 앙상블 기법으로 합치기로 해보자.

0개의 댓글