#------------------------------------矩形标注增强---------------------------------------
from PIL import Image
import os
import json

def rotate_images_and_jsons(input_folder):
    output_folder = os.path.join(input_folder, "rotated")
    os.makedirs(output_folder, exist_ok=True)

    for filename in os.listdir(input_folder):
        if filename.endswith(".jpg"):
            image_path = os.path.join(input_folder, filename)
            json_path = os.path.join(input_folder, filename.replace(".jpg", ".json"))

            # Load JSON data
            with open(json_path, 'r', encoding='utf-8') as json_file:
                json_data = json.load(json_file)

            image = Image.open(image_path)

            # Rotate the image by 90, 180, and 270 degrees
            for angle in [90, 180, 270]:
                rotated_image = image.rotate(angle, expand=True)

                # Save rotated images with a default quality value
                rotated_image_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}_{angle}.jpg")
                rotated_image.save(rotated_image_path, quality=100)

                # Update JSON data for rotated image
                rotated_json_data = update_json_for_rotation(json_data, image.size, angle, os.path.basename(rotated_image_path))
                rotated_json_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}_{angle}.json")
                with open(rotated_json_path, 'w', encoding='utf-8') as rotated_json_file:
                    json.dump(rotated_json_data, rotated_json_file, indent=4, ensure_ascii=False)

def update_json_for_rotation(json_data, image_size, angle, image_path):
    width, height = image_size
    rotated_json_data = []

    for shape in json_data['shapes']:
        points = shape['points']
        rotated_points = []

        # 提取坐标信息
        point1_x, point1_y = points[0]
        point2_x, point2_y = points[1]
        if angle == 90:
            rotated_point1 = [point1_y, width - point1_x]
            rotated_point2 = [point2_y, width - point2_x]
        elif angle == 180:
            rotated_point1 = [width - point1_x, height - point1_y]
            rotated_point2 = [width - point2_x, height - point2_y]
        elif angle == 270:
            rotated_point1 = [height - point1_y, point1_x]
            rotated_point2 = [height - point2_y, point2_x]

        rotated_points.append(rotated_point1)
        rotated_points.append(rotated_point2)
        rotated_shape = {
            'label': shape['label'],
            'points': rotated_points,
            'group_id': shape['group_id'],
            'shape_type': shape['shape_type'],
            'flags': shape['flags']
        }

        rotated_json_data.append(rotated_shape)
    if angle == 90 or angle == 270:
        return {'version': json_data['version'], 'flags': json_data['flags'], 'shapes': rotated_json_data, 'imagePath': image_path, 'imageData': json_data['imageData'], 'imageHeight': json_data['imageWidth'], 'imageWidth': json_data['imageHeight']}
    else:
        return {'version': json_data['version'], 'flags': json_data['flags'], 'shapes': rotated_json_data, 'imagePath': image_path, 'imageData': json_data['imageData'], 'imageHeight': json_data['imageHeight'], 'imageWidth': json_data['imageWidth']}
# Replace 'input_folder' with your folder containing images and JSON files
rotate_images_and_jsons('./rectangles')