Yolov7 커스텀 데이터셋에 학습시키기

 

https://blog.paperspace.com/yolov7/

How to train and use a custom YOLOv7 model

다음 포스팅을 보고 요약한 게시물입니다.

 

 

 

 

1. 우선 yolov7을 클론한다.

git clone https://github.com/WongKinYiu/yolov7.git

 

 

2. 다운로드한 폴더로 들어가서 파이썬 라이브러리들을 설치해준다.

pip install -r requirements.txt
pip install setuptools==59.5.0
pip install torchvision==0.11.3+cu111 -f https://download.pytorch.org/whl/cu111/torch_stable.html

코드를 실행하다보면 추가로 요구되는 라이브러리들이 있을 수 있는데 오류에 맞게 다운로드 해주면 된다.

 

 

3. 모델 베이스라인을 다운로드 해준다.

wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7_training.pt

 

4. 데이터셋을 다운로드 해준다.

roboflow에서 데이터를 라벨링하면 터미널에서 다운로드 할 수 있는 코드를 준다.

다음은 위의 링크의 포스트에 올라와 있던 농구 플레이어 데이터셋이다.

curl -L "https://app.roboflow.com/ds/4E12DR2cRc?key=LxK5FENSbU" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

 

 

5. 위의 데이터셋을 다운로드 하면, test, train, valid 세 개의 폴더가 생기는데 이를 v-test라는 폴더에 넣고 한 번에 관리해 줄 것이다.

mkdir v-test
mv train/ v-test/
mv valid/ v-test/

또는 다음 명령어를 사용하지 않고 그냥 v-test라는 이름의 폴더를 생성하고 train, test, valid라는 세 개의 폴더를 넣어주면 된다.

 

 

6. Roboflow에서 추출해온 파일 이름에는 추가적인 문자들이 붙어있는 경우가 있다. 이를 제거해준다.

import os 

# remove roboflow extra junk

count = 0
for i in sorted(os.listdir('v-test/train/labels')):
    if count >=3:
        count = 0
    count += 1
    if i[0] == '.':
        continue
    j = i.split('_')
    dict1 = {1:'a', 2:'b', 3:'c'}
    source = 'v-test/train/labels/'+i
    dest = 'v-test/train/labels/'+j[0]+dict1[count]+'.txt'
    os.rename(source, dest)
    
count = 0
for i in sorted(os.listdir('v-test/train/images')):
    if count >=3:
        count = 0
    count += 1
    if i[0] == '.':
        continue
    j = i.split('_')
    dict1 = {1:'a', 2:'b', 3:'c'}
    source = 'v-test/train/images/'+i
    dest = 'v-test/train/images/'+j[0]+dict1[count]+'.jpg'
    os.rename(source, dest)
    
for i in sorted(os.listdir('v-test/valid/labels')):
    if i[0] == '.':
        continue
    j = i.split('_')
    source = 'v-test/valid/labels/'+i
    dest = 'v-test/valid/labels/'+j[0]+'.txt'
    os.rename(source, dest)
    
for i in sorted(os.listdir('v-test/valid/images')):
    if i[0] == '.':
        continue
    j = i.split('_')
    source = 'v-test/valid/images/'+i
    dest = 'v-test/valid/images/'+j[0]+'.jpg'
    os.rename(source, dest)
for i in sorted(os.listdir('v-test/test/labels')):
    if i[0] == '.':
        continue
    j = i.split('_')
    source = 'v-test/test/labels/'+i
    dest = 'v-test/test/labels/'+j[0]+'.txt'
    os.rename(source, dest)
    
for i in sorted(os.listdir('v-test/test/images')):
    if i[0] == '.':
        continue
    j = i.split('_')
    source = 'v-test/test/images/'+i
    dest = 'v-test/test/images/'+j[0]+'.jpg'
    os.rename(source, dest)

다음의 코드를 실행하면 되는데, 한 번만 실행해야지 제대로 된 파일명이 나온다. 여러번 실행하면 txt 확장자가 여러번 붙는다 ...

 

 

7. data/coco.yaml 파일을 수정해준다.

train: ./v-test/train/images/
val: ./v-test/valid/images/
test: ./v-test/test/images/

nc: 2
names: ['ball-handler', 'player']

train, val, test는 각각의 폴더의 위치를 나타내며, nc는 number of class, 클래스 수를 의미한다.

roboflow에서 데이터를 가져오는 경우, data.yaml이라는 파일로 자동으로 파일이 생성된다.

 

 

8. 이제 모든 준비가 끝났다. 다음의 명령어를 실행해서 학습시키면 된다.

# Train on single GPU
!python train.py --workers 8 --device 0 --batch-size 8 --data data/coco.yaml --img 1280 720 --cfg cfg/training/yolov7.yaml --weights yolov7_training.pt --name yolov7-ballhandler --hyp data/hyp.scratch.custom.yaml --epochs 50

# Train on 2 GPUs
!python -m torch.distributed.launch --nproc_per_node 2 --master_port 9527 train.py --workers 16 --device 0,1 --sync-bn --batch-size 8 --data data/coco.yaml --img 1280 720 --cfg cfg/training/yolov7.yaml --weights yolov7_training.pt --name yolov7-ballhandler --hyp data/hyp.scratch.custom.yaml --epochs 50

 

 

9. 테스트

--source 뒤에 테스트하려는 이미지/비디오의 위치를 적어주면 된다.

!python detect.py --weights runs/train/yolov7-ballhandler/weights/best.pt --conf 0.25 --img-size 1280 --source video.mp4 --name test

runs/detect/test 폴더에 검출된 이미지가 다운로드 된다.

 

10. 추가적으로 생기는 오류

 

나의 경우에는 RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu) 다음과 같은 오류가 떴다. 다음을 참고하여 해결 할 수 있었다. utils/loss.py의 742 번째 줄을 다음으로 바꾸었더니 해결되었다.

matching_matrix = torch.zeros_like(cost, device="cpu")

 

https://github.com/WongKinYiu/yolov7/issues/1101

RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu) · Issue #1101 · WongKinYiu/yol

 

또한, numpy 버전이 맞지 않아 생기는 오류도 있었는데, numpy 버전을 다운그레이드 하여 진행하였더니 해결되었다.

https://github.com/WongKinYiu/yolov7/issues/1280

Issue while training: AttributeError: module 'numpy' has no attribute 'int' · Issue #1280 · WongKinYiu/yolov7