import altair as alt
import pandas as pd
import streamlit as st
from PIL import Image
import cv2
import numpy as np
import io
import zipfile

st.title("Image Augmentation App")
uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])

if uploaded_image:
    img = np.array(Image.open(uploaded_image))
    st.image(img, caption="Uploaded Image", use_column_width=True)
    
    st.write("**Applying augmentations...**")
    
    transformations = {
        "Translated": cv2.warpAffine(img, np.float32([[1, 0, 50], [0, 1, 50]]), (img.shape[1], img.shape[0])),
        "Scaled": cv2.resize(img, None, fx=1.2, fy=1.2),
        "Rotated": cv2.warpAffine(img, cv2.getRotationMatrix2D((img.shape[1] // 2, img.shape[0] // 2), 45, 1), (img.shape[1], img.shape[0])),
        "Cropped": img[50:img.shape[0] - 50, 50:img.shape[1] - 50]
    }
    
    augmented_images = []
    zip_buffer = io.BytesIO()
    
    with zipfile.ZipFile(zip_buffer, "w") as zipf:
        for i, (name, transformed_img) in enumerate(transformations.items()):
            img_pil = Image.fromarray(cv2.cvtColor(transformed_img, cv2.COLOR_BGR2RGB))
            img_bytes = io.BytesIO()
            img_pil.save(img_bytes, format="JPEG")
            zipf.writestr(f"aug_{i + 1}_{name}.jpg", img_bytes.getvalue())
            augmented_images.append(img_pil)
    
    st.write("**Augmentation completed!**")
    st.image(augmented_images, width=150)
    
    zip_buffer.seek(0)
    st.download_button("Download Augmented Images", data=zip_buffer, file_name="augmented_images.zip", mime="application/zip")