import albumentations as A
import cv2
from matplotlib import pyplot as plt

def visualize(image):
    plt.figure(figsize=(10, 5))
    plt.axis("off")
    plt.imshow(image)

img_path = "../images/alina_rossoshanska.jpeg"
img = cv2.imread(img_path, cv2.IMREAD_COLOR_RGB)

visualize(img)

transform = A.Compose(
    [A.ChromaticAberration(mode="red_blue", primary_distortion_limit=0.5, secondary_distortion_limit=0.1, p=1)],
    seed=137,
    strict=True,
)

plt.figure(figsize=(15, 10))

num_images = 12

# Loop through the list of images and plot them with subplot
for i in range(num_images):
    transformed_image = transform(image=img)["image"]
    plt.subplot(4, 3, i + 1)
    plt.imshow(transformed_image)
    plt.axis("off")

plt.tight_layout()
plt.show()