import streamlit as st
from PIL import Image
import random
import os
import zipfile
import io

# Directory to save images
SAVE_DIR = r"C:\Users\syam0\Downloads\Image App"

# Ensure the directory exists
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR)

# Augmentation functions with user inputs
def scale_image(image, count, scale_factor=None):
    if scale_factor is None:
        scale_factor = random.uniform(0.8, 1.2)
    return [image.resize((int(image.width * scale_factor), int(image.height * scale_factor))) for _ in range(count)]

def translate_image(image, count, axis, translation_range=(-40, 40)):
    translations = []
    for _ in range(count):
        if axis == "x-axis":
            translation_x = random.randint(*translation_range)
            translation_y = 0
        elif axis == "y-axis":
            translation_x = 0
            translation_y = random.randint(*translation_range)
        else:
            translation_x = random.randint(*translation_range)
            translation_y = random.randint(*translation_range)
        translated_image = image.transform(image.size, Image.AFFINE, (1, 0, translation_x, 0, 1, translation_y))
        translations.append(translated_image)
    return translations

def crop_image(image, count, crop_margin=None):
    if crop_margin is None:
        crop_margin = random.randint(5, 20)
    return [image.crop((crop_margin, crop_margin, image.width - crop_margin, image.height - crop_margin)) for _ in range(count)]

def rotate_image(image, count, angle=None):
    if angle is None:
        angle = random.randint(0, 360)
    return [image.rotate(angle) for _ in range(count)]

def shear_image(image, count, shear_factor=None):
    if shear_factor is None:
        shear_factor = random.uniform(-0.5, 0.5)
    return [
        image.transform(image.size, Image.AFFINE, (1, shear_factor, 0, shear_factor, 1, 0)) 
        for _ in range(count)
    ]

# Function to save images to the specified directory
def save_images(images, prefix):
    saved_files = []
    for i, img in enumerate(images):
        filename = os.path.join(SAVE_DIR, f"{prefix}_image_{i+1}.png")
        img.save(filename, format="PNG")
        saved_files.append(filename)
    return saved_files

# Function to zip the images
def create_zip(saved_files):
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w") as zipf:
        for file in saved_files:
            zipf.write(file, os.path.basename(file))
    zip_buffer.seek(0)
    return zip_buffer

# Custom CSS for background and styling
st.markdown(
    """
    <style>
    body {
        background: linear-gradient(#ffaa80, #ff5500, #b33c00);  /* Gradient background from yellow to red */
        font-family: 'Arial', sans-serif;
        color: #333;
    }
    .stApp {
        background: linear-gradient(#66d9ff, #d966ff, #66ff66);  /* Gradient background for the entire app */
    }
    .stFileUploader {
        display: flex;
        justify-content: center;
        align-items: center;
        padding: 30px;
        border: 2px dashed #4CAF50;  /* Green dashed border for a fresh look */
        background-color: rgba(76, 175, 80, 0.1);  /* Light green background on hover */
        border-radius: 8px;
        cursor: pointer;
    }
    .stFileUploader:hover {
        background-color: rgba(76, 175, 80, 0.2);  /* Slightly darker green on hover */
    }
    </style>
    """,
    unsafe_allow_html=True
)

# Streamlit app
st.title("Image Augmentation App with Gradient Background")

# Image uploader with custom styling
uploaded_image = st.file_uploader("Drag and Drop Your Image", type=["jpg", "jpeg", "png"], label_visibility="collapsed")

if uploaded_image:
    image = Image.open(uploaded_image)
    st.image(image, caption="Uploaded Image", use_container_width=True)  # Fixed the deprecation warning by using `use_container_width=True`

    # User inputs for augmentation
    st.write("### Generate Augmented Images")
    num_augments = st.number_input("How many augmented images would you like to generate?", min_value=1, max_value=20, value=5)

    # Additional inputs for specific techniques
    rotate_angle = st.number_input("Enter the angle for rotation (0-360):", min_value=0, max_value=360, value=random.randint(0, 360))
    scale_factor = st.number_input("Enter the scaling factor (0.8-1.2):", min_value=0.8, max_value=1.2, value=random.uniform(0.8, 1.2))
    crop_margin = st.number_input("Enter the crop margin (5-20):", min_value=5, max_value=20, value=random.randint(5, 20))
    shear_factor = st.number_input("Enter the shear factor (-0.5 to 0.5):", min_value=-0.5, max_value=0.5, value=random.uniform(-0.5, 0.5))
    translate_axis = st.selectbox("Select axis for translation:", ["x-axis", "y-axis", "both axes"])

    # Generate and save all augmentations
    if st.button("Generate All Techniques"):
        all_augmentations = scale_image(image, num_augments, scale_factor) + \
                            translate_image(image, num_augments, translate_axis) + \
                            crop_image(image, num_augments, crop_margin) + \
                            rotate_image(image, num_augments, rotate_angle) + \
                            shear_image(image, num_augments, shear_factor)
        
        # Save images to the specified folder
        st.write("Generating and Saving Images...")
        saved_files = save_images(all_augmentations, "combined")
        
        # Display generated images
        st.write("Generated Augmented Images:")
        for img in all_augmentations:
            st.image(img, use_container_width=True)

        # Create a zip file
        zip_buffer = create_zip(saved_files)
        
        # Provide a download button for the zip file
        st.download_button(
            label="Download Augmented Images (ZIP)",
            data=zip_buffer,
            file_name="augmented_images.zip",
            mime="application/zip"
        )

    # Generate and save technique-specific augmentations
    st.write("### Select Specific Technique")
    technique = st.selectbox("Choose an augmentation technique:", 
                              ["Scale", "Translate", "Crop", "Rotate", "Shear"])
    technique_count = st.number_input(f"How many images do you want for {technique}?", min_value=1, max_value=20, value=5)

    if st.button(f"Generate and Save {technique} Images"):
        if technique == "Scale":
            specific_augmentations = scale_image(image, technique_count, scale_factor)
        elif technique == "Translate":
            specific_augmentations = translate_image(image, technique_count, translate_axis)
        elif technique == "Crop":
            specific_augmentations = crop_image(image, technique_count, crop_margin)
        elif technique == "Rotate":
            specific_augmentations = rotate_image(image, technique_count, rotate_angle)
        elif technique == "Shear":
            specific_augmentations = shear_image(image, technique_count, shear_factor)
        
        # Save images to the specified folder
        st.write(f"Generating and Saving {technique} Images...")
        saved_files = save_images(specific_augmentations, technique.lower())
        
        # Display generated images
        st.write(f"Generated {technique} Augmented Images:")
        for img in specific_augmentations:
            st.image(img, use_container_width=True)
        
        # Create a zip file
        zip_buffer = create_zip(saved_files)
        
        # Provide a download button for the zip file
        st.download_button(
            label=f"Download {technique} Augmented Images (ZIP)",
            data=zip_buffer,
            file_name=f"{technique.lower()}_augmented_images.zip",
            mime="application/zip"
        )