Streamlit web app with example images.

.gitignore

@@ -0,0 +1,2 @@
+serving
+tensorflow

Makefile

@@ -1,2 +1,5 @@
-start: style
+model-server: style
 	docker run --gpus all -p 8500:8500 --mount type=bind,source=/home/albert/github/neural-style/style,target=/models/style -e MODEL_NAME=style -t tensorflow/serving:latest-gpu
+
+app:
+	exec streamlit run streamlit_app.py

README.md

@@ -7,4 +7,5 @@
 - Go
 - gocv
 - tensorflow and tensorflow/serving repositories (for golang protobufs)
-- golang protobuf library
+- golang protobuf library
+- protobuf compiler

client.go

@@ -1,19 +0,0 @@
-package main
-
-import (
-	"flag"
-	"fmt"
-	"gocv.io/x/gocv"
-	pb "tensorflow_serving/apis"
-)
-
-func main() {
-	request := &pb.PredictRequest{}
-	webcam, _ := gocv.VideoCaptureDevice(0)
-	window := gocv.NewWindow("hello")
-	for {
-		webcam.Read(&img)
-		window.IMShow(img)
-		window.WaitKey(1)
-	}
-}

client.py

@@ -1,12 +1,28 @@
-import os
 import cv2
 import grpc
 import tensorflow as tf
 import numpy as np
-import matplotlib.pyplot as plt
 from tensorflow_serving.apis import predict_pb2, prediction_service_pb2_grpc
-import time
 
+def style_transfer_serving(stub, content, style, resize=None):
+    content = np.array(content, dtype=np.float32)
+    style = np.array(style, dtype=np.float32)
+
+    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.)
+
+    request = predict_pb2.PredictRequest()
+    request.model_spec.name = 'style'
+    request.model_spec.signature_name = 'serving_default'
+    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]
+    return stylized_image
 
 if __name__ == "__main__":
     options = [
@@ -20,27 +36,7 @@ if __name__ == "__main__":
     file = tf.io.read_file('/home/albert/Downloads/pebbles.jpg')
     style = tf.io.decode_image(file)
 
-    style_image = cv2.resize(np.array(style, dtype=np.float32), (64, 64))[np.newaxis, ...] / 255.
-    style_proto = tf.make_tensor_proto(np.array(style, dtype=np.float32)[np.newaxis, ...] / 255.)
-
-    def style_transfer(stub, image):
-        request.model_spec.name = 'style'
-        request.model_spec.signature_name = 'serving_default'
-        image = cv2.resize(np.array(image, dtype=np.float32), (512, 512))
-        image_proto = tf.make_tensor_proto(image[np.newaxis, ...] / 255.)
-
-        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]
-        return stylized_image
-
-    video = cv2.VideoCapture('/home/albert/Downloads/cat_yelling.mp4')
-    while video.isOpened():
-        ret, frame = video.read()
-        styled_image = style_transfer(stub, frame)
-        cv2.imshow('cheese', styled_image)
-        if cv2.waitKey(1) & 0xFF == ord('q'):
-            break
-    video.release()
-    cv2.destroyAllWindows()
+    file = tf.io.read_file('/home/albert/Downloads/sam_and_nyx/sam_faces/sam_kitchen.jpg')
+    content = tf.io.decode_image(file)
+
+    style_transfer_serving(stub, content, style)

streamlit_app.py

@@ -0,0 +1,55 @@
+import os
+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
+
+
+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)
+
+    st.title("Neural Style-Transfer App")
+    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('template_styles'))
+    col1.subheader('Content Image')
+    col2.subheader('Style Image')
+    show_image = col1.empty()
+    show_style = col2.empty()
+
+    st.subheader('Style Transfer')
+    show_transfer = st.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()
+        show_style.image(style, use_column_width=True)
+    elif style_options is not None:
+        with open(os.path.join('template_styles', style_options), 'rb') as f:
+            style = f.read()
+        show_style.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 = style_transfer_serving(stub, content_image, style_image)
+            show_transfer.image(style_transfer, use_column_width=True)
+
+
+if __name__ == "__main__":
+    main()

webcam_stream.py

@@ -0,0 +1,56 @@
+import cv2
+import tensorflow as tf
+import numpy as np
+from threading import Thread
+
+
+class WebcamVideoStream:
+    def __init__(self, src=0):
+        self.stream = cv2.VideoCapture(src)
+        self.grabbed, self.frame = self.stream.read()
+        self.stopped = False
+
+    def start(self):
+        Thread(target=self.update, args=()).start()
+        return self
+
+    def update(self):
+        while True:
+            if self.stopped:
+                return
+
+            self.grabbed, self.frame = self.stream.read()
+
+    def read(self):
+        return self.frame
+
+    def stop(self):
+        self.stopped = True
+
+
+def style_transfer_direct(image, model, resize=None):
+    if resize:
+        image = cv2.resize(np.array(image, dtype=np.float32), (512, 512))
+    else:
+        image = np.array(image, dtype=np.float32)
+    resp = model(tf.convert_to_tensor(image[np.newaxis, ...] / 255.), tf.convert_to_tensor(style_image))
+    stylized_image = resp[0]
+    return stylized_image.numpy()[0]
+
+
+def main():
+    model = tf.saved_model.load('style/1')
+    video_capture = WebcamVideoStream(0)
+    video_capture.start()
+    while True:
+        content_image = video_capture.read()
+        transfer = style_transfer_direct(content_image, model)
+        cv2.imshow('style camera', transfer)
+        if cv2.waitKey(1) & 0xFF == ord('q'):
+            break
+    cv2.destroyAllWindows()
+    video_capture.stop()
+
+
+if __name__ == "__main__":
+    main()