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

def remove_background(image):
    return remove(image)

def main():
    st.title("Background Removal App")
    st.write("Upload an image to remove its background")

    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        input_image = Image.open(uploaded_file)
        st.image(input_image, caption='Uploaded Image', use_column_width=True)

        if st.button('Remove Background'):
            with st.spinner('Processing...'):
                output_image = remove_background(input_image)
                st.image(output_image, caption='Output Image', use_column_width=True)
                buf = io.BytesIO()
                output_image.save(buf, format='PNG')
                byte_im = buf.getvalue()

                st.download_button(
                    label="Download Image",
                    data=byte_im,
                    file_name="output_image.png",
                    mime="image/png"
                )

if __name__ == "__main__":
    main()