Ahmad Ayad
added passing for target codec for both synthetics and organics (#94)
4e4aa87 11 months ago
26.1 kB
'''
This script performs video encoding using various codecs, potentially optimized
for scene content type.
'''
import ffmpeg
import re
import time
import traceback
import numpy as np
from .encoder_configs import ENCODER_SETTINGS, SCENE_SPECIFIC_PARAMS, MODEL_CQ_REFERENCE_CODEC, QUALITY_MAPPING_ANCHORS

def cleanup_quality_params(settings, keep_param=None):
    """
    Remove conflicting quality parameters (CRF, CQ, QP).
    Args:
        settings (dict): Current settings dictionary
        keep_param (str, optional): Which parameter to keep ('crf', 'cq', or 'qp')
                                   If None, removes all quality parameters
    Returns:
        dict: Updated settings with cleaned quality parameters
    """
    quality_params = ['crf', 'cq', 'qp']

    if keep_param:
        # Remove all except the one we want to keep
        for param in quality_params:
            if param != keep_param:
                settings.pop(param, None)
    else:
        # Remove all quality parameters
        for param in quality_params:
            settings.pop(param, None)

    return settings


def apply_rate_mapping(codec, rate, current_settings, logging_enabled=True):
    """
    Apply rate (CQ/CRF) mapping to current settings based on codec type.
    Args:
        codec (str): Video codec name
        rate (int/float): Model-predicted CQ value
        current_settings (dict): Current encoder settings
        logging_enabled (bool): Enable logging
    Returns:
        dict: Updated settings with rate parameter applied
    """
    model_predicted_ref_cq = int(rate)

    if codec == MODEL_CQ_REFERENCE_CODEC:
        # Use the predicted CQ directly for the reference codec
        if 'cq' in current_settings:
            current_settings['cq'] = model_predicted_ref_cq
            cleanup_quality_params(current_settings, keep_param='cq')
            if logging_enabled:
                print(f"Applying model CQ directly for {MODEL_CQ_REFERENCE_CODEC}: {current_settings['cq']}")
        elif 'crf' in current_settings:
            current_settings['crf'] = model_predicted_ref_cq
            cleanup_quality_params(current_settings, keep_param='crf')
            if logging_enabled:
                print(f"Applying model CQ as CRF for {MODEL_CQ_REFERENCE_CODEC}: {current_settings['crf']}")
        else:
            if logging_enabled:
                print(f"Warning: Neither 'cq' nor 'crf' in base settings for {MODEL_CQ_REFERENCE_CODEC}. Using CRF fallback.")
            current_settings['crf'] = model_predicted_ref_cq

    elif codec in QUALITY_MAPPING_ANCHORS:
        # Use quality mapping for non-reference codecs
        mapping_config = QUALITY_MAPPING_ANCHORS[codec]

        model_anchor_cqs = [p[0] for p in mapping_config['anchor_points']]
        target_anchor_params = [p[1] for p in mapping_config['anchor_points']]

        # Interpolate with clamping
        clamped_model_cq = np.clip(model_predicted_ref_cq,
                                 mapping_config['model_ref_cq_range'][0],
                                 mapping_config['model_ref_cq_range'][1])
        if logging_enabled and clamped_model_cq != model_predicted_ref_cq:
            print(f"Clamped model_predicted_ref_cq from {model_predicted_ref_cq} to {clamped_model_cq}")

        mapped_param_float = np.interp(clamped_model_cq, model_anchor_cqs, target_anchor_params)

        # Clamp to target parameter range
        min_target_param, max_target_param = mapping_config['target_param_range']
        mapped_param_clamped = np.clip(mapped_param_float, min_target_param, max_target_param)
        mapped_param_int = int(round(mapped_param_clamped))

        target_param_type = mapping_config['target_param_type']
        if target_param_type == 'cq':
            current_settings['cq'] = mapped_param_int
            cleanup_quality_params(current_settings, keep_param='cq')
            if logging_enabled:
                print(f"Applying mapped CQ for {codec} from model ref CQ {model_predicted_ref_cq}: {current_settings['cq']}")
        elif target_param_type == 'crf':
            current_settings['crf'] = mapped_param_int
            cleanup_quality_params(current_settings, keep_param='crf')
            if logging_enabled:
                print(f"Applying mapped CRF for {codec} from model ref CQ {model_predicted_ref_cq}: {current_settings['crf']}")
        else:
            if logging_enabled:
                print(f"Warning: Unknown target_param_type '{target_param_type}' for {codec}. Using fallback.")
            if 'cq' in current_settings:
                current_settings['cq'] = model_predicted_ref_cq
            elif 'crf' in current_settings:
                current_settings['crf'] = model_predicted_ref_cq

    else:
        # Fallback for codecs without mapping
        if logging_enabled:
            print(f"Warning: No quality mapping found for {codec}. Applying model ref CQ {model_predicted_ref_cq} directly.")
        if 'cq' in current_settings:
            current_settings['cq'] = model_predicted_ref_cq
            cleanup_quality_params(current_settings, keep_param='cq')
        elif 'crf' in current_settings:
            current_settings['crf'] = model_predicted_ref_cq
            cleanup_quality_params(current_settings, keep_param='crf')
        else:
            current_settings['crf'] = model_predicted_ref_cq
        if logging_enabled:
            print(f"Applied direct rate {model_predicted_ref_cq} to {codec}.")

    return current_settings


def get_contrast_optimized_params(scene_type, contrast_value, codec):
    """
    Get contrast-optimized encoding parameters for a specific scene type and codec.
    
    Args:
        scene_type (str): The classified scene type
        contrast_value (float): The calculated perceptual contrast (0.0-1.0)
        codec (str): The codec being used
        
    Returns:
        dict: Additional encoding parameters optimized for contrast
    """
    params = {}
    
    if contrast_value > 0.7:
        contrast_category = "high"
    elif contrast_value < 0.3:
        contrast_category = "low"
    else:
        contrast_category = "medium"
    
    # --- NVENC-based codecs (AV1, HEVC, H264) ---
    if codec in ["av1_nvenc", "hevc_nvenc", "h264_nvenc"]:      
        # Set spatial AQ strength based on contrast
        if contrast_category == "high":
            params['spatial-aq'] = 1
            params['aq-strength'] = 8  # Higher strength for high contrast (NVENC 1-15)
        elif contrast_category == "low":
            params['spatial-aq'] = 1
            params['aq-strength'] = 4  # Lower strength for low contrast
        else:
            params['spatial-aq'] = 1
            params['aq-strength'] = 6  # Medium strength
        
        # Adjust temporal AQ based on scene type and contrast
        if scene_type == 'Faces / People' and contrast_category == "high":
            params['temporal-aq'] = 1  # Enable temporal AQ for high contrast faces
        else:
            params['temporal-aq'] = 0  # Default
    
    # --- x264 ---
    elif codec == "libx264":  # ✅ FIXED: Use standardized codec name
        # AQ mode and strength adjustments
        if contrast_category == "high":
            params['aq-mode'] = 2  # Variance AQ with auto-variance
            params['aq-strength'] = 1.2  # Higher strength for high contrast
        elif contrast_category == "low":
            params['aq-mode'] = 1  # Standard AQ
            params['aq-strength'] = 0.8  # Lower strength for low contrast
        else:
            params['aq-mode'] = 1
            params['aq-strength'] = 1.0  # Default
    
    # --- x265 ---
    elif codec == "libx265":  # ✅ FIXED: Use standardized codec name
        # AQ mode and strength adjustments
        if contrast_category == "high":
            params['aq-mode'] = 3  # Auto-variance AQ with more aggressive bias
            params['aq-strength'] = 1.2  # Higher strength
        elif contrast_category == "low":
            params['aq-mode'] = 2  # Standard AQ
            params['aq-strength'] = 0.8  # Lower strength
        else:
            params['aq-mode'] = 2
            params['aq-strength'] = 1.0  # Default
    
    # --- SVT-AV1 ---
    elif codec == "libsvtav1":  # ✅ FIXED: Use standardized codec name
        # SVT-AV1 doesn't have enable-hdr or aq-mode options
        # Use CRF adjustments and available parameters instead
        if contrast_category == "high":
            # For high contrast, we can rely on the encoder's default behavior
            # or use dolbyvision if HDR content is detected
            pass  # SVT-AV1 handles high contrast well by default
        elif contrast_category == "low":
            # For low contrast, we might want a slightly lower CRF for better quality
            pass  # This would be handled at the CRF level, not here
        else:
            pass  # Default SVT-AV1 behavior is generally good
    
    # --- VP9 ---
    elif codec == "libvpx_vp9":  # ✅ FIXED: Use standardized codec name (underscore)
        # VP9 adaptive quantization and quality parameters
        if contrast_category == "high":
            params['aq-mode'] = 3  # Complexity-based AQ for high contrast
            params['arnr-maxframes'] = 7  # More noise reduction frames
            params['arnr-strength'] = 5  # Higher noise reduction strength
        elif contrast_category == "low":
            params['aq-mode'] = 2  # Variance-based AQ for low contrast  
            params['arnr-maxframes'] = 5  # Standard noise reduction
            params['arnr-strength'] = 3  # Moderate noise reduction
        else:
            params['aq-mode'] = 1  # Variance-based AQ (default)
            params['arnr-maxframes'] = 7  # Standard noise reduction
            params['arnr-strength'] = 4  # Standard noise reduction strength
        
        # VP9-specific quality settings
        params['tune'] = 'psnr'  # Optimize for quality
        params['auto-alt-ref'] = 1  # Enable alternate reference frames
    
    # --- Scene-specific contrast adjustments ---
    # For text content, adjust parameters further based on contrast
    if scene_type == 'Screen Content / Text':
        if contrast_category == "high":
            params['sharpness'] = 0  # Preserve sharpness
            if "_nvenc" in codec:  # ✅ FIXED: lowercase nvenc
                params['rc-lookahead'] = 20  # More lookahead for complex text
        elif contrast_category == "low":
            params['sharpness'] = 1  # Some sharpening for low contrast text
            if "_nvenc" in codec:  # ✅ FIXED: lowercase nvenc
                params['rc-lookahead'] = 8  # Less lookahead needed
    
    # For faces, special handling based on contrast
    elif scene_type == 'Faces / People':
        if contrast_category == "high":
            if "_nvenc" in codec:  # ✅ FIXED: lowercase nvenc
                params['rc-lookahead'] = 20  # More lookahead for dramatic lighting
        else:
            if "_nvenc" in codec:  # ✅ FIXED: lowercase nvenc
                params['rc-lookahead'] = 15  # Standard lookahead
    
    return params

def encode_video(input_path, output_path, codec, rate=None, preset=None, scene_type=None, contrast_value=None, codec_mode=None, target_bitrate=None, logging_enabled=True):
    """
    Encodes a video using specified codec settings, optimized for scene type and contrast.
    Audio is copied without re-encoding for efficiency.
    Note: encoding decisions should be made BEFORE calling this function.
    Args:
        input_path (str): Path to input video file
        output_path (str): Path for output video file
        codec (str): Video codec to use (e.g., 'av1_nvenc', 'libx264')
        rate (int/float, optional): Quality parameter (CQ value from model)
        preset (str, optional): Encoder preset override
        scene_type (str, optional): Scene classification for optimization
        contrast_value (float, optional): Perceptual contrast (0.0-1.0)
        codec_mode (str, optional): Encoding mode - 'CRF', 'CBR', or 'VBR'
        target_bitrate (float, optional): Target bitrate in Mbps (for CBR/VBR modes)
        logging_enabled (bool): Enable detailed logging
    Returns:
        tuple: (encoding_results_log, encoding_time) or (None, None) on failure
    """
    
    if logging_enabled:
        print(f"Encoding video using codec: {codec}, scene: {scene_type}, model_predicted_rate: {rate}")
        if contrast_value is not None:
            print(f"Using contrast value: {contrast_value:.2f}")

    # Get base encoder settings
    base_settings = ENCODER_SETTINGS.get(codec)
    if not base_settings:
        print(f"Error: Codec settings for '{codec}' not found in ENCODER_SETTINGS.")
        return None, None

    # --- Parameter Prioritization ---
    # 1. Start with base settings
    current_settings = base_settings.copy()

    # 2. Apply scene-specific overrides if scene_type is provided
    if scene_type:
        codec_scene_params = SCENE_SPECIFIC_PARAMS.get(codec, {})
        scene_params = codec_scene_params.get(scene_type, codec_scene_params.get('other', {}))
        if scene_params:
            if logging_enabled:
                print(f"Applying scene-specific params for '{scene_type}': {scene_params}")
            current_settings.update(scene_params)

    # 3. Apply contrast-specific overrides if contrast_value is provided
    if contrast_value is not None:
        contrast_params = get_contrast_optimized_params(scene_type, contrast_value, codec)
        if contrast_params:
            if logging_enabled:
                print(f"Applying contrast-specific params: {contrast_params}")
            current_settings.update(contrast_params)

    # 4. Apply codec_mode and target_bitrate if provided (BEFORE rate mapping)
    # This determines whether we should apply CQ/CRF or use bitrate control
    skip_rate_application = False
    if codec_mode and target_bitrate:
        if logging_enabled:
            print(f"Applying codec_mode='{codec_mode}' with target_bitrate={target_bitrate} Mbps")

        if codec_mode.upper() == 'CBR':
            # Constant Bitrate Mode - fixed bitrate throughout
            skip_rate_application = True  # Don't apply CQ/CRF in CBR mode
            bitrate_kbps = int(target_bitrate * 1000)  # Convert Mbps to kbps
            bitrate_value = f"{bitrate_kbps}k"

            current_settings['bitrate'] = bitrate_value
            current_settings['maxrate'] = bitrate_value
            current_settings['bufsize'] = f"{bitrate_kbps * 2}k"  # Buffer = 2x bitrate

            # Remove quality parameters (incompatible with CBR)
            cleanup_quality_params(current_settings)

            # Set rate control mode for NVENC/QSV
            if codec.endswith('_nvenc'):
                current_settings['rc'] = 'cbr'
                if logging_enabled:
                    print(f"Set NVENC rate control to CBR")
            elif codec.endswith('_qsv'):
                current_settings['rc'] = 'cbr'
            elif 'libvpx' in codec:
                # VP9/VP8 CBR mode
                current_settings['minrate'] = bitrate_value
                if logging_enabled:
                    print(f"Set libvpx rate control to CBR with minrate={bitrate_value}")

            if logging_enabled:
                print(f"CBR mode: bitrate={bitrate_value}, maxrate={bitrate_value}, bufsize={bitrate_kbps * 2}k")

        elif codec_mode.upper() == 'VBR':
            # Variable Bitrate Mode - allows bitrate to vary but caps at target
            # We will still apply CQ/CRF for quality, but add maxrate constraint
            bitrate_kbps = int(target_bitrate * 1000)
            bitrate_value = f"{bitrate_kbps}k"

            current_settings['maxrate'] = bitrate_value
            current_settings['bufsize'] = f"{bitrate_kbps * 2}k"

            # Set rate control mode for NVENC/QSV
            if codec.endswith('_nvenc') or codec.endswith('_qsv'):
                current_settings['rc'] = 'vbr'
                if logging_enabled:
                    print(f"Set rate control to VBR")
            elif 'libvpx' in codec:
                # VP9/VP8 VBR mode - set target bitrate for VBR with CRF
                current_settings['bitrate'] = bitrate_value
                if logging_enabled:
                    print(f"Set libvpx VBR mode with target bitrate={bitrate_value}")

            if logging_enabled:
                print(f"VBR mode: maxrate={bitrate_value}, bufsize={bitrate_kbps * 2}k, will apply CQ/CRF for quality")

        elif codec_mode.upper() == 'CRF':
            # CRF mode is the default - will apply rate parameter below
            if logging_enabled:
                print(f"CRF mode: Will apply quality-based encoding (rate={rate})")

        else:
            if logging_enabled:
                print(f"Warning: Unknown codec_mode '{codec_mode}', defaulting to CRF behavior")

    # 5. Apply rate (CQ) parameter with codec-specific handling (skip if CBR mode)
    if rate is not None and not skip_rate_application:
        current_settings = apply_rate_mapping(codec, rate, current_settings, logging_enabled)

    # 6. Apply preset override if provided
    if preset is not None:
        current_settings['preset'] = preset
        if logging_enabled:
            print(f"Applying preset override: {preset}")

    # 7. Handle CRF usage - disable constqp when CRF is used
    try:
        uses_crf = 'crf' in current_settings
        has_maxrate = 'maxrate' in current_settings

        if uses_crf:
            # Remove constqp explicitly when CRF is used
            if str(current_settings.get('rc', '')).lower() == 'constqp':
                if logging_enabled:
                    print("Disabling 'constqp' because CRF is in use")
                del current_settings['rc']

            # For NVENC/QSV with CRF + maxrate, use VBR (only if rc not already set by codec_mode)
            if (codec.endswith('_nvenc') or codec.endswith('_qsv')) and has_maxrate and 'rc' not in current_settings:
                current_settings['rc'] = 'vbr'
                if logging_enabled:
                    print(f"Using 'vbr' with CRF due to specified maxrate")
    except Exception:
        pass

    # 8. CQ policy for NVENC/QSV - use VBR with CQ (only if rc not already set by codec_mode)
    try:
        has_cq = 'cq' in current_settings and isinstance(current_settings.get('cq'), (int, float))
        has_maxrate = 'maxrate' in current_settings
        rc_already_set = 'rc' in current_settings

        if (codec.endswith('_nvenc') or codec.endswith('_qsv')) and has_cq and not rc_already_set:
            # Standard CQ encoding with VBR (fallback when codec_mode not specified)
            current_settings['rc'] = 'vbr'
            cleanup_quality_params(current_settings, keep_param='cq')
            if logging_enabled:
                maxrate_info = " and maxrate" if has_maxrate else ""
                print(f"Using VBR with CQ{maxrate_info} for {codec}")
    except Exception:
        pass

    # --- Build FFmpeg arguments ---
    output_args = {}
    
    # Map internal setting names to FFmpeg argument names
    key_map = {
        'keyint': 'g',
        'bitrate': 'b:v',
        'codec': 'vcodec'
    }

    # For NVENC in constqp mode, translate 'cq' to 'qp'
    try:
        if codec.endswith('_nvenc') and str(current_settings.get('rc', '')).lower() == 'constqp':
            if 'cq' in current_settings and 'qp' not in current_settings:
                current_settings['qp'] = int(current_settings.pop('cq'))
                if logging_enabled:
                    print("Translating NVENC constqp: using -qp instead of -cq")
    except Exception:
        pass

    # Add all current settings to output_args
    for key, value in current_settings.items():
        ffmpeg_key = key_map.get(key, key)
        output_args[ffmpeg_key] = value

    # Ensure vcodec is set
    if 'vcodec' not in output_args and 'codec' in current_settings:
        output_args['vcodec'] = current_settings['codec']

    # Remove internal 'codec' key if different from 'vcodec'
    if 'codec' in output_args and output_args.get('codec') != output_args.get('vcodec'):
        del output_args['codec']

    # --- Audio handling ---
    has_audio = check_audio_stream(input_path)
    
    if has_audio:
        output_args['acodec'] = 'copy'
        output_args['map'] = '0'  # Map all streams
    else:
        output_args['map'] = '0:v:0'  # Only video stream
    
    if logging_enabled:
        print(f"Audio detected: {has_audio}")
        print("Final FFmpeg output args:", output_args)

    # --- Execute FFmpeg ---
    try:
        start_time = time.time()
        
        input_stream = ffmpeg.input(input_path)
        output_stream = ffmpeg.output(input_stream, output_path, **output_args)
        
        result = output_stream.run(
            overwrite_output=True,
            capture_stdout=True,
            capture_stderr=True,
            quiet=(not logging_enabled)
        )
        
        end_time = time.time()
        encoding_time_calculated = round(end_time - start_time, 2)
        stderr = result[1].decode("utf-8") if result[1] else ""

        # Extract final encoding log line for debugging
        encoding_results_log = None
        if logging_enabled:
            for line in stderr.splitlines():
                if "time=" in line:
                    encoding_results_log = line

        if logging_enabled:
            print(f"Successfully encoded using {codec} scene '{scene_type}': {output_path}")
        return encoding_results_log, encoding_time_calculated

    except ffmpeg.Error as e:
        err_text = e.stderr.decode('utf8') if getattr(e, 'stderr', None) else str(e)
        if logging_enabled:
            print(f"Error encoding with {codec} scene '{scene_type}': {e}")
            print(f"FFmpeg stderr: {err_text}")
            print("Trying with libsvtav1 encoder...")

        try:
            # switch codec from gpu to cpu
            output_args["vcodec"] = "libsvtav1"

            # remove nvenc-specific options that might cause fallback errors
            for key in ["spatial-aq", "temporal-aq", "aq-strength", "rc-lookahead", "preset"]:
                output_args.pop(key, None)

            start_time = time.time()
            input_stream = ffmpeg.input(input_path)
            output_stream = ffmpeg.output(input_stream, output_path, **output_args)

            result = output_stream.run(
                overwrite_output=True,
                capture_stdout=True,
                capture_stderr=True,
                quiet=(not logging_enabled)
            )

            end_time = time.time()
            encoding_time_calculated = round(end_time - start_time, 2)
            stderr = result[1].decode("utf-8") if result[1] else ""

            encoding_results_log = None
            if logging_enabled:
                for line in stderr.splitlines():
                    if "time=" in line:
                        encoding_results_log = line

            if logging_enabled:
                print(f"Successfully encoded using libsvtav1 fallback for scene '{scene_type}': {output_path}")
            return encoding_results_log, encoding_time_calculated

        except Exception as fallback_error:
            if logging_enabled:
                print(f"Fallback encoding with libsvtav1 also failed: {fallback_error}")
            return None, None
            
def encode_lossless_video(input_path, output_path, logging_enabled=True):
    """
    Encode video using FFV1 lossless codec with optimized settings.
    This function is specifically for lossless encoding and doesn't use the main
    encode_video function to avoid codec mapping issues.
    """
    import subprocess
    import time
    
    if logging_enabled:
        print(f"🎥 Encoding lossless video with FFV1 codec...")
    
    # FFV1 lossless encoding command
    cmd = [
        'ffmpeg',
        '-i', input_path,
        '-c:v', 'ffv1',         
        '-level', '3',          
        '-coder', '1',          
        '-context', '1',        
        '-g', '1',              
        '-slices', '4',         
        '-slicecrc', '1',       
        '-pix_fmt', 'yuv420p',  
        '-c:a', 'copy',         
        '-y',                   
        output_path
    ]
    
    try:
        start_time = time.time()
        
        if logging_enabled:
            print(f"🔧 FFmpeg command: {' '.join(cmd)}")
        
        # Run FFmpeg command
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True
        )
        
        end_time = time.time()
        encoding_time = round(end_time - start_time, 2)
        
        if logging_enabled:
            print(f"✅ FFV1 lossless encoding completed in {encoding_time:.1f}s")
        
        return "FFV1 lossless encoding successful", encoding_time
        
    except subprocess.CalledProcessError as e:
        if logging_enabled:
            print(f"❌ FFmpeg error: {e}")
            print(f"   stdout: {e.stdout}")
            print(f"   stderr: {e.stderr}")
        return None, None
        
    except Exception as e:
        if logging_enabled:
            print(f"❌ Unexpected error during lossless encoding: {e}")
        return None, None

def check_audio_stream(input_path):
    """Check if the input video has an audio stream."""
    try:
        import subprocess
        import json
        
        cmd = [
            'ffprobe', '-v', 'quiet', '-print_format', 'json', 
            '-show_streams', input_path
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        data = json.loads(result.stdout)
        
        # Check if any stream is audio
        has_audio = any(stream.get('codec_type') == 'audio' for stream in data.get('streams', []))
        return has_audio
    except Exception as e:
        # If we can't detect, assume no audio to be safe
        return False