图像旋转增广后JSON标注与Mask位置偏移问题求助

我将图像分为三类：

• A类（标注为Acircle）：近乎完美的圆形
• B类（标注为Bcircle）：接近圆形但存在变形或缺陷
• C类（标注为Ccircle）：形状与圆形差距较大

因数据量不足，对1980x1080的原始图像执行旋转增广，旋转角度包含30°、60°、90°、120°、150°、180°、210°、240°、270°、300°、330°。目前JPG图像旋转结果准确，但JSON标注文件和Mask文件中的类别位置无法与图像对齐，常出现偏离中心或完全错位的情况，已调试代码7天仍未解决，相关代码如下：

import os
import json
import numpy as np
import cv2
from PIL import Image
import matplotlib.pyplot as plt

# 이미지 리사이즈 함수
def resize_image(image, target_size=(1024, 1024)):
    """Resizes the image to the target size."""
    return image.resize(target_size, Image.Resampling.LANCZOS)

# 좌표 스케일링 함수
def scale_points(points, original_size, new_size):
    """Scales the points to match the resized image."""
    ox, oy = original_size  # 원본 이미지 크기
    nx, ny = new_size  # 새 이미지 크기
    scaled_points = [(x * nx / ox, y * ny / oy) for x, y in points]
    return scaled_points

# 이미지 회전 함수
def rotate_image(image, angle):
    """Rotates the image by the given angle."""
    return image.rotate(angle, expand=True)

# JSON 파일의 점 좌표를 회전하는 함수 (이미지 중심을 기준으로 좌표 변환)
def rotate_points(points, angle, image_size):
    """Rotates points in JSON to match image rotation based on image center."""
    cx, cy = image_size[0] / 2, image_size[1] / 2  # 이미지 중심
    angle_rad = np.radians(angle)
    rotation_matrix = np.array([[np.cos(angle_rad), -np.sin(angle_rad)],
                                [np.sin(angle_rad), np.cos(angle_rad)]])
    rotated_points = []
    for point in points:
        px, py = point
        # 이미지 중심을 기준으로 좌표 변환
        translated_point = np.array([px - cx, py - cy])
        rotated_point = np.dot(rotation_matrix, translated_point) + np.array([cx, cy])
        rotated_points.append(rotated_point.tolist())
    return rotated_points

# 마스크 생성 함수
def create_mask_from_json(json_file, image_size, class_mapping):
    """Creates a mask from a JSON file containing polygon points."""
    mask = np.zeros(image_size[::-1], dtype=np.uint8)  # 마스크는 기본적으로 0으로 채워짐 (배경)
    
    with open(json_file, 'r') as f:
        data = json.load(f)

    for shape in data['shapes']:
        points = np.array(shape['points'], dtype=np.int32)
        label = shape['label']

        # 라벨에 해당하는 클래스 ID로 마스크 채우기
        class_id = class_mapping.get(label, 0)  # label이 class_mapping에 없으면 0으로 설정 (배경)
        
        cv2.fillPoly(mask, [points], class_id)

    return mask

