From 4f635b68cc4c09cb9f409a6727456e38dc5f2ae7 Mon Sep 17 00:00:00 2001
From: Jacob Zelek <jacob.zelek@me.com>
Date: Mon, 2 Nov 2020 22:02:20 -0800
Subject: [PATCH] Parallel processing

---
 README.md |  1 +
 generator | 51 +++++++++++++++++++++++++++++++++------------------
 2 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index aa2c829..6642ee4 100644
--- a/README.md
+++ b/README.md
@@ -45,6 +45,7 @@ Options:
   <height>        Height of each thumbnail.
   <columns>       Total number of thumbnails per line.
   <output>        Output.
+  <parallelism>   Number of files to process in parallel
 ```
 
 ## Example
diff --git a/generator b/generator
index 164862c..5b2044a 100755
--- a/generator
+++ b/generator
@@ -16,6 +16,7 @@ Options:
   <height>        Height of each thumbnail.
   <columns>       Total number of thumbnails per line.
   <output>        Output.
+  <parallelism>   Number of files to process in parallel
 """
 
 from docopt import docopt
@@ -23,6 +24,7 @@ from moviepy.editor import VideoFileClip
 from PIL import Image
 from click import progressbar
 from collections import namedtuple
+from multiprocessing import Pool, cpu_count
 import glob
 import os
 import random
@@ -39,11 +41,11 @@ def generate_video_thumbnails(args):
     input_path = args['<video>']
     interval = int(args['<interval>'])
     size = (int(args['<width>']), int(args['<height>']))
-    output_prefix = get_output_prefix()
     columns = int(args['<columns>'])
     output_path = args['<output>']
+    parallelism = args.get('<parallelism>', cpu_count()*2-1)
 
-    file_paths = set()
+    work_units = []
 
     if os.path.isdir(input_path):
         # Ensure output path is also directory
@@ -70,26 +72,38 @@ def generate_video_thumbnails(args):
                         seperator=os.sep,
                         file_name=os.path.basename(file_path)
                     )
-                file_paths.add((file_path, single_output_path,))
+                work_units.append((file_path, single_output_path,
+                                   interval, size, columns,))
     else:
-        file_paths.add((input_path, output_path,))
+        work_units.append((input_path, output_path, interval, size, columns,))
 
-    # Process all files sequentially
-    for file_path in file_paths:
-        video_file_clip = VideoFileClip(file_path[0])
-        generate_frames(video_file_clip, interval, output_prefix, size)
-        generate_sprite_from_frames(output_prefix, columns, size, file_path[1])
+    # Limit the number of parallel jobs if lower number of files
+    parallelism = min(parallelism, len(work_units))
+
+    # Process all files in parallel
+    with Pool(parallelism) as p:
+        p.map(process_file, work_units)
 
 
-def generate_frames(video_file_clip, interval, output_prefix, size):
+def process_file(work_unit):
+    input_file, output_file, interval, size, columns = work_unit
+    video_file_clip = VideoFileClip(input_file)
+    output_prefix = get_output_prefix()
+    file_name = os.path.basename(input_file)
+    generate_frames(file_name, video_file_clip, interval, output_prefix, size)
+    generate_sprite_from_frames(output_prefix, columns, size, output_file)
+
+
+def generate_frames(file_name, video_file_clip, interval, output_prefix, size):
     duration = video_file_clip.duration
-    print("Extracting", int(duration / interval), "frames")
     frame_count = 0
-    with progressbar(range(0, int(duration), interval)) as items:
-        for i in items:
-            extract_frame(video_file_clip, i, output_prefix, size, frame_count)
-            frame_count += 1
-    print("Frames extracted.")
+    total_frames = int(duration / interval)
+    for i in range(0, int(duration), interval):
+        print("[{file_name}] Extracting frame {current}/{total}".
+              format(file_name=file_name, current=frame_count+1,
+                     total=total_frames+1))
+        extract_frame(video_file_clip, i, output_prefix, size, frame_count)
+        frame_count += 1
 
 
 def extract_frame(video_file_clip, moment, output_prefix, size, frame_count):
@@ -138,8 +152,9 @@ def generate_sprite_from_frames(frames_path, columns, size, output):
                 column = 0
 
     final_image.save(output)
-    shutil.rmtree(TMP_FRAMES_PATH, ignore_errors=True)
-    print("Saved!")
+    shutil.rmtree(frames_path, ignore_errors=True)
+    output_file = os.path.basename(output)
+    print("[{output_file}] Saved".format(output_file=output_file))
 
 
 def get_output_prefix():