Add base pipeline

CartoonGAN/network/Transformer.py

@@ -0,0 +1,177 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Transformer(nn.Module):
+    def __init__(self):
+        super(Transformer, self).__init__()
+        #
+        self.refpad01_1 = nn.ReflectionPad2d(3)
+        self.conv01_1 = nn.Conv2d(3, 64, 7)
+        self.in01_1 = InstanceNormalization(64)
+        # relu
+        self.conv02_1 = nn.Conv2d(64, 128, 3, 2, 1)
+        self.conv02_2 = nn.Conv2d(128, 128, 3, 1, 1)
+        self.in02_1 = InstanceNormalization(128)
+        # relu
+        self.conv03_1 = nn.Conv2d(128, 256, 3, 2, 1)
+        self.conv03_2 = nn.Conv2d(256, 256, 3, 1, 1)   
+        self.in03_1 = InstanceNormalization(256)    
+        # relu
+
+        ## res block 1
+        self.refpad04_1 = nn.ReflectionPad2d(1)
+        self.conv04_1 = nn.Conv2d(256, 256, 3)
+        self.in04_1 = InstanceNormalization(256)
+        # relu
+        self.refpad04_2 = nn.ReflectionPad2d(1)
+        self.conv04_2 = nn.Conv2d(256, 256, 3)
+        self.in04_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 2
+        self.refpad05_1 = nn.ReflectionPad2d(1)
+        self.conv05_1 = nn.Conv2d(256, 256, 3)
+        self.in05_1 = InstanceNormalization(256)
+        # relu
+        self.refpad05_2 = nn.ReflectionPad2d(1)
+        self.conv05_2 = nn.Conv2d(256, 256, 3)
+        self.in05_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 3
+        self.refpad06_1 = nn.ReflectionPad2d(1)
+        self.conv06_1 = nn.Conv2d(256, 256, 3)
+        self.in06_1 = InstanceNormalization(256)
+        # relu
+        self.refpad06_2 = nn.ReflectionPad2d(1)
+        self.conv06_2 = nn.Conv2d(256, 256, 3)
+        self.in06_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 4
+        self.refpad07_1 = nn.ReflectionPad2d(1)
+        self.conv07_1 = nn.Conv2d(256, 256, 3)
+        self.in07_1 = InstanceNormalization(256)
+        # relu
+        self.refpad07_2 = nn.ReflectionPad2d(1)
+        self.conv07_2 = nn.Conv2d(256, 256, 3)
+        self.in07_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 5
+        self.refpad08_1 = nn.ReflectionPad2d(1)
+        self.conv08_1 = nn.Conv2d(256, 256, 3)
+        self.in08_1 = InstanceNormalization(256)
+        # relu
+        self.refpad08_2 = nn.ReflectionPad2d(1)
+        self.conv08_2 = nn.Conv2d(256, 256, 3)
+        self.in08_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 6
+        self.refpad09_1 = nn.ReflectionPad2d(1)
+        self.conv09_1 = nn.Conv2d(256, 256, 3)
+        self.in09_1 = InstanceNormalization(256)
+        # relu
+        self.refpad09_2 = nn.ReflectionPad2d(1)
+        self.conv09_2 = nn.Conv2d(256, 256, 3)
+        self.in09_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 7
+        self.refpad10_1 = nn.ReflectionPad2d(1)
+        self.conv10_1 = nn.Conv2d(256, 256, 3)
+        self.in10_1 = InstanceNormalization(256)
+        # relu
+        self.refpad10_2 = nn.ReflectionPad2d(1)
+        self.conv10_2 = nn.Conv2d(256, 256, 3)
+        self.in10_2 = InstanceNormalization(256)
+        # + input
+
+        ## res block 8
+        self.refpad11_1 = nn.ReflectionPad2d(1)
+        self.conv11_1 = nn.Conv2d(256, 256, 3)
+        self.in11_1 = InstanceNormalization(256)
+        # relu
+        self.refpad11_2 = nn.ReflectionPad2d(1)
+        self.conv11_2 = nn.Conv2d(256, 256, 3)
+        self.in11_2 = InstanceNormalization(256)
+        # + input
+
+        ##------------------------------------##
+        self.deconv01_1 = nn.ConvTranspose2d(256, 128, 3, 2, 1, 1)
+        self.deconv01_2 = nn.Conv2d(128, 128, 3, 1, 1)
+        self.in12_1 = InstanceNormalization(128)
+        # relu
+        self.deconv02_1 = nn.ConvTranspose2d(128, 64, 3, 2, 1, 1)
+        self.deconv02_2 = nn.Conv2d(64, 64, 3, 1, 1)
+        self.in13_1 = InstanceNormalization(64)
+        # relu
+        self.refpad12_1 = nn.ReflectionPad2d(3)
+        self.deconv03_1 = nn.Conv2d(64, 3, 7)
+        # tanh
+
+    def forward(self, x):
+        y = F.relu(self.in01_1(self.conv01_1(self.refpad01_1(x))))
+        y = F.relu(self.in02_1(self.conv02_2(self.conv02_1(y))))
+        t04 = F.relu(self.in03_1(self.conv03_2(self.conv03_1(y))))
+
+        ##
+        y = F.relu(self.in04_1(self.conv04_1(self.refpad04_1(t04))))
+        t05 = self.in04_2(self.conv04_2(self.refpad04_2(y))) + t04
+
+        y = F.relu(self.in05_1(self.conv05_1(self.refpad05_1(t05))))
+        t06 = self.in05_2(self.conv05_2(self.refpad05_2(y))) + t05
+
+        y = F.relu(self.in06_1(self.conv06_1(self.refpad06_1(t06))))
+        t07 = self.in06_2(self.conv06_2(self.refpad06_2(y))) + t06
+
+        y = F.relu(self.in07_1(self.conv07_1(self.refpad07_1(t07))))
+        t08 = self.in07_2(self.conv07_2(self.refpad07_2(y))) + t07
+
+        y = F.relu(self.in08_1(self.conv08_1(self.refpad08_1(t08))))
+        t09 = self.in08_2(self.conv08_2(self.refpad08_2(y))) + t08
+
+        y = F.relu(self.in09_1(self.conv09_1(self.refpad09_1(t09))))
+        t10 = self.in09_2(self.conv09_2(self.refpad09_2(y))) + t09
+
+        y = F.relu(self.in10_1(self.conv10_1(self.refpad10_1(t10))))
+        t11 = self.in10_2(self.conv10_2(self.refpad10_2(y))) + t10
+
+        y = F.relu(self.in11_1(self.conv11_1(self.refpad11_1(t11))))
+        y = self.in11_2(self.conv11_2(self.refpad11_2(y))) + t11
+        ##
+
+        y = F.relu(self.in12_1(self.deconv01_2(self.deconv01_1(y))))
+        y = F.relu(self.in13_1(self.deconv02_2(self.deconv02_1(y))))
+        y = F.tanh(self.deconv03_1(self.refpad12_1(y)))
+
+        return y
+
+
+class InstanceNormalization(nn.Module):
+    def __init__(self, dim, eps=1e-9):
+        super(InstanceNormalization, self).__init__()
+        self.scale = nn.Parameter(torch.FloatTensor(dim))
+        self.shift = nn.Parameter(torch.FloatTensor(dim))
+        self.eps = eps
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        self.scale.data.uniform_()
+        self.shift.data.zero_()
+
+    def __call__(self, x):
+        n = x.size(2) * x.size(3)
+        t = x.view(x.size(0), x.size(1), n)
+        mean = torch.mean(t, 2).unsqueeze(2).unsqueeze(3).expand_as(x)
+        # Calculate the biased var. torch.var returns unbiased var
+        var = torch.var(t, 2).unsqueeze(2).unsqueeze(3).expand_as(x) * ((n - 1) / float(n))
+        scale_broadcast = self.scale.unsqueeze(1).unsqueeze(1).unsqueeze(0)
+        scale_broadcast = scale_broadcast.expand_as(x)
+        shift_broadcast = self.shift.unsqueeze(1).unsqueeze(1).unsqueeze(0)
+        shift_broadcast = shift_broadcast.expand_as(x)
+        out = (x - mean) / torch.sqrt(var + self.eps)
+        out = out * scale_broadcast + shift_broadcast
+        return out

