Spaces:
Runtime error
Runtime error
File size: 1,847 Bytes
6cba02d 2b1436e 6cba02d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
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")
|