pytorch-YOLOv4 커스텀 데이터 사용법

 

 

1. 서론

YOLOv4는 Darknet 기반으로 구현했으나 필자는 코드를 이해하기 어려워 Pytorch로 커스터마이징 된 코드를 사용하였다.

 

2. 설치

 

1. git을 사용하여 코드를 다운로드한다.

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

2. pytorch 설치, cuda, cuDNN 설치

windows10: https://jeo96.tistory.com/entry/Pytorch-CUDA-cuDNN-%EC%84%A4%EC%B9%98-windows-10

Pytorch, CUDA, cuDNN 설치 [windows 10]

3. 데이터셋 준비

데이터셋은 roboflow에서 적당한 예재로 예를 들겠습니다.

 

1. roboflow 접속

https://roboflow.com/

Roboflow: Give your software the power to see objects in images and video

2. Public Datasets 클릭

3. 원하는 데이터셋 선택한 후 아래 이미지처럼 YOLOv4-Pytorch를 선택하여 다운로드합니다

4. 사용법

1. 다운로드한 파일을 압축 해제하고 다음 코드를 사용하여 경로 정보가 담긴 txt 파일을 생성해 줍니다.

import os


# 파일 리스트 얻기
def get_file_list(dir_path):
    file_list = os.listdir(dir_path)
    file_list.sort()
    for f in range(len(file_list)):
        file_list[f] = os.path.join(dir_path, file_list[f])
    return file_list


# 이미지 경로를 담고있는 txt 생성
def make_txt_file(file_list, txt_name):
    with open(txt_name, 'w') as f:
        for line in file_list:
            f.write(line + '\n')


if __name__ == '__main__':
    file_list = get_file_list(r"학습데이터 이미지 경로 ")  # 데이터셋 이미지 리스트 얻기
    make_txt_file(file_list, 'train_Hard_Hat.txt')  # txt 파일 생성
    file_list = get_file_list(r"평가데이터 이미지 경로")
    make_txt_file(file_list, 'val_Hard_Hat.txt')

2. 위 코드를 사용하면 아래 그림처럼 이미지의 절대 경로가 담긴 txt 파일을 생성할 수 있습니다.

3. yaml 파일 편집

PyTorch_YOLOv4/data/coco.yaml을 복사한다.

train에는 위에서 생성한 train.txt 경로를 입력, val에도 val.txt 경로 입력, test는 val.py로 테스트 데이터셋을 사용하지 않으면 명시하지 않아도 된다. nc는 클래스 개수, names는 클래스 이름을 클래스 번호에 맞게 순서대로 나열한다.

4. cfg 파일 편집

cfg 파일 변경하는 방법은 Darknet으로 구현된 YOLOv4에 명시되어있다.

https://github.com/AlexeyAB/darknet.git

GitHub - AlexeyAB/darknet: YOLOv4 / Scaled-YOLOv4 / YOLO - Neural Networks for Object Detection (Windows and Linux version of Da

PyTorch_YOLOv4/cfg/yolov4.cfg를 복사한다.

yolov4.cfg를 맨 밑으로 내리면 [yolo] 밑에 classes를 본인이 사용하는 데이터셋의 클래스 개수로 변경한다.

max_batches = classes * 2000

steps = max_batches * 0.8, max_batches * 0.9

[yolo] 앞에 있는 [convolutional] 안에 filters를 수정한다. (3군데 있음)

filters = (classes + 5) * 3

 

이제 train.py를 --cfg, --data 등 parser를 사용하여 실행하면 된다.