260 lines
11 KiB
Python
260 lines
11 KiB
Python
#!/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() |