import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import cv2
import albumentations as A

def visualize(ax, img, bboxes, grid=False):
    box_colors = ["grey", "black"]
    cmap = plt.cm.get_cmap("tab10")
    ax.imshow(img)
    for i, bbox in enumerate(bboxes):
        ax.add_patch(Rectangle(bbox[:2], bbox[2], bbox[3], fill=False, lw=2, color=box_colors[i]))
    ax.grid(grid)
    return ax

def get_data():    
    ## prepare sample image and bbox
    img = np.full((300, 300, 3), 255, dtype=np.uint8)
    cx, cy = (150, 150)
    rx = 50
    ry = 100
    cv2.rectangle(img, (cx-rx, cy-ry), (cx+rx, cy+ry), (255, 127, 127), -1)
    cv2.ellipse(img, (cx, cy), (rx, ry), 0, 0, 360, (127, 255, 127), -1)
    bboxes = [[cx-rx, cy-ry, 2 * rx, 2 * ry, "label"]]
    return img, bboxes

def get_affine_aug(rot=0, shear=(0, 0), rotate_method="largest_box"):
    shear_x, shear_y = shear
    aug = A.Compose([
        A.Affine(
            p=1, 
            scale=1,
            translate_percent=0,
            rotate=rot,
            shear={"x": (shear_x, shear_x), "y": (shear_y, shear_y)},
            rotate_method=rotate_method,
            mode=cv2.BORDER_CONSTANT,
            cval=(255, 255, 255)
        ),
    ], bbox_params=A.BboxParams(format='coco'))
    return aug

fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharey=True)    

img, bboxes = get_data()
visualize(axes[0], img, bboxes)
aug1 = get_affine_aug(rot=15, shear=(50, 0), rotate_method="largest_box")
aug2 = get_affine_aug(rot=15, shear=(50, 0), rotate_method="ellipse")

data1 = aug1(image=img, bboxes=bboxes)
data2 = aug2(image=img, bboxes=bboxes)

img1 = data1["image"]
bboxes1 = data1["bboxes"]
bboxes2 = data2["bboxes"]

visualize(axes[1], img1, bboxes1 + bboxes2)
fig.savefig("./affine_ellipse.jpg", bbox_inches="tight")