Refactor & Rewrite

utils/upcale.py

@@ -1,38 +1,92 @@
 import os
+import sys
 import subprocess
 
 from utils.mediainfo import get_framerate
 
-def re_encode(input_file, upscale_out, temp_out, max_width, input_aspect):
-    if os.path.exists(upscale_out):
-        print('Skipped Pre-Encode')
-        return
+MAX_INPUT_WIDTH  = '720'
+
+def _re_encode(
+    source_video: str, 
+    temp_out_video: str,
+    input_aspect: str = "16:9"
+):
+    """
+    Re-Encodes the source video to avoid nasty video bugs
     
-    command = (
-        f'ffmpeg -v quiet -stats -i "{input_file}" '
-        '-c:v ffv1 -level 3 '
-        f'-vf "fps={get_framerate(input_file)},scale=-1:\'min({max_width},ih)\'" -aspect {input_aspect} '
-        '-pix_fmt yuv420p -color_primaries 1 -color_trc 1 -colorspace 1 '
-        '-an -sn -map_metadata -1 '
-        f'"{temp_out}"'
-    )
+    :param source_video: Video Input
+    :type source_video: str
+    :param input_aspect: Aspect Ratio of Video
+    :type input_aspect: str
+    """
 
-    subprocess.call(command, shell=True)
+    cmd = [
+        "ffmpeg", "-v", "quiet", "-stats",
+        "-i", source_video,
+        "-c:v", "ffv1",
+        "-level", "3",
+        "-vf", f"fps={get_framerate(source_video)},scale=-1:\'min({MAX_INPUT_WIDTH},ih)\'",
+        "-aspect", input_aspect,
+        "-pix_fmt", "yuv420p",
+        "-color_primaries", "1",
+        "-color_trc", "1",
+        "-colorspace", "1",
+        "-an", 
+        "-sn",
+        "-map_metadata", "-1",
+        temp_out_video
+    ]
 
-def upscale(input_file, upscale_out, max_width, input_aspect):
-    temp_out = os.path.join('1-Temp', 'source.mkv')
-    vsgan = os.path.join('utils', 'vs-realesrgan.vpy')
+    try:
+        subprocess.run(cmd, shell=True, check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"\nffmpeg failed with error code {e.returncode} at _re_encode()", file=sys.stderr)
+        sys.exit(e.returncode)
 
-    # Re-Encode to fix issues
-    re_encode(input_file, upscale_out, temp_out, max_width, input_aspect)
+def _upscale(
+    upscale_output: str,
+    input_aspect: str = "16:9"
+):
+    print('Started Upscale')
 
-    if os.path.exists(upscale_out):
+    vapoursynth_script = os.path.join('utils', 'vs-realesrgan.vpy')
+
+    cmd = [
+        "vspipe", 
+        "-c", "y4m",
+        vapoursynth_script,
+        "-", # Video output to pipe
+        "|", # Pipe
+        "ffmpeg", "-v", "quiet", "-stats",
+        "-f", "yuv4mpegpipe",
+        "-i", "-", # Pipe Video Input
+        "-c:v", "hevc_nvenc",
+        "-qp", "5",
+        "-aspect", input_aspect,
+        upscale_output
+    ]
+
+    try:
+        subprocess.run(cmd, shell=True, check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"\nffmpeg failed with error code {e.returncode}", file=sys.stderr)
+        sys.exit(e.returncode)
+
+
+def upscale(
+    source_video: str, 
+    upscaled_video_output: str, 
+    input_aspect: str,
+):
+    if os.path.exists(upscaled_video_output):
         print('Skipped Upscale')
         return
+    
+    temp_out_video = os.path.join('1-Temp', 'source.mkv')
 
-    print('Started Upscale')
-    subprocess.call(f'vspipe -c y4m {vsgan} - | ffmpeg -v quiet -stats -f yuv4mpegpipe -i - -c:v hevc_nvenc -qp 5 -aspect {input_aspect} "{upscale_out}"', shell=True)
-
+    _re_encode(source_video, temp_out_video, input_aspect)
+    _upscale(upscaled_video_output, input_aspect)
+    
     # Remove Temp Files
-    os.remove(temp_out)
-    os.remove(temp_out + '.ffindex')
+    os.remove(temp_out_video)
+    os.remove(f'{temp_out_video}.ffindex')