[colab] yolov6 커스텀 데이터 셋으로 훈련시키기

데이터셋 구성하기

Labelimg 다운

https://github.com/heartexlabs/labelImg

GitHub - heartexlabs/labelImg: LabelImg is now part of the Label Studio community. The popular image annotation tool created by

설치 명령어

pip install PyQt5
pip install lxml
pyrcc5 -o libs/resources.py resources.qrc
pip3 install labelImg

실행 명령어

python labelImg.py

 

Labelimg 사용법

루트 폴더의 data/predefined_classes.txt를 열어서 라벨 이름을 적어준다.

gumban
mylabel
youlable
#줄바꿈해서 아래에 계속 적어주면 된다

 

좌측하단의 Save 밑의 아이콧이 YOLO가 될때까지 클릭한다. 

이미지 폴더를 선택한다. 

w 버튼을 누른 상태에서 드래그해서 사각형을 선택하고 선택한 사물의 이름을 적어준다.

저장하고 다음 버튼을 누른다. 

 

만약 2개 이상의 물체를 라벨링 할거면 순서를 정하고 순서대로 라벨링을 해야한다. 

ex. 알, 애벌레, 성충 순으로 저장

안그러면 성능이 저하될 수 있다.

저장하면 선택한 폴더에 classes.txt 파일이 생성되는데 사물 순서를 라벨링할 당시 정했던 순서대로 적어준다. 

 

1개의 물체만 라벨링 하는 경우 UseDefault Label에 체크해두면 편하다.

 

이미지 크기가 너무 크면 학습시키는데 자원이 많이 들어가기 때문에 이미지 크기를 적절히 줄여준다. 보통 200~600px 사이로 줄이는 것 같다. 

 

이미지 리사이즈 python 스크립트

from PIL import Image
import os

image_dir = [원본 이미지 경로]

save_dir = [저장할 이미지 경로]
images = os.listdir(image_dir)

for img in images:
    im = Image.open(image_dir + "/" +img)
    re = im.resize((332, 187)) #줄일 사이즈
    re.save(save_dir + "/" + img)

 

데이터셋 폴더 구조

images
ㄴtest
    ㄴimage1.jpg
    ㄴimage2.jpg
    ㄴimage3.jpg
    ㄴ...
ㄴtrain
    ㄴimage1.jpg
    ㄴimage2.jpg
    ㄴimage3.jpg
    ㄴ...
ㄴval
    ㄴimage1.jpg
    ㄴimage2.jpg
    ㄴimage3.jpg
    ㄴ...
lables
ㄴtest
    ㄴimage1.txt
    ㄴimage2.txt
    ㄴimage3.txt
    ㄴ...
ㄴtrain
    ㄴimage1.txt
    ㄴimage2.txt
    ㄴimage3.txt
    ㄴ...
ㄴval
    ㄴimage1.txt
    ㄴimage2.txt
    ㄴimage3.txt
    ㄴ...
 ㄴclasses.txt #labelImg로 라벨링을 한 후 저장하면 선택한 폴더에서 자동으로 생성이된다. labels 폴더로 옮겨준다.

 

데이터셋 설정 파일 만들기

train: ./resized_custom_dataset/images/train #학습 이미지 경로
val: ./resized_custom_dataset/images/val #검증 이미지 경로
test: ./resized_custom_dataset/images/test #테스트 이미지경로

nc: 1 #탐지할 사물 개수

is_coco: False 
names: ['gumban'] #라벨링한 사물이름

 

환경 설정 

파이썬 3.8을 설치하고 가상환경을 생성한다.

python -m venv [가상환경이름]
cd Scripts
activate
cd ..

 

Yolov6 다운받기 

#yolov6 다운로드
git clone https://github.com/meituan/YOLOv6
cd YOLOv6

#의존 패키지 설치
pip install -r requirements.txt

 

모델 다운받기 

import torch
#링크는 나중에 변경될 수 있다.
torch.hub.download_url_to_file('https://github.com/meituan/YOLOv6/releases/download/0.4.0/yolov6n.pt', 'yolov6n.pt')

 

훈련시키기

훈련 명령어 

python tools/train.py 
    --batch 256  
    --conf configs/yolov6n.py  #백지상태에서 훈련
    # --conf configs/yolov6n_finetune.py  #기존 학습된 모델에 추가로 학습시키기
    --data data/custom_dataset.yaml #데이터셋 설정 파일 경로
    --epochs 100 #훈련 횟수: 100번
    --name custom 모델이 저장될 폴더 이름
    --img-size 640 #훈련 시 이미지 크기

 

검증 명령어

python tools/eval.py 
    --weight runs/train/custom/weights/best_ckpt.pt  #훈련 모델 경로
    --data data/custom_dataset.yaml #데이터셋 설정파일 경로
    --img-size 640 #훈련 시 이미지 크기

훈련 모델 테스트 명령어

python tools/infer.py 
    --weight runs/train/custom/weights/best_ckpt.pt 
    --source custom_dataset/image.jpg #테스트할 이미지 파일 경로
    --yaml  data/custom_dataset.yaml #데이터셋 설정파일 경로
    --img-size 640 #훈련 시 이미지 크기