Refactor application to just use the model locally and not need gRPC.

Makefile

@@ -7,7 +7,7 @@ model-server-cpu: style
 	docker run -p 8500:8500 --mount type=bind,source=${CURRENT_DIR}/style,target=/models/style -e MODEL_NAME=style -t tensorflow/serving:2.4.0
 
 app:
-	exec streamlit run streamlit_app.py
+	exec streamlit run app.py
 
 webcam:
 	exec streamlit run webcam_stream.py

streamlit_app.py → app.py

@@ -1,21 +1,34 @@
 import os
+import cv2
+import numpy as np
 import streamlit as st
 import tensorflow as tf
-from tensorflow_serving.apis import prediction_service_pb2_grpc
 
-from client import style_transfer_serving
-import grpc
+@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():
-    options = [
-        ('grpc.max_send_message_length', 200 * 1024 * 1024),
-        ('grpc.max_receive_message_length', 200 * 1024 * 1024)
-    ]
-    channel = grpc.insecure_channel('localhost:8500', options=options)
-    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
+    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'])
@@ -26,7 +39,6 @@ def main():
     show_image = col1.empty()
     show_style = col2.empty()
 
-
     style = None
     content = None
 
@@ -46,7 +58,7 @@ def main():
         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 = style_transfer_serving(stub, content_image, style_image)
+            style_transfer = apply_style_transfer(model, content_image, style_image)
             show_style.image(style_transfer, use_column_width=True)
 
 

client.py

@@ -1,52 +0,0 @@
-import cv2
-import grpc
-import tensorflow as tf
-import tensorflow_hub as hub
-import numpy as np
-from tensorflow_serving.apis import predict_pb2, prediction_service_pb2_grpc
-
-hub_module = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
-
-def style_transfer_serving(stub, 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))
-
-
-    image_proto = tf.make_tensor_proto(content[np.newaxis, ...] / 255.)
-    style_proto = tf.make_tensor_proto(style[np.newaxis, ...] / 255.)
-
-    stylized_image = hub_module(tf.constant(content[np.newaxis, ...]), tf.constant(style[np.newaxis, ...]))
-    # request = predict_pb2.PredictRequest()
-    # request.model_spec.name = 'style'
-    # request.inputs['placeholder'].CopyFrom(image_proto)
-    # request.inputs['placeholder_1'].CopyFrom(style_proto)
-    # resp = stub.Predict(request)
-    # stylized_image = tf.make_ndarray(resp.outputs['output_0'])[0]
-    stylized_image = stylized_image[0] * 255
-    stylized_image = np.array(stylized_image, dtype=np.uint8)
-    stylized_image = stylized_image
-    return stylized_image
-
-if __name__ == "__main__":
-    options = [
-        ('grpc.max_send_message_length', 200 * 1024 * 1024),
-        ('grpc.max_receive_message_length', 200 * 1024 * 1024)
-    ]
-    # channel = grpc.insecure_channel('localhost:8500', options=options)
-    # stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
-
-    file = tf.io.read_file('/home/albert/github/neural-style/assets/template_styles/pebbles.jpg')
-    style = tf.io.decode_image(file)
-
-    file = tf.io.read_file('/home/albert/Downloads/sam_and_nyx/sam_stairs.jpeg')
-    content = tf.io.decode_image(file)
-
-    stub = None
-    result = style_transfer_serving(stub, content, style)
-    import matplotlib.pyplot as plt
-    plt.imshow(result[0])
-    plt.show()

requirements.txt

@@ -0,0 +1,5 @@
+opencv-python==4.4.0.46
+tensorflow==2.4.1
+streamlit==0.70.0
+numpy==1.19.5
+grpcio==1.32.0

setup.py

@@ -1,21 +1,16 @@
-from setuptools import setup
+from setuptools import setup, find_packages
 
-setup(
-    name='neural-style',
-    version='',
-    packages=[],
-    install_requires=[
-        'opencv-python==4.4.0.46',
-        'tensorflow==2.4.1',
-        'streamlit==0.70.0',
-        'numpy==1.19.5',
-        'grpcio==1.32.0',
-        'tensorflow-serving-api==2.4.1',
+with open("requirements.txt") as f:
+    requirements = f.readlines()
 
-    ],
-    url='',
-    license='',
-    author='albert',
-    author_email='',
-    description=''
+setup(
+    name="neural-style",
+    version="",
+    packages=find_packages(),
+    install_requires=requirements,
+    url="",
+    license="",
+    author="albert",
+    author_email="",
+    description=""
 )