Add downloading videos from YouTube (#7)

- Add downloading videos from YouTube
- Minor cleanup

README.md

@@ -2,7 +2,8 @@
 
 #TODO:
 - [ ] Prepare readme
-- [ ] Fix model loading to memory
+- [x] Download videos directly from YouTube
+- [x] Fix model loading to memory
 - [ ] Replace ffmpeg with some keyframes extractor
 - [ ] Prepare own ComixGAN
 - [ ] Prepare some more intelligent method for comic layout

api/models.py

@@ -2,15 +2,43 @@ import os
 import uuid
 
 import cv2
+import pafy
 from django.conf import settings
 from django.core.files import File
 from django.db import models
 
+from api.exceptions import TooLargeFile
+from comic_layout.comic_layout import LayoutGenerator
+from keyframes.keyframes import KeyFramesExtractor
+from style_transfer.style_transfer import StyleTransfer
+from utils import jj
+
 
 class Video(models.Model):
     file = models.FileField(blank=False, null=False, upload_to="raw_videos")
     timestamp = models.DateTimeField(auto_now_add=True)
 
+    def download_from_youtube(self, yt_url):
+        yt_pafy = pafy.new(yt_url)
+
+        # Use the biggest possible quality with file size < MAX_FILE_SIZE and resolution <= 480px
+        for stream in yt_pafy.videostreams:
+            if stream.get_filesize() < settings.MAX_FILE_SIZE and int(stream.quality.split("x")[1]) <= 480:
+                tmp_name = uuid.uuid4().hex + ".mp4"
+                relative_path = jj('raw_videos', tmp_name)
+                full_path = jj(settings.MEDIA_ROOT, relative_path)
+                stream.download(full_path)
+                self.file.name = relative_path
+                break
+        else:
+            raise TooLargeFile()
+
+    def create_comic(self):
+        keyframes = KeyFramesExtractor.get_keyframes(video=self)
+        stylized_keyframes = StyleTransfer.get_stylized_frames(frames=keyframes)
+        comic_image = LayoutGenerator.get_layout(frames=stylized_keyframes)
+        return comic_image
+
 
 class Comic(models.Model):
     file = models.FileField(blank=False, null=False, upload_to="comic")
@@ -20,10 +48,10 @@ class Comic(models.Model):
     def create_from_nparray(cls, nparray_file, video):
         if nparray_file.max() <= 1:
             nparray_file = (nparray_file * 255).astype(int)
-        tmp_name = uuid.uuid4()
+        tmp_name = uuid.uuid4().hex + ".png"
-        cv2.imwrite(f"{settings.TMP_DIR}{tmp_name}.png", nparray_file)
+        cv2.imwrite(jj(settings.TMP_DIR, tmp_name), nparray_file)
-        with open(f"{settings.TMP_DIR}{tmp_name}.png", mode="rb") as tmp_file:
+        with open(jj(settings.TMP_DIR, tmp_name), mode="rb") as tmp_file:
-            comic_image = File(tmp_file, name=f"{tmp_name}.png")
+            comic_image = File(tmp_file, name=tmp_name)
             comic = Comic.objects.create(file=comic_image, video=video)
-        os.remove(f"{settings.TMP_DIR}{tmp_name}.png")
+        os.remove(jj(settings.TMP_DIR, tmp_name))
         return comic

api/serializers.py

@@ -17,3 +17,7 @@ class VideoSerializer(serializers.ModelSerializer):
         if file.size > settings.MAX_FILE_SIZE:
             raise TooLargeFile
         return attrs
+
+
+class YouTubeDownloadSerializer(serializers.Serializer):
+    url = serializers.URLField()

api/urls.py

@@ -1,8 +1,8 @@
-from django.conf.urls import url
+from django.urls import path
-
-from .views import Comixify
 
+from .views import Comixify, ComixifyFromYoutube
 
 urlpatterns = [
-    url(r'^$', Comixify.as_view(), name='annotate'),
+    path(r'', Comixify.as_view(), name='comixify'),
+    path(r'from_yt/', ComixifyFromYoutube.as_view(), name='comixify_from_yt'),
 ]

api/views.py

@@ -2,11 +2,8 @@ from rest_framework.parsers import FormParser, MultiPartParser
 from rest_framework.response import Response
 from rest_framework.views import APIView
 
-from comic_layout.comic_layout import LayoutGenerator
-from keyframes.keyframes import KeyFramesExtractor
-from style_transfer.style_transfer import StyleTransfer
 from .models import Video, Comic
-from .serializers import VideoSerializer
+from .serializers import VideoSerializer, YouTubeDownloadSerializer
 
 
 class Comixify(APIView):
@@ -22,12 +19,35 @@ class Comixify(APIView):
 
         video_file = serializer.validated_data["file"]
         video = Video.objects.create(file=video_file)
-
+        comic_image = video.create_comic()
-        keyframes = KeyFramesExtractor.get_keyframes(video=video)
-        stylized_keyframes = StyleTransfer.get_stylized_frames(frames=keyframes)
-        comic_image = LayoutGenerator.get_layout(frames=stylized_keyframes)
-
         comic = Comic.create_from_nparray(comic_image, video)
+
+        response = {
+            "status_message": "ok",
+            "comic": comic.file.url,
+        }
+        # Remove to spare storage
+        video.file.delete()
+        return Response(response)
+
+
+class ComixifyFromYoutube(APIView):
+
+    def post(self, request):
+        """
+        Receives video, and returns comic image
+        """
+
+        serializer = YouTubeDownloadSerializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        yt_url = serializer.validated_data["url"]
+
+        video = Video()
+        video.download_from_youtube(yt_url)
+        video.save()
+        comic_image = video.create_comic()
+        comic = Comic.create_from_nparray(comic_image, video)
+
         response = {
             "status_message": "ok",
             "comic": comic.file.url,

keyframes/keyframes.py

@@ -1,11 +1,12 @@
 import os
 import uuid
+
 import numpy as np
 import torch
 import torch.nn as nn
-os.environ['GLOG_minloglevel'] = '2' # Prevent caffe shell loging
+
+os.environ['GLOG_minloglevel'] = '2'  # Prevent caffe shell loging
 import caffe
-from datetime import datetime
 from subprocess import call
 from math import ceil
 from sklearn.preprocessing import normalize
@@ -35,12 +36,12 @@ class KeyFramesExtractor:
 
     @staticmethod
     def _get_all_frames(video):
-        all_frames_tmp_dir = uuid.uuid4()
+        all_frames_tmp_dir = uuid.uuid4().hex
-        os.mkdir(jj(f"{settings.TMP_DIR}", f"{all_frames_tmp_dir}"))
+        os.mkdir(jj(settings.TMP_DIR, all_frames_tmp_dir))
-        call(["ffmpeg", "-i", f"{video.file.path}", "-vf", "select=not(mod(n\\,15))", "-vsync", "vfr", "-q:v", "2",
+        call(["ffmpeg", "-i", video.file.path, "-vf", "select=not(mod(n\\,15))", "-vsync", "vfr", "-q:v", "2",
-            jj(f"{settings.TMP_DIR}", f"{all_frames_tmp_dir}", "%06d.jpeg")])
+              jj(settings.TMP_DIR, 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 dirname, dirnames, filenames in os.walk(jj(settings.TMP_DIR, all_frames_tmp_dir)):
             for filename in filenames:
                 frames_paths.append(jj(dirname, filename))
         return sorted(frames_paths), all_frames_tmp_dir
@@ -89,7 +90,7 @@ class KeyFramesExtractor:
             features[idx_batch * batch_size:idx_batch * batch_size + n_batch] = temp
         normalize(features, copy=False)
         return features.astype(np.float32)
-        
+
     @staticmethod
     def _get_probs(features, gpu=True):
         model_cache_key = "keyframes_rl_model_cache"

nginx.conf

@@ -22,6 +22,10 @@ server {
         proxy_set_header Host $host;
         proxy_redirect off;
         proxy_pass http://hello_server;
+        proxy_connect_timeout 300;
+        proxy_send_timeout 300;
+        proxy_read_timeout 300;
+        send_timeout 300;
     }
     location ^~ /.well-known {
         allow all;

requirements.txt

@@ -4,10 +4,12 @@ djangorestframework==3.8.2
 gunicorn==19.9.0
 numpy==1.14.5
 opencv-python==3.4.2.17
+pafy==0.5.4
 Pillow==5.2.0
 psycopg2==2.7.5
 pytz==2018.5
 six==1.11.0
 torch==0.4.1
 torchvision==0.2.1
-scikit-learn==0.19.2
+scikit-learn==0.19.2
+youtube-dl==2018.9.18