File size: 2,568 Bytes
daf7bc7
15c5016
 
daf7bc7
 
 
15c5016
 
 
 
daf7bc7
15c5016
 
 
 
 
 
 
 
 
 
 
 
 
 
daf7bc7
 
15c5016
daf7bc7
 
15c5016
daf7bc7
 
 
7165e24
daf7bc7
c8cd80b
 
daf7bc7
80446b6
daf7bc7
 
 
 
 
 
 
 
 
 
c8cd80b
daf7bc7
7165e24
daf7bc7
c8cd80b
daf7bc7
 
 
 
 
15c5016
c8cd80b
daf7bc7
 
 
 
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
import os
import cv2
import numpy as np
import streamlit as st
import tensorflow as tf

@st.cache(suppress_st_warning=True)
def load_model():
    model = tf.keras.models.load_model("style/1")
    return model

@st.cache(suppress_st_warning=True)
def apply_style_transfer(model, content, style, resize=None):
    content = np.array(content, dtype=np.float32) / 255.
    style = np.array(style, dtype=np.float32) / 255.

    if resize:
        content = cv2.resize(content, (512, 512))
        style = cv2.resize(style, (512, 512))

    stylized_image = model(tf.constant(content[np.newaxis, ...]), tf.constant(style[np.newaxis, ...]))
    stylized_image = stylized_image[0] * 255
    stylized_image = np.array(stylized_image, dtype=np.uint8)
    stylized_image = stylized_image
    return stylized_image

def main():
    model = load_model()

    st.title("Neural Style-Transfer App")
    st.write("`neural-style` is a pre-trained model from Tensorflow-Hub that allows you to apply styles to images and create pretty art. This app allows you to upload your own content or style images to create some funky effects. We provide some example styles which you can use.")
    col1, col2 = st.beta_columns(2)
    content_file = st.sidebar.file_uploader('Upload Image', type=['jpg', 'jpeg', 'png'])
    style_file = st.sidebar.file_uploader('Upload Style', type=['jpg', 'jpeg', 'png'])
    style_options = st.sidebar.selectbox(label='Example Styles', options=os.listdir('assets/template_styles'))
    col1.subheader('Content Image')
    col2.subheader('Style Transfer')
    st.sidebar.subheader('Style Image')
    show_image = col1.empty()
    show_style = col2.empty()

    style = None
    content = None

    if content_file:
        content = content_file.getvalue()
        show_image.image(content, use_column_width=True)

    if style_file:
        style = style_file.getvalue()
        st.sidebar.image(style, use_column_width=True)
    elif style_options is not None:
        with open(os.path.join('assets/template_styles', style_options), 'rb') as f:
            style = f.read()
        st.sidebar.image(style, use_column_width=True)

    if content is not None and style is not None:
        content_image = tf.io.decode_image(content)
        style_image = tf.image.resize(tf.io.decode_image(style), (256, 256))
        with st.spinner('Generating style transfer...'):
            style_transfer = apply_style_transfer(model, content_image, style_image)
            show_style.image(style_transfer, use_column_width=True)


if __name__ == "__main__":
    main()