Merge pull request #97 from n9Mtq4/feature/firstpassreuse

Reuse first pass from AOM keyframe split on the chunks

av1an.py

@@ -45,6 +45,10 @@ class Av1an:
                 print(f'No such model: {Path(self.vmaf_path).as_posix()}')
                 terminate()
 
+        if self.reuse_first_pass and self.encoder != 'aom' and self.split_method != 'aom_keyframes':
+            print('Reusing the first pass is only supported with the aom encoder and aom_keyframes split method.')
+            terminate()
+
         if self.video_params is None:
             self.video_params = get_default_params_for_encoder(self.encoder)
 
@@ -130,7 +134,7 @@ class Av1an:
                 tg_cq = self.target_vmaf(source)
                 cm1 = man_cq(commands[0], tg_cq)
 
-                if self.passes == 2:
+                if self.passes == 2 and (not self.reuse_first_pass):
                     cm2 = man_cq(commands[1], tg_cq)
                     commands = (cm1, cm2) + commands[2:]
                 else:
@@ -138,7 +142,7 @@ class Av1an:
 
             # Boost
             if self.boost:
-                commands = boosting(self.boost_limit, self.boost_range, source, commands, self.passes)
+                commands = boosting(self.boost_limit, self.boost_range, source, commands, self.passes, self.reuse_first_pass)
 
             log(f'Enc: {source.name}, {frame_probe_source} fr\n')
 
@@ -223,13 +227,16 @@ class Av1an:
             if self.extra_split:
                 framenums = extra_splits(self.input, framenums, self.extra_split)
 
+            if self.reuse_first_pass:
+                segment_first_pass(self.temp, framenums)
+
             segment(self.input, self.temp, framenums)
             extract_audio(self.input, self.temp,  self.audio_params)
 
         chunk = get_video_queue(self.temp,  self.resume)
 
         # Make encode queue
-        commands = compose_encoding_queue(chunk, self.temp, self.encoder, self.video_params, self.ffmpeg_pipe, self.passes)
+        commands = compose_encoding_queue(chunk, self.temp, self.encoder, self.video_params, self.ffmpeg_pipe, self.passes, self.reuse_first_pass)
         log(f'Encoding Queue Composed\n'
             f'Encoder: {self.encoder.upper()} Queue Size: {len(commands)} Passes: {self.passes}\n'
             f'Params: {self.video_params}\n\n')

utils/__init__.py

@@ -7,8 +7,9 @@ from .boost import *
 from .compose import *
 from .dynamic_progress_bar import *
 from .ffmpeg import *
+from .firstpassreuse import *
 from .pyscenedetect import *
 from .utils import *
 from .vmaf import *
 from .setup import *
-from .split import *
+from .split import *

utils/arg_parse.py

@@ -22,6 +22,9 @@ def arg_parsing():
     parser.add_argument('--scenes', '-s', type=str, default=None, help='File location for scenes')
     parser.add_argument('--threshold', '-tr', type=float, default=50, help='PySceneDetect Threshold')
 
+    # AOM Keyframe split
+    parser.add_argument('--reuse_first_pass', help='Reuse the first pass from aom_keyframes split on the chunks', action='store_true')
+
     # Encoding
     parser.add_argument('--passes', '-p', type=int, default=2, help='Specify encoding passes')
     parser.add_argument('--video_params', '-v', type=str, default=None, help='encoding settings')

utils/boost.py

@@ -26,16 +26,16 @@ def boost(command: str, brightness, b_limit, b_range, new_cq=0):
         print(f'Error in encoding loop {e}\nAt line {exc_tb.tb_lineno}')
 
 
-def boosting(bl, br, source, commands, passes):
+def boosting(bl, br, source, commands, passes, reuse_first_pass):
     try:
         brightness = get_brightness(source.absolute().as_posix())
         com0, cq = boost(commands[0], brightness, bl, br )
 
-        if passes == 2:
+        if passes == 2 and (not reuse_first_pass):
             com1, _ = boost(commands[1], brightness, bl, br, new_cq=cq)
             commands = (com0, com1) + commands[2:]
         else:
-            commands = com0 + commands[1:]
+            commands = (com0,) + commands[1:]
         log(f'{source.name}\n[Boost]\nAvg brightness: {br}\nAdjusted CQ: {cq}\n\n')
         return commands
 

utils/compose.py

@@ -90,7 +90,7 @@ def svt_av1_encode(inputs, passes, pipe, params):
     return commands
 
 
-def aom_vpx_encode(inputs, enc, passes, pipe, params):
+def aom_vpx_encode(inputs, enc, passes, pipe, params, skip_first_pass):
     """
     Generates commands for AOM, VPX encoders
 
@@ -103,23 +103,26 @@ def aom_vpx_encode(inputs, enc, passes, pipe, params):
     single_p = f'{enc} --passes=1 '
     two_p_1 = f'{enc} --passes=2 --pass=1'
     two_p_2 = f'{enc} --passes=2 --pass=2'
-    commands = []
 
     if passes == 1:
-        commands = [
+        return [
             (f'-i {file[0]} {pipe} {single_p} {params} -o {file[1].with_suffix(".ivf")} - ',
              (file[0], file[1].with_suffix('.ivf')))
             for file in inputs]
 
+    if passes == 2 and skip_first_pass:
+        return [
+            (f'-i {file[0]} {pipe} {two_p_2} {params} --fpf={file[0].with_suffix(".log")} -o {file[1].with_suffix(".ivf")} - ',
+             (file[0], file[1].with_suffix('.ivf')))
+            for file in inputs]
+
     if passes == 2:
-        commands = [
+        return [
             (f'-i {file[0]} {pipe} {two_p_1} {params} --fpf={file[0].with_suffix(".log")} -o {os.devnull} - ',
              f'-i {file[0]} {pipe} {two_p_2} {params} --fpf={file[0].with_suffix(".log")} -o {file[1].with_suffix(".ivf")} - ',
              (file[0], file[1].with_suffix('.ivf')))
             for file in inputs]
 
-    return commands
-
 
 def rav1e_encode(inputs, passes, pipe, params):
     """
