【Python-实操】LabelMe to YOLOv8 Converter

LabelMe to YOLOv8 Converter

这是一个 Python 脚本，用于将 LabelMe 标注工具导出的 JSON 文件转换为 YOLOv8 格式的标注文件，并同时在图像上绘制标注的多边形。

功能

• 读取 LabelMe JSON 文件。
• 解码并显示图像。
• 从 classes.txt 文件加载类别标签。
• 将多边形标注转换为 YOLOv8 的格式。
• 在图像上绘制多边形标注。
• 将原始图像和带有标注的图像拼接在一起并保存。
• 保存 YOLOv8 格式的标注文件。

安装

确保安装了以下依赖包：

• OpenCV (pip install opencv-python)
• NumPy (pip install numpy)

使用方法

1. 将此脚本放置在包含 LabelMe JSON 文件的目录内。
2. 在同一目录下创建一个 classes.txt 文件，每行一个类别标签。
3. 运行脚本，并指定 JSON 文件所在的目录作为命令行参数。

命令行参数

• /path/to/json/files: 包含 LabelMe JSON 文件的目录路径。

示例

python script.py /path/to/json/files

代码

import os
import json
import base64
import cv2
import numpy as np


def draw_polygon(image, points, color=(0, 255, 0), thickness=2):
    """
    在给定图像上绘制一个多边形。

    :param image: 待绘制多边形的图像（numpy数组）
    :param points: 多边形顶点坐标列表
    :param color: 多边形的颜色 (B, G, R)
    :param thickness: 边缘线条厚度
    """
    # 绘制多边形轮廓
    cv2.polylines(image, [np.int32(points)], isClosed=True, color=color, thickness=thickness)
    # 绘制多边形内部
    cv2.fillPoly(image, [np.int32(points)], color=color)


def concat_images(image1, image2):
    """
    拼接两个图像为一个垂直堆叠的图像。

    :param image1: 第一张图像（numpy数组）
    :param image2: 第二张图像（numpy数组）
    :return: 垂直拼接后的图像（numpy数组）
    """
    # 确保两个图像的宽度相同
    max_height = max(image1.shape[0], image2.shape[0])
    max_width = max(image1.shape[1], image2.shape[1])

    # 调整图像大小以匹配最大宽度
    if image1.shape[1] < max_width:
        image1 = cv2.copyMakeBorder(image1, 0, 0, 0, max_width - image1.shape[1], cv2.BORDER_CONSTANT, value=[255, 255, 255])
    if image2.shape[1] < max_width:
        image2 = cv2.copyMakeBorder(image2, 0, 0, 0, max_width - image2.shape[1], cv2.BORDER_CONSTANT, value=[255, 255, 255])

    # 如果需要，调整图像高度以匹配最大高度
    if image1.shape[0] < max_height:
        image1 = cv2.copyMakeBorder(image1, 0, max_height - image1.shape[0], 0, 0, cv2.BORDER_CONSTANT, value=[255, 255, 255])
    if image2.shape[0] < max_height:
        image2 = cv2.copyMakeBorder(image2, 0, max_height - image2.shape[0], 0, 0, cv2.BORDER_CONSTANT, value=[255, 255, 255])

    # 拼接图像
    return np.vstack((image1, image2))