CartoonGAN/pretrained_model/download_pth.sh

@@ -0,0 +1,8 @@
+cd pretrained_model
+
+wget -c http://vllab1.ucmerced.edu/~yli62/CartoonGAN/pytorch_pth/Hayao_net_G_float.pth
+wget -c http://vllab1.ucmerced.edu/~yli62/CartoonGAN/pytorch_pth/Hosoda_net_G_float.pth
+wget -c http://vllab1.ucmerced.edu/~yli62/CartoonGAN/pytorch_pth/Paprika_net_G_float.pth	
+wget -c http://vllab1.ucmerced.edu/~yli62/CartoonGAN/pytorch_pth/Shinkai_net_G_float.pth
+
+cd ..

CartoonGAN/pretrained_model/download_t7.sh

@@ -0,0 +1,8 @@
+cd pretrained_model
+
+wget -c http://vllab1.ucmerced.edu/~yli62/CartoonGAN/torch_t7/Hayao_net_G_float.t7
+wget -c http://vllab1.ucmerced.edu/~yli62/CartoonGAN/torch_t7/Hosoda_net_G_float.t7
+wget -c http://vllab1.ucmerced.edu/~yli62/CartoonGAN/torch_t7/Paprika_net_G_float.t7	
+wget -c http://vllab1.ucmerced.edu/~yli62/CartoonGAN/torch_t7/Shinkai_net_G_float.t7
+
+cd ..

