import os

import streamlit as st
import base64
from io import BytesIO
from PIL import Image
import os
import rembg

try :
    from rembg import remove ,new_session
except ImportError:
    os.system("pip install rembg")

# local_img ={"zebra":os.path.join(os.getcwd(),"zebra.jpg"),"wallaby":os.path.join(os.getcwd(),"wallaby.png")}
local_img ={"zebra":"zebra.jpg","wallaby":"wallaby.png"}
models_name = ["u2net","u2netp","u2net_human_seg","u2net_cloth_seg","silueta",'isnet-general-use','sam ']


# Download the fixed image
def convert_image(img):
    buf = BytesIO()
    img.save(buf, format="PNG")
    byte_im = buf.getvalue()
    return byte_im

st.set_page_config(layout="wide", page_title="Image Background Remover")

st.write("## Remove background from your image")
st.write(
    ":dog: Try uploading an image to watch the background magically removed. Full quality images can be downloaded from the sidebar. This code is open source and available [here](https://github.com/tyler-simons/BackgroundRemoval) on GitHub. Special thanks to the [rembg library](https://github.com/danielgatis/rembg) :grin:"
)
col1, col2 = st.columns(2)

uploaded_img = col1.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
custom_img = col1.selectbox('Or pick custom image',local_img.keys())
model_name = col1.selectbox('Or pick custom image',models_name)

session = new_session(model_name=model_name)
if uploaded_img:
    image = Image.open(uploaded_img)
    col1.write("Original Image :camera:")
    col1.image(image)
else :

    image = Image.open(local_img[custom_img])
    col1.write("Custom Image :camera:")
    col1.image(image)

fixed = remove(image,session =session)
col2.write("Fixed Image :wrench:")
col2.image(fixed)
col2.markdown("\n")
col2.download_button("Download fixed image", convert_image(fixed), "fixed.png", "image/png")