|
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 |
|
|
|
|
|
|
|
|
|
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" |
|
) |
