diff --git a/CartoonGAN/network/Transformer.py b/CartoonGAN/network/Transformer.py new file mode 100644 index 0000000..02623a1 --- /dev/null +++ b/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 diff --git a/CartoonGAN/network/__init__.py b/CartoonGAN/network/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/CartoonGAN/pretrained_model/Hayao_net_G_float.pth b/CartoonGAN/pretrained_model/Hayao_net_G_float.pth new file mode 100644 index 0000000..fb83ab0 Binary files /dev/null and b/CartoonGAN/pretrained_model/Hayao_net_G_float.pth differ diff --git a/CartoonGAN/pretrained_model/Hayao_net_G_float.t7 b/CartoonGAN/pretrained_model/Hayao_net_G_float.t7 new file mode 100644 index 0000000..bc32954 Binary files /dev/null and b/CartoonGAN/pretrained_model/Hayao_net_G_float.t7 differ diff --git a/CartoonGAN/pretrained_model/Hosoda_net_G_float.pth b/CartoonGAN/pretrained_model/Hosoda_net_G_float.pth new file mode 100644 index 0000000..27fdcd9 Binary files /dev/null and b/CartoonGAN/pretrained_model/Hosoda_net_G_float.pth differ diff --git a/CartoonGAN/pretrained_model/Hosoda_net_G_float.t7 b/CartoonGAN/pretrained_model/Hosoda_net_G_float.t7 new file mode 100644 index 0000000..104b70f Binary files /dev/null and b/CartoonGAN/pretrained_model/Hosoda_net_G_float.t7 differ diff --git a/CartoonGAN/pretrained_model/Paprika_net_G_float.pth b/CartoonGAN/pretrained_model/Paprika_net_G_float.pth new file mode 100644 index 0000000..fb5976b Binary files /dev/null and b/CartoonGAN/pretrained_model/Paprika_net_G_float.pth differ diff --git a/CartoonGAN/pretrained_model/Paprika_net_G_float.t7 b/CartoonGAN/pretrained_model/Paprika_net_G_float.t7 new file mode 100644 index 0000000..45c833f Binary files /dev/null and b/CartoonGAN/pretrained_model/Paprika_net_G_float.t7 differ diff --git a/CartoonGAN/pretrained_model/Shinkai_net_G_float.pth b/CartoonGAN/pretrained_model/Shinkai_net_G_float.pth new file mode 100644 index 0000000..1900bab Binary files /dev/null and b/CartoonGAN/pretrained_model/Shinkai_net_G_float.pth differ diff --git a/CartoonGAN/pretrained_model/Shinkai_net_G_float.t7 b/CartoonGAN/pretrained_model/Shinkai_net_G_float.t7 new file mode 100644 index 0000000..f26db3e Binary files /dev/null and b/CartoonGAN/pretrained_model/Shinkai_net_G_float.t7 differ diff --git a/CartoonGAN/pretrained_model/download_pth.sh b/CartoonGAN/pretrained_model/download_pth.sh new file mode 100644 index 0000000..5876d8a --- /dev/null +++ b/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 .. diff --git a/CartoonGAN/pretrained_model/download_t7.sh b/CartoonGAN/pretrained_model/download_t7.sh new file mode 100644 index 0000000..8f3e801 --- /dev/null +++ b/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 .. diff --git a/CartoonGAN/src/InstanceNormalization.lua b/CartoonGAN/src/InstanceNormalization.lua new file mode 100644 index 0000000..35d17eb --- /dev/null +++ b/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 diff --git a/CartoonGAN/src/util.lua b/CartoonGAN/src/util.lua new file mode 100755 index 0000000..759877f --- /dev/null +++ b/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 diff --git a/api/models.py b/api/models.py index 2a538fa..12afb59 100644 --- a/api/models.py +++ b/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 diff --git a/api/serializers.py b/api/serializers.py index 31f912b..2bebcef 100644 --- a/api/serializers.py +++ b/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 diff --git a/api/views.py b/api/views.py index e524f7f..9c1247d 100644 --- a/api/views.py +++ b/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) diff --git a/comic_layout/comic_layout.py b/comic_layout/comic_layout.py index c34e5af..9a9842e 100644 --- a/comic_layout/comic_layout.py +++ b/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 diff --git a/dockerfile b/dockerfile index 7a1d775..b081dec 100644 --- a/dockerfile +++ b/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'] \ No newline at end of file +CMD ['start'] diff --git a/entrypoint.sh b/entrypoint.sh index 3d5f215..4007f19 100644 --- a/entrypoint.sh +++ b/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 diff --git a/keyframes/keyframes.py b/keyframes/keyframes.py index 4e8d46a..d12695b 100644 --- a/keyframes/keyframes.py +++ b/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 diff --git a/requirements.txt b/requirements.txt index 6a99bc6..44b32ff 100644 --- a/requirements.txt +++ b/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 diff --git a/settings/settings.py b/settings/settings.py index 7601f6d..bd3ff6c 100644 --- a/settings/settings.py +++ b/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 diff --git a/style_transfer/style_transfer.py b/style_transfer/style_transfer.py index 0f34a89..6c30aff 100644 --- a/style_transfer/style_transfer.py +++ b/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 diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..c2a0df0 --- /dev/null +++ b/utils.py @@ -0,0 +1,5 @@ +import os + + +def jj(*args): + return os.path.join(*args)