932 lines
28 KiB
Python
932 lines
28 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, get_ident
|
|
|
|
# --- 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
|
|
|
|
# Hardware encoder quality compensation (GPU needs higher VMAF target to match CPU quality)
|
|
HW_ENCODER_VMAF_OFFSET = 2.0
|
|
HW_ENCODERS = {"av1_amf", "av1_nvenc", "av1_qsv", "av1_vaapi"}
|
|
|
|
# Global state for resume capability
|
|
_processed_files = set()
|
|
_lock = Lock()
|
|
_shutdown_requested = False
|
|
_AB_AV1_HELP_CACHE = {}
|
|
_hw_worker_id = None # Thread ID of the designated hardware worker
|
|
|
|
|
|
def claim_hardware_worker():
|
|
"""Attempt to claim hardware worker status for this thread. Returns True if claimed."""
|
|
global _hw_worker_id
|
|
with _lock:
|
|
thread_id = get_ident()
|
|
if _hw_worker_id is None:
|
|
_hw_worker_id = thread_id
|
|
return True
|
|
return _hw_worker_id == thread_id
|
|
|
|
|
|
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),
|
|
]
|
|
# FIX: Added encoding="utf-8" and errors="ignore"
|
|
res = subprocess.run(
|
|
cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
check=True,
|
|
encoding="utf-8",
|
|
errors="ignore"
|
|
)
|
|
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,
|
|
encoder="svt-av1",
|
|
use_hw=False,
|
|
hwaccel=None,
|
|
):
|
|
"""Run CRF search for a specific VMAF target"""
|
|
is_hw_encoder = encoder in HW_ENCODERS
|
|
|
|
# Apply VMAF offset for hardware encoders to match software quality
|
|
effective_vmaf = (
|
|
target_vmaf + HW_ENCODER_VMAF_OFFSET if is_hw_encoder else target_vmaf
|
|
)
|
|
|
|
# --- Build the Command (Used for both running and parsing) ---
|
|
cmd = [
|
|
"ab-av1",
|
|
"crf-search",
|
|
"-i",
|
|
str(filepath),
|
|
"--min-vmaf",
|
|
str(effective_vmaf),
|
|
"--preset",
|
|
str(preset),
|
|
"--max-encoded-percent",
|
|
"100",
|
|
"--temp-dir",
|
|
temp_dir,
|
|
"--samples",
|
|
"4",
|
|
]
|
|
|
|
# FORCE use of the specified encoder (Fixes the CPU default crash)
|
|
if ab_av1_supports("crf-search", "--encoder"):
|
|
cmd.extend(["--encoder", encoder])
|
|
|
|
# Hardware decode acceleration
|
|
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"])
|
|
|
|
vmaf_label = f"VMAF {effective_vmaf}" if is_hw_encoder else f"VMAF {target_vmaf}"
|
|
print(f" - Searching for CRF to hit {vmaf_label}...")
|
|
|
|
# 1. First Run: Stream output to console
|
|
returncode = run_command_streaming(cmd, f"crf-search {vmaf_label}")
|
|
|
|
if returncode == 0:
|
|
# 2. Second Run: Capture output to parse the CRF
|
|
# FIX: Reuse 'cmd' and add encoding parameters to prevent charmap crashes
|
|
res = subprocess.run(
|
|
cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
encoding="utf-8",
|
|
errors="ignore"
|
|
)
|
|
|
|
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,
|
|
encoder="svt-av1",
|
|
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, encoder, 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, encoder="svt-av1", 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",
|
|
]
|
|
|
|
if encoder != "svt-av1":
|
|
if ab_av1_supports("encode", "--encoder"):
|
|
cmd.extend(["--encoder", encoder])
|
|
|
|
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,
|
|
hw_encoder="svt-av1",
|
|
use_hw_mode=False,
|
|
hwaccel=None,
|
|
):
|
|
"""Process a single video file with intelligent VMAF targeting"""
|
|
global _shutdown_requested
|
|
|
|
# --- LOGIC FIX START ---
|
|
# If "use_hw_mode" (Hybrid) is FALSE, use the requested encoder for EVERYONE.
|
|
if not use_hw_mode:
|
|
encoder = hw_encoder
|
|
use_hw = (encoder in HW_ENCODERS)
|
|
else:
|
|
# Hybrid Mode: Only 1 worker gets HW, rest get CPU
|
|
use_hw = False
|
|
if hwaccel and hw_encoder in HW_ENCODERS:
|
|
use_hw = claim_hardware_worker()
|
|
encoder = hw_encoder if use_hw else "svt-av1"
|
|
# --- LOGIC FIX END ---
|
|
|
|
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} ---")
|
|
# Added 'errors="ignore"' to print to avoid Unicode crashes on console
|
|
try:
|
|
print(
|
|
f" Source: {stats_before['codec']} @ {stats_before['bitrate']}k, {stats_before['size'] / (1024**3):.2f} GB"
|
|
)
|
|
except UnicodeEncodeError:
|
|
print(f" Source: {stats_before['codec']} @ {stats_before['bitrate']}k")
|
|
|
|
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), encoder, 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), encoder, use_hw, hwaccel
|
|
)
|
|
if not search_result_94:
|
|
search_result_94 = run_crf_search(
|
|
filepath, 92.0, preset, str(temp_dir), encoder, use_hw, hwaccel
|
|
)
|
|
if not search_result_94:
|
|
search_result_94 = run_crf_search(
|
|
filepath, 90.0, preset, str(temp_dir), encoder, 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), encoder, 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), encoder, 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,
|
|
encoder,
|
|
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(
|
|
"--encoder",
|
|
default="svt-av1",
|
|
help=(
|
|
"Video encoder to use. Default: svt-av1 (CPU). "
|
|
"Hardware encoders: av1_amf (AMD), av1_nvenc (NVIDIA), av1_qsv (Intel)."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--use-hardware-worker",
|
|
action="store_true",
|
|
help="Use 1 hardware encoding worker + rest CPU workers (requires --encoder with HW encoder)",
|
|
)
|
|
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
|
|
# HW worker uses specified encoder; CPU workers use svt-av1
|
|
hw_encoder = args.encoder if args.encoder in HW_ENCODERS else None
|
|
use_hw_primary = args.use_hardware_worker and hw_encoder 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,
|
|
args.encoder,
|
|
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
|
|
|
|
# All workers try to claim HW; only the first thread succeeds
|
|
# and will use HW for ALL its tasks
|
|
future = executor.submit(
|
|
process_file,
|
|
file_path,
|
|
args.log_dir,
|
|
log_name,
|
|
args.preset,
|
|
args.encoder,
|
|
use_hw_primary,
|
|
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()
|