@@ -160,7 +163,7 @@ def rav1e_encode(inputs, passes, pipe, params):
     return commands
 
 
-def compose_encoding_queue(files, temp, encoder, params, pipe, passes):
+def compose_encoding_queue(files, temp, encoder, params, pipe, passes, reuse_first_pass):
     """
     Composing encoding queue with split videos.
     :param files: List of files that need to be encoded
@@ -180,7 +183,7 @@ def compose_encoding_queue(files, temp, encoder, params, pipe, passes):
                file) for file in files]
 
     if encoder in ('aom', 'vpx'):
-        queue = aom_vpx_encode(inputs, enc_exe, passes, pipe, params)
+        queue = aom_vpx_encode(inputs, enc_exe, passes, pipe, params, reuse_first_pass)
 
     elif encoder == 'rav1e':
         queue = rav1e_encode(inputs, passes, pipe, params)

utils/firstpassreuse.py

@@ -0,0 +1,99 @@
+#!/bin/env python
+import os
+import struct
+from typing import List, Dict
+
+from .aom_keyframes import fields
+
+
+def read_first_pass(log_path):
+    """
+    Reads libaom first pass log into a list of dictionaries.
+
+    :param log_path: the path to the log file
+    :return: A list of dictionaries. The keys are the fields from aom_keyframes.py
+    """
+    frame_stats = []
+    with open(log_path, 'rb') as file:
+        frame_buf = file.read(208)
+        while len(frame_buf) > 0:
+            stats = struct.unpack('d' * 26, frame_buf)
+            p = dict(zip(fields, stats))
+            frame_stats.append(p)
+            frame_buf = file.read(208)
+    return frame_stats
+
+
+def write_first_pass_log(log_path, frm_lst: List[Dict]):
+    """
+    Writes a libaom compatible first pass log from a list of dictionaries containing frame stats.
+
+    :param log_path: the path of the ouput file
+    :param frm_lst: the list of dictionaries of the frame stats + eos stat
+    :return: None
+    """
+    with open(log_path, 'wb') as file:
+        for frm in frm_lst:
+            frm_bin = struct.pack('d' * 26, *frm.values())
+            file.write(frm_bin)
+
+
+def reindex_chunk(chunk_stats: List[Dict]):
+    """
+    The stats for each frame includes its frame number. This will reindex them to start at 0 in place.
+
+    :param chunk_stats: the list of stats for just this chunk
+    :return: None
+    """
+    for i, frm_stats in enumerate(chunk_stats):
+        frm_stats['frame'] = i
+
+
+def compute_eos_stats(chunk_stats: List[Dict], old_eos: Dict):
+    """
+    The end of sequence stat is a final packet at the end of the log. It contains the sum of all the previous
+    frame packets. When we split the log file, we need to sum up just the included frames as a new EOS packet.
+
+    :param chunk_stats: the list of stats for just this chunk
+    :param old_eos: the old eos stat packet
+    :return: A dict for the new eos packet
+    """
+    eos = old_eos.copy()
+    for key in eos.keys():
+        eos[key] = sum([d[key] for d in chunk_stats])
+        # eos[key] = (old_eos[key] / old_eos['count']) * len(chunk_stats)  # TODO(n9Mtq4): I think this will work well for VBR encodes
+    return eos
+
+
+def segment_first_pass(temp, framenums):
+    """
+    Segments the first pass file in temp/keyframes.log into individual log files for each chunk.
+    Looks at the len of framenums to determine file names for the chunks.
+
+    :param temp: the temp directory Path
+    :param framenums: a list of frame numbers along the split boundaries
+    :return: None
+    """
+    stat_file = temp / 'keyframes.log'  # TODO(n9Mtq4): makes this a constant for use here and w/ aom_keyframes.py
+    stats = read_first_pass(stat_file)
+
+    # special case for only 1 scene
+    # we don't need to do anything with the log
+    if len(framenums) == 0:
+        write_first_pass_log(os.path.join(temp, "split", "0.log"), stats)
+        return
+
+    eos_stats = stats[-1]  # EOS stats is the last one
+    split_names = [str(i).zfill(5) for i in range(len(framenums) + 1)]
+    frm_split = [0] + framenums + [len(stats) - 1]
+
+    for i in range(0, len(frm_split) - 1):
+        frm_start_idx = frm_split[i]
+        frm_end_idx = frm_split[i + 1]
+        log_name = split_names[i] + '.log'
+
+        chunk_stats = stats[frm_start_idx:frm_end_idx]
+        reindex_chunk(chunk_stats)
+        chunk_stats = chunk_stats + [compute_eos_stats(chunk_stats, eos_stats)]
+
+        write_first_pass_log(os.path.join(temp, "split", log_name), chunk_stats)