def convert_labelme_to_yolov8(json_dir):
    """
    将 LabelMe 格式的标注转换为 YOLOv8 格式，并绘制多边形到图像上。

    :param json_dir: 包含 LabelMe JSON 文件的目录路径
    """
    # 生成颜色列表
    color_list = [
        (0, 0, 255),  # Red
        (0, 255, 0),  # Green
        (255, 0, 0),  # Blue
        (0, 255, 255),  # Yellow
        (255, 255, 0),  # Cyan
        (255, 0, 255),  # Magenta
        (0, 165, 255),  # Orange
        (203, 192, 255),  # Pink
        (42, 42, 165),  # Brown
        (0, 128, 128),  # Olive
        (128, 128, 0),  # Teal
        (238, 130, 238),  # Violet
        (128, 128, 128),  # Gray
        (192, 192, 192),  # Silver
        (0, 0, 128),  # Maroon
        (128, 0, 128),  # Purple
        (0, 0, 128),  # Navy
        (0, 255, 0),  # Lime
        (0, 255, 255),  # Aqua
        (255, 0, 255),  # Fuchsia
        (255, 255, 255),  # White
        (0, 0, 0),  # Black
        (235, 206, 135),  # Light Blue
        (144, 238, 144),  # Light Green
        (193, 182, 255),  # Light Pink
        (224, 255, 255),  # Light Yellow
        (216, 191, 216),  # Light Purple
        (0, 128, 128),  # Light Olive
        (30, 105, 210),  # Light Brown
        (211, 211, 211)  # Light Gray
    ]

    # 加载类别文件
    classes_file = os.path.join(json_dir, 'classes.txt')
    
    if not os.path.exists(classes_file):
        print("Error: Could not find 'classes.txt' in the specified directory.")
        exit(1)

    with open(classes_file, 'r') as f:
        class_names = [line.strip() for line in f.readlines()]

    # 获取 JSON 文件列表
    json_files = [f for f in os.listdir(json_dir) if f.endswith('.json')]

    for json_file in json_files:
        json_file_path = os.path.join(json_dir, json_file)
        
        # 输出文件名
        output_file_name = json_file.replace('.json', '.txt')
        output_file_path = os.path.join(json_dir, output_file_name)

        # 读取 JSON 文件
        with open(json_file_path, 'r') as f:
            data = json.load(f)

        image_width = data['imageWidth']
        image_height = data['imageHeight']

        # 解码图像数据
        imageData = data.get('imageData')
        if imageData is not None:
            image_data = base64.b64decode(imageData)
            image_np = np.frombuffer(image_data, dtype=np.uint8)
            image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
            
            # 创建一个副本用于绘制标注
            annotated_image = image.copy()

            # 绘制标注
            for i, shape in enumerate(data['shapes']):
                if shape['shape_type'] == 'polygon':
                    points = np.array(shape['points'], dtype=np.int32)
                    label = shape['label']
                    class_index = class_names.index(label)
                    color = color_list[class_index % len(color_list)]
                    draw_polygon(annotated_image, points, color=color)

            # 保存原始图像
            image_file_name = json_file.replace('.json', '.jpg')
            image_file_path = os.path.join(json_dir, image_file_name)
            cv2.imwrite(image_file_path, image)

            # 将原始图像和带有标注的图像上下拼接
            concatenated_image = concat_images(image, annotated_image)

            # 保存拼接后的图像
            concatenated_image_file_name = json_file.replace('.json', '_check.jpg')
            concatenated_image_file_path = os.path.join(json_dir, concatenated_image_file_name)
            cv2.imwrite(concatenated_image_file_path, concatenated_image)
            
            # 显示带有标注的图像
            cv2.imshow('Annotated Image', concatenated_image)
            cv2.waitKey(0)
            cv2.destroyAllWindows()

        # 开始写入 YOLOv8 格式的文本文件
        with open(output_file_path, 'w') as f:
            for shape in data['shapes']:
                if shape['shape_type'] == 'polygon':
                    label = shape['label']
                    points = shape['points']
                    
                    # 获取类别索引
                    try:
                        class_index = class_names.index(label)
                    except ValueError:
                        print(f"Warning: Label '{label}' not found in 'classes.txt'. Skipping this label.")
                        continue
                    
                    # 归一化坐标
                    normalized_points = []
                    for point in points:
                        x = point[0] / image_width
                        y = point[1] / image_height
                        normalized_points.extend([x, y])

                    # 写入 YOLOv8 格式的行
                    line = f"{class_index} {' '.join(map(str, normalized_points))}\n"
                    f.write(line)

if __name__ == '__main__':
    import sys

    # 从命令行参数获取 JSON 目录
    if len(sys.argv) != 2:
        print("Usage: python script.py /path/to/json/files")
        exit(1)

    json_dir = sys.argv[1]
    convert_labelme_to_yolov8(json_dir)

代码结构

draw_polygon 函数

该函数在给定图像上绘制一个多边形，并填充颜色。

concat_images 函数

该函数将两张图像拼接为一个垂直堆叠的图像。

convert_labelme_to_yolov8 函数

该函数执行以下操作：

• 读取 classes.txt 文件中的类别标签。
• 遍历目录中的所有 JSON 文件。
• 对每个 JSON 文件执行以下操作：
  • 解码并显示图像。
  • 读取 JSON 文件的内容。
  • 在图像上绘制标注的多边形。
  • 将原始图像与带标注的图像拼接。
  • 保存拼接后的图像。
  • 将多边形标注转换为 YOLOv8 格式，并保存为 .txt 文件。

注意事项

• 确保 classes.txt 文件正确无误地列出了所有的类别标签。
• 请确保脚本有足够的权限来读写文件。