GH-6: Fix the memory leaking issue (GH-16)

 - Optimize the code - remove the redundant snippets
 - Replace lists with tuples
 - Optimize imports (scanned with cProfile)
 - Change the `interval` default to `1`
 - Optimize the memory usage by not keeping the frames in dynamic memory
 - Implement `compress` argument usage

Author: ArtyomVancyan

README.md

@@ -44,7 +44,7 @@ thumbnails [options]
 |-------------|:-------:|-----------------------------------------------------------------------------------------------|
 | `parallel`  | `false` | Process in parallel to make it faster.                                                        |
 | `overwrite` | `false` | Overwrite the existing files with new ones.                                                   |
-| `interval`  |   0.5   | The interval between two thumbnails in seconds.                                               |
+| `interval`  |    1    | The interval between two thumbnails in seconds.                                               |
 | `compress`  |    1    | The image compression coefficient is a number from 0 to 1 where 1 means _no compression_.     |
 | `as`        |  `vtt`  | Output as either `VTT` or `JSON` format.                                                      |
 | `basepath`  |    -    | In case the relative path does not work, the prefix of the thumbnails path can be customized. |

ffmpeg/ffmpeg.py

@@ -1,27 +1,54 @@
+import concurrent.futures
+import glob
+import math
 import os
 import re
 import subprocess
-from functools import reduce
+from datetime import timedelta
+from tempfile import TemporaryDirectory
 
-import imageio_ffmpeg
-import numpy
 from PIL import Image
 from imageio.v3 import immeta
+from imageio_ffmpeg import get_ffmpeg_exe
 
-FFMPEG_BINARY = imageio_ffmpeg.get_ffmpeg_exe()
+FFMPEG_BINARY = get_ffmpeg_exe()
 
 
 class FFMpeg:
     def __init__(self, filename):
-        duration, size = self.parse_metadata(filename)
-        self.area = reduce(int.__mul__, size)
-        self.bytes = self.area * 4
+        self.__compress = 1
+        self.__interval = 1
+        duration, size = self._parse_metadata(filename)
+        self.tempdir = TemporaryDirectory()
         self.filename = filename
         self.duration = duration
         self.size = size
 
+    def get_compress(self):
+        return self.__compress
+
+    def set_compress(self, compress):
+        if type(compress) not in (int, float):
+            raise TypeError("Compress must be a number.")
+        self.__compress = compress
+
+    def get_interval(self):
+        return self.__interval
+
+    def set_interval(self, interval):
+        if not isinstance(interval, int):
+            raise TypeError("Interval must be an integer.")
+        self.__interval = interval
+
     @staticmethod
