Add profiling (#11)

* Add profiling

* Fix decorators order

api/models.py

@@ -11,13 +11,14 @@ 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
+from utils import jj, profile
 
 
 class Video(models.Model):
     file = models.FileField(blank=False, null=False, upload_to="raw_videos")
     timestamp = models.DateTimeField(auto_now_add=True)
 
+    @profile
     def download_from_youtube(self, yt_url):
         yt_pafy = pafy.new(yt_url)
 
@@ -34,10 +35,16 @@ class Video(models.Model):
             raise TooLargeFile()
 
     def create_comic(self, frames_mode=0, rl_mode=0):
-        keyframes = KeyFramesExtractor.get_keyframes(video=self, frames_mode=frames_mode, rl_mode=rl_mode)
-        stylized_keyframes = StyleTransfer.get_stylized_frames(frames=keyframes)
-        comic_image = LayoutGenerator.get_layout(frames=stylized_keyframes)
-        return comic_image
+        keyframes, keyframes_extraction_time = KeyFramesExtractor.get_keyframes(video=self,
+                                                                                frames_mode=frames_mode,
+                                                                                rl_mode=rl_mode)
+        stylized_keyframes, stylization_time = StyleTransfer.get_stylized_frames(frames=keyframes)
+        comic_image, layout_generation_time = LayoutGenerator.get_layout(frames=stylized_keyframes)
+
+        timings = {'keyframes_extraction_time': keyframes_extraction_time,
+                   'stylization_time': stylization_time,
+                   'layout_generation_time': layout_generation_time}
+        return comic_image, timings
 
 
 class Comic(models.Model):
@@ -45,6 +52,7 @@ class Comic(models.Model):
     video = models.ForeignKey(Video, on_delete=models.CASCADE, related_name="comic")
 
     @classmethod
+    @profile
     def create_from_nparray(cls, nparray_file, video):
         if nparray_file.max() <= 1:
             nparray_file = (nparray_file * 255).astype(int)

api/views.py

@@ -19,15 +19,17 @@ class Comixify(APIView):
 
         video_file = serializer.validated_data["file"]
         video = Video.objects.create(file=video_file)
-        comic_image = video.create_comic(
+        comic_image, timings = video.create_comic(
             frames_mode=serializer.validated_data["frames_mode"],
             rl_mode=serializer.validated_data["rl_mode"]
         )
-        comic = Comic.create_from_nparray(comic_image, video)
+        comic, from_nparray_time = Comic.create_from_nparray(comic_image, video)
+        timings['from_nparray_time'] = from_nparray_time
 
         response = {
             "status_message": "ok",
             "comic": comic.file.url,
+            "timings": timings,
         }
         # Remove to spare storage
         video.file.delete()
@@ -46,17 +48,20 @@ class ComixifyFromYoutube(APIView):
         yt_url = serializer.validated_data["url"]
 
         video = Video()
-        video.download_from_youtube(yt_url)
+        _, yt_download_time = video.download_from_youtube(yt_url)
         video.save()
-        comic_image = video.create_comic(
+        comic_image, timings = video.create_comic(
             frames_mode=serializer.validated_data["frames_mode"],
             rl_mode=serializer.validated_data["rl_mode"]
         )
-        comic = Comic.create_from_nparray(comic_image, video)
+        comic, from_nparray_time = Comic.create_from_nparray(comic_image, video)
+        timings['from_nparray_time'] = from_nparray_time
+        timings['yt_download_time'] = yt_download_time
 
         response = {
             "status_message": "ok",
             "comic": comic.file.url,
+            "timings": timings,
         }
         # Remove to spare storage
         video.file.delete()

comic_layout/comic_layout.py

@@ -1,9 +1,12 @@
 import cv2
 import numpy as np
 
+from utils import profile
+
 
 class LayoutGenerator():
     @classmethod
+    @profile
     def get_layout(cls, frames):
         result_imgs = cls._pad_images(frames)
 

keyframes/keyframes.py

@@ -15,7 +15,7 @@ from django.core.cache import cache
 from skimage import img_as_ubyte
 import logging
 
-from utils import jj
+from utils import jj, profile
 from keyframes_rl.models import DSN
 from popularity.models import PopularityPredictor
 from keyframes.kts import cpd_auto
@@ -26,6 +26,7 @@ logger = logging.getLogger(__name__)
 
 class KeyFramesExtractor:
     @classmethod
+    @profile
     def get_keyframes(cls, video, gpu=settings.GPU, features_batch_size=settings.FEATURE_BATCH_SIZE,
                       frames_mode=0, rl_mode=0):
         frames_paths, all_frames_tmp_dir = cls._get_all_frames(video, mode=frames_mode)

style_transfer/style_transfer.py

@@ -9,10 +9,12 @@ from django.core.cache import cache
 from torch.autograd import Variable
 
 from CartoonGAN.network.Transformer import Transformer
+from utils import profile
 
 
 class StyleTransfer():
     @classmethod
+    @profile
     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)

utils.py

@@ -1,5 +1,32 @@
 import os
+import time
+from functools import wraps
 
 
 def jj(*args):
     return os.path.join(*args)
+
+
+class Timer(object):
+    def __init__(self, verbose=False):
+        self.verbose = verbose
+
+    def __enter__(self):
+        self.start = time.time()
+        return self
+
+    def __exit__(self, *args):
+        self.end = time.time()
+        self.secs = self.end - self.start
+        self.msecs = self.secs * 1000  # millisecs
+
+
+def profile(fn):
+    @wraps(fn)
+    def with_profiling(*args, **kwargs):
+        with Timer() as t:
+            ret = fn(*args, **kwargs)
+
+        return ret, t.secs
+
+    return with_profiling