chunk_encoder/vmaf_encoder.py

#!/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()