GH-23: Add support for JSON thumbnail generation (GH-24)

* Fix the time format at VTT for floating intervals
* Remove options' setters and provide through the constructor
* Implement Factory for creating `ThumbnailFormat`
* Implement the strategy for `VTT` generation
* Implement the strategy for `JSON` generation
* Change output file format to fix the order on sort
* Update mock data of the JSON demo

Author: ArtyomVancyan

requirements.txt

@@ -1,3 +1,4 @@
+click>=8.0.3
 imageio-ffmpeg>=0.4.7
 imageio>=2.23.0
 pillow>=8.4.0

thumbnails.py

@@ -7,22 +7,19 @@
 interval = .5
 basepath = "/stc/"
 
-files = ["valerian-1080p.avi", "valerian-1080p.mkv", "valerian-1080p.mov", "valerian-1080p.mp4",
-         "valerian-1080p.webm", "valerian-1080p.wmv", "valerian-1080p.mpeg", "valerian-1080p.mpg", "valerian-1080p.ogv"]
+files = ["valerian-1080p.avi"] #, "valerian-1080p.mkv", "valerian-1080p.mov", "valerian-1080p.mp4",
+         # "valerian-1080p.webm", "valerian-1080p.wmv", "valerian-1080p.mpeg", "valerian-1080p.mpg", "valerian-1080p.ogv"]
 
 
 def worker(video):
-    video.compress = compress
-    video.interval = interval
-    video.basepath = basepath
     video.extract_frames()
     video.join_frames()
     video.to_vtt()
 
 
 def main():
     with concurrent.futures.ProcessPoolExecutor() as executor:
-        executor.map(worker, map(Thumbnails, files))
+        executor.map(worker, (Thumbnails(file, compress, interval, basepath) for file in files))
 
 
 if __name__ == "__main__":

thumbnails/__init__.py

@@ -1 +1,95 @@
+import json
+import os
+import shutil
+from datetime import timedelta
+from distutils.dir_util import copy_tree
+
+from PIL import Image
+
+from .formatter import ThumbnailFactory
+from .formatter import ThumbnailFormat
 from .thumbnails import Thumbnails
+from .thumbnails import arange
+
+DEFAULT_AS = "vtt"
+DEFAULT_COMPRESS = 1.
+DEFAULT_INTERVAL = 1.
+DEFAULT_BASEPATH = ""
+
+
+def register_format(typename):
+    def _registrator(cls: ThumbnailFormat):
+        cls.extension = typename
+        ThumbnailFactory.thumbnails[typename] = cls
+
+    return _registrator
+
+
+@register_format("vtt")
+class VTT(ThumbnailFormat):
+    def __init__(self, video):
+        super().__init__(video)
+        self._master_name = self.filename + ".png"
+
+    def prepare_thumbnails(self):
+        _thumbnails = self.video.thumbnails(True)
+        master = Image.new(mode="RGBA", size=next(_thumbnails))
+
+        for frame, start, end, x, y in self.video.thumbnails():
+            with Image.open(frame) as image:
+                image = image.resize((self.width, self.height), Image.ANTIALIAS)
+                master.paste(image, (x, y))
+
+        master.save(self._master_name)
+        self.tempdir.cleanup()
+
+    def generate(self):
+        def _format_time(secs):
+            delta = timedelta(seconds=secs)
+            return ("0%s.000" % delta)[:12]
+
+        _lines = ["WEBVTT\n\n"]
+        _img_src = self.basepath + self._master_name
+
+        for frame, start, end, x, y in self.video.thumbnails():
+            _thumbnail = "%s --> %s\n%s#xywh=%d,%d,%d,%d\n\n" % (
+                _format_time(start), _format_time(end),
+                _img_src, x, y, self.width, self.height
+            )
+            _lines.append(_thumbnail)
+
+        with open(self.output_format, "w") as fp:
+            fp.writelines(_lines)
+
+
+@register_format("json")
+class JSON(ThumbnailFormat):
+    def __init__(self, video):
+        super().__init__(video)
+        self._outdir = "outdir"  # temp dirname
+
+    def prepare_thumbnails(self):
+        if os.path.isdir(self._outdir):
+            shutil.rmtree(self._outdir)
+        copy_tree(self.tempdir.name, self._outdir)
+        self.tempdir.cleanup()
+
+    def generate(self):
+        _content = {}
+
+        for frame, start, end, x, y in self.video.thumbnails():
+            frame = self._outdir + os.sep + os.path.split(frame)[1]
+            with Image.open(frame) as image:
+                image.resize((self.width, self.height), Image.ANTIALIAS).save(frame)
+                _thumbnail = {
+                    "src": self.basepath + frame,
+                    "width": "%spx" % self.width,
+                }
+                _content[int(start)] = _thumbnail
+
+        with open(self.output_format, "w") as fp:
+            json.dump(_content, fp, indent=2)
+
+
+__version__ = "v1.0"
+__all__ = (Thumbnails,)

