Custom Dataset으로 YOLOv5/YOLOv7 학습하기 (22.12.07 Update)

(22.12.07) YOLOv7 추가

Custom Dataset 준비

자신이 가지고 있는 데이터를 사용하고 싶다! → labelImg

labeling이 귀찮다! → 존재하는 dataset을 사용한다.

 

1. LabelImg

LabelImg를 사용해 데이터셋을 생성하는 방법은 다음을 참조한다.

LabelImg 설치 및 사용방법

 

이 글에서는 roboflow에서 제공하는 dataset을 사용해 학습할 것이다!

 

 

2. roboflow

데이터셋은 roboflow에서 제공하는 North American Mushrooms Dataset을 사용한다. 링크로 이동 후 데이터셋을 다운로드 받자.

 

학습 시간을 줄이기 위해 51장의 이미지를 다운 받는다. (둘 중 어느 것을 다운 받아도 상관 없음)

 

빨간 박스 부분을 눌러보면 다양한 label format이 존재하는데, 사용할 format에 맞게 다운로드 하면 된다. 이 글에서는 YOLOv5를 사용할 것이므로 YOLO v5 PyTorch 선택 후 다운로드 한다.

 

(22.12.07 추가)

YOLOv7의 경우 다음과 같이 YOLO v7 PyTorch 선택 후 다운로드 한다.

 

 

데이터셋의 폴더 구성은 다음과 같다.

 

North_American_Mushrooms 폴더 내의 data.yaml 파일을 열어서 다음과 같이 데이터 경로를 수정 해준다.

train: North_American_Mushrooms/train/images
val: North_American_Mushrooms/valid/images

nc: 2
names: ['CoW', 'chanterelle']

• train : train dataset 경로.
• val : validation dataset 경로.
• nc : class 개수.
• names : class 목록. class_id를 기준으로 왼쪽부터 오름차순 정렬

 

YOLOv5 불러오기

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

 

(22.12.07 추가)

YOLOv7

git clone https://github.com/WongKinYiu/yolov7.git
cd yolov7
pip install -r requirements.txt

 

yolov5 폴더 아래에 다운 받은 데이터셋을 위치시킨다. 폴더 구성은 다음과 같다.

 

Train

다음 명령어로 모델을 학습 시킬 수 있다.

python train.py --data North_American_Mushrooms/data.yaml

 

사용할 수 있는 주요 옵션들은 다음과 같다. (자세한 내용은 train.py의 line 476~525 참고)

 

• --weights : 사전 학습 된 모델 weight를 불러온다. (.pt 파일)
• --cfg : 모델 config. 모델의 parameter나 클래스 개수 등을 설정 해준다. ex) models/yolov5s.yaml
• --data : dataset 기본 정보를 담고 있다. 클래스 목록/개수 등의 정보가 들어 있다.
• --epochs : 에폭을 설정할 수 있다.
• --batch_size : 배치 사이즈를 설정할 수 있다.
• --imgsz, --img, --img_size : input image size를 설정한다.

 

명령어를 실행하면 학습 환경이나 사용되는 parameter 등이 출력되며, 다음과 같은 화면이 출력된다면 성공적으로 학습을 진행중인 것이다.

 

학습이 완료되면 runs/train/ 경로에 학습 결과가 저장된다.

첫 학습 시 exp, 두 번째 학습 부터 exp1, exp2... 와 같은 형태로 네이밍 된다.

 

runs/train/exp/weights/best.pt 파일이 한 번의 학습(exp)에서 성능이 가장 좋은 model weights가 되며, 이를 이용해 모델 검증, 예측이 가능하다.

 

모델 검증

모델 학습과 비슷한 방법으로 모델 검증이 가능하다. 데이터 경로(--data)와 모델 가중치(--weights)를 옵션으로 넣어준다.

python val.py --data North_American_Mushrooms/data.yaml --weights runs/train/exp/weights/best.pt

 

(22.12.07 추가)

YOLOv7는 val.py 존재 하지 않음.

 

 

검증 결과 역시 모델 학습과 비슷한 화면이 출력 되며, runs/val/exp 경로에 결과가 저장된다.

 

모델 예측

모델 예측에는 detect.py 파일을 사용한다. 모델 가중치(--weights)와 test data 경로(--source)를 옵션으로 넣어준다.

python detect.py --weights runs/train/exp/weights/best.pt --source North_American_Mushrooms/test/images

 

--source 옵션을 수정해주면 이미지 뿐만 아니라 웹캠, 비디오 등에서도 예측이 가능하다.

python detect.py --source 0  # webcam
python detect.py --source image.jpg  # image
python detect.py --source video.mp4  # video
python detect.py --source path/  # directory
python detect.py --source path/*.jpg  # glob
python detect.py --source "YouTube Link"  # YouTube
python detect.py --source "RTSP stream"  # RTSP

 

학습, 검증과 동일하게 예측 결과는 runs/detect/exp에 저장된다. 해당 경로를 탐색 해보면 예측 결과 이미지가 저장 되어 있는 것을 확인할 수 있다.