src/streamlit_app.py

import streamlit as st
from PIL import Image
import numpy as np
import cv2
import random
import io
import base64

st.title(":red[Image Data Augmentation]")

# Explanation of Image Augmentation
st.write("""
## :blue[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.
### :blue[Common Types of Image Augmentation]:
- **Flipping**: Horizontally or vertically flipping an image.
- **Rotation**: Rotating the image by a certain angle.
- **Scaling**: Zooming in or out on the image.
- **Translation**: Shifting the image in the x or y direction.
- **Shearing**: Distorting the image along one axis.
- **Color Jittering**: Randomly changing the brightness, contrast, saturation, and hue of the image.
- **Cropping**: Randomly cropping a portion of the image.
- **Affine Transformations**: Applying geometric transformations like scaling, rotation, and translation together.
""")

# 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)
    img = img.convert("RGB")  # Convert to RGB mode
    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:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image.', use_container_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")
    
    

    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)


    # Convert numpy array back to image
    augmented_image = Image.fromarray(img_array)
    st.image(augmented_image, caption='Augmented Image.', use_container_width=True)