[AI 부트캠프] DAY 120 - 파이널 프로젝트 11

[오늘의 일지]

프로젝트 - YOLO 모델을 위한 데이터셋 알아보기

[상세 내용]

YOLO 모델을 위한 데이터셋

- 다른 모델과는 다르게 딥러닝 모델들은 각 모델에 맞는 데이터셋을 잘 구축하는 것이 중요합니다. 제대로 된 데이터셋을 구축하지 않으면 모델이 제대로 학습되지 않기 때문입니다. 그래서 오늘은 YOLO 모델에 맞는 데이터셋을 구축하는 과정에 대해서 정리해 보겠습니다. 오늘은 아마도 일부만 정리하고 나머지는 내일 정리할 것 같습니다. 

 

특징

- 여러 가지 사이트를 검색해 본 결과 YOLO 모델은 예를 들면 하나의 이미지가 존재한다면 하나의 annotation 파일이 존재해야지 학습이 가능한 것으로 파악을 했습니다. annotation이라는 단어는 그 자체로 주석이라는 뜻을 가지는데 YOLO 모델의 annotation 안에는 다음과 같은 순서로 값이 존재해야 한다는 것을 알게 되었습니다.

# annotation 내용

<클래스> <x_중심> <y_중심> <너비> <높이>

 

 

데이터셋 구축 코드 예시 (데이콘 대회에 올라온 코드 참고)

- 여기서는 annotation 파일을 txt형식으로 저장해서 사용한 것으로 파악했습니다. 다른 형식으로 저장해도 되는지는 저도 알아봐야 할 것 같습니다.

def make_yolo_dataset(image_paths, txt_paths, type="train"):
    for image_path, txt_path in tqdm(zip(image_paths, txt_paths if not type == "test" else image_paths), total=len(image_paths)):
        source_image = cv2.imread(image_path, cv2.IMREAD_COLOR)        
        image_height, image_width, _ = source_image.shape
        
        target_image_path = f"../data/yolo/{type}/{os.path.basename(image_path)}"
        cv2.imwrite(target_image_path, source_image)
        
        if type == "test":
            continue
        
        with open(txt_path, "r") as reader:
            yolo_labels = []
            for line in reader.readlines():
                line = list(map(float, line.strip().split(" ")))
                class_name = int(line[0])
                x_min, y_min = float(min(line[5], line[7])), float(min(line[6], line[8]))
                x_max, y_max = float(max(line[1], line[3])), float(max(line[2], line[4]))
                x, y = float(((x_min + x_max) / 2) / image_width), float(((y_min + y_max) / 2) / image_height)
                w, h = abs(x_max - x_min) / image_width, abs(y_max - y_min) / image_height
                yolo_labels.append(f"{class_name} {x} {y} {w} {h}")
            
        target_label_txt = f"../data/yolo/{type}/{os.path.basename(txt_path)}"      
        with open(target_label_txt, "w") as writer:
            for yolo_label in yolo_labels:
                writer.write(f"{yolo_label}\n")

 

[마무리]

 오늘은 YOLO 모델의 데이터셋을 구성하는 방법에 대해서 알아보았습니다. 내용은 짧은 것 같아도 자세하게 파악하기까지 시간이 생각보다 오래 걸렸던 것 같습니다. 그냥 코드만 보고 파악하느라 더 오래 걸린 것 같습니다. 내일은 데이터셋을 구성해서 모델을 간단하게 학습하는 것까지 해볼 수 있으면 좋을 거 같습니다.