File size: 2,664 Bytes
206ce41 |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import streamlit as st
import os
from datetime import datetime
from PIL import Image
from io import BytesIO
from src.utils import change_background, matte
from src.st_style import apply_prod_style
# apply_prod_style(st) # NOTE: Uncomment this for production!
def image_download_button(pil_image, filename: str, fmt: str, label="Download"):
if fmt not in ["jpg", "png"]:
raise Exception(f"Unknown image format (Available: {fmt} - case sensitive)")
pil_format = "JPEG" if fmt == "jpg" else "PNG"
file_format = "jpg" if fmt == "jpg" else "png"
mime = "image/jpeg" if fmt == "jpg" else "image/png"
buf = BytesIO()
pil_image.save(buf, format=pil_format)
return st.download_button(
label=label,
data=buf.getvalue(),
file_name=f'{filename}.{file_format}',
mime=mime,
)
st.title("AI Photo Background Removal")
st.image(Image.open("assets/demo.jpg"))
st.write(
"""
You want to remove your photo background, but don't have the time and effort to learn photo editing skills?
**This app will change or remove your photo background, in seconds.**
"""
)
uploaded_file = st.file_uploader(
label="Upload your photo here",
accept_multiple_files=False, type=["png", "jpg", "jpeg"],
)
if uploaded_file is not None:
with st.expander("Original photo", expanded=True):
if uploaded_file is not None:
st.image(uploaded_file)
else:
st.warning("You haven't uploaded any photo yet")
in_mode = st.selectbox("Choose background color", ["Transparent (PNG)", "White", "Black", "Green", "Red", "Blue"])
in_submit = st.button("Submit")
if uploaded_file is not None and in_submit:
img_input = Image.open(uploaded_file)
with st.spinner("AI is doing magic to your photo. Please wait..."):
hexmap = {
"Transparent (PNG)": "#000000",
"Black": "#000000",
"White": "#FFFFFF",
"Green": "#22EE22",
"Red": "#EE2222",
"Blue": "#2222EE",
}
alpha = 0.0 if in_mode == "Transparent (PNG)" else 1.0
img_matte = matte(img_input)
img_output = change_background(img_input, img_matte, background_alpha=alpha, background_hex=hexmap[in_mode])
with st.expander("Success!", expanded=True):
st.image(img_output)
uploaded_name = os.path.splitext(uploaded_file.name)[0]
image_download_button(
pil_image=img_output,
filename=uploaded_name,
fmt="png"
) |