Update app.py

app.py

@@ -1,61 +1,50 @@
-# Define basic float and int types to patch sctypes without numpy
-if not hasattr(__builtins__, 'sctypes'):
-    sctypes = {
-        "float": [float],
-        "int": [int],
-        "uint": [int]  # Handle 'uint' as normal 'int' since it's not natively supported without numpy
-    }
-
-# Now, import imgaug after patching
-import imgaug.augmenters as iaa
-import cv2
-import matplotlib.pyplot as plt
 import gradio as gr
+import albumentations as A
+import cv2
+import numpy as np
 
-def augment_image(image, flip, rotate, brightness, noise_scale, elastic_alpha, elastic_sigma):
-    # Ensure the image is a NumPy array (required by imgaug)
+# Function for image augmentation
+def augment_image(image, flip=False, rotate=0, brightness=1.0, noise=0, elastic=False):
     image = np.array(image)
     
-    # Apply augmentations based on user inputs
+    aug_list = []
     if flip:
-        flip_aug = iaa.Fliplr(1.0)  # flips horizontally 100% of the time
-        image = flip_aug.augment_image(image)
-
-    rotate_aug = iaa.Affine(rotate=rotate)  # Rotate by specified degrees
-    image = rotate_aug.augment_image(image)
-
-    brightness_aug = iaa.Multiply(brightness)  # Adjust brightness
-    image = brightness_aug.augment_image(image)
-
-    noise_aug = iaa.AdditiveGaussianNoise(scale=(noise_scale))  # Gaussian noise
-    image = noise_aug.augment_image(image)
-
-    elastic_aug = iaa.ElasticTransformation(alpha=elastic_alpha, sigma=elastic_sigma)  # Elastic transformation
-    image = elastic_aug.augment_image(image)
+        aug_list.append(A.HorizontalFlip(p=1.0))
     
-    return image
+    if rotate:
+        aug_list.append(A.Rotate(limit=rotate, p=1.0))
+    
+    if brightness != 1.0:
+        aug_list.append(A.RandomBrightnessContrast(brightness_limit=(brightness-1, brightness-1), p=1.0))
+    
+    if noise > 0:
+        aug_list.append(A.GaussNoise(var_limit=(noise, noise), p=1.0))
+    
+    if elastic:
+        aug_list.append(A.ElasticTransform(alpha=1, sigma=50, alpha_affine=50, p=1.0))
+    
+    aug = A.Compose(aug_list)
+    augmented = aug(image=image)
+    
+    return augmented["image"]
 
-def gradio_interface(image, flip, rotate, brightness, noise_scale, elastic_alpha, elastic_sigma):
-    augmented_image = augment_image(image, flip, rotate, brightness, noise_scale, elastic_alpha, elastic_sigma)
+# Gradio Interface
+def image_augmentor_interface(image, flip, rotate, brightness, noise, elastic):
+    augmented_image = augment_image(image, flip, rotate, brightness, noise, elastic)
     return augmented_image
 
-# Gradio UI
-inputs = [
-    gr.Image(type="pil"),  # Image input
-    gr.Checkbox(label="Flip Image Horizontally"),  # Flip input
-    gr.Slider(minimum=-180, maximum=180, step=1, value=0, label="Rotate Image (degrees)"),  # Rotation input
-    gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Adjust Brightness"),  # Brightness input
-    gr.Slider(minimum=0, maximum=100, step=1, value=10, label="Gaussian Noise Scale"),  # Noise input
-    gr.Slider(minimum=0, maximum=200, step=10, value=100, label="Elastic Transformation Alpha"),  # Elastic Alpha input
-    gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=3.0, label="Elastic Transformation Sigma")  # Elastic Sigma input
-]
-
 iface = gr.Interface(
-    fn=gradio_interface, 
-    inputs=inputs, 
-    outputs=gr.Image(type="numpy"),
-    title="Image Augmentation Demo",
-    description="Try out different data augmentation techniques on your image.",
+    fn=image_augmentor_interface,
+    inputs=[
+        gr.Image(type="numpy"),
+        gr.Checkbox(label="Flip Horizontally"),
+        gr.Slider(0, 360, label="Rotate Degrees"),
+        gr.Slider(0.5, 1.5, step=0.1, label="Brightness Adjustment"),
+        gr.Slider(0, 100, step=1, label="Noise Scale"),
+        gr.Checkbox(label="Elastic Distortion")
+    ],
+    outputs="image",
+    title="Image Augmentation with Gradio"
 )
 
 iface.launch()