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