최근에 EfficientDet이 발표되고나서, 많은 사람들이 이 기술을 Custom dataset에 적용하려고 한다.
보통 Object Detection에는 COCO와 VOC 데이터 셋을 많이 사용하고는 하는 데, 데이터 포멧때문에 애를 먹는 경우가 많다. 귀찮거나 변환을 해야하기때문에...
일반적으로 1 stage detector를 적용할 경우, yolov3를 사용하여 적용하는 경우가 많은 데, yolo의 데이터 라벨링은 이에 비해 엄청 간단하다. 그저 클래스, 중앙 x, 중앙 y, width, height만 입력하면 되기 때문이다. ( ex : 1 120 240 10 10 )
그래서 이 format을 토대로 yolov3등을 (요즘에는 yolov4까지 나왔지만.. ) 사용하는 경우가 많은 데, 문제는 포멧을 쉽게 변환하고 싶은 데, 그게 매우 귀찮음을 동반하면서 누군가에게는 어려울 수도 있다는 것이다.
특히, coco나 voc를 yolo로 변환시켜주는 코드는 많다. 그리고 voc를 coco로, 또는 coco를 voc로 변환시켜주는 것은 많은 데, 이상하게 yolo코드를 coco로 변환시켜주는 코드는 내 능력 부족일 수도 있겠지만 찾아보니까 없었다.
그래서 그냥 직접 짜보았고, 이 귀찮음을 느끼는 분들을 해소하기 위해 코드를 공개한다. 별거 아니긴 하지만...
사용법은 간단하다.
각 번호마다 주의사항을 반드시 읽고, 다른 부분이 있다면 본인이 직접 수정해주자!
1. train 혹은 valid 의 이미지 데이터가 있는 곳을 image_dir_path 에 넣어준다.
2. train 혹은 valid 의 라벨 데이터가 있는 곳을 label_dir_path에 넣어준다.
** 이 코드는 이미지 데이터와 라벨 데이터가 동일하다는 전제하에 돌아가는 코드이기 때문에 혹시라도 라벨과 이미지의 이름이 다르다면, (그런 경우는 많이는 없지만), 라벨에 해당하는 코드 부분을 수정하길 바란다.
일치하지 않으면, 에러나게 코딩 되어있음.
3. save_file_name 에 저장할 이름.json 으로 써주자. 나의 경우에는 saving.json 으로 하였다.
** 반드시 뒤에 .json 을 붙여줘야한다.
4. is_normalized에 True 나 False를 넣어준다.
Yolo 코드들 중에서 어떤 것들은 normalized 된 것들이 있고, 어떤 것들은 없다.
예시 (1) normalized 되지 않은 경우의 라벨 파일
(이미지 가로 세로가 모두 1000이라고 가정)
0 100 100 10 10
1 200 200 30 30
3 300 300 15 15
예시 (2) normalized 된 경우의 라벨 파일
(이미지 가로 세로가 모두 1000이라고 가정)
0 0.1 0.1 0.01 0.01
1 0.2 0.2 0.03 0.03
3 0.3 0.3 0.15 0.15
예시 2처럼 Normalized 된 파일은 True를, 예시 1처럼 안되어 있다면 False를 넣어주자.
5. 라이센스나 클래스 추가
내 경우에는, 학습하는 데이터에 탐지할 것이 1개여서 코드 안에 딱 1개의 클래스만 설정되어 있다.
만일 학습할 클래스가 3개라면
# make categories <= 이부분 밑에
class 를 3개 추가해주면된다.
내가 주석으로 예시를 넣었으니, 그대로 추가시키면된다.
라이센스도 마찬가지로 #make licenses를 보면된다.
보통 라이센스는 별로 건드릴 일이 없고, 클래스만 설정해주면 된다.
예시) 클래스 3개이고, 이름이 cow, cat, coco일 경우,
# make categories
category_list = []
class_0 = {
'id': 1,
'name' : 'cow',
'supercategory' : 'None'
}
category_list.append(class_0)
class_1 = {
'id': 2,
'name' : 'cat',
'supercategory' : 'None'
}
category_list.append(class_1)
class_2 = {
'id': 3,
'name' : 'coco',
'supercategory' : 'None'
}
category_list.append(class_2)
import json
import os
import cv2
def yolo2coco(image_dir_path, label_dir_path,save_file_name , is_normalized):
total = {}
# make info
info = {
'description' : '',
'url' : '',
'version' : '',
'year' : 2020,
'contributor' : '',
'data_created' : '2020-04-14 01:45:18.567988'
}
total['info'] = info
# make licenses
licenses_list = []
licenses_0= {
'id' : '1',
'name' : 'your_name',
'url' : 'your_name'
}
licenses_list.append(licenses_0)
''' if you want to add licenses, copy this code
licenses_1 = {
'id': '2',
'name': 'your_name',
'url': 'your_name'
}
licenses_list.append(licenses_1)
'''
total['licenses'] = licenses_list
# make categories
category_list = []
class_0 = {
'id': 1,
'name' : 'defect',
'supercategory' : 'None'
}
category_list.append(class_0)
'''
# if you want to add class
class_1 = {
'id': 2,
'name' : 'defect',
'supercategory' : 'None'
}
category_list.append(class_1)
'''
total['categories'] = category_list
# make yolo to coco format
# get images
image_list = os.listdir(image_dir_path)
print('image length : ', len(image_list))
label_list = os.listdir(label_dir_path)
print('label length : ',len(label_list))
image_dict_list = []
count = 0
for image_name in image_list :
img = cv2.imread(image_dir_path+image_name)
image_dict = {
'id' : count,
'file_name' : image_name,
'width' : img.shape[1],
'height' : img.shape[0],
'date_captured' : '2020-04-14 -1:45:18.567975',
'license' : 1, # put correct license
'coco_url' : '',
'flickr_url' : ''
}
image_dict_list.append(image_dict)
count += 1
total['images'] = image_dict_list
# make yolo annotation to coco format
label_dict_list = []
image_count = 0
label_count = 0
for image_name in image_list :
img = cv2.imread(image_dir_path+image_name)
label = open(label_dir_path+image_name[0:-4] + '.txt','r')
if not os.path.isfile(label_dir_path + image_name[0:-4] + '.txt'): # debug code
print('there is no label match with ',image_dir_path + image_name)
return
while True:
line = label.readline()
if not line:
break
class_number, center_x,center_y,box_width,box_height = line.split()
# should put bbox x,y,width,height
# bbox x,y is top left
if is_normalized :
center_x = int(float(center_x) * int(img.shape[1]))
center_y = int(float(center_y) * int(img.shape[0]))
box_width = int(float(box_width) * int(img.shape[1]))
box_height = int(float(box_height) * int(img.shape[0]))
top_left_x = center_x - int(box_width/2)
top_left_y = center_y - int(box_height/2)
if not is_normalized :
center_x = float(center_x)
center_y = float(center_y)
box_width = float(box_width)
box_height = float(box_height)
top_left_x = center_x - int(box_width / 2)
top_left_y = center_y - int(box_height / 2)
bbox_dict = []
bbox_dict.append(top_left_x)
bbox_dict.append(top_left_y)
bbox_dict.append(box_width)
bbox_dict.append(box_height)
# segmetation dict : 8 points to fill, x1,y1,x2,y2,x3,y3,x4,y4
segmentation_list_list = []
segmentation_list= []
segmentation_list.append(bbox_dict[0])
segmentation_list.append(bbox_dict[1])
segmentation_list.append(bbox_dict[0] + bbox_dict[2])
segmentation_list.append(bbox_dict[1])
segmentation_list.append(bbox_dict[0]+bbox_dict[2])
segmentation_list.append(bbox_dict[1]+bbox_dict[3])
segmentation_list.append(bbox_dict[0])
segmentation_list.append(bbox_dict[1] + bbox_dict[3])
segmentation_list_list.append(segmentation_list)
label_dict = {
'id' : label_count,
'image_id' : image_count,
'category_id' : int(class_number)+1,
'iscrowd' : 0,
'area' : int(bbox_dict[2] * bbox_dict[3]),
'bbox' : bbox_dict,
'segmentation' : segmentation_list_list
}
label_dict_list.append(label_dict)
label_count += 1
label.close()
image_count += 1
total['annotations'] = label_dict_list
with open(save_file_name,'w',encoding='utf-8') as make_file :
json.dump(total,make_file, ensure_ascii=False,indent='\t')
if __name__ == '__main__':
image_dir_path = './images/'
label_dir_path = './labels/'
save_file_name = 'saving.json'
is_normalized = False
# if you want to add more licenses or classes
# add in code
yolo2coco(image_dir_path, label_dir_path, save_file_name,is_normalized)
이렇게 할 경우, json 파일 뷰어로 보면, 정상적으로 생성되었음을 알 수 있다.
혹시 틀리거나 수정할 부분이 있다면 댓글로 남겨주면 수정을 하겠습니다.
** 5.27 수정
내가 잘못 한 부분이 좀 많아서 다시 올린다.
** 2020.10.7
완성본인데
혹시라도 잘못된 점 있으면 알려주시면 감사드리겠습니다.