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

st.set_page_config(layout="wide")
st.title("Image Augmentation Tool")

# Set max file size limit in bytes (e.g., 5MB)
MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB

uploaded_files = st.file_uploader(
    "Upload one or more images (Max 5MB each)",
    type=["jpg", "jpeg", "png"],
    accept_multiple_files=True
)



# Function to apply all transformations
def apply_all_transformations(img_array):
    results = []

    gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
    results.append(("Grayscale", gray))

    rotated = cv2.rotate(img_array, cv2.ROTATE_90_CLOCKWISE)
    results.append(("Rotate 90°", rotated))

    flip_h = cv2.flip(img_array, 1)
    results.append(("Flip Horizontally", flip_h))

    flip_v = cv2.flip(img_array, 0)
    results.append(("Flip Vertically", flip_v))

    rows, cols = img_array.shape[:2]
    shear_M = np.float32([[1, 0.5, 0], [0.2, 1, 0]])
    sheared = cv2.warpAffine(img_array, shear_M, (cols, rows))
    results.append(("Shearing", sheared))

    trans_M = np.float32([[1, 0, 50], [0, 1, 30]])
    translated = cv2.warpAffine(img_array, trans_M, (cols, rows))
    results.append(("Translation", translated))

    h, w = img_array.shape[:2]
    x_start = w // 4
    x_end = 3 * x_start
    y_start = h // 4
    y_end = 3 * y_start
    cropped = img_array[y_start:y_end, x_start:x_end]
    results.append(("Crop (center 50%)", cropped, (x_start, x_end, y_start, y_end)))

    return results

if uploaded_files:
    for uploaded_file in uploaded_files:
        # Check file size limit
        if uploaded_file.size > MAX_FILE_SIZE:
            st.error(f"❌ File {uploaded_file.name} is too large. Max size allowed is 5MB.")
            continue

        st.markdown("---")
        st.subheader(f"📷 Original: {uploaded_file.name}")

        image = Image.open(uploaded_file).convert("RGB")
        img_array = np.array(image)

        # Resize original image if width exceeds 400px
        max_display_width = 400
        orig_w, orig_h = image.size
        if orig_w > max_display_width:
            new_h = int(orig_h * max_display_width / orig_w)
            image_resized = image.resize((max_display_width, new_h))
        else:
            image_resized = image

        # Show resized original image
        st.image(image_resized, caption="Original Image", use_container_width=False)

        # Apply all transformations
        transformations = apply_all_transformations(img_array)

        st.markdown("### 🧪 Transformed Versions")

        # Create columns dynamically (3 per row)
        cols = st.columns(3)
        for idx, item in enumerate(transformations):
            name = item[0]
            transformed = item[1]
            coords = item[2] if len(item) > 2 else None

            col = cols[idx % 3]
            with col:
                st.markdown(f"**{name}**")

                if len(transformed.shape) == 2:
                    st.image(transformed, caption=name, use_container_width=True, channels="GRAY")
                    save_img = cv2.cvtColor(transformed, cv2.COLOR_GRAY2RGB)
                else:
                    st.image(transformed, caption=name, use_container_width=True)
                    save_img = transformed

               

                # Save transformed image to bytes
                result_img = Image.fromarray(save_img)
                buf = io.BytesIO()
                result_img.save(buf, format="PNG")
                byte_data = buf.getvalue()

                st.download_button(
                    label="⬇️ Download",
                    data=byte_data,
                    file_name=f"{uploaded_file.name.split('.')[0]}_{name.replace(' ', '_').lower()}.png",
                    mime="image/png"
                )