diff --git a/ComixGAN/generator_model.h5 b/ComixGAN/generator_model.h5 new file mode 100644 index 0000000..ebe0ee1 Binary files /dev/null and b/ComixGAN/generator_model.h5 differ diff --git a/dockerfile b/dockerfile index 48429f1..4cb1229 100644 --- a/dockerfile +++ b/dockerfile @@ -42,7 +42,7 @@ RUN echo "$CAFFE_ROOT/build/lib" >> /etc/ld.so.conf.d/caffe.conf && ldconfig && WORKDIR /comixify COPY . /comixify RUN unzip popularity/pretrained_model/svr_test_11.10.sk.zip -d popularity/pretrained_model/ && \ - python3.6 -m pip install -r requirements.txt + python3.6 -m pip install -r requirements.txt && \ python3.6 -m pip install git+https://www.github.com/keras-team/keras-contrib.git diff --git a/requirements.txt b/requirements.txt index 7f4cfc2..a367be5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,8 @@ -absl-py==0.6.1 astor==0.7.1 +cloudpickle==0.6.1 +cycler==0.10.0 +dask==0.20.0 +decorator==4.3.0 Django==2.0.7 django-rest-framework==0.1.0 djangorestframework==3.8.2 @@ -9,22 +12,32 @@ gunicorn==19.9.0 h5py==2.8.0 Keras==2.2.2 Keras-Applications==1.0.6 -keras-contrib==2.0.8 Keras-Preprocessing==1.0.5 +kiwisolver==1.0.1 Markdown==3.0.1 +matplotlib==3.0.1 +networkx==2.2 numpy==1.14.5 +opencv-contrib-python==3.4.3.18 opencv-python==3.4.2.17 pafy==0.5.4 Pillow==5.2.0 protobuf==3.6.1 psycopg2==2.7.5 +pyparsing==2.3.0 +python-dateutil==2.7.5 pytz==2018.5 +PyWavelets==1.0.1 PyYAML==3.13 +scikit-image==0.14.1 +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 termcolor==1.1.0 +toolz==0.9.0 torch==0.4.1 torchvision==0.2.1 Werkzeug==0.14.1 diff --git a/settings/settings.py b/settings/settings.py index 658ff10..8ab2442 100644 --- a/settings/settings.py +++ b/settings/settings.py @@ -148,3 +148,4 @@ DEFAULT_IMAGE_ASSESSMENT_MODE = 0 DEFAULT_STYLE_TRANSFER_MODE = 0 COMIX_GAN_WEIGHTS_PATH = os.path.join(BASE_DIR, 'ComixGAN', 'G_weights.best.hdf5') +COMIX_GAN_MODEL_PATH = os.path.join(BASE_DIR, 'ComixGAN', 'generator_model.h5') diff --git a/style_transfer/style_transfer.py b/style_transfer/style_transfer.py index 3e1135a..b00af5d 100644 --- a/style_transfer/style_transfer.py +++ b/style_transfer/style_transfer.py @@ -6,6 +6,7 @@ import torch import torchvision.transforms as transforms from django.conf import settings from django.core.cache import cache +from keras.models import load_model from torch.autograd import Variable from CartoonGAN.network.Transformer import Transformer @@ -43,19 +44,19 @@ class StyleTransfer(): @classmethod def _comix_gan_stylize(cls, frames): - model_cache_key = 'model_cache' - model = cache.get(model_cache_key) # get model from cache + comixGAN_cache_key = 'comixGAN_model_cache' + comixGAN_model = cache.get(comixGAN_cache_key) # get model from cache - if model is None: + if comixGAN_model is None: # load pretrained model - model = ComixGAN() - model.load_weights(settings.COMIX_GAN_WEIGHTS_PATH) - cache.set(model_cache_key, model, None) # None is the timeout parameter. It means cache forever + comixGAN_model = ComixGAN() + comixGAN_model.generator = load_model(settings.COMIX_GAN_MODEL_PATH) + 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 / frames.max() - stylized_imgs = model.generator.predict(frames) + stylized_imgs = comixGAN_model.generator.predict(frames) return list(255 * stylized_imgs) @classmethod