Update pages/4_Image Augmentation.py

pages/4_Image Augmentation.py

@@ -1,33 +1,105 @@
 import streamlit as st
-from PIL import Image, ImageEnhance
+from PIL import Image
+import numpy as np
+import cv2
+import random
+import io
 
-# Title of the app
-st.title("Image Data Augmentation")
+st.title("Interactive Image Augmentation with Streamlit")
 
-# File uploader to upload images
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
+# Explanation of Image Augmentation
+st.write("""
+## What is Image Augmentation?
+Image augmentation refers to a set of techniques used to increase the diversity of training images available for a machine learning model without actually collecting new images. It's commonly used in computer vision tasks to improve the performance of models by making them more robust to variations in the data.
 
+### Common Types of Image Augmentation:
+1. **Flipping**: Horizontally or vertically flipping an image.
+2. **Rotation**: Rotating the image by a certain angle.
+3. **Scaling**: Zooming in or out on the image.
+4. **Translation**: Shifting the image in the x or y direction.
+5. **Shearing**: Distorting the image along one axis.
+6. **Color Jittering**: Randomly changing the brightness, contrast, saturation, and hue of the image.
+7. **Noise Addition**: Adding random noise to the image to simulate variations.
+8. **Cropping**: Randomly cropping a portion of the image.
+9. **Blurring**: Applying different types of blurs (e.g., Gaussian blur).
+10. **Affine Transformations**: Applying geometric transformations like scaling, rotation, and translation together.
+11. **Elastic Transformations**: Randomly deforming the image using displacement fields.
+12. **Cutout**: Randomly masking out a section of the image.
+""")
+
+# Function to convert numpy array to image and provide download link
+def get_image_download_link(img_array, filename, text):
+    img = Image.fromarray(img_array)
+    buffered = io.BytesIO()
+    img.save(buffered, format="JPEG")
+    buffered.seek(0)
+    b64 = base64.b64encode(buffered.read()).decode()
+    href = f'<a href="data:file/jpg;base64,{b64}" download="{filename}">{text}</a>'
+    return href
+
+# Upload image
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 if uploaded_file is not None:
-    # Open the uploaded image
     image = Image.open(uploaded_file)
-    st.image(image, caption='Uploaded Image', use_column_width=True)
-    
-    st.write("Augment your image below:")
-
-    # Slider for adjusting brightness
-    brightness = st.slider("Brightness", 0.5, 3.0, 1.0)
-    enhancer = ImageEnhance.Brightness(image)
-    img_brightness = enhancer.enhance(brightness)
-    st.image(img_brightness, caption='Brightness Adjusted', use_column_width=True)
-
-    # Slider for adjusting contrast
-    contrast = st.slider("Contrast", 0.5, 3.0, 1.0)
-    enhancer = ImageEnhance.Contrast(image)
-    img_contrast = enhancer.enhance(contrast)
-    st.image(img_contrast, caption='Contrast Adjusted', use_column_width=True)
-
-    # Slider for adjusting sharpness
-    sharpness = st.slider("Sharpness", 0.5, 3.0, 1.0)
-    enhancer = ImageEnhance.Sharpness(image)
-    img_sharpness = enhancer.enhance(sharpness)
-    st.image(img_sharpness, caption='Sharpness Adjusted', use_column_width=True)
+    st.image(image, caption='Uploaded Image.', use_column_width=True)
+    st.write("")
+    st.write("Select the augmentation operations you want to apply:")
+
+    # Convert image to numpy array
+    img_array = np.array(image)
+
+    # Create checkboxes for each augmentation operation
+    flip_horizontally = st.checkbox("Flip Horizontally")
+    flip_vertically = st.checkbox("Flip Vertically")
+    rotate = st.checkbox("Rotate 90 degrees")
+    scale = st.checkbox("Scale")
+    translate = st.checkbox("Translate")
+    shear = st.checkbox("Shear")
+    color_jitter = st.checkbox("Color Jittering")
+    add_noise = st.checkbox("Add Noise")
+
+    if flip_horizontally:
+        img_array = cv2.flip(img_array, 1)
+        st.markdown(get_image_download_link(img_array, "flipped_horizontally.jpg", "Download Horizontally Flipped Image"), unsafe_allow_html=True)
+
+    if flip_vertically:
+        img_array = cv2.flip(img_array, 0)
+        st.markdown(get_image_download_link(img_array, "flipped_vertically.jpg", "Download Vertically Flipped Image"), unsafe_allow_html=True)
+
+    if rotate:
+        img_array = cv2.rotate(img_array, cv2.ROTATE_90_CLOCKWISE)
+        st.markdown(get_image_download_link(img_array, "rotated.jpg", "Download Rotated Image"), unsafe_allow_html=True)
+
+    if scale:
+        scale_factor = st.slider("Scale Factor", 0.5, 2.0, 1.0)
+        img_array = cv2.resize(img_array, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR)
+        st.markdown(get_image_download_link(img_array, "scaled.jpg", "Download Scaled Image"), unsafe_allow_html=True)
+
+    if translate:
+        translate_x = st.slider("Translate X", -50, 50, 0)
+        translate_y = st.slider("Translate Y", -50, 50, 0)
+        M = np.float32([[1, 0, translate_x], [0, 1, translate_y]])
+        img_array = cv2.warpAffine(img_array, M, (img_array.shape[1], img_array.shape[0]))
+        st.markdown(get_image_download_link(img_array, "translated.jpg", "Download Translated Image"), unsafe_allow_html=True)
+
+    if shear:
+        shear_factor = st.slider("Shear Factor", -0.5, 0.5, 0.0)
+        M_shear = np.float32([[1, shear_factor, 0], [shear_factor, 1, 0]])
+        img_array = cv2.warpAffine(img_array, M_shear, (img_array.shape[1], img_array.shape[0]))
+        st.markdown(get_image_download_link(img_array, "sheared.jpg", "Download Sheared Image"), unsafe_allow_html=True)
+
+    if color_jitter:
+        brightness = st.slider("Brightness", 0.5, 1.5, 1.0)
+        contrast = st.slider("Contrast", 0.5, 1.5, 1.0)
+        img_array = cv2.convertScaleAbs(img_array, alpha=contrast, beta=brightness * 127)
+        st.markdown(get_image_download_link(img_array, "color_jittered.jpg", "Download Color Jittered Image"), unsafe_allow_html=True)
+
+    if add_noise:
+        noise = np.random.randn(*img_array.shape) * 10
+        img_array = img_array + noise
+        img_array = np.clip(img_array, 0, 255).astype(np.uint8)
+        st.markdown(get_image_download_link(img_array, "noisy.jpg", "Download Noisy Image"), unsafe_allow_html=True)
+
+    # Convert numpy array back to image
+    augmented_image = Image.fromarray(img_array)
+    st.image(augmented_image, caption='Augmented Image.', use_column_width=True)