thumbnails/__main__.py

@@ -0,0 +1,68 @@
+import concurrent.futures
+import functools
+import os
+
+import click
+
+from . import DEFAULT_AS
+from . import DEFAULT_BASEPATH
+from . import DEFAULT_COMPRESS
+from . import DEFAULT_INTERVAL
+from . import ThumbnailFactory
+from . import Thumbnails
+from . import __version__
+
+
+def worker(video, as_):
+    video.extract_frames()
+    formatter = ThumbnailFactory.get_thumbnail(as_, video)
+    formatter.prepare_thumbnails()
+    formatter.generate()
+
+
+class _ThumbnailsCLI(click.Command):
+    def format_usage(self, ctx, formatter):
+        usages = (
+            "[OPTIONS] INPUT_DIR OUTPUT_DIR",
+            "[OPTIONS] INPUT_FILE OUTPUT_FILE",
+            "[OPTIONS] INPUT_FILES... OUTPUT_DIR",
+        )
+        formatter.write_usage(ctx.command_path, "\n\t\t\t".join(usages), prefix="Usages: ")
+
+
+@click.command(cls=_ThumbnailsCLI)
+@click.option("--as", "-F", default=DEFAULT_AS, help="Output format. Default is %s." % DEFAULT_AS,
+              type=click.Choice(ThumbnailFactory.thumbnails.keys(), case_sensitive=False))
+@click.option("--compress", "-C", default=DEFAULT_COMPRESS, help="The image scale coefficient. A number from 0 to 1.")
+@click.option("--interval", "-I", default=DEFAULT_INTERVAL, help="The interval between neighbor thumbnails in seconds.")
+@click.option("--basepath", "-B", default=DEFAULT_BASEPATH, help="The prefix of the thumbnails path can be customized.")
+@click.argument("inputs", required=True, type=click.Path(), nargs=-1)
+@click.argument("output", required=True, type=click.Path(), nargs=1)
+@click.version_option(__version__)
+def thumbnails_cli(compress, interval, basepath, inputs, output, **kwargs):
+    """TODO: Add more description about particular usages."""
+    as_ = kwargs.pop("as")
+    output_is_directory = all((len(inputs) > 1, *map(os.path.isfile, inputs))) or os.path.isdir(inputs[0])
+
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        videos = executor.map(
+            functools.partial(
+                Thumbnails,
+                compress=compress,
+                interval=interval,
+                basepath=basepath
+            ),
+            inputs,
+        )
+
+    with concurrent.futures.ProcessPoolExecutor() as executor:
+        executor.map(functools.partial(worker, as_=as_), videos)
+
+
+# @click.confirmation_option("--overwrite", "-y", prompt="Are you sure you want to overwrite the existing output files?")
+# def overwrite():
+#     print("overwritten")
+
+
+if __name__ == "__main__":
+    thumbnails_cli()

thumbnails/formatter.py

@@ -0,0 +1,31 @@
+class ThumbnailFormat:
+    extension = None
+
+    def __init__(self, video):
+        self.video = video
+
+    def __getattr__(self, item):
+        return getattr(self.video, item)
+
+    @property
+    def output_format(self):
+        return "%s.%s" % (self.filename, self.extension)
+
+    def prepare_thumbnails(self):
+        """Prepare the thumbnails before generating the output."""
+        raise NotImplementedError
+
+    def generate(self):
+        """Generate the thumbnails for the given video."""
+        raise NotImplementedError
+
+
+class ThumbnailFactory:
+    thumbnails = {}
+
+    @classmethod
+    def get_thumbnail(cls, typename, *args, **kwargs) -> ThumbnailFormat:
+        try:
+            return cls.thumbnails[typename](*args, **kwargs)
+        except KeyError:
+            raise ValueError("Thumbnail type '%s' is not supported." % typename)

thumbnails/thumbnails.py

@@ -6,7 +6,6 @@
 from datetime import timedelta
 from tempfile import TemporaryDirectory
 
-from PIL import Image
 from imageio_ffmpeg import get_ffmpeg_exe
 
 from .ffmpeg import _FFMpeg