CartoonGAN/src/InstanceNormalization.lua

@@ -0,0 +1,98 @@
+require 'nn'
+
+_ = [[
+   An implementation for https://arxiv.org/abs/1607.08022
+]]
+
+local InstanceNormalization, parent = torch.class('nn.InstanceNormalization', 'nn.Module')
+
+function InstanceNormalization:__init(nOutput, eps, momentum, affine)
+   parent.__init(self)
+   self.running_mean = torch.zeros(nOutput)
+   self.running_var = torch.ones(nOutput)
+
+   self.eps = eps or 1e-5
+   self.momentum = momentum or 0.0
+   if affine ~= nil then
+      assert(type(affine) == 'boolean', 'affine has to be true/false')
+      self.affine = affine
+   else
+      self.affine = true
+   end
+   
+   self.nOutput = nOutput
+   self.prev_batch_size = -1
+
+   if self.affine then 
+      self.weight = torch.Tensor(nOutput):uniform()
+      self.bias = torch.Tensor(nOutput):zero()
+      self.gradWeight = torch.Tensor(nOutput)
+      self.gradBias = torch.Tensor(nOutput)
+   end 
+end
+
+function InstanceNormalization:updateOutput(input)
+   self.output = self.output or input.new()
+   assert(input:size(2) == self.nOutput)
+
+   local batch_size = input:size(1)
+   
+   if batch_size ~= self.prev_batch_size or (self.bn and self:type() ~= self.bn:type())  then
+      self.bn = nn.SpatialBatchNormalization(input:size(1)*input:size(2), self.eps, self.momentum, self.affine)
+      self.bn:type(self:type())
+      self.bn.running_mean:copy(self.running_mean:repeatTensor(batch_size))
+      self.bn.running_var:copy(self.running_var:repeatTensor(batch_size))
+
+      self.prev_batch_size = input:size(1)
+   end
+
+   -- Get statistics
+   self.running_mean:copy(self.bn.running_mean:view(input:size(1),self.nOutput):mean(1))
+   self.running_var:copy(self.bn.running_var:view(input:size(1),self.nOutput):mean(1))
+
+   -- Set params for BN
+   if self.affine then
+      self.bn.weight:copy(self.weight:repeatTensor(batch_size))
+      self.bn.bias:copy(self.bias:repeatTensor(batch_size))
+   end
+
+   local input_1obj = input:contiguous():view(1,input:size(1)*input:size(2),input:size(3),input:size(4))
+   self.output = self.bn:forward(input_1obj):viewAs(input)
+   
+   return self.output
+end
+
+function InstanceNormalization:updateGradInput(input, gradOutput)
+   self.gradInput = self.gradInput or gradOutput.new()
+
+   assert(self.bn)
+
+   local input_1obj = input:contiguous():view(1,input:size(1)*input:size(2),input:size(3),input:size(4)) 
+   local gradOutput_1obj = gradOutput:contiguous():view(1,input:size(1)*input:size(2),input:size(3),input:size(4)) 
+   
+   if self.affine then
+      self.bn.gradWeight:zero()
+      self.bn.gradBias:zero()
+   end
+
+   self.gradInput = self.bn:backward(input_1obj, gradOutput_1obj):viewAs(input)
+
+   if self.affine then
+      self.gradWeight:add(self.bn.gradWeight:view(input:size(1),self.nOutput):sum(1))
+      self.gradBias:add(self.bn.gradBias:view(input:size(1),self.nOutput):sum(1))
+   end
+   return self.gradInput
+end
+
+function InstanceNormalization:clearState()
+   self.output = self.output.new()
+   self.gradInput = self.gradInput.new()
+   
+   self.bn:clearState()
+end
+
+function InstanceNormalization:evaluate()
+end
+
+function InstanceNormalization:training()
+end

