将labelme标注的实例分割数据转换为yolo_obb模型需要的txt数据

分享一个大家可能会用到的东西：将labelme标注的实例分割数据转换为yolo_obb模型需要的txt数据

import os
import json
import shutil
import numpy as np
import cv2
import matplotlib.pyplot as plt
from tqdm import tqdm


def get_file_list(path, file_list):
    """
    递归地获取指定路径下所有.json文件的路径列表。

    本函数会遍历指定的目录，寻找所有的.json文件，并将它们的路径添加到一个列表中。
    如果找到了对应的图片文件（通过替换.json为.png得到图片文件名），则只添加.json文件的路径到列表中；
    否则，会打印出那些没有对应图片文件的.json文件的路径。

    参数:
    path (str): 要遍历的目录路径。
    file_list (list): 用于存储找到的.json文件路径的列表。

    返回:
    list: 包含所有找到的.json文件路径的列表。
    """
    # 获取当前路径下的所有文件和文件夹
    file_paths = os.listdir(path)

    # 遍历当前路径下的所有文件和文件夹
    for file in file_paths:
        # 构成文件或文件夹的完整路径
        file_path = os.path.join(path, file)

        # 如果是文件夹，则递归调用本函数
        if os.path.isdir(file_path):
            get_file_list(file_path, file_list)
        else:
            # 如果是.json文件，则进行进一步检查
            if file_path.endswith('.json'):
                # 通过替换.json为.png得到对应图片文件的名称
                img_name = file_path.replace('.json', '.png')

                # 如果对应的图片文件存在于文件列表中，则添加.json文件的路径到列表中
                if os.path.basename(img_name) in file_paths:
                    file_list.append(file_path)
                else:
                    # 如果没有找到对应的图片文件，则打印.json文件的路径
                    print(file_path)

    # 返回包含所有找到的.json文件路径的列表
    return file_list


def labelme_to_obb(img_path: str, label_path: str, save_dir: str, if_test: bool=False):
    """
    将LabelMe标注工具生成的实例分割标签文件转换为yolo obb检测标签格式，并保存为TXT文件。

    参数:
    img_path: 图像文件路径。
    label_path: LabelMe生成的JSON标签文件路径。
    save_dir: 转换后的TXT标签文件保存目录。
    if_test: 是否生成测试图像，默认为False。

    返回:
    成功转换并保存标签文件时返回True，否则返回False。
    """
    # 定义类别名称，包括数字1-24和特定的字母
    classe_names = [str(i) for i in range(1, 25)] + ['X', 'Y', 'x', 'y']

    # 读取JSON标签文件
    with open(label_path, 'r') as f:
        label_data = json.load(f)
        shapes = label_data['shapes']
        img_w = label_data['imageWidth']
        img_h = label_data['imageHeight']
        img_w_h = np.array([[img_w, img_h]])

    label_list = []
    for shape in shapes:
        points = np.array(shape['points'], dtype=np.int32)
        # 检查标签是否在定义的类别名称中
        if shape['label'] not in classe_names:
            print('label %s not in classes' % shape['label'], label_path)
            label_list = None
            break
        # 根据标签分配类别ID
        if shape['label'] == 'X' or shape['label'] == 'x':
            label = 22
        elif shape['label'] == 'Y' or shape['label'] == 'y':
            label = 23
        else:
            label = int(shape['label']) - 1

        # 计算最小面积外接矩形
        rect = cv2.minAreaRect(points)
        box = cv2.boxPoints(rect)
        obb = box / img_w_h
        shape['points'] = obb.tolist()
        rex = np.round(obb, 5)
        # 将二维数组转为字符串
        rex = rex.reshape(-1)
        points_nor_str = ' '.join(map(str, rex))
        label_str = str(label) + ' ' + points_nor_str + '\n'
        label_list.append(label_str)

    # 如果标签列表不为空，则保存标签文件
    if label_list is not None:
        label_path = os.path.join(save_dir, os.path.basename(label_path).replace('.json', '.txt'))
        with open(label_path, 'w') as txt_file:
            for i in label_list:
                txt_file.writelines(i)
        # 如果需要，生成测试图像
        if if_test:
            if not os.path.exists(save_dir+'_test'):
                os.makedirs(save_dir+'_test')
            obb_label_to_img(img_path, label_path, os.path.join(r'E:\DeepLearning_data\yolo_obb\test',
                                                                os.path.basename(label_path).replace('.txt', '.png')))
        return True
    else:
        return False


def obb_label_to_img(img_path, label_path, save_path):
    """
    将旋转框标签映射到图像上。

    读取图像和标签文件，根据标签信息在图像上绘制旋转框，并保存结果图像。

    参数:
    img_path (str): 图像文件路径。
    label_path (str): 标签文件路径，每行包含一个旋转框的信息。
    save_path (str): 绘制旋转框后图像的保存路径。
    """
    # 读取图像
    img = cv2.imread(img_path)
    # 获取图像的高、宽
    height, width, _ = img.shape
    # 打开标签文件
    with open(label_path, 'r') as file_handle:
        # 逐行读取标签信息
        cnt_info = file_handle.readlines()
        # 解析标签信息，去除换行符并分割
        new_cnt_info = [line_str.replace("\n", "").split(" ") for line_str in cnt_info]

    # 遍历每个标签信息
    for new_info in new_cnt_info:
        points = []
        # 提取坐标点，转换为图像坐标
        for i in range(1, len(new_info), 2):
            x_y = [float(tmp) for tmp in new_info[i:i + 2]]
            points.append([int(x_y[0] * width), int(x_y[1] * height)])
        # 在图像上绘制多边形
        cv2.polylines(img, [np.array(points, np.int32)], True, (0, 255, 255))
    # 保存绘制后的图像
    plt.imsave(save_path, img)


if __name__ == '__main__':
    path_img = [r'E:\DeepLearning_data\in_seg_data\train', r'E:\DeepLearning_data\in_seg_data\val']
    img_save_path = [r'E:\DeepLearning_data\yolo_obb\images\train', r'E:\DeepLearning_data\yolo_obb\images\val']
    label_save_path = [r'E:\DeepLearning_data\yolo_obb\labels\train', r'E:\DeepLearning_data\yolo_obb\labels\val']
    for i in range(len(path_img)):
        for j in tqdm(os.listdir(path_img[i])):
            if j.endswith('.json'):
                continue
            else:
                if_copy = labelme_to_obb(os.path.join(path_img[i], j), os.path.join(path_img[i], j)[:-4] + '.json', label_save_path[i])
                if if_copy:
                    shutil.copyfile(os.path.join(path_img[i], j), os.path.join(img_save_path[i], j))