import cv2
import os

def rotate_image(image, angle):
    """Rotate an image by a specific angle."""
    (h, w) = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv2.warpAffine(image, M, (w, h))
    return rotated

def main():
    folder_path = r"E:\AP\img0204\P"
    angles_offsets = [(90, 25), (180, 50), (270, 75)]  # Rotation angles and filename offsets

    # Ensure the folder exists
    if not os.path.exists(folder_path):
        print(f"Folder '{folder_path}' does not exist.")
        return

    for filename in os.listdir(folder_path):
        if filename.endswith(".png"):
            base_name = os.path.splitext(filename)[0]
            try:
                image_idx = int(base_name)
            except ValueError:
                print(f"Skipping invalid filename: {filename}")
                continue

            # Load the image
            file_path = os.path.join(folder_path, filename)
            image = cv2.imread(file_path)
            if image is None:
                print(f"Failed to read image: {file_path}")
                continue

            # Rotate and save images
            for angle, offset in angles_offsets:
                rotated_image = rotate_image(image, angle)
                new_name = f"{image_idx + offset}.png"
                new_path = os.path.join(folder_path, new_name)

                if os.path.exists(new_path):
                    print(f"File {new_path} already exists. Skipping.")
                    continue

                cv2.imwrite(new_path, rotated_image)
                print(f"Saved rotated image: {new_path}")

if __name__ == "__main__":
    main()