refactoring, reducing complexity

Av1an/aom_kf.py

@@ -12,7 +12,8 @@ from tqdm import tqdm
 
 from .compose import compose_aomsplit_first_pass_command
 from .logger import log
-from .utils import frame_probe, terminate
+from .utils import  terminate
+from .ffmpeg import frame_probe
 
 # This is a script that returns a list of keyframes that aom would likely place. Port of aom's C code.
 # It requires an aom first-pass stats file as input. FFMPEG first-pass file is not OK. Default filename is stats.bin.

Av1an/encode.py

@@ -4,7 +4,7 @@ import json
 from pathlib import Path
 from .target_vmaf import target_vmaf
 from .boost import boosting
-from .utils import frame_probe, frame_check, frame_probe_fast, terminate
+from .utils import  frame_probe_cv2, terminate, process_inputs
 from .fp_reuse import remove_first_pass_from_commands
 from .utils import man_q
 from .logger import log
@@ -16,11 +16,10 @@ import concurrent
 import concurrent.futures
 from .logger import log, set_log
 from .config  import conf
-from .compose import (compose_encoding_queue, get_video_queue)
-from .ffmpeg import concatenate_video, extract_audio
+from .compose import compose_encoding_queue, get_video_queue
+from .ffmpeg import concatenate_video, extract_audio, frame_probe, frame_check
 from .fp_reuse import segment_first_pass
 from .split import extra_splits, segment, split_routine
-from .utils import (frame_probe, frame_probe_fast,process_inputs, terminate)
 from .vmaf import plot_vmaf
 import shutil
 
@@ -75,6 +74,7 @@ def video_encoding(args):
 def main_queue(args):
     # Todo: Redo Queue
     try:
+        conf(args)
         tm = time.time()
 
         args.queue = process_inputs(args.input)
@@ -112,7 +112,7 @@ def encoding_loop(args, commands):
                 log(f'Resumed with {done} encoded clips done\n\n')
             else:
                 initial = 0
-                total = frame_probe_fast(args.input)
+                total = frame_probe_cv2(args.input)
 
                 if total < 1:
                     total = frame_probe(args.input)

Av1an/ffmpeg.py

@@ -1,11 +1,87 @@
 #!/bin/env python
 
-
+import json
 import subprocess
 from pathlib import Path
 from subprocess import PIPE, STDOUT
-
+from .utils import terminate
 from .logger import log
