[Yolov5] Custom Dataset으로 Yolov5 돌려보기 : 데이터셋 준비

이전에 Custom Dataset을 통해 yolo를 학습하는 글을 쓴 적이 있는데

다시 보니 미숙함이 너무 잘 보여서 처음부터 차근차근 써보고자 한다.

 

https://github.com/ultralytics/yolov5

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

 

- 시작하기 전에 Object Detection이 뭔지, Yolo가 어떤 모델인지 알고 가면 좋다! *준비중

- 사실은 아래 링크에 모든 것이 나와있어요

https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

 

---

데이터셋 준비하기

당연하게도 모델을 학습하기 위해서는 데이터셋이 필요하다.

yolo를 학습하기 위한 데이터셋 형태는 다음과 같다.

---

학습을 위한 이미지 파일(들)

개별 이미지의 yolo 형식 annotation 텍스트 파일

---

 

이 때 Annotation 파일들은 개별 이미지 파일과 같은 이름을 가져야 하고,

이미지 폴더가 위치한 디렉토리 내의 /labels 폴더에 저장되어야 한다.

따라서 yolo 데이터들은 다음과 같은 구조로 저장되어야 한다.

 

root 디렉토리

ㄴ 이미지 폴더

ㄴ labels

 

yolo bounding box format은 다음과 같다.

 

출처 : yolov5 Github

 

bounding box는 [x, y, w, h] 형태이고 이는 이미지 크기를 0-1로 Normalize했을 때의 값이다.

또한 개별 annotation txt 파일은 다음과 같이 클래스(인덱스), x, y, w, h 형태로 쓰여 있어야 한다.

 

 

0, x1, y1, w1, h1

1, x2, y2, w2, h2,

...

 

 

운이 좋다면 여러분이 하고자 하는 작업에 알맞은 데이터셋이 존재할 것이고 그렇지 않다면 직접 데이터를 수집해야 하는 수고로움이 필요하다. 또는 이미 존재하는 label을 yolo 형식으로 잘 변환해야 할 수도 있을 것이다...

 

이 글에서는 이런 방법들을 최대한 많이 소개해볼 것이다. (왜냐? 필자가 거의 다 해 봤음..)

---

밑바닥부터 데이터셋 구축하기

여기서 말하는 "밑바닥부터"는 모든 이미지 데이터를 수집하고 annotation 또한 손수 달아야 함을 뜻한다.

다양한 노가다 방법들이 있겠지만 유용한 사이트를 하나 소개하고자 한다.

 

https://www.makesense.ai/

Make Sense

 

 

해당 사이트에 들어가면 이렇게 쉽게! bounding box를 class별로 그리고 label을 yolo format으로 쉽게 저장할 수 있다.

 

---

xml 형식을 yolo 형식으로 변환하기

 

아래 사이트를 매우 매우 많이 참고하였다. (정말정말 감사합니다...)

https://towardsdatascience.com/convert-pascal-voc-xml-to-yolo-for-object-detection-f969811ccba5

Convert PASCAL VOC XML to YOLO for Object Detection

 

파일 경로는 다음과 같다고 가정하자.

 

dataset

ㄴ annotations (개별 xml 파일이 존재하는)

ㄴ images (이미지 파일이 존재하는)

ㄴ labels (새 yolo label이 들어갈 폴더)

 

# 경로를 지정해주자.

path = "./dataset"
annot_path = os.path.join(path,"annotations")
img_path = os.path.join(path,"images")
label_path = os.path.join(path,"labels")

import xml.etree.ElementTree as ET
import glob
import os
import json

# xml bbox 형식을 yolo bbox 형태로 변환하는 함수

def xml_to_yolo_bbox(bbox, w, h):
    # xmin, ymin, xmax, ymax
    x_center = ((bbox[2] + bbox[0]) / 2) / w
    y_center = ((bbox[3] + bbox[1]) / 2) / h
    width = (bbox[2] - bbox[0]) / w
    height = (bbox[3] - bbox[1]) / h
    return [x_center, y_center, width, height]
    
 classes = []

from tqdm import tqdm

files = glob.glob(os.path.join(annot_path, '*.xml'))
for fil in tqdm(files):
    
    basename = os.path.basename(fil)
    filename = os.path.splitext(basename)[0]
    
    result = []
    
    tree = ET.parse(fil)
    root = tree.getroot()
    width = int(root.find("size").find("width").text)
    height = int(root.find("size").find("height").text)
    for obj in root.findall('object'):
        label = obj.find("name").text
        if label not in classes:
            classes.append(label)
        index = classes.index(label)
        pil_bbox = [int(x.text) for x in obj.find("bndbox")]
        yolo_bbox = xml_to_yolo_bbox(pil_bbox, width, height)
        bbox_string = " ".join([str(x) for x in yolo_bbox])
        result.append(f"{index} {bbox_string}")
    if result:
        with open(os.path.join(label_path, f"{filename}.txt"), "w", encoding="utf-8") as f:
            f.write("\n".join(result))

 

코드가 돌아가면 labels 폴더 안에 yolo format의 annotation 텍스트 파일들이 예쁘게 저장되어 있을 것이다.

 

---

처음부터 yolo 형식으로 저장된 데이터셋 구하기

https://www.kaggle.com/datasets/alvarobasily/road-damage

Road Damage

 

운이 정말정말 좋다면 이렇게 처음부터 yolo format으로 annotation 된 데이터셋을 구할 수도 있다!