import streamlit as st
from rembg import remove
from PIL import Image
import io

# Function to remove background from image
def remove_background(image_bytes):
    input_image = Image.open(io.BytesIO(image_bytes))
    # Rembg expects the image to be in the correct mode, typically RGBA or RGB
    input_image = input_image.convert("RGBA")
    output_image = remove(input_image)
    return output_image

# Streamlit UI setup
st.set_page_config(page_title="Background Remover", layout="wide")

# Title and Description
st.title("Background Removal Tool")
st.markdown("### Upload an image, and we will remove its background for you!")

# Create columns for layout
col1, col2 = st.columns(2)

# Column 1: Image upload
with col1:
    uploaded_file = st.file_uploader("Choose an image file", type=["png", "jpg", "jpeg"])

# Column 2: Instructions
with col2:
    st.markdown("""
    #### Instructions:
    - Upload an image (PNG, JPG, or JPEG) using the button above.
    - The background will be removed automatically.
    - Download the result below once it's ready.
    """)

# When an image is uploaded
if uploaded_file is not None:
    # Show the uploaded image
    st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)

    # Remove the background
    with st.spinner("Removing background..."):
        output_image = remove_background(uploaded_file.read())

    # Show the processed image
    st.image(output_image, caption="Image with Background Removed", use_column_width=True)

    # Allow the user to download the result
    output_image_pil = Image.open(io.BytesIO(output_image))
    buf = io.BytesIO()
    output_image_pil.save(buf, format="PNG")
    byte_im = buf.getvalue()

    st.download_button(
        label="Download Image with Removed Background",
        data=byte_im,
        file_name="image_no_bg.png",
        mime="image/png"
    )

# Footer
st.markdown("---")
st.markdown("Created by [Your Name].")