yolo v3 형식 data를 yolo v5 형식 data로 변경하기

프로젝트 진행 중 yolo v3를 학습시킨 데이터를 이용해 yolo v5에 학습을 시켜야하는 상황이 생겼습니다.

yolov3의 데이터 형식으로 csv를 사용하였는데, csv는 "파일경로 obj1의 xmin, ymin, xmax, ymax, class_id obj2의 xmin, ymin, xmax, ymax ......"의 형식으로 이뤄져 있습니다.

 

yolov5의 데이터 형식은 튜토리얼에 따르면 다음과 같습니다.

사진과 같이 yolov5 디렉토리 옆에 dataset directory가 존재하고, 데이터 디렉토리는 image, labels로 나뉘고 그 안에서 train set과 validation set으로 나누어지는 트리구조를 가졌습니다.

 

그리고 yaml파일로 데이터를 정의해주는데요, 

사진과 같이 trainset, validation set의 폴더 경로(이미지)와 class 개수, class명 list로 이루어져 있습니다.

그리고 labels 폴더 안에는 이미지명과 동일한 명을 가진 txt파일이 존재하고, 그 txt파일은 

다음과 같이 object별로 줄바꿈을 해주고, class_id, 중앙x좌표, 중앙y좌표, w, h 의 형식이지만 좌표는 모두 normalized되어 있습니다.

 

다음은 v3의 csv 파일을 v5에 맞추어 변경하는 코드입니다.

v5모델 자체에 대해서는 다음 글에서 다루겠습니다.

import cv2

with open('C:/Users/playcarnival/Desktop/train_data/total_texts.txt') as f:
    lines = f.readlines()
    print("***lines***\n", lines)
    for line in lines:
        objects = line.split(' ')
        img_path = objects[0]
        name = img_path.split('/')[-1].split('.')[0]
        objects = objects[1:]
        print('name:', name, 'objects:', objects)
        text_list = []
        for obj in objects:
            if len(obj) > 2:
                obj = obj.strip()
                xy_list = obj.split(',')
                print('xy_list:', xy_list)
                xy_list = list(map(int, xy_list))
                xmin, ymin, xmax, ymax, class_id = xy_list[0:]
                img = cv2.imread(img_path)
                w, h = img.shape[:2]
                # center, w, h, : Normalized
                xc, yc, w, h = round((xmin+xmax)/2/w, 6), round((ymin+ymax)/2/h, 6), round((xmax-xmin)/w, 6), round((ymax-ymin)/h, 6) 
                text = f'{class_id} {xc} {yc} {w} {h}\n'
                print('text:', text)
                text_list.append(text)
                
            else:
                continue
            
        with open('C:/Users/playcarnival/Desktop/train_data/text_v5/'+ name + '.txt', 'w') as t:
                for text in text_list:
                    t.write(text)
                t.close()