File size: 2,579 Bytes
da406d5
 
 
 
 
 
 
 
 
d7f113c
da406d5
26ef69b
 
 
9d62619
67f9880
164ffb9
 
9d62619
164ffb9
 
 
9d62619
164ffb9
8d3aa08
164ffb9
9d62619
8d3aa08
 
 
 
 
9d62619
 
8d3aa08
 
9d62619
 
 
 
 
 
da406d5
 
 
 
 
 
 
 
f3f08a7
da406d5
 
 
9d62619
 
da406d5
 
 
 
 
 
 
 
 
 
 
 
9d62619
da406d5
9d62619
da406d5
8d3aa08
164ffb9
 
 
79690b2
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
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 V_SPACE(lines):
    for _ in range(lines):
        st.write(' ')
        
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/?target=external" if fmt == "jpg" else "image/png/?target=external"
    
    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,
        on_click=open_in_new_tab,
        args=(buf.getvalue(),)
    )

def open_in_new_tab(file_content):
    file_ = BytesIO(file_content)
    file_.seek(0)
    b64_img = base64.b64encode(file_.read()).decode()
    href = f'<a href="data:image/png;base64,{b64_img}" target="_blank" rel="noopener noreferrer">Open image in new tab</a>'
    st.sidebar.markdown(href, unsafe_allow_html=True)

uploaded_file = st.file_uploader(
    label="Upload your photo here", 
    accept_multiple_files=False, type=["png", "jpg", "jpeg"],
)

if uploaded_file is not None:

    in_mode = "Transparent (PNG)"
    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"
            )