#!/usr/bin/env python3 import subprocess import json import os import sys import argparse import urllib.request import shutil def run_vmaf_check(reference_path, distorted_path, model_filename): """ Compares two video files using FFmpeg's internal libvmaf filter and returns the VMAF score. This version uses a local 'temp' directory and changes the CWD to ensure all paths are relative. """ print(f" - Running VMAF check against '{os.path.basename(reference_path)}'...") temp_dir = "temp" log_filename = 'vmaf_log.json' filter_script_filename = 'filter_script.txt' # The 'model' option for libvmaf requires its own key=value pair, # so we must specify 'model=path=...' to load from a file. filter_content = f"[0:v][1:v]libvmaf=log_fmt=json:log_path={log_filename}:model=path={model_filename}" filter_script_path = os.path.join(temp_dir, filter_script_filename) with open(filter_script_path, 'w') as f: f.write(filter_content) # We need absolute paths for the inputs before we change directory abs_distorted_path = os.path.abspath(distorted_path) abs_reference_path = os.path.abspath(reference_path) # Build the FFmpeg command using -filter_complex_script ffmpeg_command = [ 'ffmpeg', '-hide_banner', '-loglevel', 'error', '-i', abs_distorted_path, '-i', abs_reference_path, '-filter_complex_script', filter_script_filename, '-f', 'null', '-' ] original_dir = os.getcwd() try: # Execute FFmpeg from within the temp directory so filter script paths are relative os.chdir(temp_dir) subprocess.run(ffmpeg_command, check=True, capture_output=True) # Return to original directory before reading the log file os.chdir(original_dir) # Parse the VMAF results log_path = os.path.join(temp_dir, log_filename) with open(log_path, 'r') as f: vmaf_data = json.load(f) score = vmaf_data['pooled_metrics']['vmaf']['mean'] return float(score) except (subprocess.CalledProcessError, FileNotFoundError, json.JSONDecodeError, KeyError) as e: print(f" Error during VMAF calculation: {e}") if isinstance(e, subprocess.CalledProcessError): print(f" FFmpeg stderr: {e.stderr.decode().strip()}") return 0.0 finally: # Crucial: ensure we always change back to the original directory if os.getcwd() != original_dir: os.chdir(original_dir) def get_video_resolution(video_path): """Gets video resolution using ffprobe.""" command = [ 'ffprobe', '-v', 'error', '-select_streams', 'v:0', '-show_entries', 'stream=width,height', '-of', 'csv=s=x:p=0', video_path ] try: result = subprocess.run(command, capture_output=True, text=True, check=True) width, height = map(int, result.stdout.strip().split('x')) return width, height except (subprocess.CalledProcessError, ValueError) as e: print(f" Error getting video resolution: {e}") return None, None def encode_and_verify_segment(segment_path, target_vmaf, start_crf, min_crf, crf_step, model_filename): """ Encodes a single segment, iteratively adjusting CRF to meet the target VMAF score. Uses ffmpeg to pipe raw video to the external SvtAv1EncApp.exe encoder. """ output_dir = "segments" base_name = os.path.basename(segment_path) final_output_path = os.path.join(output_dir, base_name) # Get video resolution to pass to the encoder width, height = get_video_resolution(segment_path) if not width or not height: print(f" Could not determine resolution for '{base_name}'. Skipping segment.") return False current_crf = start_crf while True: temp_output_path = os.path.join(output_dir, f"temp_{base_name}") print(f" -> Encoding with CRF {current_crf} using SvtAv1EncApp.exe...") # Command to use ffmpeg as a frameserver, decoding the UTVideo segment # and piping raw 10-bit y4m video to stdout. ffmpeg_command = [ 'ffmpeg', '-hide_banner', '-loglevel', 'error', '-i', segment_path, '-pix_fmt', 'yuv420p10le', # Force 10-bit pixel format for the pipe '-f', 'yuv4mpegpipe', '-strict', '-1', # Allow non-standard pixel format in y4m pipe '-' # to stdout ] # Command for the external SVT-AV1-PSY encoder, reading from stdin svt_command = [ 'SvtAv1EncApp.exe', '-i', 'stdin', '--width', str(width), '--height', str(height), '--progress', '2', # Show aomenc-style progress '--preset', '4', # Temporarily changed from 2 for speed '--input-depth', '10', # Explicitly set 10-bit encoding '--crf', str(current_crf), '--film-grain', '8', '--tune', '2', '--keyint', '-1', '--color-primaries', '1', '--transfer-characteristics', '1', '--matrix-coefficients', '1', '-b', temp_output_path ] ffmpeg_process = None svt_process = None try: # Start the ffmpeg frameserver process ffmpeg_process = subprocess.Popen(ffmpeg_command, stdout=subprocess.PIPE) # Start the SVT encoder process, taking stdin from ffmpeg's stdout # We pass sys.stderr to the encoder so progress is shown in real-time svt_process = subprocess.Popen(svt_command, stdin=ffmpeg_process.stdout, stdout=subprocess.PIPE, stderr=sys.stderr) # This allows ffmpeg to receive a SIGPIPE if svt_process exits before ffmpeg is done. if ffmpeg_process.stdout: ffmpeg_process.stdout.close() # Wait for the encoder to finish. stdout is captured, stderr is shown live. svt_stdout, _ = svt_process.communicate() ffmpeg_process.wait() if svt_process.returncode != 0: # Manually raise an error to be caught by the except block raise subprocess.CalledProcessError(svt_process.returncode, svt_command) except (subprocess.CalledProcessError, FileNotFoundError) as e: print(f"\n Encoding failed with CRF {current_crf}.") if isinstance(e, subprocess.CalledProcessError): print(f" Encoder returned non-zero exit code {e.returncode}. See encoder output above for details.") else: # FileNotFoundError print(f" Error: '{e.filename}' not found. Please ensure SvtAv1EncApp.exe is in your PATH.") if os.path.exists(temp_output_path): os.remove(temp_output_path) current_crf -= crf_step if current_crf < min_crf: print(f" Could not encode segment even at min CRF. Skipping.") return False continue vmaf_score = run_vmaf_check(segment_path, temp_output_path, model_filename) print(f" - VMAF Score: {vmaf_score:.2f} (Target: {target_vmaf})") if vmaf_score >= target_vmaf or current_crf <= min_crf: if vmaf_score < target_vmaf: print(f" - Warning: Target VMAF not met, but hit min CRF {min_crf}. Keeping this version.") os.rename(temp_output_path, final_output_path) print(f" -> Success! Final CRF {current_crf} for '{os.path.basename(final_output_path)}'") return True else: os.remove(temp_output_path) current_crf -= crf_step def main(): parser = argparse.ArgumentParser( description="Encodes video segments from the 'cuts' directory to AV1, targeting a specific VMAF score.", formatter_class=argparse.RawTextHelpFormatter ) parser.add_argument("--target-vmaf", type=float, default=96.0, help="The target VMAF score to achieve. Default: 96.0") parser.add_argument("--start-crf", type=int, default=35, help="The CRF to start encoding with (higher is lower quality). Default: 35") parser.add_argument("--min-crf", type=int, default=17, help="The lowest CRF to allow (lower is higher quality). Default: 17") parser.add_argument("--crf-step", type=int, default=2, help="The value to decrease CRF by on each attempt. Default: 2") parser.add_argument("--cleanup", action='store_true', help="Remove the 'temp' directory after completion.") args = parser.parse_args() # --- Pre-flight check for executables --- if not shutil.which("SvtAv1EncApp.exe"): print("Error: 'SvtAv1EncApp.exe' not found. Please ensure it is in your system's PATH.") sys.exit(1) if not shutil.which("ffprobe"): print("Error: 'ffprobe' not found. Please ensure it is in your system's PATH.") sys.exit(1) # --- Setup Local Temp Directory (per user request) --- temp_dir = "temp" os.makedirs(temp_dir, exist_ok=True) # --- VMAF Model Check & Download --- model_filename = "vmaf_v0.6.1.json" model_path = os.path.join(temp_dir, model_filename) model_url = "https://raw.githubusercontent.com/Netflix/vmaf/master/model/vmaf_v0.6.1.json" if not os.path.exists(model_path): print(f"VMAF model '{model_filename}' not found in '{temp_dir}'.") print(f"Downloading from {model_url}...") try: urllib.request.urlretrieve(model_url, model_path) print("Download complete.") except Exception as e: print(f"Error: Failed to download VMAF model: {e}") sys.exit(1) # --- Pre-flight Checks --- input_dir = "cuts" output_dir = "segments" if not os.path.isdir(input_dir): print(f"Error: Input directory '{input_dir}' not found. Please run the scene cutter script first.") sys.exit(1) os.makedirs(output_dir, exist_ok=True) # --- Process Segments --- segments = sorted([f for f in os.listdir(input_dir) if f.endswith('.mkv')]) total_segments = len(segments) print(f"Found {total_segments} segments to process from '{input_dir}'.") print(f"Final files will be saved in '{output_dir}'.") try: for i, segment_file in enumerate(segments): print(f"\n--- Processing segment {i+1}/{total_segments}: {segment_file} ---") segment_path = os.path.join(input_dir, segment_file) encode_and_verify_segment( segment_path, args.target_vmaf, args.start_crf, args.min_crf, args.crf_step, model_filename # Pass only the filename, as CWD will be 'temp' ) finally: if args.cleanup: print(f"\n--- Cleaning up temporary directory '{temp_dir}'... ---") shutil.rmtree(temp_dir) print("Cleanup complete.") print("\n--- All segments processed. ---") if __name__ == "__main__": main()