@@ -56,15 +55,16 @@ def height(self):
 
 
 class Thumbnails(_ThumbnailMixin, _FFMpeg):
-    def __init__(self, filename):
-        self.__compress = 1.
-        self.__interval = 1.
-        self.__basepath = ""
-        self.thumbnails = []
-        self.tempdir = TemporaryDirectory()
+    def __init__(self, filename, compress, interval, basepath):
+        self.__compress = float(compress)
+        self.__interval = float(interval)
+        self.__basepath = basepath
+
+        if self.__compress <= 0 or self.__compress > 1:
+            raise ValueError("Compress must be between 0 and 1.")
+
         self.filename = filename
-        self._vtt_name = filename + ".vtt"
-        self._image_name = filename + ".png"
+        self.tempdir = TemporaryDirectory()
 
         _FFMpeg.__init__(self, filename)
         _ThumbnailMixin.__init__(self, self.size)
@@ -73,43 +73,25 @@ def __init__(self, filename):
     def compress(self):
         return self.__compress
 
-    @compress.setter
-    def compress(self, value):
-        try:
-            self.__compress = float(value)
-        except ValueError:
-            raise ValueError("Compress must be a number.")
-
     @property
     def interval(self):
         return self.__interval
 
-    @interval.setter
-    def interval(self, value):
-        try:
-            self.__interval = float(value)
-        except ValueError:
-            raise ValueError("Interval must be a number.")
-
     @property
     def basepath(self):
         return self.__basepath
 
-    @basepath.setter
-    def basepath(self, value):
-        self.__basepath = value
-
     @staticmethod
-    def _calc_columns(frames_count, width, height):
+    def calc_columns(frames_count, width, height):
         ratio = 16 / 9
         for col in range(1, frames_count):
             if (col * width) / (frames_count // col * height) > ratio:
                 return col
 
     def _extract_frame(self, start_time):
         _input_file = self.filename
-        _output_file = "%s/%s-%s.png" % (self.tempdir.name, start_time, self.filename)
         _timestamp = str(timedelta(seconds=start_time))
+        _output_file = "%s/%s-%s.png" % (self.tempdir.name, _timestamp, self.filename)
 
         cmd = (
             ffmpeg_bin,
@@ -128,48 +110,24 @@ def extract_frames(self):
         with concurrent.futures.ThreadPoolExecutor() as executor:
             executor.map(self._extract_frame, _intervals)
 
-    def join_frames(self):
+    def thumbnails(self, master_size=False):
         line, column = 0, 0
         frames = sorted(glob.glob(self.tempdir.name + os.sep + "*.png"))
         frames_count = len(arange(0, self.duration, self.interval))
-        columns = self._calc_columns(frames_count, self.width, self.height)
-        master_height = self.height * math.ceil(frames_count / columns)
-        master = Image.new(mode="RGBA", size=(self.width * columns, master_height))
+        columns = self.calc_columns(frames_count, self.width, self.height)
 
-        for n, frame in enumerate(frames):
-            with Image.open(frame) as image:
-                x, y = self.width * column, self.height * line
-
-                start = n * self.interval
-                end = (n + 1) * self.interval
-                self.thumbnails.append((start, end, x, y))
-
-                image = image.resize((self.width, self.height), Image.ANTIALIAS)
-                master.paste(image, (x, y))
-
-                column += 1
+        if master_size:
+            yield self.width * columns, self.height * math.ceil(frames_count / columns)
 
-                if column == columns:
-                    line += 1
-                    column = 0
-
-        master.save(self._image_name)
-        self.tempdir.cleanup()
-
-    def to_vtt(self):
-        def _format_time(secs):
-            delta = timedelta(seconds=secs)
-            return ("0%s.000" % delta)[:12]
+        for n, frame in enumerate(frames):
+            x, y = self.width * column, self.height * line
 
-        _lines = ["WEBVTT\n\n"]
-        _img_src = self.basepath + self._image_name
+            start = n * self.interval
+            end = (n + 1) * self.interval
+            yield frame, start, end, x, y
 
-        for start, end, x, y in self.thumbnails:
-            _thumbnail = "%s --> %s\n%s#xywh=%d,%d,%d,%d\n\n" % (
-                _format_time(start), _format_time(end),
-                _img_src, x, y, self.width, self.height
-            )
-            _lines.append(_thumbnail)
+            column += 1
 
-        with open(self._vtt_name, "w") as vtt:
-            vtt.writelines(_lines)
+            if column == columns:
+                line += 1
+                column = 0