CartoonGAN/src/util.lua

@@ -0,0 +1,90 @@
+--
+-- code derived from https://github.com/soumith/dcgan.torch
+--
+
+local util = {}
+
+require 'torch'
+require 'nn'
+require 'lfs'
+
+-- Preprocesses an image before passing it to a net
+-- Converts from RGB to BGR and rescales from [0,1] to [-1,1]
+function util.preprocess(img)
+    -- RGB to BGR
+    local perm = torch.LongTensor{3, 2, 1}
+    img = img:index(1, perm)
+    
+    -- [0,1] to [-1,1]
+    img = img:mul(2):add(-1)
+    
+    -- check that input is in expected range
+    assert(img:max()<=1,"badly scaled inputs")
+    assert(img:min()>=-1,"badly scaled inputs")
+    
+    return img
+end
+
+-- Undo the above preprocessing.
+function util.deprocess(img)
+    -- BGR to RGB
+    local perm = torch.LongTensor{3, 2, 1}
+    img = img:index(1, perm)
+    
+    -- [-1,1] to [0,1]
+    
+    img = img:add(1):div(2)
+    
+    return img
+end
+
+function util.preprocess_batch(batch)
+  for i = 1, batch:size(1) do
+    batch[i] = util.preprocess(batch[i]:squeeze())
+  end
+  return batch
+end
+
+function util.deprocess_batch(batch)
+  for i = 1, batch:size(1) do
+   batch[i] = util.deprocess(batch[i]:squeeze())
+  end
+return batch
+end
+
+--
+-- code derived from AdaIN https://github.com/xunhuang1995/AdaIN-style
+--
+
+function util.extractImageNamesRecursive(dir)
+    local files = {}
+    print("Extracting image paths: " .. dir)
+  
+    local function browseFolder(root, pathTable)
+        for entity in lfs.dir(root) do
+            if entity~="." and entity~=".." then
+                local fullPath=root..'/'..entity
+                local mode=lfs.attributes(fullPath,"mode")
+                if mode=="file" then
+                    local filepath = paths.concat(root, entity)
+  
+                    if string.find(filepath, 'jpg$')
+                    or string.find(filepath, 'png$')
+                    or string.find(filepath, 'jpeg$')
+                    or string.find(filepath, 'JPEG$')
+                    or string.find(filepath, 'ppm$') then
+                        table.insert(pathTable, filepath)
+                    end
+                elseif mode=="directory" then
+                    browseFolder(fullPath, pathTable);
+                end
+            end
+        end
+    end
+
+    browseFolder(dir, files)
+    return files
+end
+
+
+return util

api/models.py

@@ -2,27 +2,28 @@ import os
 import uuid
 
 import cv2
