将多边形标注数据转为YOLO的txt格式

格式：一对一对的浮点数，表示多边形顶点的 x 和 y 坐标。坐标是归一化的，即值范围在 [0, 1] 之间，分别表示相对于图像的宽度和高度。

import json
import os

'''
任务：实例分割，labelme的json文件, 转txt文件
Ultralytics YOLO format
     ...  
'''

# 类别映射表，定义每个类别对应的ID
label_to_class_id = {
    "corn": 0,

    # 根据需要添加更多类别
}


# json转txt
def convert_labelme_json_to_yolo(json_file, output_dir, img_width, img_height):
    with open(json_file, 'r') as f:
        labelme_data = json.load(f)

    # 获取文件名（不含扩展名）
    file_name = os.path.splitext(os.path.basename(json_file))[0]

    # 输出的txt文件路径
    txt_file_path = os.path.join(output_dir, f"{file_name}.txt")

    with open(txt_file_path, 'w') as txt_file:
        for shape in labelme_data['shapes']:
            label = shape['label']
            points = shape['points']

            # 根据类别映射表获取类别ID，如果类别不在映射表中，跳过该标签
            class_id = label_to_class_id.get(label)
            if class_id is None:
                print(f"Warning: Label '{label}' not found in class mapping. Skipping.")
                continue

            # 将点的坐标归一化到0-1范围
            normalized_points = [(x / img_width, y / img_height) for x, y in points]

            # 写入类别ID
            txt_file.write(f"{class_id}")

            # 写入多边形掩膜的所有归一化顶点坐标
            for point in normalized_points:
                txt_file.write(f" {point[0]:.6f} {point[1]:.6f}")
            txt_file.write("\n")


if __name__ == "__main__":
    json_dir = "D:/草/25.4.19 17点-19点 晴天/玉米224/近154/biaozhu"
    output_dir = "D:/草/25.4.19 17点-19点 晴天/玉米224/近154/biaozhu"

    # 调试信息
    print(f"Input directory: {json_dir}")
    print(f"Files in input directory: {os.listdir(json_dir)}")

    # 创建输出文件夹
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"Created output directory: {output_dir}")
    else:
        print(f"Output directory already exists: {output_dir}")

    processed_files = 0

    # 批量处理所有json文件
    for json_file in os.listdir(json_dir):
        if json_file.endswith(".json"):
            json_path = os.path.join(json_dir, json_file)
            print(f"Processing: {json_path}")

            # 从JSON文件中读取图像尺寸
            with open(json_path, 'r') as f:
                labelme_data = json.load(f)
                img_height = labelme_data['imageHeight']
                img_width = labelme_data['imageWidth']

            convert_labelme_json_to_yolo(json_path, output_dir, img_width, img_height)
            processed_files += 1

    print(f"Processing complete. Converted {processed_files} files.")