Improve audio concat; cleanup temp files; add -T(temp_dir) and -P (video_pool_size) args

main.py

@@ -8,7 +8,7 @@
 from multiprocessing import Pool
 from pathlib import Path
 from tempfile import NamedTemporaryFile, TemporaryDirectory
-from typing import List
+from typing import List, Tuple
 
 import ffmpeg
 from pydub import AudioSegment, silence
@@ -76,11 +76,11 @@ def concatenate_videos(temp_dir, chunks: List[str], outfile: str):
         subprocess.check_output(concat_args)
 
 
-def trim_silence(infile: str, outfile: str, min_silence_len: int, silence_thresh: int, margin: int) -> bool:
+def trim_silence(infile: str, min_silence_len: int, silence_thresh: int, margin: int) -> Tuple[str, str, List[Tuple[int, int]]]:
     """
     Returns
     -------
-    True trimming was successful and resulted in nonempty outfile
+    Video output, audio output and a list of nonsilent intervals in form of (start_ms, stop_ms)
     """
     infile_video = infile
     infile_audio = str(Path(infile_video).with_suffix('.mp3'))
@@ -103,7 +103,7 @@ def trim_silence(infile: str, outfile: str, min_silence_len: int, silence_thresh
     logging.info(parts)
 
     if not parts:
-        return False
+        return infile_video, infile_audio, []
 
     logging.info("Trimming audio")
     segments = AudioSegment._sync(*[audio[start:stop] for start, stop in parts])
@@ -112,32 +112,36 @@ def trim_silence(infile: str, outfile: str, min_silence_len: int, silence_thresh
 
     logging.info(f"Writing trimmed audio into {infile_audio} with duration {trimmed_audio.duration_seconds} s.")
     trimmed_audio.export(infile_audio)
+    return infile_video, infile_audio, parts
 
+
+def process_chunk(workdir: str, args: argparse.Namespace, total_segments: int, id_: int, segment: str):
+    logging.info(f"Processing chunk {id_ + 1}/{total_segments}: {segment} -> {Path(segment).with_suffix('.mp3')}")
+    res = trim_silence(segment,
+                       args.min_silence_len,
+                       args.silence_thresh,
+                       args.margin)
+    logging.info(f"Done processing audio chunk {id_ + 1}/{total_segments}")
+    return res
+
+
+def trim_video(infile_video: str, infile_audio: str, parts: List[Tuple[int, int]]) -> str:
+    in_path = Path(infile_video)
+    outfile = str(in_path.with_stem(in_path.stem + '_cropped'))
     parts = [(start / 1000, stop / 1000) for (start, stop) in parts]
 
     in_file = ffmpeg.input(infile_video)
-
     joined = ffmpeg.concat(
         ffmpeg.concat(
             *[in_file.trim(start=start, end=stop).setpts('PTS-STARTPTS')
               for start, stop in parts]),
         ffmpeg.input(infile_audio),
         v=1,
         a=1).node
-    ffmpeg.output(joined[0], joined[1], outfile).run(quiet=True, overwrite_output=True)
-    return True
-
 
-def process_chunk(workdir: str, args: argparse.Namespace, total_segments: int, id_: int, segment: str):
-    outfile = os.path.join(workdir, BASENAME + f"_cropped_{id_}" + Path(args.infile).suffix)
-    logging.info(f"Processing chunk {id_ + 1}/{total_segments}: {segment} -> {outfile}")
-    if not trim_silence(segment,
-                        outfile,
-                        args.min_silence_len,
-                        args.silence_thresh,
-                        args.margin):
-        outfile = None
-    logging.info(f"Done processing chunk {id_ + 1}/{total_segments}")
+    logging.info(f'Trimming {infile_video} -> {outfile}')
+    ffmpeg.output(joined[0], joined[1], outfile).run(quiet=True, overwrite_output=True)
+    logging.info(f'Done trimming {infile_video}')
     return outfile
 
 
@@ -153,8 +157,12 @@ def main():
                         help='margin (ms)', default=100)
     parser.add_argument('-n', dest='n_segments', type=int,
                         help='number of chunks to split input file to be processed independently', default=10)
-    parser.add_argument('-p', dest='pool_size', type=int,
-                        help='number of chunks to be processed concurrently', default=1)
+    parser.add_argument('-p', dest='audio_pool_size', type=int,
+                        help='number of audio chunks to be processed concurrently', default=1)
+    parser.add_argument('-P', dest='video_pool_size', type=int,
+                        help='number of video chunks to be processed concurrently. '
+                             'FFmpeg (without hwaccel) uses optimal number of threads by default, '
+                             'so values greater than 1 usually won\'t improve performance', default=1)
     parser.add_argument('-T', dest='temp_dir',
                         help='directory for temporary files')
     parser.add_argument('-d', dest='log_level', action='store_const', const=logging.DEBUG, help='print debugging info',
@@ -165,11 +173,16 @@ def main():
 
     with TemporaryDirectory(dir=args.temp_dir) as workdir:
         segments = split_video(args.infile, workdir, BASENAME, args.n_segments)
-        cropped_segments = []
-        with Pool(processes=args.pool_size) as pool:
+
+        with Pool(processes=args.audio_pool_size) as pool:
             cropped_segments = pool.starmap(functools.partial(process_chunk, workdir, args, len(segments)), enumerate(segments))
-        cropped_segments = [segment for segment in cropped_segments if segment is not None]
-        concatenate_videos(args.temp_dir, cropped_segments, args.outfile)
+        cropped_segments = [(video, audio, parts) for video, audio, parts in cropped_segments if parts]
+
+        logging.info('Trimming video chunks')
+        with Pool(processes=args.video_pool_size) as pool:
+            resulting_segments = pool.starmap(trim_video, cropped_segments)
+
+        concatenate_videos(args.temp_dir, resulting_segments, args.outfile)
 
 
 if __name__ == '__main__':