import json
import os
from glob import glob
def labelme_to_yolo(labelme_json_path, output_dir, class_list):
    """
    将LabelMe标注的JSON文件转换为YOLO格式的txt文件
    参数:
        labelme_json_path: LabelMe JSON文件路径或包含多个JSON文件的目录
        output_dir: 输出YOLO格式文件的目录
        class_list: 类别列表，确定类别索引
    """
    # 如果输入是目录，处理所有JSON文件
    if os.path.isdir(labelme_json_path):
        json_files = glob(os.path.join(labelme_json_path, '*.json'))
    else:
        json_files = [labelme_json_path]
    print(len(json_files))
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    # 处理每个JSON文件
    for json_file in json_files:
        with open(json_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
        # 获取图像尺寸
        img_width = data['imageWidth']
        img_height = data['imageHeight']
        # 准备YOLO格式内容
        yolo_lines = []
        for shape in data['shapes']:
            label = shape['label']
            points = shape['points']
            # 确保类别在class_list中
            if label not in class_list:
                class_list.append(label)
            class_id = class_list.index(label)
            # 转换多边形或矩形到YOLO格式
            if shape['shape_type'] in ['rectangle', 'polygon']:
                # 获取所有点的x,y坐标
                x_coords = [p[0] for p in points]
                y_coords = [p[1] for p in points]
                # 计算边界框
                x_min = min(x_coords)
                x_max = max(x_coords)
                y_min = min(y_coords)
                y_max = max(y_coords)
                # 计算中心点、宽度和高度（归一化）
                x_center = (x_min + x_max) / 2 / img_width
                y_center = (y_min + y_max) / 2 / img_height
                width = (x_max - x_min) / img_width
                height = (y_max - y_min) / img_height
                # 确保值在0-1范围内
                x_center = max(0, min(1, x_center))
                y_center = max(0, min(1, y_center))
                width = max(0, min(1, width))
                height = max(0, min(1, height))
                yolo_lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
        # 写入YOLO格式文件
        output_filename = os.path.splitext(os.path.basename(json_file))[0] + '.txt'
        output_path = os.path.join(output_dir, output_filename)
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(yolo_lines))
    # 保存类别文件
    #with open(os.path.join(output_dir, 'classes.txt'), 'w', encoding='utf-8') as f:
    #    f.write('\n'.join(class_list))
    print(class_list)
    print(f"转换完成！共处理 {len(json_files)} 个文件。")
    #print(f"类别列表已保存到: {os.path.join(output_dir, 'classes.txt')}")
if __name__ == '__main__':
    # 使用示例
    labelme_path = './num_yolo/labels/val'  # 替换为你的LabelMe JSON文件或目录路径
    output_directory = './num_yolo/labels_txt/val'       # 替换为输出目录
    classes = ["0", "1", "2", "3"]                  # 替换为你的类别列表，可以留空[]自动从JSON中收集
    labelme_to_yolo(labelme_path, output_directory, classes)