YOLO label to Labelme JSON label 변환

YOLO txt 형식 label 데이터를 라벨링 툴인 Labelme에서 읽을 수 있는 JSON 형식으로 바꿔 준다.

이미지와 txt 라벨이 같은 디렉터리에 있을 경우 그대로 사용 가능하다.

다른 디렉터리일 경우 코드 수정이 필요하다.

import os
import glob
import json
import time
import datetime
import shutil
import cv2

def calculate_points(image_width: int, image_height: int, x_center_scaling: float, y_center_scaling:float, 
                     w_scaling:float, h_scaling:float):
    w = w_scaling * image_width
    h = h_scaling * image_height
    
    xmin = (x_center_scaling * image_width) - int(w / 2)
    ymin = (y_center_scaling * image_height) - int(h / 2)
    xmax = (x_center_scaling * image_width) + int(w / 2)
    ymax = (y_center_scaling * image_height) + int(h / 2)
    
    return xmin, ymin, xmax, ymax

labelme_version = "5.0.1"
flages = {}

class_names = ["person", "car"]

# 이미지, 라벨이 다른 디렉터리일 경우 코드 수정 필요
# 서로 다른 디렉터리 pair 확인 시 파일이름으로 정렬 필요
origin_image_and_yolo_txt_label_path = "./pair_data_img_label/"
dest_json_label_path = "./json_label/"

origin_image_and_yolo_txt_label_path_list = glob.glob(os.path.join(origin_image_and_yolo_txt_label_path, "*"))

pair_count = 0
pair_len = int(len(origin_image_and_yolo_txt_label_path_list) / 2)

if len(origin_image_and_yolo_txt_label_path_list) % 2 != 0:
    print("No pair image and label data")
    exit()

start_time = time.time()

for i in range(1, len(origin_image_and_yolo_txt_label_path_list), 2):
    file_name_1 = os.path.splitext(origin_image_and_yolo_txt_label_path_list[i - 1])[0]
    file_name_1 = os.path.split(file_name_1)[1]
    file_name_1_ext = os.path.splitext(origin_image_and_yolo_txt_label_path_list[i - 1])[1]
    
    file_name_2 = os.path.splitext(origin_image_and_yolo_txt_label_path_list[i])[0]
    file_name_2 = os.path.split(file_name_2)[1]
    file_name_2_ext = os.path.splitext(origin_image_and_yolo_txt_label_path_list[i])[1]
    
    if file_name_1 == file_name_2:
        pair_count += 1
    else:
        print(f"No pair image and label: ({origin_image_and_yolo_txt_label_path_list[i - 1]}) != ({origin_image_and_yolo_txt_label_path_list[i]})")
        break
    
    if pair_count % 100 == 0:
        print(f"{pair_count} / {pair_len}")
    
    text_label_path = None
    text_label_file_name = None
    image_file_path = None
    
    if file_name_1_ext == ".txt":
        text_label_path = origin_image_and_yolo_txt_label_path_list[i - 1]
        text_label_file_name = file_name_1
        image_file_path = origin_image_and_yolo_txt_label_path_list[i]
    else:
        text_label_path = origin_image_and_yolo_txt_label_path_list[i]
        text_label_file_name = file_name_2
        image_file_path = origin_image_and_yolo_txt_label_path_list[i - 1]
    
    with open (text_label_path, "r") as f:
        lines = f.readlines()

        if len(lines) == 0 or lines[0].replace("\n", "") == "":
            print("Empty label data: ", text_label_path)
            
        with open(dest_json_label_path + text_label_file_name + ".json", 'w') as json_file:
            json_data = {}
            json_data["version"] = labelme_version
            json_data["flags"] = flages
            json_data["shapes"] = []
            
            image = cv2.imread(image_file_path)
            h, w, c = image.shape
            
            for line in lines:
                object_data = line.replace("\n", "").split(" ")
                shape = {}
                
                shape["label"] = class_names[int(object_data[0])]
                
                xmin, ymin, xmax, ymax = calculate_points(w, h, float(object_data[1]), float(object_data[2]), float(object_data[3]), float(object_data[4]))
                shape["points"] = []
                shape["points"].append([xmin, ymin])
                shape["points"].append([xmax, ymax])
                shape["group_id"] = None
                shape["shape_type"] = "rectangle"
                shape["flages"] = {}
                
                json_data["shapes"].append(shape)
                
                
            json_data["imagePath"] = os.path.split(image_file_path)[1]
            json_data["imageData"] = None
            json_data["imageHeight"] = h
            json_data["imageWidth"] = w
            
            json.dump(json_data, json_file, indent=2)
            shutil.copyfile(image_file_path, dest_json_label_path + os.path.split(image_file_path)[1])
        
print(f"Total Data: {len(origin_image_and_yolo_txt_label_path_list)} Total Pair: {pair_count}")
print(f"Total Time: {datetime.timedelta(seconds=int(time.time() - start_time))}")