+from django.conf import settings
 from django.core.files import File
 from django.db import models
 
 
 class Video(models.Model):
-    file = models.FileField(blank=False, null=False, upload_to='raw_videos')
+    file = models.FileField(blank=False, null=False, upload_to="raw_videos")
     timestamp = models.DateTimeField(auto_now_add=True)
 
 
 class Comic(models.Model):
-    file = models.FileField(blank=False, null=False, upload_to='comic')
-    video = models.ForeignKey(Video, on_delete=models.CASCADE, related_name='comic')
+    file = models.FileField(blank=False, null=False, upload_to="comic")
+    video = models.ForeignKey(Video, on_delete=models.CASCADE, related_name="comic")
 
     @classmethod
     def create_from_nparray(cls, nparray_file, video):
         tmp_name = uuid.uuid4()
-        if not os.path.exists('tmp/'):
-            os.makedirs('tmp/')
-        cv2.imwrite(f'tmp/{tmp_name}.png', nparray_file)
-        with open(f'tmp/{tmp_name}.png', mode='rb') as tmp_file:
-            comic_image = File(tmp_file, name=f'{tmp_name}.png')
+        if not os.path.exists(f"{settings.TMP_DIR}"):
+            os.makedirs(f"{settings.TMP_DIR}")
+        cv2.imwrite(f"{settings.TMP_DIR}{tmp_name}.png", nparray_file)
+        with open(f"{settings.TMP_DIR}{tmp_name}.png", mode="rb") as tmp_file:
+            comic_image = File(tmp_file, name=f"{tmp_name}.png")
             comic = Comic.objects.create(file=comic_image, video=video)
-        os.remove(f'tmp/{tmp_name}.png')
+        os.remove(f"{settings.TMP_DIR}{tmp_name}.png")
         return comic

api/serializers.py

@@ -8,11 +8,11 @@ from .models import Video
 class VideoSerializer(serializers.ModelSerializer):
     class Meta:
         model = Video
-        fields = ('file', 'timestamp')
+        fields = ("file", "timestamp")
 
     def validate(self, attrs):
-        file = attrs.get('file')
-        if file.name.split('.')[-1] not in settings.PERMITTED_VIDEO_EXTENSIONS:
+        file = attrs.get("file")
+        if file.name.split(".")[-1] not in settings.PERMITTED_VIDEO_EXTENSIONS:
             raise FileExtensionError
         if file.size > settings.MAX_FILE_SIZE:
             raise TooLargeFile

api/views.py

@@ -20,16 +20,18 @@ class Comixify(APIView):
         serializer = VideoSerializer(data=request.data)
         serializer.is_valid(raise_exception=True)
 
-        video_file = serializer.validated_data['file']
+        video_file = serializer.validated_data["file"]
+        video = Video.objects.create(file=video_file)
 
-        keyframes = KeyFramesExtractor.get_keyframes(video=video_file)
+        keyframes = KeyFramesExtractor.get_keyframes(video=video)
         stylized_keyframes = StyleTransfer.get_stylized_frames(frames=keyframes)
         comic_image = LayoutGenerator.get_layout(frames=stylized_keyframes)
 
-        video = Video.objects.create(file=video_file)
         comic = Comic.create_from_nparray(comic_image, video)
         response = {
-            'status_message': 'ok',
-            'comic': comic.file.url,
+            "status_message": "ok",
+            "comic": comic.file.url,
         }
+        # Remove to spare storage
+        video.file.delete()
         return Response(response)

comic_layout/comic_layout.py

@@ -1,7 +1,30 @@
-from PIL import Image
 import cv2
+import numpy as np
+
 
 class LayoutGenerator():
+    @classmethod
+    def get_layout(cls, frames):
+        result_imgs = cls._pad_images(frames)
+
+        first_row = np.hstack(result_imgs[:2])
+        second_row = np.hstack(result_imgs[2:5])
+        third_row = np.hstack(result_imgs[5:7])
+        fourth_row = np.hstack(result_imgs[7:10])
+
+        second_row = cv2.resize(second_row,
+                                (first_row.shape[1],
+                                 int(second_row.shape[0] * first_row.shape[1] / second_row.shape[1])))
+        fourth_row = cv2.resize(fourth_row,
+                                (first_row.shape[1],
+                                 int(fourth_row.shape[0] * first_row.shape[1] / fourth_row.shape[1])))
+
+        return np.vstack([first_row, second_row, third_row, fourth_row])
+
     @staticmethod