# 이미지와 JSON 파일을 회전 및 리사이즈하여 저장하는 함수
def augment_and_transform(image_dir, json_dir, save_image_dir, save_json_dir, save_mask_dir, target_size=(1024, 1024), angles=[30, 60, 90]):
    os.makedirs(save_image_dir, exist_ok=True)
    os.makedirs(save_json_dir, exist_ok=True)
    os.makedirs(save_mask_dir, exist_ok=True)

    image_files = [f for f in os.listdir(image_dir) if f.endswith('.jpg')]
    
    for image_file in image_files:
        image_path = os.path.join(image_dir, image_file)
        json_path = os.path.join(json_dir, image_file.replace('.jpg', '.json'))

        # 이미지 및 JSON 파일 열기
        original_image = Image.open(image_path)
        with open(json_path, 'r') as f:
            data = json.load(f)

        original_size = original_image.size
        resized_image = resize_image(original_image, target_size)  # 이미지 리사이즈

        for angle in angles:
            rotated_image = rotate_image(resized_image, angle)  # 이미지 회전
            rotated_size = rotated_image.size  # 获取旋转后的实际尺寸

            rotated_shapes = []
            for shape in data['shapes']:
                points = shape['points']

                # 좌표 스케일링
                scaled_points = scale_points(points, original_size, target_size)

                # 先将缩放后的坐标转换为旋转后图像的坐标系（因为expand=True会改变尺寸，需要平移到新中心）
                # 计算原缩放图像中心到旋转后图像中心的偏移
                offset_x = (rotated_size[0] - target_size[0]) / 2
                offset_y = (rotated_size[1] - target_size[1]) / 2
                translated_to_rotated = [(x + offset_x, y + offset_y) for x, y in scaled_points]

                # 修正旋转矩阵方向：PIL是逆时针旋转，原矩阵是顺时针，所以取负角度
                rotated_points = rotate_points(translated_to_rotated, -angle, rotated_size)

                rotated_shapes.append({
                    "label": shape['label'],
                    "points": rotated_points,
                    "shape_type": shape["shape_type"]
                })

            # 회전된 이미지 저장
            rotated_image_name = f"{image_file.replace('.jpg', '')}_rot{angle}.jpg"
            rotated_image.save(os.path.join(save_image_dir, rotated_image_name))

            # 회전된 JSON 파일 저장
            rotated_json_name = f"{image_file.replace('.jpg', '')}_rot{angle}.json"
            rotated_data = data.copy()
            rotated_data['shapes'] = rotated_shapes
            rotated_data['imageWidth'], rotated_data['imageHeight'] = rotated_size  # 更新为旋转后的实际尺寸
            with open(os.path.join(save_json_dir, rotated_json_name), 'w') as f:
                json.dump(rotated_data, f, indent=4)

            # 마스크 생성：使用旋转后的图像尺寸
            mask = create_mask_from_json(os.path.join(save_json_dir, rotated_json_name), rotated_size, class_mapping)

            # 마스크 저장
            mask_filename = rotated_image_name.replace('.jpg', '_mask.png')
            mask_save_path = os.path.join(save_mask_dir, mask_filename)
            cv2.imwrite(mask_save_path, mask)

            # 시각화
            visualize_mask(rotated_image, mask, rotated_image_name)

# 마스크 시각화 함수
def visualize_mask(image, mask, filename):
    """Displays the original image and corresponding mask."""
    plt.figure(figsize=(10, 5))

    # 원본 이미지 시각화
    plt.subplot(1, 2, 1)
    plt.imshow(image)
    plt.title(f"Original Image: {filename}")
    plt.axis('off')

    # 마스크 시각화
    plt.subplot(1, 2, 2)
    plt.imshow(mask, cmap='jet', alpha=0.5)
    plt.title(f"Generated Mask: {filename}")
    plt.axis('off')

    plt.show()

# 클래스 ID 매핑 (라벨을 클래스 ID로 변환)
class_mapping = {
    'Acircle': 1,
    'Bcircle': 2,
    'Ccircle': 3
}

# 경로 설정
train_image_dir = 'C:/New/resize_train_jpg'
train_json_dir = 'C:/New/resize_train_json'
save_image_dir = 'C:/New/augmented_train_jpg'
save_json_dir = 'C:/New/augmented_train_json'
save_mask_dir = 'C:/New/augmented_train_mask'

# 각도 설정 (30도씩 회전)
angles = [30, 60, 90, 120, 150, 180, 210, 240, 270, 300, 330]

# 증강 및 변환 실행
augment_and_transform(train_image_dir, train_json_dir, save_image_dir, save_json_dir, save_mask_dir, angles=angles)

关键修改点 #

1. 同步旋转后的图像尺寸：获取rotated_image.size作为实际尺寸，更新JSON中的imageWidth和imageHeight，确保标注坐标基于正确的图像尺寸。
2. 修正坐标偏移：因expand=True导致旋转后图像尺寸变大，将缩放后的坐标先平移到新图像的中心位置，再执行旋转。
3. 匹配PIL旋转方向：PIL默认逆时针旋转，原旋转矩阵是顺时针计算，因此传入负角度修正方向，保证坐标旋转和图像旋转一致。
4. 调整Mask尺寸：创建Mask时使用旋转后的图像尺寸，解决Mask与图像尺寸不匹配的问题。
5. 修正Mask维度：原代码中mask = np.zeros(image_size, dtype=np.uint8)维度错误，改为image_size[::-1]（cv2的尺寸是(height, width)，而PIL是(width, height)）。

内容的提问来源于stack exchange，提问作者이민수