chunk_encoder/static_encoder.py

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