Converting to tensorflow serving gpu container.

Makefile

@@ -1,2 +1,2 @@
 start: style
-	docker run -p 8500:8500 --mount type=bind,source=C:/Users/doge/github/neural-style/style,target=/models/style -e MODEL_NAME=style -t tensorflow/serving
+	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

client.py

@@ -1,25 +1,61 @@
-import grpc
-import tensorflow as tf
-import numpy as np
-from tensorflow_serving.apis import predict_pb2, prediction_service_pb2_grpc
-
-if __name__ == "__main__":
-    options = [('grpc.max_message_length', 100 * 1024 * 1024)]
-    channel = grpc.insecure_channel('localhost:8500', options=options)
-    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
-    request = predict_pb2.PredictRequest()
-
-    file = tf.io.read_file('C:\\Users\\doge\\Downloads\\sam.jpg')
-    image = tf.io.decode_image(file)
-
-    request.model_spec.name = 'style'
-    request.model_spec.signature_name = 'serving_default'
-    image_proto = tf.make_tensor_proto(np.array(image, dtype=np.float32)[np.newaxis, ...])
-    request.inputs['placeholder'].CopyFrom(image_proto)
-    request.inputs['placeholder_1'].CopyFrom(image_proto)
-    resp = stub.Predict(request)
-    print(resp)
-
-
-
-
+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
+
+
+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)
+    request = predict_pb2.PredictRequest()
+
+    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()
+
+    # styled_plot = plt.imshow(styled_image)
+    # plt.axis('off')
+    # plt.show()
+    #
+    # IMAGES_PATH = '/home/albert/Downloads/sam'
+    # for path in os.listdir(IMAGES_PATH):
+    #     start = time.time()
+    #     file = tf.io.read_file(os.path.join(IMAGES_PATH, path))
+    #     image = tf.io.decode_image(file)
+    #     output = style_transfer(stub, image)
+    #     styled_plot.set_data(output)
+    #     plt.draw()
+    #     end = time.time()
+    #     print(f'Time taken to predict: {end - start}s')

main.py

@@ -23,11 +23,12 @@ content_path = tf.keras.utils.get_file('YellowLabradorLooking_new.jpg',
 style_path = tf.keras.utils.get_file('kandinsky5.jpg',
                                      'https://storage.googleapis.com/download.tensorflow.org/example_images/Vassily_Kandinsky%2C_1913_-_Composition_7.jpg')
 
+
 def load_img(path_to_img):
     max_dim = 512
     img = tf.io.read_file(path_to_img)
     img = tf.image.decode_image(img, channels=3)
-    img = tf.image.convert_image_dtype(img ,tf.float32)
+    img = tf.image.convert_image_dtype(img, tf.float32)
 
     shape = tf.cast(tf.shape(img)[:-1], tf.float32)
     long_dim = max(shape)
@@ -39,6 +40,7 @@ def load_img(path_to_img):
     img = img[tf.newaxis, ...]
     return img
 
+
 def imshow(image, title=None):
     if np.ndim(image) > 3:
         image = tf.squeeze(image, axis=0)
@@ -47,6 +49,7 @@ def imshow(image, title=None):
     if title:
         plt.title(title)
 
+
 content_image = load_img(content_path)
 style_image = load_img(style_path)
 
@@ -56,10 +59,10 @@ plt.subplot(1, 2, 2)
 imshow(style_image, 'Style Image')
 
 import tensorflow_hub as hub
+
 hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
 stylized_image = hub_model(tf.constant(content_image), tf.constant(style_image))[0]
 tensor_to_image(stylized_image)
 print(type(hub_model))
 
-
 tf.saved_model.save(hub_model, 'style')