From 22a5df4332b00c4afd83bf6588b7b2ee98bef529 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maciej=20P=C4=99=C5=9Bko?= <mpesko3031@gmail.com>
Date: Fri, 2 Nov 2018 18:26:07 +0100
Subject: [PATCH] Add more fixes

---
 ComixGAN/model.py                | 61 --------------------------------
 prepare_setup.sh                 |  4 +++
 requirements.txt                 |  4 +--
 style_transfer/style_transfer.py | 11 +++---
 4 files changed, 11 insertions(+), 69 deletions(-)
 delete mode 100644 ComixGAN/model.py

diff --git a/ComixGAN/model.py b/ComixGAN/model.py
deleted file mode 100644
index 7cca114..0000000
--- a/ComixGAN/model.py
+++ /dev/null
@@ -1,61 +0,0 @@
-from keras import layers, Model
-from keras_contrib.layers import InstanceNormalization
-
-
-class ComixGAN():
-    def __init__(self):
-        # Build and compile the generator
-        self.generator = self.build_generator()
-
-    def build_generator(self):
-        def residual_block(input_tensor):
-            c1 = layers.Conv2D(256, (3, 3), strides=(1, 1), padding='same', use_bias=False)(input_tensor)
-            bn1 = InstanceNormalization(axis=3)(c1)
-            a1 = layers.Activation('relu')(bn1)
-            c2 = layers.Conv2D(256, (3, 3), strides=(1, 1), padding='same', use_bias=False)(a1)
-            bn2 = InstanceNormalization(axis=3)(c2)
-            add1 = layers.add([bn2, input_tensor])
-
-            return add1
-
-        inp = layers.Input((None, None, 3))
-
-        c1 = layers.Conv2D(64, (7, 7), strides=(1, 1), padding='same', use_bias=False)(inp)
-        bn1 = InstanceNormalization(axis=3)(c1)
-        a1 = layers.Activation('relu')(bn1)
-
-        c2 = layers.Conv2D(128, (3, 3), strides=(2, 2), padding='same')(a1)
-        c3 = layers.Conv2D(128, (3, 3), strides=(1, 1), padding='same', use_bias=False)(c2)
-        bn2 = InstanceNormalization(axis=3)(c3)
-        a2 = layers.Activation('relu')(bn2)
-
-        c4 = layers.Conv2D(256, (3, 3), strides=(2, 2), padding='same')(a2)
-        c5 = layers.Conv2D(256, (3, 3), strides=(1, 1), padding='same', use_bias=False)(c4)
-        bn3 = InstanceNormalization(axis=3)(c5)
-        a3 = layers.Activation('relu')(bn3)
-
-        r1 = residual_block(a3)
-        r2 = residual_block(r1)
-        r3 = residual_block(r2)
-        r4 = residual_block(r3)
-        r5 = residual_block(r4)
-        r6 = residual_block(r5)
-        r7 = residual_block(r6)
-        r8 = residual_block(r7)
-
-        u1 = layers.UpSampling2D(size=(2, 2))(r8)
-        c6 = layers.Conv2D(128, (3, 3), strides=(1, 1), padding='same', use_bias=False)(u1)
-        bn4 = InstanceNormalization(axis=3)(c6)
-        a4 = layers.Activation('relu')(bn4)
-
-        u2 = layers.UpSampling2D(size=(2, 2))(a4)
-        c7 = layers.Conv2D(64, (3, 3), strides=(1, 1), padding='same', use_bias=False)(u2)
-        bn5 = InstanceNormalization(axis=3)(c7)
-        a5 = layers.Activation('relu')(bn5)
-
-        output = layers.Conv2D(3, (7, 7), strides=(1, 1), activation='sigmoid', padding='same')(a5)
-
-        return Model(inputs=[inp], outputs=[output])
-
-
-comixGAN = ComixGAN()
diff --git a/prepare_setup.sh b/prepare_setup.sh
index c34709d..b196920 100644
--- a/prepare_setup.sh
+++ b/prepare_setup.sh
@@ -72,6 +72,10 @@ echo 'export LD_LIBRARY_PATH=/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
 source ~/.bashrc
 
 
+# CHECK CUDNN VERSION
+cat /usr/include/x86_64-linux-gnu/cudnn_v*.h | grep CUDNN_MAJOR -A 2
+
+
 # INSTALL PACKAGES
 export LC_ALL="en_US.UTF-8"
 export LC_CTYPE="en_US.UTF-8"
diff --git a/requirements.txt b/requirements.txt
index 3efd366..8c58b1d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -34,8 +34,8 @@ scikit-learn==0.20.0
 scipy==1.1.0
 six==1.11.0
 sklearn==0.0
-tensorboard==1.11.0
-tensorflow-gpu==1.11.0
+tensorboard==1.10.0
+tensorflow-gpu==1.10.1
 termcolor==1.1.0
 toolz==0.9.0
 torch==0.4.1
diff --git a/style_transfer/style_transfer.py b/style_transfer/style_transfer.py
index 9b8c104..92b7020 100644
--- a/style_transfer/style_transfer.py
+++ b/style_transfer/style_transfer.py
@@ -8,10 +8,9 @@ from django.conf import settings
 from django.core.cache import cache
 from keras.models import load_model
 from torch.autograd import Variable
-
+from keras_contrib.layers import InstanceNormalization
 from CartoonGAN.network.Transformer import Transformer
 from utils import profile
-from ComixGAN.model import ComixGAN
 
 
 class StyleTransfer():
@@ -40,7 +39,7 @@ class StyleTransfer():
                 h = int(w * ratio)
             resized_img = cv2.resize(img, (w, h))
             resized_images.append(resized_img)
-            return resized_images
+        return resized_images
 
     @classmethod
     def _comix_gan_stylize(cls, frames):
@@ -49,14 +48,14 @@ class StyleTransfer():
 
         if comixGAN_model is None:
             # load pretrained model
-            comixGAN_model = ComixGAN()
-            comixGAN_model.generator = load_model(settings.COMIX_GAN_MODEL_PATH)
+            comixGAN_model = load_model(settings.COMIX_GAN_MODEL_PATH,
+                                                  custom_objects={'InstanceNormalization': InstanceNormalization})
             cache.set(comixGAN_cache_key, comixGAN_model, None)  # None is the timeout parameter. It means cache forever
 
         frames = cls._resize_images(frames, size=384)
         frames = np.stack(frames)
         frames = frames / 255
-        stylized_imgs = comixGAN_model.generator.predict(frames)
+        stylized_imgs = comixGAN_model.predict(frames)
         return list(255 * stylized_imgs)
 
     @classmethod