[Python] YOLOv5 Custom dataset 으로 학습하기

CNN을 사용해 제품 인식을 구현하다보니,

사무실 바닥같은 복잡한 패턴을 가진 이미지에서는 정상적으로 인식하지 못했다.

또한, 이미지에 제품이 없는 경우에도 제품이 없다는 걸 알지 못하고 그 이미지가 닮은 클래스를 찾아서 출력한다..

 

OpenCV로 배경을 찾아서 지우는 방법도 해봤는데

이미지마다 임계치가 달라서 배경을 제대로 지우지 못하는 경우도 있었다.

 

그러던 중, 객체 인식 알고리즘 YOLO가 있어서 한번 사용해 본다.

YOLO도 여러가지 버전이 있는데, v5가 올해 5월에 나온 최신 버전이여서 v5로 선택했다.

(새로 구현을 한다던가 하는 것은 아니고, GitHub에 있는 코드를 그대로 사용할 것임)

 

출처 : GitHub

ultralytics/yolov5

---

1. 요구 사항

Python 최신버전(작성일 기준 3.8.5)과 Anaconda 환경에서 실행한다.

가장 먼저 Anaconda Prompt를 열고 아래 명령어를 실행한다.

conda install torch
conda install -c conda-forge pytorch
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
 
pip install --upgrade pip
pip install tensorboard
pip install opencv-python
 
https://aka.ms/vs/16/release/vc_redist.x64.exe

 

vc_redist.x64는 visual c++ 패키지라는데,

나는 계속 오류가 떠서 그냥 비쥬얼 스튜디오 C++를 통째로 다운받았다... 

 

 

2. 학습 데이터 만들기

그 다음으로 학습 데이터를 만든다.

labelImg를 사용해서 만드는데, 왼쪽 메뉴에서 형식을 YOLO로 바꾸는 걸 잊지말자.

(까먹고 300장째쯤에 알아서 처음부터 다시한적 있음..)

꼭 바꾸자

학습할 데이터의 클래스명을 predefined_classes.txt파일에 추가하고

Open Dir 버튼을 클릭해 데이터셋이 저장된 폴더를 선택한다.

 

w키를 눌러 학습할 이미지 범위를 선택하고, 해당 범위의 클래스를 선택한 뒤 ctrl+s를 눌러 저장한다.

a키는 이전 이미지, d키는 다음 이미지를 선택할 수 있음.

원하는 데이터 전부를 이렇게 해줘야한다.

하고 나면 원본 이미지 파일이랑

그 이미지 파일 이름이랑 같은 txt 파일이 생성되는데,

이 txt 파일에는 아마도 내가 설정한 라벨의 위치가 들어가는 것 같다. (첫번째로 어떤 클래스인지도 들어감)

 

하나하나 전부 해줘야한다.

 

 

3. 학습 해보기

이제 만든 데이터셋으로 학습을 해본다.

yolov5 폴더에 coco128 이라는 폴더를 만들고, (다른 이름으로도 할 수 있지만 굳이 바꾸지 말자..)

그 폴더 안에 images 폴더와 labels 폴더를 만든 뒤

images 폴더에 또 train2017 폴더를 만들고 안에는 원본 이미지를,

마찬가지로 labels 폴더에도 train2017 폴더를 만들고 txt 파일을 넣는다.

 

폴더명이나 이런건 자유롭게 바꿔도 되는데, 바꿨다가 이름 다르다고 오류나는게 귀찮아서

원래 코드에 있는 coco128에 데이터셋만 바꿨다.

 

다 만들고 data폴더에 coco128.yaml 파일을 연다.

보면 다운로드 하는 부분은 필요없으니까 주석처리하고

train과 val의 경로를 저렇게 바꾼다.

 

number of classes도 자기가 분류하고 싶은 클래스의 개수로 바꾼다.

class names도 위에서 설정한 number of classes 의 개수에 맞춰서 임의로 넣는다. 

(클래스 이름을 한글로 하면 오류남!!!!)

(알파벳이나 숫자로 넣고 나중에 detect.py에서 dictionary로 매칭해서 한글로 바꿨다)

 

데이터셋이 적은 경우에는 Pretrained된 모델을 사용하는게 더 좋을수도 있다.

 

    Model            APval            APtest           AP50          SpeedGPU       FPSGPU         params         FLOPS

YOLOv5s	37.0	37.0	56.2	2.4ms	416	7.5M	13.2B
YOLOv5m	44.3	44.3	63.2	3.4ms	294	21.8M	39.4B
YOLOv5l	47.7	47.7	66.5	4.4ms	227	47.8M	88.1B
YOLOv5x	49.2	49.2	67.7	6.9ms	145	89.0M	166.4B
YOLOv5x + TTA	50.8	50.8	68.9	25.5ms	39	89.0M	354.3B
YOLOv3-SPP	45.6	45.5	65.2	4.5ms	222	63.0M	118.0B

다운받아서 weights 폴더에 넣자.

weights 폴더에 넣고 class 개수도 내가 설정한 class의 개수로 바꿔야한다.

 

마지막으로 Anaconda Prompt를 켜서 cd 명령어로 yolo v5 폴더로 이동하고,

python train.py --img 640 --batch 16 --epochs 5 --data data/coco128.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt

를 입력하면 

 

학습이 시작된다.

학습 결과는 runs/exp .. 폴더에서 확인할 수 있다.

 

4. 객체 인식

가장 먼저 객체 인식하고 싶은 이미지를

inference/images 폴더에 넣고,

똑같이 Anaconda Prompt에서 yolov5 폴더로 이동한 뒤

python detect.py --source ./inference/images/ --weights runs/exp()/weights/best.pt --conf 0.4

를 실행한다.

(conf를 낮게하면 0.1, 0.2라고 판단한 객체도 출력한다)

 

 

5. 결과

inference/output 에서 확인할 수 있다.