diff --git a/api/models.py b/api/models.py index 52a47a4..ea5e219 100644 --- a/api/models.py +++ b/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) diff --git a/api/views.py b/api/views.py index 03a9e8a..54343bd 100644 --- a/api/views.py +++ b/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() diff --git a/comic_layout/comic_layout.py b/comic_layout/comic_layout.py index af478a1..27a9dad 100644 --- a/comic_layout/comic_layout.py +++ b/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) diff --git a/keyframes/keyframes.py b/keyframes/keyframes.py index 72e7088..caf8635 100644 --- a/keyframes/keyframes.py +++ b/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) diff --git a/style_transfer/style_transfer.py b/style_transfer/style_transfer.py index a61a7f5..e4b5c45 100644 --- a/style_transfer/style_transfer.py +++ b/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) diff --git a/utils.py b/utils.py index c2a0df0..cb738f7 100644 --- a/utils.py +++ b/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