first changes to original scripts
This commit is contained in:
260
old/vmaf_encoder.py
Normal file
260
old/vmaf_encoder.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user