From d486fbe2ed065f84245f3e907f8e5117335da508 Mon Sep 17 00:00:00 2001 From: pat-e Date: Sun, 20 Jul 2025 09:01:49 +0200 Subject: [PATCH] first commit --- scene_cutter.py | 215 ++++++++++++++++++++++++++++++++++++++ segment_muxer.py | 93 +++++++++++++++++ static_encoder.py | 150 ++++++++++++++++++++++++++ vmaf_encoder.py | 260 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 718 insertions(+) create mode 100644 scene_cutter.py create mode 100644 segment_muxer.py create mode 100644 static_encoder.py create mode 100644 vmaf_encoder.py diff --git a/scene_cutter.py b/scene_cutter.py new file mode 100644 index 0000000..12920b9 --- /dev/null +++ b/scene_cutter.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +import subprocess +import json +import os +import sys +import argparse + +# --- Utility Functions (from previous scripts) --- + +def get_best_hwaccel(): + """ + Checks for available FFmpeg hardware acceleration methods and returns the best one + based on the current operating system. + """ + # Determine the priority list based on the operating system + if sys.platform == "win32": + # Windows: CUDA (Nvidia) > QSV (Intel) > D3D11VA (Modern, AMD/All) > DXVA2 (Legacy) + priority = ['cuda', 'qsv', 'd3d11va', 'dxva2'] + elif sys.platform == "linux": + # Linux: CUDA (Nvidia) > VAAPI (Intel/AMD) + priority = ['cuda', 'vaapi'] + elif sys.platform == "darwin": + # macOS: VideoToolbox is the native framework for Apple Silicon + priority = ['videotoolbox'] + else: + # Fallback for other operating systems (e.g., BSD) + priority = [] + + if not priority: + print(f"No hardware acceleration priority list for this OS ({sys.platform}). Using software.") + return None + + print(f"Checking for available hardware acceleration on {sys.platform}...") + try: + result = subprocess.run( + ['ffmpeg', '-hide_banner', '-hwaccels'], + capture_output=True, text=True, encoding='utf-8', check=False + ) + if result.returncode != 0: + print("Warning: Could not query FFmpeg. Using software decoding.") + return None + + available_methods = result.stdout.strip().split('\n') + if len(available_methods) > 1: + available_methods = available_methods[1:] + + for method in priority: + if method in available_methods: + print(f"Found best available hardware acceleration: {method}") + return method + + print("No high-priority hardware acceleration found. Using software decoding.") + return None + except (FileNotFoundError, Exception): + return None + +def get_video_duration(video_path): + """Gets the duration of a video file in seconds using ffprobe.""" + command = ['ffprobe', '-v', 'error', '-show_entries', 'format=duration', '-of', 'default=noprint_wrappers=1:nokey=1', video_path] + try: + result = subprocess.run(command, capture_output=True, text=True, check=True) + return float(result.stdout.strip()) + except (FileNotFoundError, subprocess.CalledProcessError, ValueError) as e: + print(f"\nError getting video duration: {e}") + return None + +# --- Core Logic Functions --- + +def detect_scenes(video_path, json_output_path, hwaccel=None, threshold=0.4): + """Uses FFmpeg to detect scene changes and saves timestamps to a JSON file.""" + print(f"\nStarting scene detection for: {os.path.basename(video_path)}") + command = ['ffmpeg', '-hide_banner'] + if hwaccel: + print(f"Attempting to use hardware acceleration: {hwaccel}") + command.extend(['-hwaccel', hwaccel]) + + filter_string = f"select='gt(scene,{threshold})',showinfo" + command.extend(['-i', video_path, '-vf', filter_string, '-f', 'null', '-']) + + try: + process = subprocess.Popen(command, stderr=subprocess.PIPE, text=True, encoding='utf-8') + scene_timestamps = [] + for line in iter(process.stderr.readline, ''): + if 'pts_time:' in line: + timestamp_str = line.split('pts_time:')[1].split()[0] + scene_timestamps.append(float(timestamp_str)) + elif line.strip().startswith('frame='): + sys.stderr.write(f'\r{line.strip()}') + sys.stderr.flush() + + process.wait() + sys.stderr.write('\n') + sys.stderr.flush() + + if process.returncode != 0: + print(f"\nWarning: FFmpeg may have encountered an error (exit code: {process.returncode}).") + + with open(json_output_path, 'w') as json_file: + json.dump(scene_timestamps, json_file, indent=4) + + print(f"Scene detection completed. {len(scene_timestamps)} scenes found.") + return True + + except (FileNotFoundError, Exception) as e: + print(f"\nAn error occurred during scene detection: {e}") + return False + +def cut_video_into_scenes(video_path, json_path, max_segment_length, hwaccel=None): + """Cuts a video into segments, ensuring no segment exceeds a maximum length.""" + print(f"\nStarting segment cutting for: {os.path.basename(video_path)}") + try: + with open(json_path, 'r') as f: + scene_timestamps = json.load(f) + except (FileNotFoundError, json.JSONDecodeError) as e: + print(f"\nError reading scene file '{json_path}': {e}") + return + + video_duration = get_video_duration(video_path) + if video_duration is None: return + + output_dir = "cuts" + os.makedirs(output_dir, exist_ok=True) + print(f"Output will be saved to the '{output_dir}' directory.") + + print(f"Enforcing maximum segment length of {max_segment_length} seconds...") + segment_boundaries = [0.0] + sorted(list(set(scene_timestamps))) + [video_duration] + final_cut_points = [] + for i in range(len(segment_boundaries) - 1): + start, end = segment_boundaries[i], segment_boundaries[i+1] + current_time = start + while (end - current_time) > max_segment_length: + current_time += max_segment_length + final_cut_points.append(current_time) + final_cut_points.append(end) + + final_cut_points = sorted(list(set(final_cut_points))) + if final_cut_points and final_cut_points[-1] >= video_duration - 0.1: + final_cut_points.pop() + + print(f"Original scenes: {len(scene_timestamps)}. Total segments after splitting: {len(final_cut_points) + 1}.") + + segment_times_str = ",".join(map(str, final_cut_points)) + base_name = os.path.splitext(os.path.basename(video_path))[0] + output_pattern = os.path.join(output_dir, f"{base_name}_segment%03d.mkv") + + # Add -loglevel error to hide info messages and -stats to show progress + command = ['ffmpeg', '-hide_banner', '-loglevel', 'error', '-stats'] + if hwaccel: + command.extend(['-hwaccel', hwaccel]) + + command.extend(['-i', video_path, '-c:v', 'utvideo', '-an', '-sn', '-dn', '-map_metadata', '-1', '-map_chapters', '-1', '-f', 'segment', '-segment_times', segment_times_str, '-segment_start_number', '1', '-reset_timestamps', '1', output_pattern]) + + print("\nStarting FFmpeg to cut all segments in a single pass...") + try: + # -stats will print progress to stderr, which subprocess.run will display + subprocess.run(command, check=True) + total_segments = len(final_cut_points) + 1 + print(f"Successfully created {total_segments} segments in the '{output_dir}' directory.") + except (FileNotFoundError, subprocess.CalledProcessError, Exception) as e: + print(f"\nAn error occurred during cutting: {e}") + +# --- Main Orchestrator --- + +def main(): + parser = argparse.ArgumentParser( + description="A comprehensive video processing script to detect scenes and cut segments.", + formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument("video_path", help="Path to the input video file.") + parser.add_argument( + "--so", "--sceneonly", + action='store_true', + dest='sceneonly', # Explicitly set the destination attribute name + help="Only run scene detection and create the .scenes.json file." + ) + parser.add_argument( + "--segtime", + type=int, + default=10, + help="Maximum length of any cut segment in seconds. Default: 10" + ) + parser.add_argument( + "-t", "--threshold", + type=float, + default=0.4, + help="Scene detection threshold (0.0 to 1.0). Lower is more sensitive. Default: 0.4" + ) + args = parser.parse_args() + + if not os.path.isfile(args.video_path): + print(f"Error: Input file not found: '{args.video_path}'") + sys.exit(1) + + base_name = os.path.splitext(args.video_path)[0] + json_path = f"{base_name}.scenes.json" + + hwaccel_method = get_best_hwaccel() + + if args.sceneonly: + detect_scenes(args.video_path, json_path, hwaccel_method, args.threshold) + sys.exit(0) + + # --- Full Workflow (Detect if needed, then Cut) --- + if not os.path.isfile(json_path): + print("--- Scene file not found, running detection first ---") + if not detect_scenes(args.video_path, json_path, hwaccel_method, args.threshold): + print("\nScene detection failed. Aborting process.") + sys.exit(1) + else: + print(f"--- Found existing scene file: {os.path.basename(json_path)} ---") + + cut_video_into_scenes(args.video_path, json_path, args.segtime, hwaccel_method) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/segment_muxer.py b/segment_muxer.py new file mode 100644 index 0000000..adbfcd6 --- /dev/null +++ b/segment_muxer.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +import os +import sys +import subprocess +import argparse +import shutil + +def main(): + """ + Finds all encoded segments in the 'segments' directory and uses mkvmerge + to mux them into a single, final video file. + """ + parser = argparse.ArgumentParser( + description="Muxes encoded segments from the 'segments' directory into a final file using mkvmerge.", + formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument("--cleanup", action='store_true', help="Remove the 'segments' and 'cuts' directories after completion.") + args = parser.parse_args() + + segments_dir = "segments" + cuts_dir = "cuts" + + # --- Pre-flight Checks --- + if not shutil.which("mkvmerge"): + print("Error: 'mkvmerge' not found. Please ensure it is installed and in your system's PATH.") + sys.exit(1) + + if not os.path.isdir(segments_dir): + print(f"Error: Input directory '{segments_dir}' not found. Please run the encoder script first.") + sys.exit(1) + + segment_files = sorted([f for f in os.listdir(segments_dir) if f.endswith('.mkv')]) + if not segment_files: + print(f"Error: No .mkv files found in the '{segments_dir}' directory.") + sys.exit(1) + + print(f"Found {len(segment_files)} segments to mux.") + + # --- Determine Output Filename --- + # Infer the base name from the first segment file. + # e.g., "My Video_segment001.mkv" -> "My Video" + try: + base_name = segment_files[0].rsplit('_segment', 1)[0] + output_filename = f"temp_{base_name}.mkv" + except IndexError: + print("Error: Could not determine a base name from the segment files.") + print("Files should be named like '..._segmentXXX.mkv'.") + sys.exit(1) + + print(f"Output file will be: {output_filename}") + + # --- Build mkvmerge Command --- + # mkvmerge -o output.mkv segment1.mkv + segment2.mkv + segment3.mkv ... + command = ['mkvmerge', '-o', output_filename] + + # Add the first file + command.append(os.path.join(segments_dir, segment_files[0])) + + # Add the subsequent files with the '+' append operator + for segment in segment_files[1:]: + command.append('+') + command.append(os.path.join(segments_dir, segment)) + + # --- Execute Muxing --- + print("Running mkvmerge...") + try: + process = subprocess.run(command, check=True, capture_output=True, text=True) + print("mkvmerge output:") + print(process.stdout) + print(f"\nSuccess! Segments muxed into '{output_filename}'.") + except subprocess.CalledProcessError as e: + print("\nError: mkvmerge failed.") + print(f"Command: {' '.join(e.cmd)}") + print(f"Return Code: {e.returncode}") + print(f"Stderr:\n{e.stderr}") + sys.exit(1) + + # --- Cleanup --- + if args.cleanup: + print("\n--- Cleaning up temporary directories... ---") + try: + if os.path.isdir(segments_dir): + shutil.rmtree(segments_dir) + print(f"Removed '{segments_dir}' directory.") + if os.path.isdir(cuts_dir): + shutil.rmtree(cuts_dir) + print(f"Removed '{cuts_dir}' directory.") + print("Cleanup complete.") + except OSError as e: + print(f"Error during cleanup: {e}") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/static_encoder.py b/static_encoder.py new file mode 100644 index 0000000..f02e010 --- /dev/null +++ b/static_encoder.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +import subprocess +import os +import sys +import argparse +import shutil + +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 for '{video_path}': {e}") + return None, None + +def encode_segment(segment_path, crf): + """ + Encodes a single segment with a static CRF value. + 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 + + print(f" -> Encoding with static CRF {crf} using SvtAv1EncApp.exe...") + + # Command to use ffmpeg as a frameserver, decoding the 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', + '--preset', '2', + '--input-depth', '10', + '--crf', str(crf), + '--film-grain', '8', + '--tune', '2', + '--keyint', '-1', + '--color-primaries', '1', + '--transfer-characteristics', '1', + '--matrix-coefficients', '1', + '-b', final_output_path # Encode directly to the final 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. + 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) + + print(f" -> Success! Finished encoding '{os.path.basename(final_output_path)}'") + return True + + except (subprocess.CalledProcessError, FileNotFoundError) as e: + print(f"\n Encoding failed for CRF {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.") + + # Clean up partially created file on failure + if os.path.exists(final_output_path): + os.remove(final_output_path) + return False + +def main(): + parser = argparse.ArgumentParser( + description="Encodes video segments from the 'cuts' directory to AV1 using a static CRF value.", + formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument("--crf", type=int, default=27, help="The static CRF value to use for all segments. Default: 27") + + args = parser.parse_args() + + # --- Pre-flight check for executables --- + for exe in ["SvtAv1EncApp.exe", "ffprobe"]: + if not shutil.which(exe): + print(f"Error: '{exe}' not found. Please ensure it is in your system's PATH.") + sys.exit(1) + + # --- Setup Directories --- + 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) + if total_segments == 0: + print(f"No segments found in '{input_dir}'.") + sys.exit(0) + + print(f"Found {total_segments} segments to process from '{input_dir}'.") + print(f"Encoding with static CRF {args.crf}.") + print(f"Final files will be saved in '{output_dir}'.") + + 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_segment(segment_path, args.crf) + + print("\n--- All segments processed. ---") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/vmaf_encoder.py b/vmaf_encoder.py new file mode 100644 index 0000000..80706a4 --- /dev/null +++ b/vmaf_encoder.py @@ -0,0 +1,260 @@ +#!/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() \ No newline at end of file