+from threading import Lock
+import re
+import sys
+
+
+
+def frame_probe(source: Path):
+    """Get frame count."""
+    cmd = ["ffmpeg", "-hide_banner", "-i", source.as_posix(), "-map", "0:v:0", "-f", "null", "-"]
+    r = subprocess.run(cmd, stdout=PIPE, stderr=PIPE)
+    matches = re.findall(r"frame=\s*([0-9]+)\s", r.stderr.decode("utf-8") + r.stdout.decode("utf-8"))
+    return int(matches[-1])
+
+def get_keyframes(file: Path):
+    """
+    Read file info and return list of all keyframes
+
+    :param file: Path for input file
+    :return: list with frame numbers of keyframes
+    """
+
+    keyframes = []
+
+    ff = ["ffmpeg", "-hide_banner", "-i", file.as_posix(),
+    "-vf", r"select=eq(pict_type\,PICT_TYPE_I)",
+    "-f", "null", "-loglevel", "debug", "-"]
+
+    pipe = subprocess.Popen(ff, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+
+    while True:
+        line = pipe.stdout.readline().strip().decode("utf-8")
+
+        if len(line) == 0 and pipe.poll() is not None:
+            break
+
+        match = re.search(r"n:([0-9]+)\.[0-9]+ pts:.+key:1", line)
+        if match:
+            keyframe = int(match.group(1))
+            keyframes.append(keyframe)
+
+    return keyframes
+
+doneFileLock = Lock()
+def write_progress_file(file, chunk, frames):
+    doneFileLock.acquire()
+    with file.open() as f:
+        d = json.load(f)
+    d['done'][chunk.name] = frames
+    with file.open('w') as f:
+        json.dump(d, f)
+
+
+def frame_check(source: Path, encoded: Path, temp, check):
+    """Checking is source and encoded video frame count match."""
+    try:
+        status_file = Path(temp / 'done.json')
+
+        if check:
+            s1 = frame_probe(source)
+            write_progress_file(status_file, source, s1)
+        else:
+            s1, s2 = [frame_probe(i) for i in (source, encoded)]
+            if s1 == s2:
+                write_progress_file(status_file, source, s1)
+            else:
+                print(f'Frame Count Differ for Source {source.name}: {s2}/{s1}')
+
+    except IndexError:
+        print('Encoding failed, check validity of your encoding settings/commands and start again')
+        terminate()
+    except Exception as e:
+        _, _, exc_tb = sys.exc_info()
+        print(f'\nError frame_check: {e}\nAt line: {exc_tb.tb_lineno}\n')
+    finally:
+        if doneFileLock.locked():
+            doneFileLock.release()
 
 
 def concatenate_video(temp, output, encoder):

Av1an/split.py

@@ -7,7 +7,7 @@ from ast import literal_eval
 from pathlib import Path
 from subprocess import PIPE, STDOUT
 
-from .utils import frame_probe, get_keyframes
+from .ffmpeg import frame_probe, get_keyframes
 from .aom_kf import aom_keyframes
 from .logger import log
 from .pyscene import pyscene

Av1an/target_vmaf.py

@@ -1,6 +1,7 @@
 #!/bin/env python
 
-from .utils import terminate, frame_probe
+from .utils import terminate
+from .ffmpeg import frame_probe
 from .vmaf import call_vmaf, read_vmaf_xml
 from scipy import interpolate
 from pathlib import Path

Av1an/utils.py

@@ -7,7 +7,7 @@ import subprocess
 import sys
 from pathlib import Path
 from subprocess import PIPE
-from threading import Lock
+
 
 import cv2
 import numpy as np
@@ -33,35 +33,6 @@ def process_inputs(inputs):
 
     return inputs
 
-def get_keyframes(file: Path):
-    """
-    Read file info and return list of all keyframes
-
-    :param file: Path for input file
-    :return: list with frame numbers of keyframes
-    """
-
-    keyframes = []
-
-    ff = ["ffmpeg", "-hide_banner", "-i", file.as_posix(),
-    "-vf", "select=eq(pict_type\,PICT_TYPE_I)",
-    "-f", "null", "-loglevel", "debug", "-"]
-
-    pipe = subprocess.Popen(ff, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
-
-    while True:
-        line = pipe.stdout.readline().strip().decode("utf-8")
-
-        if len(line) == 0 and pipe.poll() is not None:
-            break
-
-        match = re.search(r"n:([0-9]+)\.[0-9]+ pts:.+key:1", line)
-        if match:
-            keyframe = int(match.group(1))
-            keyframes.append(keyframe)
-
-    return keyframes
-
 
 def get_cq(command):
     """
@@ -86,57 +57,12 @@ def man_q(command: str, q: int):
     return cmd
 
 
-
-def frame_probe_fast(source: Path):
+def frame_probe_cv2(source: Path):
     video = cv2.VideoCapture(source.as_posix())
     total = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
     return total
 
 
-def frame_probe(source: Path):
-    """Get frame count."""
-    cmd = ["ffmpeg", "-hide_banner", "-i", source.as_posix(), "-map", "0:v:0", "-f", "null", "-"]
-    r = subprocess.run(cmd, stdout=PIPE, stderr=PIPE)
-    matches = re.findall(r"frame=\s*([0-9]+)\s", r.stderr.decode("utf-8") + r.stdout.decode("utf-8"))
-    return int(matches[-1])
-
-
-doneFileLock = Lock()
-def frame_check(source: Path, encoded: Path, temp, check):
-    """Checking is source and encoded video frame count match."""
-    try:
-        status_file = Path(temp / 'done.json')
-
-        if check:
-            s1 = frame_probe(source)
-            doneFileLock.acquire()
-            with status_file.open() as f:
-                d = json.load(f)
-            d['done'][source.name] = s1
-            with status_file.open('w') as f:
-                json.dump(d, f)
-        else:
-            s1, s2 = [frame_probe(i) for i in (source, encoded)]
-            if s1 == s2:
-                doneFileLock.acquire()
-                with status_file.open() as f:
-                    d = json.load(f)
-                d['done'][source.name] = s1
-                with status_file.open('w') as f:
-                    json.dump(d, f)
-            else:
-                print(f'Frame Count Differ for Source {source.name}: {s2}/{s1}')
-    except IndexError:
-        print('Encoding failed, check validity of your encoding settings/commands and start again')
-        terminate()
-    except Exception as e:
-        _, _, exc_tb = sys.exc_info()
-        print(f'\nError frame_check: {e}\nAt line: {exc_tb.tb_lineno}\n')
-    finally:
-        if doneFileLock.locked():
-            doneFileLock.release()
-
-
 def get_brightness(video):
     """Getting average brightness value for single video."""
     brightness = []

av1an.py

@@ -11,7 +11,6 @@ class Av1an:
     def main_thread(self):
         """Main."""
         startup_check(self.args)
-        conf(self.args)
         main_queue(self.args)