点击查看代码
import json
import os
import base64
import numpy as np
import cv2
from math import cos, sin, radians
import argparsedef 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_abs_path = os.path.splitext(json_path)[0] + '.jpg'# 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] += dxM[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.jsonrelative_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_abs_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编码# 保存旋转后的图像和JSONoutput_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 = 0success_count = 0fail_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 += 1json_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 += 1else: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', default=r'trains\seg_gelinmei_yicikeli\labeled', help='输入根目录（会递归遍历所有子目录的JSON）')# parser.add_argument('output_dir', default=r'seg_gelinmei_yicikeli\augumented',help='输出根目录（保持与输入一致的目录结构）')# parser.add_argument('angle', type=float, default='3', help='旋转角度（度，逆时针为正，例：90/180/-90）')# input_dir = r'trains\seg_gelinmei_yicikeli\labeled'# output_dir = r'trains\seg_gelinmei_yicikeli\augumented'input_dir = r'trains\seg_hongmoshui\labeled'output_dir = r'trains\seg_hongmoshui\augumented'angle = 3# args = parser.parse_args()# 启动批量处理batch_rotate_labelme(input_dir, output_dir, angle)
