labelme标注后的json文件和原图同步按角度旋转

点击查看代码

import json
import os
import base64
import numpy as np
import cv2
from math import cos, sin, radians
import argparse


def rotate_point_opencv_style(point, rotation_matrix):
    """用OpenCV旋转矩阵计算标注点，确保与图像旋转逻辑一致"""
    point_homo = np.array([point[0], point[1], 1], dtype=np.float32)
    rotated_point = np.dot(rotation_matrix, point_homo)
    return (round(rotated_point[0]), round(rotated_point[1]))


def image_to_base64(img):
    """OpenCV图像转Base64编码（LabelMe imageData字段需求）"""
    ret, buffer = cv2.imencode('.jpg', img)
    if not ret:
        raise ValueError("无法编码图像为JPEG格式")
    return base64.b64encode(buffer).decode('utf-8')


def rotate_single_labelme(json_path, output_root_dir, angle):
    """
    单文件LabelMe旋转处理（核心功能，复用原逻辑）
    :param json_path: 单个JSON文件路径
    :param output_root_dir: 批量输出的根目录（会保持原目录结构）
    :param angle: 旋转角度（度，逆时针为正）
    """
    try:
        # 1. 读取LabelMe数据和原始图像
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)

        # 拼接原始图像路径（LabelMe的imagePath是相对JSON的路径）
        img_origin_rel_path = data['imagePath']
        img_origin_abs_path = os.path.join(os.path.dirname(json_path), img_origin_rel_path)
        img_origin = cv2.imread(img_origin_abs_path)
        if img_origin is None:
            raise FileNotFoundError(f"图像文件不存在：{img_origin_abs_path}")

        # 2. 图像维度与旋转中心（严格对齐LabelMe与OpenCV坐标）
        h_origin, w_origin = img_origin.shape[:2]  # OpenCV：h=行（y），w=列（x）
        center_origin = (w_origin // 2, h_origin // 2)  # 旋转中心：(x=列, y=行)

        # 3. 计算旋转后图像尺寸（避免裁剪，完整保留旋转内容）
        angle_rad = radians(angle)
        w_rot = int(h_origin * abs(sin(angle_rad)) + w_origin * abs(cos(angle_rad)))
        h_rot = int(w_origin * abs(sin(angle_rad)) + h_origin * abs(cos(angle_rad)))

        # 4. 生成OpenCV旋转矩阵（图像与标注共用，强制同步）
        M = cv2.getRotationMatrix2D(center=center_origin, angle=angle, scale=1.0)
        # 偏移量：让旋转后的图像居中显示
        dx = (w_rot // 2) - center_origin[0]
        dy = (h_rot // 2) - center_origin[1]
        M[0, 2] += dx
        M[1, 2] += dy

        # 5. 旋转图像
        img_rot = cv2.warpAffine(
            src=img_origin,
            M=M,
            dsize=(w_rot, h_rot),
            flags=cv2.INTER_LINEAR,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=(255, 255, 255)  # 白色背景（可自定义）
        )

        # 6. 旋转标注点（用同一个矩阵M，无偏差）
        for shape in data['shapes']:
            rotated_points = [rotate_point_opencv_style(p, M) for p in shape['points']]
            shape['points'] = [list(p) for p in rotated_points]  # 转为List格式（LabelMe要求）

        # 7. 构建输出路径（保持原目录结构，避免文件混乱）
        # 例：原路径 ./data/label1.json → 输出 ./output/data/label1_rotated_90deg.json
        relative_json_path = os.path.relpath(json_path, start=os.path.dirname(output_root_dir))
        output_json_dir = os.path.join(output_root_dir, os.path.dirname(relative_json_path))
        os.makedirs(output_json_dir, exist_ok=True)

        # 8. 更新LabelMe JSON字段并保存
        # 处理图像文件名（添加旋转标识）
        img_filename = os.path.basename(img_origin_rel_path)
        img_name, img_ext = os.path.splitext(img_filename)
        new_img_name = f"{img_name}_rotated_{angle}deg{img_ext}"
        # 处理JSON文件名
        json_filename = os.path.basename(json_path)
        json_name, json_ext = os.path.splitext(json_filename)
        new_json_name = f"{json_name}_rotated_{angle}deg{json_ext}"

        # 更新JSON内容
        data['imagePath'] = new_img_name  # 关联旋转后的图像
        data['imageWidth'] = w_rot  # 旋转后图像宽度（列数）
        data['imageHeight'] = h_rot  # 旋转后图像高度（行数）
        data['imageData'] = image_to_base64(img_rot)  # 更新Base64编码

        # 保存旋转后的图像和JSON
        output_img_path = os.path.join(output_json_dir, new_img_name)
        output_json_path = os.path.join(output_json_dir, new_json_name)
        cv2.imwrite(output_img_path, img_rot)
        with open(output_json_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

        return True, f"成功：{os.path.basename(json_path)} → {new_json_name}"

    except Exception as e:
        return False, f"失败：{os.path.basename(json_path)} → 错误原因：{str(e)}"


def batch_rotate_labelme(input_root_dir, output_root_dir, angle):
    """
    批量旋转目录中所有LabelMe JSON文件
    :param input_root_dir: 输入根目录（会递归遍历所有子目录）
    :param output_root_dir: 输出根目录（保持与输入一致的目录结构）
    :param angle: 旋转角度（度，逆时针为正）
    """
    # 校验输入目录是否存在
    if not os.path.isdir(input_root_dir):
        raise NotADirectoryError(f"输入目录不存在：{input_root_dir}")
    os.makedirs(output_root_dir, exist_ok=True)

    # 统计变量
    total_count = 0
    success_count = 0
    fail_list = []

    # 递归遍历所有子目录，寻找JSON文件
    for root, dirs, files in os.walk(input_root_dir):
        for file in files:
            # 只处理LabelMe生成的JSON文件（避免非标注JSON）
            if file.endswith('.json'):
                total_count += 1
                json_abs_path = os.path.join(root, file)
                # 调用单文件处理函数
                success, msg = rotate_single_labelme(json_abs_path, output_root_dir, angle)
                print(msg)  # 实时打印处理进度
                if success:
                    success_count += 1
                else:
                    fail_list.append(msg)

    # 输出批量处理总结
    print("\n" + "=" * 50)
    print(f"批量处理完成！")
    print(f"总文件数：{total_count}")
    print(f"成功数：{success_count}")
    print(f"失败数：{len(fail_list)}")
    if fail_list:
        print("失败详情：")
        for fail_msg in fail_list:
            print(f"  - {fail_msg}")


if __name__ == "__main__":
    # 命令行参数配置（支持批量处理）
    parser = argparse.ArgumentParser(description='批量旋转LabelMe标注数据（图像与标注100%对齐）')
    parser.add_argument('input_dir', help='输入根目录（会递归遍历所有子目录的JSON）')
    parser.add_argument('output_dir', help='输出根目录（保持与输入一致的目录结构）')
    parser.add_argument('angle', type=float, help='旋转角度（度，逆时针为正，例：90/180/-90）')

    args = parser.parse_args()
    # 启动批量处理
    batch_rotate_labelme(args.input_dir, args.output_dir, args.angle)