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