[AI] Pytorch YoloV5을 이용한 이미지 탐색(Feat..소주병 분류)

Pytroch 에서 Yolo 모델을 사용하여 객체 탐지하는 모델을 만드는 과정입니다.

만들면서 과정을 기록해놔서, 기록 해놓는겸 yolov5 공부하시는 분들에게 도움이 되고자 정리해서 올립니다.

개인적으로 Yolo 모델이 제일 사용하기 쉬운 것 같더라구요.. Detectron2 모델도 사용해봤는데 Yolo가 사용하기 제일 쉬웠습니다.

일단 저는 Yolov5 모델을 사용하여 최대한 간단하게(성능이 쫌 떨어집니다..) 소주병을 segementation 할 수 있도록 학습 시켰습니다.

Yolov5 공식 홈페이지 입니다.
https://github.com/ultralytics/yolov5

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

 

개발 환경

• OS: Windows 10
• Pycharm
• anaconda 사용
  • Torch : 1.19
  • Cuda : 11.1
  • Python : 3.8

데이터 셋

라벨링 된 소주병 이미지 4500장

종류 : 참이슬, 처음처럼, 진로, 청하

학습 결과

results (2).csv

0.02MB

 

추론 결과

개발 순서 (위에서 명시한 환경에서 진행했습니다.)

• Anaconda 가상 환경 만들기

conda create -n mec_yolo5_v1 python=3.8

• Pycharm 기준 Conda 설정 하기
  • Setting - python interpreter - Base interpreter 설정 - anaconda install path - envs - 위에서 만든 가상 환경의 python.exe 선택

 

• pytorch 설치
  • PYTORCH 1.9 설치 명령어 (CUDA 11.1 기반) 

pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html

• CUDA Version이 다를 경우 참고 링크

https://pytorch.org/get-started/previous-versions/

PyTorch

• yolov5 설치

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

• dataset 생성 및 Soju.yaml 파일 생성

#E:\yolov5\soju_dataset

path: /yolov5/soju_dataset

train: train/images

val: valid/images

test: test/images #optinal

# Classes

nc: 4 # number of classes

names: ["cheoeumcheoreom", "chamiseul", "jinro", "chungha"] # 진로, 참이슬, 처음처럼,청하

• traing.py 실행 
  • 제가 사용한 실행 옵션입니다. 

- python train.py --imgsz 768 --batch 32 --epochs 200 --optimizer Adam --data "data/Soju.yaml" --device 0 --weights yolov5x.pt --workers 7  --cache ram

- python train.py --imgsz 320 --batch 32 --epochs 200 --optimizer Adam --data "data/Soju.yaml" --device 0 --weights yolov5s.pt --workers 4 --cache ram

• GPU 설정이 잘 안 되시는 분들은 CPU를 통해 진행하시면 됩니다.(엄청 오래걸려요)

- python train.py --imgsz 320 --batch 64 --epochs 10 --optimizer Adam --data "data/Soju.yaml" --device cpu --weights yolov5s.pt