import streamlit as st
import os
from rembg import remove
from PIL import Image
from werkzeug.utils import secure_filename

UPLOAD_FOLDER = 'static/uploads'
ALLOWED_EXTENSIONS = set(['png', 'jpg', 'jpeg', 'webp'])

if 'static' not in os.listdir('.'):
    os.mkdir('static')

if 'uploads' not in os.listdir('static/'):
    os.mkdir('static/uploads')

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS

def remove_background(input_path, output_path):
    input_image = Image.open(input_path)
    output_image = remove(input_image)
    output_image.save(output_path)

def main():
    st.title("Remove Background")

    uploaded_file = st.file_uploader("Choose a file", type=['png', 'jpg', 'jpeg', 'webp'])

    if uploaded_file is not None and allowed_file(uploaded_file.name):
        filename = secure_filename(uploaded_file.name)
        file_path = os.path.join(UPLOAD_FOLDER, filename)
        uploaded_file.save(file_path)

        rembg_img_name = filename.split('.')[0] + "_rembg.png"
        remove_background(file_path, os.path.join(UPLOAD_FOLDER, rembg_img_name))

        st.image(Image.open(file_path), caption='Original Image', use_column_width=True)
        st.image(Image.open(os.path.join(UPLOAD_FOLDER, rembg_img_name)),
                 caption='Image with Background Removed', use_column_width=True)

if __name__ == '__main__':
    main()