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

# Function to remove background using rembg
def remove_background(image_file):
    # Open the uploaded image using PIL
    input_image = Image.open(image_file)
    
    # Convert image to byte format for processing
    img_byte_arr = io.BytesIO()
    input_image.save(img_byte_arr, format='PNG')
    img_byte_arr = img_byte_arr.getvalue()

    # Remove background using rembg
    output_image = remove(img_byte_arr)
    
    # Convert output byte data to an image
    output_image = Image.open(io.BytesIO(output_image))
    
    return output_image

# Streamlit UI setup
st.title("Background Removal Tool")

st.markdown("""
    Upload an image, and this tool will automatically remove the background. 
    You can then download the processed image with a transparent background.
""")

# Sidebar for uploading the image
st.sidebar.header("Upload Your Image")
uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

# Columns layout for displaying images
col1, col2 = st.columns(2)

if uploaded_file is not None:
    # Show the original image in the first column
    with col1:
        st.subheader("Original Image")
        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
    
    # Remove background and display the processed image in the second column
    with col2:
        st.subheader("Processed Image")
        result_image = remove_background(uploaded_file)
        st.image(result_image, caption="Processed Image", use_column_width=True)
        
        # Provide a download button for the processed image
        result_image_bytes = io.BytesIO()
        result_image.save(result_image_bytes, format='PNG')
        result_image_bytes.seek(0)
        
        st.download_button(
            label="Download Image with Removed Background",
            data=result_image_bytes,
            file_name="output_image.png",
            mime="image/png"
        )
else:
    st.info("Please upload an image file to remove the background.")