-    def get_layout(frames):
-        return cv2.imread('tmp/test.jpg')
+    def _pad_images(frames):
+        padded_result_imgs = []
+        for img in frames:
+            padded_img = cv2.copyMakeBorder(img, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=(255, 255, 255))
+            padded_result_imgs.append(padded_img)
+        return padded_result_imgs

dockerfile

@@ -1,12 +1,22 @@
-FROM python:3.6
+FROM nvidia/cuda:9.0-cudnn7-runtime
+
+RUN apt-get update && apt-get install -y apt-utils software-properties-common && \
+    add-apt-repository ppa:jonathonf/python-3.6 && \
+    apt-get update && apt-get -y install python3 python3-pip python3.6 python3.6-dev python3-pip python3.6-venv vim ffmpeg \
+    build-essential cmake git libgtk2.0-dev pkg-config libavcodec-dev \
+    libavformat-dev libswscale-dev && \
+    python3.6 -m pip install --upgrade pip && \
+    python3.6 -m pip install jupyter ipywidgets jupyterlab && \
+    python3.6 -m pip install tensorflow-gpu h5py keras && \
+    python3.6 -m pip install scikit-image opencv-contrib-python
 
 RUN mkdir /comixify
 WORKDIR /comixify
 COPY . /comixify
-RUN pip install --upgrade pip && pip install -r requirements.txt
+RUN python3.6 -m pip install -r requirements.txt
 
 # Port to expose
 EXPOSE 8080
 
 ENTRYPOINT ["sh", "entrypoint.sh"]
-CMD ['start']
+CMD ['start']

entrypoint.sh

@@ -1,15 +1,15 @@
 #!/bin/sh
 
 start (){
-    gunicorn --bind :8080 settings.wsgi:application
+    gunicorn --bind :8080 settings.wsgi:application --timeout 300
 }
 
 migrate (){
-    python3 manage.py migrate --noinput
+    python3.6 manage.py migrate --noinput
 }
 
 collectstatic (){
-    python3 manage.py collectstatic --clear --no-input
+    python3.6 manage.py collectstatic --clear --no-input
 }
 
 migrate

keyframes/keyframes.py

@@ -1,4 +1,48 @@
+import os
+import shutil
+import uuid
+from subprocess import call
+
+import cv2
+from django.conf import settings
+
+from utils import jj
+
+
 class KeyFramesExtractor():
+    @classmethod
+    def get_keyframes(cls, video):
+        all_keyframes, all_frames_tmp_dir = cls._get_all_frames(video)
+        interval = cls._count_interval(all_keyframes)
+        chosen_frames = cls._get_frames_with_interval(interval, all_keyframes)
+
+        shutil.rmtree(jj(f"{settings.TMP_DIR}", f"{all_frames_tmp_dir}"))
+        return chosen_frames
+
     @staticmethod
-    def get_keyframes(video):
-        return []
+    def _get_all_frames(video):
+        all_frames_tmp_dir = uuid.uuid4()
+        os.mkdir(jj(f"{settings.TMP_DIR}", f"{all_frames_tmp_dir}"))
+        call(["ffmpeg", "-skip_frame", "nokey", "-i", f"{video.file.path}", "-vsync", "0", "-qscale:v", "1",
+              "-f", "image2", jj(f"{settings.TMP_DIR}", f"{all_frames_tmp_dir}", "%06d.jpeg")])
+
+        frames_paths = []
+        for dirname, dirnames, filenames in os.walk(jj(f"{settings.TMP_DIR}", f"{all_frames_tmp_dir}")):
+            for filename in filenames:
+                frames_paths.append(jj(dirname, filename))
+        return sorted(frames_paths), all_frames_tmp_dir
+
+    @staticmethod
+    def _count_interval(all_keyframes):
+        return int((len(all_keyframes) - settings.NUMBERS_OF_FRAMES_TO_SHOW) / (settings.NUMBERS_OF_FRAMES_TO_SHOW + 1))
+
+    @staticmethod
+    def _get_frames_with_interval(interval, all_keyframes):
+        chosen_frames = []
+        chosen_frames_tmp_dir = uuid.uuid4()
+        os.mkdir(jj(f"{settings.TMP_DIR}", f"{chosen_frames_tmp_dir}"))
+
+        for i in range(settings.NUMBERS_OF_FRAMES_TO_SHOW):
+            frame = cv2.imread(all_keyframes[(i + 1) * interval])
+            chosen_frames.append(frame)
+        return chosen_frames