-    def cross_platform_popen_params(bufsize=100000):
+    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
+
+    @staticmethod
+    def _cross_platform_popen_params(bufsize=100000):
         popen_params = {
             "bufsize": bufsize,
             "stdout": subprocess.PIPE,
@@ -36,7 +63,7 @@ def cross_platform_popen_params(bufsize=100000):
     def _parse_duration(stdout):
         duration_regex = r"duration[^\n]+([0-9][0-9]:[0-9][0-9]:[0-9][0-9].[0-9][0-9])"
         time = re.search(duration_regex, stdout, re.M | re.I).group(1)
-        time = [float(part.replace(",", ".")) for part in time.split(":")]
+        time = (float(part.replace(",", ".")) for part in time.split(":"))
         return sum(mult * part for mult, part in zip((3600, 60, 1), time))
 
     @staticmethod
@@ -45,67 +72,72 @@ def _parse_size(stdout):
         match_size = re.search(size_regex, stdout, re.M)
         return tuple(map(int, match_size.groups()))
 
-    def parse_metadata(self, filename):
+    def _parse_metadata(self, filename):
         meta = immeta(filename)
         duration, size = meta.get("duration"), meta.get("size")
 
-        if not all([duration, size]):
-            cmd = [FFMPEG_BINARY, "-hide_banner", "-i", filename]
+        if not all((duration, size)):
+            cmd = (FFMPEG_BINARY, "-hide_banner", "-i", filename)
 
-            popen_params = self.cross_platform_popen_params()
+            popen_params = self._cross_platform_popen_params()
             process = subprocess.Popen(cmd, **popen_params)
             _, stderr = process.communicate()
             stdout = stderr.decode("utf8", errors="ignore")
 
-            process.terminate()
-            del process
-
             duration = self._parse_duration(stdout)
             size = self._parse_size(stdout)
 
         return duration, size
 
-    @staticmethod
-    def frame_to_buffer(image):
-        image = image.astype("uint8")
-        return Image.fromarray(image)
-
-    def get_frame(self, start_time):
-        if start_time != 0:
-            offset = min(1, start_time)
-            i_arg = [
-                "-ss", "%.06f" % (start_time - offset),
-                "-i", self.filename,
-                "-ss", "%.06f" % offset,
-            ]
-        else:
-            i_arg = ["-i", self.filename]
+    def _extract_frame(self, start_time):
+        _input_file = self.filename
+        _output_file = "%s/%d-%s.png" % (self.tempdir.name, start_time, self.filename)
+        _timestamp = str(timedelta(seconds=start_time))
 
         cmd = (
-                [FFMPEG_BINARY]
-                + i_arg
-                + [
-                    "-loglevel", "error",
-                    "-f", "image2pipe",
-                    "-vf", "scale=%d:%d" % tuple(self.size),
-                    "-sws_flags", "bicubic",
-                    "-pix_fmt", "rgba",
-                    "-vcodec", "rawvideo", "-",
-                ]
+            FFMPEG_BINARY,
+            "-ss", _timestamp,
+            "-i", _input_file,
+            "-loglevel", "error",
+            "-vframes", "1",
+            _output_file,
+            "-y",
         )
 
-        popen_params = self.cross_platform_popen_params(self.bytes + 100)
-        process = subprocess.Popen(cmd, **popen_params)
-        buffer = process.stdout.read(self.bytes)
+        subprocess.Popen(cmd).wait()
+
+    def extract_frames(self):
+        _intervals = range(0, int(self.duration), self.get_interval())
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            executor.map(self._extract_frame, _intervals)
+
+    def join_frames(self):
+        width, height = self.size
+
+        _min_width = 300
+        _min_height = math.ceil(_min_width * height / width)
+
+        width = max(_min_width, width * self.__compress / 10)
+        height = max(_min_height, height * self.__compress / 10)
+
+        line, column = 0, 0
+        frames = sorted(glob.glob(self.tempdir.name + os.sep + "*.png"))
+        frames_count = len(range(0, int(self.duration), self.get_interval()))
+        columns = self.calc_columns(frames_count, width, height)
+        master_height = height * int(math.ceil(float(frames_count) / columns))
+        master = Image.new(mode="RGBA", size=(width * columns, master_height))
 
-        process.terminate()
-        del process
+        for frame in frames:
+            with Image.open(frame) as image:
+                x, y = width * column, height * line
+                image = image.resize((width, height), Image.ANTIALIAS)
+                master.paste(image, (x, y))
 
-        if hasattr(numpy, "frombuffer"):
-            result = numpy.frombuffer(buffer, dtype="uint8")
-        else:
-            result = numpy.fromstring(buffer, dtype="uint8")
+                column += 1
 
-        result.shape = (*self.size[::-1], len(buffer) // self.area)
+                if column == columns:
+                    line += 1
+                    column = 0
 
-        return result
+        master.save(self.filename + ".png")
+        self.tempdir.cleanup()

thumbnails.py

@@ -1,55 +1,26 @@
-import asyncio
 import concurrent.futures
-import math
-
-from PIL import Image
 
 from ffmpeg import FFMpeg
 
-width, height = 300, 200
+# Read from the program arguments.
+compress = 1
 interval = 20
-columns = 3
 
 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):
-    line, column = 0, 0
-    frames_count = len(range(0, int(video.duration), interval))
-    master_width = width * columns
-    master_height = height * int(math.ceil(float(frames_count) / columns))
-
-    try:
-        master = Image.new(mode="RGBA", size=(master_width, master_height))
-    except IOError:
-        master = Image.new(mode="RGB", size=(master_width, master_height))
-
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        frames = executor.map(video.get_frame, range(0, int(video.duration), interval))
-
-    for frame in frames:
-        x, y = width * column, height * line
-        image = video.frame_to_buffer(frame)
-        image = image.resize((width, height), Image.ANTIALIAS)
-        master.paste(image, (x, y))
-
-        column += 1
-
-        if column == columns:
-            line += 1
-            column = 0
-
-    master.save(video.filename + ".png")
-
+    video.set_compress(compress)
+    video.set_interval(interval)
+    video.extract_frames()
+    video.join_frames()
 
-async def main():
-    fps = [FFMpeg(fp) for fp in files]
 
+def main():
     with concurrent.futures.ProcessPoolExecutor() as executor:
-        executor.map(worker, fps)
+        executor.map(worker, map(FFMpeg, files))
 
 
-loop = asyncio.get_event_loop()
-loop.run_until_complete(main())
-loop.close()
+if __name__ == "__main__":
+    main()