import streamlit as st
from rembg import remove, new_session
from PIL import Image, ImageFilter
import io

# Load high-quality model
session = new_session("u2net")

# Streamlit App Configuration
st.set_page_config(page_title="Background Remover & Changer", layout="wide")
st.title("Background Remover & Changer")

# Upload multiple images
uploaded_files = st.file_uploader("Upload Images (Multiple Allowed)", type=["png", "jpg", "jpeg"], accept_multiple_files=True)

# Upload custom background
bg_image = st.file_uploader("Upload Custom Background (Optional)", type=["png", "jpg", "jpeg"])

# Process images
if uploaded_files:
    st.subheader("Processed Images")
    for uploaded_file in uploaded_files:
        image = Image.open(uploaded_file)
        image = image.convert("RGBA")
        
        # Remove background using high-quality model
        removed_bg = remove(image, session=session)
        
        # Apply smoothing to improve edges
        removed_bg = removed_bg.filter(ImageFilter.SMOOTH)
        
        if bg_image:
            bg = Image.open(bg_image).convert("RGBA")
            bg = bg.resize(removed_bg.size)
            final_image = Image.alpha_composite(bg, removed_bg)
        else:
            final_image = removed_bg
        
        st.image(final_image, caption="Processed Image", use_column_width=True)
        
        # Download button
        img_bytes = io.BytesIO()
        final_image.save(img_bytes, format='PNG')
        st.download_button(label="Download Image", data=img_bytes.getvalue(), file_name="processed_image.png", mime="image/png")

st.write("Upload images to remove background and optionally change it.")