requirements.txt

@@ -2,8 +2,11 @@ Django==2.0.7
 django-rest-framework==0.1.0
 djangorestframework==3.8.2
 gunicorn==19.9.0
-numpy==1.15.0
+numpy==1.14.5
 opencv-python==3.4.2.17
 Pillow==5.2.0
 psycopg2==2.7.5
 pytz==2018.5
+six==1.11.0
+torch==0.4.1
+torchvision==0.2.1

settings/settings.py

@@ -135,3 +135,6 @@ MEDIA_ROOT = os.path.join(BASE_DIR, 'media')
 
 PERMITTED_VIDEO_EXTENSIONS = ['mp4', 'avi']
 MAX_FILE_SIZE = 50000000
+NUMBERS_OF_FRAMES_TO_SHOW = 10
+TMP_DIR = 'tmp/'
+GPU = True

style_transfer/style_transfer.py

@@ -1,4 +1,64 @@
+import os
+
+import cv2
+import numpy as np
+import torch
+import torchvision.transforms as transforms
+from django.conf import settings
+from torch.autograd import Variable
+
+from CartoonGAN.network.Transformer import Transformer
+
+
 class StyleTransfer():
+    @classmethod
+    def get_stylized_frames(cls, frames, method="cartoon_gan", gpu=settings.GPU, **kwargs):
+        if method == "cartoon_gan":
+            return cls._cartoon_gan_stylize(frames, gpu=gpu, **kwargs)
+
     @staticmethod
-    def get_stylized_frames(frames):
-        return frames
+    def _cartoon_gan_stylize(frames, gpu=True, **kwargs):
+        style = kwargs.get("style", "Hosoda")
+        resize = kwargs.get("resize", 450)
+
+        # load pretrained model
+        model = Transformer()
+        model.load_state_dict(torch.load(os.path.join("CartoonGAN/pretrained_model", style + "_net_G_float.pth")))
+        model.eval()
+        model.cuda() if gpu else model.float()
+
+        stylized_imgs = []
+        for img in frames:
+            # resize image, keep aspect ratio
+            h, w, _ = img.shape
+            ratio = h * 1.0 / w
+            if ratio > 1:
+                h = resize
+                w = int(h * 1.0 / ratio)
+            else:
+                w = resize
+                h = int(w * ratio)
+            input_image = cv2.resize(img, (w, h), interpolation=cv2.INTER_CUBIC)
+
+            # RGB -> BGR
+            input_image = input_image[:, :, [2, 1, 0]]
+            input_image = transforms.ToTensor()(input_image).unsqueeze(0)
+
+            # preprocess, (-1, 1)
+            input_image = -1 + 2 * input_image
+            input_image = Variable(input_image).cuda() if gpu else Variable(input_image).float()
+
+            # forward
+            output_image = model(input_image)
+            output_image = output_image[0]
+
+            # BGR -> RGB
+            output_image = output_image[[2, 1, 0], :, :]
+
+            # deprocess, (0, 1)
+            output_image = output_image.data.cpu().float() * 0.5 + 0.5
+
+            # switch channels and append image to result images
+            stylized_imgs.append(np.rollaxis(output_image.numpy(), 0, 3))
+
+        return stylized_imgs

utils.py

@@ -0,0 +1,5 @@
+import os
+
+
+def jj(*args):
+    return os.path.join(*args)