YOLO v5 KITTI 학습 시키기

OpenJR·2022년 7월 4일

Yolo v5를 사용해 KITTI Dataset을 학습시켜보았다.

YOLO v5 받기

$ conda create -n yolo python=3.8.13
$ conda activate yolo

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

KITTI dataset 받기

http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d

첫번째 12GB 데이터와, 7번째 5MB 데이터를 다운받는다.

yolov5
└──dataset
   ├──trainning
   │  └──image_2
   │  └──label_2
   └──testing
      └──image_2

KITTI 형식 Yolo에 맞게 변경하기

gt값인 label_2의 txt파일들의 형식을 YOLO 형식에 맞게 바꿔준다. 아래의 깃허브를 통해 바꿔줄 수 있다.
https://github.com/ssaru/convert2Yolo

yolov5의 폴터 트리 구조에 맞게 폴더 이름을 변경한다.

yolov5
└──dataset
   ├──trainning
   │  └──images
   │  └──labels
   └──testing
      └──images

KITTI 데이터 파이프라인 만들기

yolov5를 사용하기 위해서 train에 쓰일 이미지들의 txt파일과 val에 쓰일 txt파일이 필요하다.

yolov5
└──dataset
   ├──trainning
   ├──testing
   ├──kitti.yaml
   ├──train.txt
   ├──val.txt
   ├──split.py

train.txt와 val.txt는 아래의 코드로 만들어 줄수 있다.

$ pip install -U scikit-learn

split.py

import glob
from sklearn.model_selection import train_test_split

img_list = sorted(glob.glob("dataset/training/images/*.png"))

train_list, val_list = train_test_split(img_list, test_size=0.2, random_state=2000)

with open("dataset/train.txt", "w") as f:
    f.write("\n".join(train_list) + "\n")

with open("dataset/val.txt", "w") as f:
    f.write("\n".join(val_list) + "\n")

두번째로 해야할 것은 kitti.yaml파일을 만드는 것이다. 아래와 같이 만들어 주면 된다.

train: /home/USER/Desktop/yolov5/dataset/train.txt
val: /home/USER/Desktop/yolov5/dataset/val.txt

# Classes
nc: 8  # number of classes
names: ['Car', 'Van', 'Truck', 'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram', 'DontCare']  # class names

디렉토리 위치나 USER는 사용자의 폴더에 맞게 바꿔주면 된다.

사용 모델의 클래스 수 바꿔주기

models/yolov5x.yaml안에 nc는 80으로 설정되어 있다.
사용하고자 하는 모델의 크기맞는 yaml파일을 하나 복사 해서 nc를 8개로 설정해준다.

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

# Parameters
nc: 8  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32
...

학습시키기

이제 모든 준비가 끝났다 아래의 명령어로 학습을 시켜준다.

python train.py --img 640 --batch 16 --epochs 5 --data dataset/kitti.yaml --cfg models/yolov5s_kitti.yaml  --weights yolov5s.pt --name results

--img: 이미지를 resize할 크기를 설정한다. kitti는 1224x370 이므로 640이 적정하다.
--batch: 학습시킬 배치 사이즈를 지정한다.
--epochs 학습 시킬 에포크 수를 지정한다.
--data: 앞서 만든 kitti.yaml로 학습시킬 데이터의 위치 정보를 가지고 있다.
--cfg: 학습시킬 모델의 크기나 설정 정보를 가지고 있다.
--weights: 미리 학습된 모델을 가져온다. 자동으로 모델을 서버에서 자동으로 다운로드 한다.
--name: 학습 결과가 저장될 폴더를 설정해준다. runs/train/results에 저장된다.