developer0hye/image-augmentation-example

1import gradio as gr2import cv23import numpy as np4import albumentations as A5import random6 7def apply_augmentations(image, flip_h, flip_v, rotate, crop, gray, scale,8                       prob_flip_h, prob_flip_v, prob_rotate, prob_crop, prob_gray, prob_scale,9                       rotation_limit):10    augmentations = []11    12    if flip_h:13        augmentations.append(A.HorizontalFlip(p=float(prob_flip_h)))14    if flip_v:15        augmentations.append(A.VerticalFlip(p=float(prob_flip_v)))16    if rotate:17        augmentations.append(A.Rotate(limit=(-rotation_limit, rotation_limit), p=float(prob_rotate)))18    if crop:19        augmentations.append(A.RandomResizedCrop(20            size=(image.shape[0], image.shape[1]),21            scale=(0.8, 1.0),22            p=float(prob_crop)23        ))24    if gray:25        augmentations.append(A.ToGray(p=float(prob_gray)))26    if scale:27        scale_factor = random.uniform(0.8, 1.2)28        augmentations.append(A.Resize(29            height=int(image.shape[0] * scale_factor),30            width=int(image.shape[1] * scale_factor),31            p=float(prob_scale)32        ))33 34    transform = A.Compose(augmentations)35    36    # Generate 9 augmented images37    augmented_images = []38    for _ in range(9):39        augmented = transform(image=image)40        augmented_images.append(augmented['image'])41    42    # Find maximum dimensions43    max_height = max(img.shape[0] for img in augmented_images)44    max_width = max(img.shape[1] for img in augmented_images)45    46    # Add padding to all images47    padded_images = []48    for img in augmented_images:49        h, w = img.shape[:2]50        pad_top = (max_height - h) // 251        pad_bottom = max_height - h - pad_top52        pad_left = (max_width - w) // 253        pad_right = max_width - w - pad_left54        55        # Handle both RGB and grayscale images56        if len(img.shape) == 3:57            padded = cv2.copyMakeBorder(img, pad_top, pad_bottom, pad_left, pad_right, 58                                      cv2.BORDER_CONSTANT, value=[128, 128, 128])59        else:60            padded = cv2.copyMakeBorder(img, pad_top, pad_bottom, pad_left, pad_right, 61                                      cv2.BORDER_CONSTANT, value=[128])62        padded_images.append(padded)63    64    # Create a 3x3 grid65    rows = []66    for i in range(0, 9, 3):67        row = np.hstack(padded_images[i:i+3])68        rows.append(row)69    grid = np.vstack(rows)70    71    return grid72 73def main():74    with gr.Blocks() as demo:75        with gr.Row():76            with gr.Column():77                input_image = gr.Image(label="Input Image")78            with gr.Column():79                output_image = gr.Image(label="Output Image (3x3 Grid)")80        81        with gr.Row():82            with gr.Column():83                flip_h = gr.Checkbox(label="Horizontal Flip")84                prob_flip_h = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")85            86            with gr.Column():87                flip_v = gr.Checkbox(label="Vertical Flip")88                prob_flip_v = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")89            90            with gr.Column():91                rotate = gr.Checkbox(label="Rotate")92                prob_rotate = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")93                rotation_limit = gr.Slider(minimum=0, maximum=90, value=10, label="Rotation Limit (±degrees)")94            95            with gr.Column():96                crop = gr.Checkbox(label="Random Crop")97                prob_crop = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")98 99            with gr.Column():100                gray = gr.Checkbox(label="Grayscale")101                prob_gray = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")102                103            with gr.Column():104                scale = gr.Checkbox(label="Random Scale (0.8-1.2x)")105                prob_scale = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")106 107        with gr.Row():108            run_button = gr.Button("Apply Augmentations")109 110        run_button.click(111            fn=apply_augmentations,112            inputs=[113                input_image,114                flip_h, flip_v, rotate, crop, gray, scale,115                prob_flip_h, prob_flip_v, prob_rotate, prob_crop, prob_gray, prob_scale,116                rotation_limit117            ],118            outputs=output_image119        )120 121    demo.launch()122 123if __name__ == "__main__":124    main()