Files
VMAFOptimiser/optimize_library.py
2025-12-31 23:13:32 +04:00

834 lines
26 KiB
Python

import os
import sys
import subprocess
import argparse
import json
import shutil
import platform
import time
import signal
from pathlib import Path
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed, ProcessPoolExecutor
from threading import Lock
# --- Configuration ---
DEFAULT_VMAF = 95.0
DEFAULT_PRESET = 6
DEFAULT_WORKERS = 1
DEFAULT_SAMPLES = 4
EXTENSIONS = {".mkv", ".mp4", ".mov", ".avi", ".ts"}
TARGETS = [94.0, 93.0, 92.0, 90.0]
MIN_SAVINGS_PERCENT = 12.0
TARGET_SAVINGS_FOR_ESTIMATE = 15.0
# Global state for resume capability
_processed_files = set()
_lock = Lock()
_shutdown_requested = False
_AB_AV1_HELP_CACHE = {}
def signal_handler(signum, frame):
"""Handle graceful shutdown"""
global _shutdown_requested
_shutdown_requested = True
print("\n\n⚠️ Shutdown requested. Finishing current tasks...")
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
def check_dependencies():
missing = []
for tool in ["ffmpeg", "ffprobe", "ab-av1"]:
if not shutil.which(tool):
missing.append(tool)
if missing:
print(f"Error: Missing required tools: {', '.join(missing)}")
print(
"Please install FFmpeg and 'ab-av1' (via cargo install ab-av1) before running."
)
sys.exit(1)
def is_wsl():
if os.environ.get("WSL_DISTRO_NAME"):
return True
try:
with open("/proc/sys/kernel/osrelease", "r", encoding="utf-8") as f:
return "microsoft" in f.read().lower()
except FileNotFoundError:
return False
def platform_label():
system = platform.system()
if system == "Linux" and is_wsl():
return "Linux (WSL)"
return system
def _ab_av1_help(subcommand):
cached = _AB_AV1_HELP_CACHE.get(subcommand)
if cached is not None:
return cached
try:
result = subprocess.run(
["ab-av1", subcommand, "--help"],
capture_output=True,
text=True,
check=False,
)
help_text = (result.stdout or "") + "\n" + (result.stderr or "")
except Exception:
help_text = ""
_AB_AV1_HELP_CACHE[subcommand] = help_text
return help_text
def ab_av1_supports(subcommand, flag):
return flag in _ab_av1_help(subcommand)
def normalize_hwaccel(value):
if value is None:
return None
v = value.strip()
if not v:
return None
v_lower = v.lower()
if v_lower in {"none", "off", "false", "0"}:
return None
if v_lower != "auto":
return v
system = platform.system()
if system == "Windows":
return "d3d11va"
if system == "Darwin":
return "videotoolbox"
return "vaapi"
def get_probe_data(filepath):
"""Get comprehensive video data using ffprobe"""
try:
cmd = [
"ffprobe",
"-v",
"quiet",
"-print_format",
"json",
"-show_streams",
"-show_format",
str(filepath),
]
res = subprocess.run(cmd, capture_output=True, text=True, check=True)
return json.loads(res.stdout)
except Exception as e:
print(f"Error probing {filepath}: {e}")
return None
def get_video_stats(data):
"""Extract video statistics from ffprobe output"""
if not data or "streams" not in data or "format" not in data:
return None
v_stream = next((s for s in data["streams"] if s["codec_type"] == "video"), None)
if not v_stream:
return None
size = int(data["format"].get("size", 0))
duration = float(data["format"].get("duration", 0))
# Calculate bitrate from file size and duration (more reliable than ffprobe's bitrate)
if size > 0 and duration > 0:
bitrate = int((size * 8) / duration / 1000)
else:
bitrate = int(data["format"].get("bitrate", 0)) // 1000
return {
"codec": v_stream.get("codec_name"),
"width": v_stream.get("width"),
"height": v_stream.get("height"),
"bitrate": bitrate,
"size": size,
"duration": duration,
}
def log_result(log_dir, log_name, data):
"""Log result to JSONL file"""
os.makedirs(log_dir, exist_ok=True)
log_file = Path(log_dir) / f"{log_name}.jsonl"
data["timestamp"] = datetime.now().isoformat()
with _lock:
with open(log_file, "a") as f:
f.write(json.dumps(data) + "\n")
def run_command_streaming(cmd, description=""):
"""Run command and stream output in real-time"""
print(f" [Running {description}]")
process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
universal_newlines=True,
bufsize=1,
)
if process.stdout:
for line in process.stdout:
if _shutdown_requested:
process.terminate()
break
print(f" {line.rstrip()}")
process.wait()
return process.returncode
def run_crf_search(filepath, target_vmaf, preset, temp_dir, use_hw=False, hwaccel=None):
"""Run CRF search for a specific VMAF target"""
cmd = [
"ab-av1",
"crf-search",
"-i",
str(filepath),
"--min-vmaf",
str(target_vmaf),
"--preset",
str(preset),
"--max-encoded-percent",
"100",
"--temp-dir",
temp_dir,
"--samples",
"4", # Use 4 samples for speed/accuracy balance
]
# Hardware encoding support
if use_hw and hwaccel:
if ab_av1_supports("crf-search", "--enc-input"):
cmd.extend(["--enc-input", f"hwaccel={hwaccel}"])
if hwaccel == "vaapi":
cmd.extend(["--enc-input", "hwaccel_output_format=vaapi"])
print(f" - Searching for CRF to hit VMAF {target_vmaf}...")
returncode = run_command_streaming(cmd, f"crf-search VMAF {target_vmaf}")
if returncode == 0:
# Parse output to find CRF and predicted size
res = subprocess.run(
[
"ab-av1",
"crf-search",
"-i",
str(filepath),
"--min-vmaf",
str(target_vmaf),
"--preset",
str(preset),
"--temp-dir",
temp_dir,
"--samples",
"4",
],
capture_output=True,
text=True,
)
lines = res.stdout.strip().split("\n")
for line in reversed(lines):
if "crf" in line.lower():
try:
parts = line.split()
crf_val = float(parts[1])
percent = 100.0
for p in parts:
if "%" in p:
percent = float(p.strip("()%"))
break
return {
"crf": crf_val,
"predicted_percent": percent,
"vmaf": target_vmaf,
}
except Exception as e:
print(f" ! Failed to parse crf-search output: {e}")
return None
def find_target_savings_params(
filepath, start_vmaf, preset, temp_dir, use_hw=False, hwaccel=None
):
"""Find VMAF target that achieves minimum savings"""
print(f"\n --- Finding VMAF for {TARGET_SAVINGS_FOR_ESTIMATE}% savings ---")
test_targets = [t for t in TARGETS if t <= start_vmaf]
for i, target in enumerate(test_targets):
if _shutdown_requested:
return None
print(
f" Testing VMAF {target} for {TARGET_SAVINGS_FOR_ESTIMATE}% target... (test {i + 1}/{len(test_targets)})"
)
result = run_crf_search(filepath, target, preset, temp_dir, use_hw, hwaccel)
if result:
predicted_savings = 100.0 - result["predicted_percent"]
quality_drop = start_vmaf - target
print(
f" ✓ VMAF {target}: CRF {result['crf']}, Savings {predicted_savings:.1f}%, Drop -{quality_drop:.0f}"
)
if predicted_savings >= TARGET_SAVINGS_FOR_ESTIMATE:
print(f"\n ✅ FOUND {TARGET_SAVINGS_FOR_ESTIMATE}%+ SAVINGS:")
print(f" Target VMAF: {target} (quality drop: -{quality_drop:.0f})")
print(f" CRF: {result['crf']}")
print(f" Predicted savings: {predicted_savings:.1f}%")
return {
"target_vmaf": target,
"crf": result["crf"],
"savings": predicted_savings,
"quality_drop": quality_drop,
"found": True,
}
else:
print(f" ✗ Could not achieve VMAF {target}")
print(f"\n 📝 COULD NOT ACHIEVE {TARGET_SAVINGS_FOR_ESTIMATE}% SAVINGS")
print(f" Tried VMAF targets: {test_targets}")
return None
def run_encode(filepath, output_path, crf, preset, use_hw=False, hwaccel=None):
"""Run full encoding with real-time output streaming"""
cmd = [
"ab-av1",
"encode",
"--input",
str(filepath),
"--output",
str(output_path),
"--crf",
str(crf),
"--preset",
str(preset),
"--acodec",
"copy",
]
# Hardware encoding support
if use_hw and hwaccel:
if ab_av1_supports("encode", "--enc-input"):
cmd.extend(["--enc-input", f"hwaccel={hwaccel}"])
if hwaccel == "vaapi":
cmd.extend(["--enc-input", "hwaccel_output_format=vaapi"])
return run_command_streaming(cmd, f"encoding (CRF {crf})")
def provide_recommendations(
stats_before, hit_vmaf, predicted_savings, target_result=None
):
"""Provide recommendations based on analysis results"""
print(f"\n --- Recommendations for {stats_before['codec']} content ---")
if target_result and target_result["found"]:
print(f" 📊 TO HIT {TARGET_SAVINGS_FOR_ESTIMATE}% SAVINGS:")
print(
f" → Target VMAF: {target_result['target_vmaf']} (drop: -{target_result['quality_drop']:.0f})"
)
print(f" → CRF: {target_result['crf']}")
print(f" → Predicted: {target_result['savings']:.1f}% savings")
print(f" → Trade-off: Quality reduction for space savings")
print()
if stats_before["bitrate"] < 2000:
print(f" → Source bitrate is low ({stats_before['bitrate']}k)")
print(f" → AV1 gains minimal on highly compressed sources")
print(f" → Recommendation: SKIP - Source already optimized")
return
if stats_before["codec"] in ["hevc", "h265", "vp9"]:
print(f" → Source already uses modern codec ({stats_before['codec']})")
print(f" → AV1 gains minimal on already-compressed content")
print(f" → Recommendation: SKIP - Codec already efficient")
return
if target_result and not target_result["found"]:
print(
f" → Could not achieve {TARGET_SAVINGS_FOR_ESTIMATE}% even at lowest VMAF"
)
print(f" → Content may not compress well with AV1")
print(f" → Recommendation: SKIP - Review manually")
def refresh_plex(plex_url, plex_token):
"""Refresh Plex library after encoding completion"""
if not plex_url or not plex_token:
return
try:
print("\n📡 Refreshing Plex library...")
cmd = [
"curl",
"-X",
"GET",
f"{plex_url}/library/sections/1/refresh",
"-H",
f"X-Plex-Token: {plex_token}",
]
subprocess.run(cmd, capture_output=True, check=False)
print(" ✓ Plex refresh triggered")
except Exception as e:
print(f" ⚠️ Failed to refresh Plex: {e}")
def process_file(filepath, log_dir, log_name, preset, use_hw=False, hwaccel=None):
"""Process a single video file with intelligent VMAF targeting"""
global _shutdown_requested
filepath = Path(filepath)
lock_file = Path(log_dir).parent / ".lock" / f"{filepath.name}.lock"
# Check lock file (multi-machine coordination)
if lock_file.exists():
print(f"Skipping (Locked): {filepath.name}")
return True
# Create lock
lock_file.parent.mkdir(parents=True, exist_ok=True)
lock_file.touch()
try:
probe_data = get_probe_data(filepath)
if not probe_data:
print(f"Skipping (Invalid probe data): {filepath.name}")
return True
stats_before = get_video_stats(probe_data)
if not stats_before or stats_before["codec"] == "av1":
print(f"Skipping (Already AV1 or invalid): {filepath.name}")
return True
# Mark as processed
file_key = str(filepath)
with _lock:
if file_key in _processed_files:
return True
_processed_files.add(file_key)
print(f"\n--- Processing: {filepath.name} ---")
print(
f" Source: {stats_before['codec']} @ {stats_before['bitrate']}k, {stats_before['size'] / (1024**3):.2f} GB"
)
if _shutdown_requested:
return False
temp_dir = Path(log_dir).parent / "tmp"
temp_dir.mkdir(exist_ok=True)
# Step 1: Try VMAF 94
print(f"\n [Step 1] Testing VMAF 94...")
search_result_94 = run_crf_search(
filepath, 94.0, preset, str(temp_dir), use_hw, hwaccel
)
if not search_result_94:
print(f" !! Could not hit VMAF 94")
search_result_94 = run_crf_search(
filepath, 93.0, preset, str(temp_dir), use_hw, hwaccel
)
if not search_result_94:
search_result_94 = run_crf_search(
filepath, 92.0, preset, str(temp_dir), use_hw, hwaccel
)
if not search_result_94:
search_result_94 = run_crf_search(
filepath, 90.0, preset, str(temp_dir), use_hw, hwaccel
)
if not search_result_94:
print(f" !! Failed all VMAF targets ({TARGETS}) for {filepath.name}")
log_result(
log_dir,
"failed_searches",
{
"file": str(filepath),
"status": "all_targets_failed",
"targets": TARGETS,
},
)
provide_recommendations(stats_before, None, 0)
return False
crf_94 = search_result_94["crf"]
predicted_savings_94 = 100.0 - search_result_94["predicted_percent"]
if predicted_savings_94 >= MIN_SAVINGS_PERCENT:
print(
f"\n ✅ VMAF 94 gives {predicted_savings_94:.1f}% savings (≥{MIN_SAVINGS_PERCENT}%)"
)
print(f" → Proceeding with VMAF 94, CRF {crf_94}")
encode_params = {
"crf": crf_94,
"vmaf": 94.0,
"predicted_percent": search_result_94["predicted_percent"],
}
else:
print(
f"\n ⚠️ VMAF 94 gives {predicted_savings_94:.1f}% savings (<{MIN_SAVINGS_PERCENT}%)"
)
search_result_93 = run_crf_search(
filepath, 93.0, preset, str(temp_dir), use_hw, hwaccel
)
if search_result_93:
predicted_savings_93 = 100.0 - search_result_93["predicted_percent"]
if predicted_savings_93 >= MIN_SAVINGS_PERCENT:
print(
f" ✅ VMAF 93 gives {predicted_savings_93:.1f}% savings (≥{MIN_SAVINGS_PERCENT}%)"
)
print(
f" → Proceeding with VMAF 93, CRF {search_result_93['crf']}"
)
encode_params = {
"crf": search_result_93["crf"],
"vmaf": 93.0,
"predicted_percent": search_result_93["predicted_percent"],
}
else:
print(
f" ⚠️ VMAF 93 gives {predicted_savings_93:.1f}% savings (also <{MIN_SAVINGS_PERCENT}%)"
)
print(
f" → Finding VMAF for {TARGET_SAVINGS_FOR_ESTIMATE}% savings..."
)
target_result = find_target_savings_params(
filepath, 93.0, preset, str(temp_dir), use_hw, hwaccel
)
provide_recommendations(
stats_before, 93.0, predicted_savings_93, target_result
)
log_result(
log_dir,
"low_savings_skips",
{
"file": str(filepath),
"vmaf_94": 94.0,
"savings_94": predicted_savings_94,
"vmaf_93": 93.0,
"savings_93": predicted_savings_93,
"target_for_15_percent": target_result,
"recommendations": "logged_for_review",
},
)
return True
else:
print(f" !! Could not achieve VMAF 93")
log_result(
log_dir,
"failed_searches",
{"file": str(filepath), "status": "vmaf_93_failed"},
)
return False
temp_output = temp_dir / f"{filepath.stem}.av1_temp.mkv"
if temp_output.exists():
temp_output.unlink()
start_time = time.time()
res = run_encode(
filepath, temp_output, encode_params["crf"], preset, use_hw, hwaccel
)
if res != 0:
print(f"\n !! Encode failed with code {res}")
if temp_output.exists():
temp_output.unlink()
log_result(
log_dir,
"failed_encodes",
{"file": str(filepath), "status": "encode_failed", "returncode": res},
)
return False
encode_duration = time.time() - start_time
print(f" ✓ Encode completed in {encode_duration:.1f}s")
probe_after = get_probe_data(temp_output)
stats_after = get_video_stats(probe_after)
if not stats_after:
print(f" !! Failed to probe encoded file")
if temp_output.exists():
temp_output.unlink()
return False
actual_savings = (1 - (stats_after["size"] / stats_before["size"])) * 100
print(f"\n Results:")
print(
f" - Before: {stats_before['size'] / (1024**3):.2f} GB @ {stats_before['bitrate']}k"
)
print(
f" - After: {stats_after['size'] / (1024**3):.2f} GB @ {stats_after['bitrate']}k"
)
print(f" - Savings: {actual_savings:.2f}%")
final_path = filepath
if filepath.suffix.lower() == ".mp4":
final_path = filepath.with_suffix(".mkv")
if final_path.exists():
final_path.unlink()
shutil.move(str(filepath), str(final_path))
shutil.move(str(temp_output), str(final_path))
print(f" ✓ Successfully optimized: {final_path.name}")
log_result(
log_dir,
log_name,
{
"file": str(final_path),
"status": "success",
"vmaf": encode_params["vmaf"],
"crf": encode_params["crf"],
"before": stats_before,
"after": stats_after,
"duration": encode_duration,
"savings": actual_savings,
},
)
return True
finally:
# Remove lock file
if lock_file.exists():
lock_file.unlink()
def scan_library(root, exclude_dirs=None):
"""Scan library for video files, excluding certain directories"""
exclude_dirs = exclude_dirs or []
files = []
for dirpath, dirnames, filenames in os.walk(root):
# Skip excluded directories
dirnames[:] = [d for d in dirnames if d not in exclude_dirs]
for filename in filenames:
if Path(filename).suffix.lower() not in EXTENSIONS:
continue
full_path = Path(dirpath) / filename
if "_av1" in full_path.stem:
continue
files.append(full_path)
return files
def main():
parser = argparse.ArgumentParser(
description="Optimize video library to AV1 using VMAF targeting."
)
parser.add_argument("directory", help="Root directory to scan")
parser.add_argument(
"--vmaf",
type=float,
default=DEFAULT_VMAF,
help=f"Target VMAF score (default: {DEFAULT_VMAF})",
)
parser.add_argument(
"--preset",
type=int,
default=DEFAULT_PRESET,
help=f"SVT-AV1 Preset (default: {DEFAULT_PRESET})",
)
parser.add_argument(
"--workers",
type=int,
default=DEFAULT_WORKERS,
help=f"Concurrent files to process (default: {DEFAULT_WORKERS})",
)
parser.add_argument(
"--samples",
type=int,
default=DEFAULT_SAMPLES,
help=f"Samples to use for CRF search if supported (default: {DEFAULT_SAMPLES})",
)
parser.add_argument(
"--hwaccel",
default=None,
help=(
"Hardware acceleration for decode. "
"Examples: auto, vaapi, d3d11va, videotoolbox. Use 'none' to disable."
),
)
parser.add_argument(
"--use-hardware-worker",
action="store_true",
help="Use 1 hardware encoding worker + rest CPU workers (requires --hwaccel)",
)
parser.add_argument(
"--plex-url",
default=None,
help="Plex server URL (e.g., http://localhost:32400)",
)
parser.add_argument(
"--plex-token",
default=None,
help="Plex auth token",
)
parser.add_argument(
"--log-dir",
default="/opt/Optmiser/logs",
help="Log directory (default: /opt/Optmiser/logs)",
)
args = parser.parse_args()
if args.workers < 1:
print("Error: --workers must be >= 1")
sys.exit(2)
check_dependencies()
root = Path(args.directory)
if not root.exists():
print(f"Directory not found: {root}")
sys.exit(1)
hwaccel = normalize_hwaccel(args.hwaccel)
print(f"Platform: {platform_label()}")
print(f"Scanning library: {root}")
print(f"VMAF targets: {TARGETS}")
print(f"Minimum savings: {MIN_SAVINGS_PERCENT}%")
print(f"Estimate target: {TARGET_SAVINGS_FOR_ESTIMATE}%")
print(f"Encoder Preset: {args.preset}")
print(f"Workers: {args.workers}")
if hwaccel:
print(f"HWAccel: {hwaccel}")
if args.use_hardware_worker:
print(f"Hardware worker: 1 HW + {args.workers - 1} CPU")
if args.plex_url:
print(f"Plex refresh: Enabled")
print("-" * 60)
# Determine log name based on directory
root_parts = str(root).lower().split("/")
if "content" in root_parts:
log_name = "content"
exclude_dirs = []
else:
log_name = "tv_movies"
exclude_dirs = ["content"]
print(f"Log file: {log_name}.jsonl")
files = scan_library(root, exclude_dirs)
if not files:
print("No media files found.")
return
print(f"Found {len(files)} files to process")
print("-" * 60)
processed_count = 0
success_count = 0
fail_count = 0
# Hardware worker configuration
use_hw_primary = args.use_hardware_worker and hwaccel is not None
if args.workers == 1:
# Single thread - just process files
for file_path in files:
if _shutdown_requested:
break
processed_count += 1
result = process_file(
file_path, args.log_dir, log_name, args.preset, use_hw_primary, hwaccel
)
if result:
success_count += 1
else:
fail_count += 1
else:
# Multi-threaded processing
with ThreadPoolExecutor(max_workers=args.workers) as executor:
futures = []
for file_path in files:
if _shutdown_requested:
break
# Use hardware for first file, CPU for rest
use_hw_for_this = use_hw_primary and len(futures) == 0
future = executor.submit(
process_file,
file_path,
args.log_dir,
log_name,
args.preset,
use_hw_for_this,
hwaccel,
)
futures.append(future)
for future in as_completed(futures):
if _shutdown_requested:
break
processed_count += 1
try:
result = future.result()
if result:
success_count += 1
else:
fail_count += 1
except Exception as e:
print(f" !! Unexpected error: {e}")
import traceback
traceback.print_exc()
fail_count += 1
print("\n" + "=" * 60)
print(f"SUMMARY: {root}")
print(f" Processed: {processed_count} files")
print(f" Success/Skip: {success_count}")
print(f" Failed: {fail_count}")
print("=" * 60)
# Refresh Plex on completion
if success_count > 0:
refresh_plex(args.plex_url, args.plex_token)
if __name__ == "__main__":
main()