495 lines
18 KiB
Python
495 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Smart GPU Encoder (Windows/AMD/NVIDIA Optimized)
|
|
------------------------------------------------
|
|
Refactored to use vmaf_common.py
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
import json
|
|
import shutil
|
|
import time
|
|
import argparse
|
|
import signal
|
|
import platform
|
|
import re
|
|
import threading
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
# --- Import Common Module ---
|
|
import vmaf_common as common
|
|
|
|
# --- Configuration ---
|
|
TARGET_VMAF_MIN = common.DEFAULT_CONFIG["target_vmaf"]
|
|
MIN_SAVINGS_PERCENT = common.DEFAULT_CONFIG["min_savings"]
|
|
MAX_JOBS = common.DEFAULT_CONFIG["gpu_jobs"]
|
|
|
|
DEFAULT_AV1_QP = 32
|
|
DEFAULT_HEVC_QP = 28
|
|
|
|
SAMPLE_DURATION = 60
|
|
SAMPLE_START_TIME = 300
|
|
|
|
TEMP_DIR = None # Will be set in main
|
|
|
|
# Tools (Local override if needed, else used from common checks)
|
|
FFMPEG_BIN = "ffmpeg"
|
|
AB_AV1_BIN = "ab-av1"
|
|
|
|
# Global state
|
|
shutdown_requested = False
|
|
active_processes = set()
|
|
upload_lock = threading.Lock()
|
|
proc_lock = threading.Lock()
|
|
debug_mode = False
|
|
|
|
def handle_sigint(signum, frame):
|
|
global shutdown_requested
|
|
print("\n\n[!] CRITICAL: Shutdown requested (Ctrl+C).")
|
|
print(" Killing active encoder processes...")
|
|
shutdown_requested = True
|
|
|
|
with proc_lock:
|
|
for proc in list(active_processes):
|
|
try:
|
|
proc.terminate()
|
|
time.sleep(0.1)
|
|
if proc.poll() is None:
|
|
proc.kill()
|
|
except:
|
|
pass
|
|
print(" Cleanup complete. Exiting.")
|
|
sys.exit(1)
|
|
|
|
signal.signal(signal.SIGINT, handle_sigint)
|
|
|
|
# --- Hardware Detection ---
|
|
def detect_hardware_encoder():
|
|
"""Detects available hardware encoders via ffmpeg"""
|
|
# Use common module detection
|
|
return common.detect_hardware_encoder()
|
|
|
|
def get_encoder_args(codec, encoder, qp):
|
|
"""Returns correct ffmpeg args for specific HW vendor"""
|
|
# Use common module args
|
|
return common.get_encoder_args(codec, encoder, qp)
|
|
|
|
# --- Helpers ---
|
|
def run_process(cmd, description="", status_callback=None):
|
|
"""Run a process with real-time output and clean shutdown tracking"""
|
|
if shutdown_requested: return False
|
|
if status_callback: status_callback(description)
|
|
|
|
try:
|
|
# Windows: Hide console window
|
|
cflags = 0x08000000 if platform.system() == 'Windows' else 0
|
|
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
universal_newlines=True,
|
|
bufsize=1,
|
|
creationflags=cflags
|
|
)
|
|
|
|
with proc_lock:
|
|
active_processes.add(proc)
|
|
|
|
if proc.stdout:
|
|
for line in proc.stdout:
|
|
if shutdown_requested:
|
|
proc.terminate()
|
|
break
|
|
line = line.strip()
|
|
if line:
|
|
if debug_mode: print(f" [Debug] {line}")
|
|
|
|
if status_callback and ("frame=" in line or "size=" in line or "time=" in line):
|
|
status_callback(line)
|
|
|
|
proc.wait()
|
|
|
|
with proc_lock:
|
|
if proc in active_processes:
|
|
active_processes.remove(proc)
|
|
|
|
return proc.returncode == 0
|
|
except Exception as e:
|
|
if status_callback: status_callback(f"Error: {e}")
|
|
return False
|
|
|
|
def run_vmaf_check(reference, distorted, status_callback=None):
|
|
"""Run ab-av1 vmaf to get score"""
|
|
# Use common dependency check to find binary if needed, but here just assume it's in path or bin
|
|
ab_exe = "ab-av1"
|
|
|
|
# Check if bundled exists
|
|
bundled = Path(__file__).parent / "bin" / "ab-av1.exe"
|
|
if bundled.exists():
|
|
ab_exe = str(bundled)
|
|
|
|
cmd = [ab_exe, "vmaf", "--reference", str(reference), "--distorted", str(distorted)]
|
|
if status_callback: status_callback("Calculating VMAF...")
|
|
|
|
try:
|
|
result = subprocess.run(cmd, capture_output=True, text=True, check=True)
|
|
for line in result.stdout.splitlines():
|
|
line = line.strip()
|
|
match = re.search(r"VMAF\s+([0-9.]+)", line)
|
|
if match: return float(match.group(1))
|
|
try:
|
|
val = float(line)
|
|
if 0 <= val <= 100: return val
|
|
except: pass
|
|
return 0.0
|
|
except Exception:
|
|
return -1.0
|
|
|
|
# --- Core Logic ---
|
|
def process_file(filepath, log_category, lock_dir, log_dir, encoders, category_root=None, worker_id=0, status_cb=None):
|
|
"""
|
|
Process a single file.
|
|
status_cb: function(worker_id, filename, status_text, color)
|
|
category_root: Root directory (tv_dir or content_dir) for relative path calculation
|
|
"""
|
|
av1_enc, hevc_enc, hw_type = encoders
|
|
filepath = Path(filepath)
|
|
filename = filepath.name
|
|
|
|
def update(msg, color="white"):
|
|
if status_cb: status_cb(worker_id, filename, msg, color)
|
|
else: print(f"[{worker_id}] {msg}")
|
|
|
|
if shutdown_requested: return
|
|
|
|
# 0. Check if already processed in a previous run
|
|
if common.is_already_processed(log_dir, filepath):
|
|
update("Already Processed (Skipping)", "dim")
|
|
return
|
|
|
|
# 1. Lock Check (Shared Storage)
|
|
lock_file = common.acquire_lock(lock_dir, filepath, category_root)
|
|
if not lock_file:
|
|
return # Locked or skipped
|
|
|
|
try:
|
|
update("Analyzing...", "blue")
|
|
|
|
# 2. Analyze Source
|
|
info = common.get_video_info(filepath)
|
|
if not info:
|
|
update("Metadata Error", "red")
|
|
return
|
|
|
|
if info["codec"] == "av1":
|
|
update("Already AV1 (Skipping)", "green")
|
|
time.sleep(1)
|
|
return
|
|
|
|
# 3. Create Samples
|
|
sample_ref = TEMP_DIR / f"{filepath.stem}_{worker_id}_ref.mkv"
|
|
sample_enc = TEMP_DIR / f"{filepath.stem}_{worker_id}_enc.mkv"
|
|
|
|
sample_start = SAMPLE_START_TIME
|
|
if info["duration"] < (SAMPLE_START_TIME + SAMPLE_DURATION):
|
|
sample_start = max(0, (info["duration"] / 2) - (SAMPLE_DURATION / 2))
|
|
|
|
update("Extracting Ref", "magenta")
|
|
cmd_ref = [
|
|
FFMPEG_BIN, "-y", "-hide_banner", "-loglevel", "error",
|
|
"-ss", str(sample_start), "-t", str(SAMPLE_DURATION),
|
|
"-i", str(filepath), "-c", "copy", "-map", "0:v:0",
|
|
str(sample_ref)
|
|
]
|
|
if not run_process(cmd_ref):
|
|
update("Extract Ref Failed", "red")
|
|
return
|
|
|
|
# TEST 1: AV1
|
|
vmaf_score = 0
|
|
savings = 0
|
|
|
|
if av1_enc:
|
|
update(f"Testing AV1 QP{DEFAULT_AV1_QP}", "yellow")
|
|
enc_args = get_encoder_args("av1", av1_enc, DEFAULT_AV1_QP)
|
|
cmd_enc = [
|
|
FFMPEG_BIN, "-y", "-hide_banner", "-loglevel", "error",
|
|
"-i", str(sample_ref), *enc_args, "-an", str(sample_enc)
|
|
]
|
|
|
|
if run_process(cmd_enc):
|
|
update("Calculating VMAF", "cyan")
|
|
vmaf_score = run_vmaf_check(sample_ref, sample_enc)
|
|
|
|
ref_size = sample_ref.stat().st_size
|
|
enc_size = sample_enc.stat().st_size
|
|
savings = (1 - (enc_size / ref_size)) * 100 if ref_size > 0 else 0
|
|
else:
|
|
update("AV1 Test Failed", "red")
|
|
|
|
chosen_codec = None
|
|
chosen_qp = 0
|
|
|
|
# Decision Logic
|
|
if vmaf_score >= TARGET_VMAF_MIN and savings >= MIN_SAVINGS_PERCENT:
|
|
chosen_codec = "av1"
|
|
chosen_qp = DEFAULT_AV1_QP
|
|
update(f"AV1 Good (VMAF {vmaf_score:.1f})", "green")
|
|
|
|
# Smart Optimization
|
|
if vmaf_score > 97.0:
|
|
update(f"Optimizing (High Quality {vmaf_score:.1f})", "yellow")
|
|
new_qp = DEFAULT_AV1_QP + 4
|
|
args_opt = get_encoder_args("av1", av1_enc, new_qp)
|
|
cmd_opt = [FFMPEG_BIN, "-y", "-hide_banner", "-loglevel", "error", "-i", str(sample_ref), *args_opt, "-an", str(sample_enc)]
|
|
|
|
if run_process(cmd_opt):
|
|
vmaf_opt = run_vmaf_check(sample_ref, sample_enc)
|
|
size_opt = sample_enc.stat().st_size
|
|
sav_opt = (1 - (size_opt / sample_ref.stat().st_size)) * 100
|
|
|
|
if vmaf_opt >= TARGET_VMAF_MIN and sav_opt > savings:
|
|
update(f"Opt Accepted (+{sav_opt - savings:.1f}%)", "green")
|
|
chosen_qp = new_qp
|
|
vmaf_score = vmaf_opt
|
|
savings = sav_opt
|
|
else:
|
|
update("Testing HEVC Fallback", "magenta")
|
|
if info["codec"] != "hevc" and hevc_enc:
|
|
hevc_args = get_encoder_args("hevc", hevc_enc, DEFAULT_HEVC_QP)
|
|
cmd_hevc = [FFMPEG_BIN, "-y", "-hide_banner", "-loglevel", "error", "-i", str(sample_ref), *hevc_args, "-an", str(sample_enc)]
|
|
run_process(cmd_hevc)
|
|
|
|
vmaf_score = run_vmaf_check(sample_ref, sample_enc)
|
|
enc_size = sample_enc.stat().st_size
|
|
savings = (1 - (enc_size / sample_ref.stat().st_size)) * 100
|
|
|
|
if vmaf_score >= TARGET_VMAF_MIN and savings >= MIN_SAVINGS_PERCENT:
|
|
update(f"HEVC Accepted (VMAF {vmaf_score:.1f})", "green")
|
|
chosen_codec = "hevc"
|
|
chosen_qp = DEFAULT_HEVC_QP
|
|
else:
|
|
update("HEVC Rejected", "red")
|
|
common.log_event(log_dir, "rejected.jsonl", {"file": str(filepath), "status": "rejected", "vmaf": vmaf_score})
|
|
else:
|
|
update("Skipping HEVC", "yellow")
|
|
|
|
# Cleanup Samples
|
|
if sample_ref.exists(): sample_ref.unlink()
|
|
if sample_enc.exists(): sample_enc.unlink()
|
|
|
|
# 4. Full Encode
|
|
if chosen_codec:
|
|
update(f"Encoding {chosen_codec.upper()} (QP {chosen_qp})", "green")
|
|
output_file = TEMP_DIR / f"{filepath.stem}.{chosen_codec}.mkv"
|
|
final_args = get_encoder_args(chosen_codec, av1_enc if chosen_codec=="av1" else hevc_enc, chosen_qp)
|
|
|
|
cmd_full = [
|
|
FFMPEG_BIN, "-y", "-hide_banner", "-loglevel", "info", "-stats",
|
|
"-i", str(filepath),
|
|
*final_args,
|
|
"-c:a", "copy", "-c:s", "copy", "-map", "0",
|
|
str(output_file)
|
|
]
|
|
|
|
def prog_cb(msg):
|
|
if "frame=" in msg:
|
|
try:
|
|
if "time=" in msg:
|
|
t_str = msg.split("time=")[1].split(" ")[0]
|
|
h, m, s = map(float, t_str.split(':'))
|
|
cur_sec = h*3600 + m*60 + s
|
|
percent = (cur_sec / info["duration"]) * 100
|
|
else:
|
|
percent = 0.0
|
|
|
|
speed = "1x"
|
|
if "speed=" in msg:
|
|
speed = msg.split("speed=")[1].split("x")[0] + "x"
|
|
|
|
update(f"Encoding {chosen_codec} | {percent:.1f}% | {speed}", "green")
|
|
except:
|
|
pass
|
|
|
|
if run_process(cmd_full, f"Full Encode ({chosen_codec.upper()} QP {chosen_qp})", status_callback=prog_cb):
|
|
final_info = common.get_video_info(output_file)
|
|
if not final_info: final_info = {"size": output_file.stat().st_size}
|
|
|
|
final_size = final_info["size"]
|
|
final_savings = (1 - (final_size / info["size"])) * 100
|
|
saved_bytes = info["size"] - final_size
|
|
|
|
update(f"Uploading (Saved {final_savings:.1f}%)", "blue")
|
|
|
|
# UPLOAD (Serialized)
|
|
with upload_lock:
|
|
update(f"Uploading...", "blue")
|
|
backup_path = filepath.with_suffix(f"{filepath.suffix}.original")
|
|
try:
|
|
shutil.move(str(filepath), str(backup_path))
|
|
shutil.copy2(str(output_file), str(filepath))
|
|
|
|
# Verify integrity
|
|
if Path(filepath).stat().st_size == final_size:
|
|
output_file.unlink()
|
|
backup_path.unlink()
|
|
|
|
# Refresh metadata for accuracy
|
|
final_info_verified = common.get_video_info(filepath)
|
|
|
|
common.log_event(log_dir, f"{log_category}.jsonl", {
|
|
"file": str(filepath),
|
|
"status": "success",
|
|
"codec": chosen_codec,
|
|
"vmaf": vmaf_score,
|
|
"savings": final_savings,
|
|
"original_metadata": info,
|
|
"encoded_metadata": final_info_verified or final_info
|
|
})
|
|
|
|
# Mark as processed to prevent re-encoding in future runs
|
|
common.mark_processed(log_dir, filepath, chosen_codec, vmaf_score, final_savings)
|
|
|
|
# Mark lock as completed (keep it for future runs)
|
|
common.mark_lock_completed(lock_file)
|
|
|
|
update("Done", "green")
|
|
if status_cb: status_cb(worker_id, filename, f"STATS:SAVED:{saved_bytes}", "green")
|
|
else:
|
|
update("Upload Failed (Size Mismatch)", "red")
|
|
shutil.move(str(backup_path), str(filepath))
|
|
except Exception as e:
|
|
update(f"Move Error: {str(e)[:20]}", "red")
|
|
if backup_path.exists(): shutil.move(str(backup_path), str(filepath))
|
|
else:
|
|
update("Encode Failed", "red")
|
|
if output_file.exists(): output_file.unlink()
|
|
|
|
except Exception as e:
|
|
update(f"Error: {str(e)[:30]}", "red")
|
|
# On error, delete lock so file can be retried
|
|
if lock_file and lock_file.exists():
|
|
lock_file.unlink()
|
|
finally:
|
|
update("Idle", "dim")
|
|
|
|
def main():
|
|
global debug_mode
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--tv-dir", default=common.DEFAULT_CONFIG["tv_dir"])
|
|
parser.add_argument("--content-dir", default=common.DEFAULT_CONFIG["content_dir"])
|
|
parser.add_argument("--jobs", type=int, default=MAX_JOBS)
|
|
parser.add_argument("--debug", action="store_true")
|
|
parser.add_argument("--skip-until", help="Skip all files alphabetically until this filename substring is found")
|
|
parser.add_argument("--cpu-only", action="store_true", help="Force software encoding (CPU only)")
|
|
parser.add_argument("--temp-dir", help="Override local temp directory")
|
|
parser.add_argument("--av1-encoder", choices=["hw", "sw", "off"], default="hw", help="AV1 encoder: hw (hardware), sw (software), off (disable)")
|
|
parser.add_argument("--hevc-encoder", choices=["hw", "sw", "off"], default="hw", help="HEVC encoder: hw (hardware), sw (software), off (disable)")
|
|
args = parser.parse_args()
|
|
|
|
if args.debug:
|
|
debug_mode = True
|
|
print("[Debug Mode Enabled]")
|
|
|
|
# 0. Check Dependencies
|
|
common.check_dependencies()
|
|
|
|
# 1. Setup Directories
|
|
lock_dir, log_dir = common.get_base_paths(args)
|
|
|
|
global TEMP_DIR
|
|
TEMP_DIR = common.get_temp_dir(args)
|
|
|
|
# 2. Detect Hardware
|
|
av1, hevc, hw = common.detect_hardware_encoder(args)
|
|
|
|
print("="*60)
|
|
print(f" SMART ENCODER | Hardware: {hw.upper()} | Jobs: {args.jobs}")
|
|
|
|
if hw == "cpu":
|
|
print(f" [!] Fallback to CPU Software Encoding (Slow)")
|
|
if av1: print(f" AV1: {av1} (libsvtav1)")
|
|
if hevc: print(f" HEVC: {hevc} (libx265)")
|
|
else:
|
|
print(f" AV1: {av1} | HEVC: {hevc}")
|
|
|
|
print(f" Locks: {lock_dir}")
|
|
print("="*60)
|
|
|
|
# 3. Scan & Queue
|
|
tasks = []
|
|
|
|
# Skip-until filtering
|
|
skip_until = args.skip_until
|
|
skipping = bool(skip_until)
|
|
skipped_count = 0
|
|
|
|
tv_path = Path(args.tv_dir)
|
|
if tv_path.exists():
|
|
print(f"Scanning TV: {tv_path}")
|
|
files = list(tv_path.rglob("*.mkv")) + list(tv_path.rglob("*.mp4"))
|
|
files.sort() # Alphabetical order for consistency across platforms
|
|
for f in files:
|
|
if skipping:
|
|
if skip_until.lower() in str(f).lower():
|
|
skipping = False
|
|
print(f" Found '{skip_until}' - resuming from here")
|
|
else:
|
|
skipped_count += 1
|
|
continue
|
|
tasks.append((f, "tv_shows"))
|
|
|
|
content_path = Path(args.content_dir)
|
|
if content_path.exists():
|
|
print(f"Scanning Content: {content_path}")
|
|
files = list(content_path.rglob("*.mkv")) + list(content_path.rglob("*.mp4"))
|
|
files.sort() # Alphabetical order for consistency across platforms
|
|
for f in files:
|
|
if skipping:
|
|
if skip_until.lower() in str(f).lower():
|
|
skipping = False
|
|
print(f" Found '{skip_until}' - resuming from here")
|
|
else:
|
|
skipped_count += 1
|
|
continue
|
|
tasks.append((f, "content"))
|
|
|
|
if skipped_count > 0:
|
|
print(f" Skipped {skipped_count} files (--skip-until)")
|
|
|
|
if not tasks:
|
|
print("No files found.")
|
|
return
|
|
|
|
# 4. Execute
|
|
print(f"\n🚀 Processing {len(tasks)} files...")
|
|
|
|
# Build category root map
|
|
category_roots = {
|
|
"tv_shows": Path(args.tv_dir),
|
|
"content": Path(args.content_dir)
|
|
}
|
|
|
|
with ThreadPoolExecutor(max_workers=args.jobs) as executor:
|
|
futures = {
|
|
executor.submit(process_file, f, cat, lock_dir, log_dir, (av1, hevc, hw), category_roots.get(cat)): f
|
|
for f, cat in tasks
|
|
}
|
|
|
|
for future in as_completed(futures):
|
|
if shutdown_requested:
|
|
executor.shutdown(wait=False, cancel_futures=True)
|
|
break
|
|
try:
|
|
future.result()
|
|
except Exception as e:
|
|
print(f"Worker Error: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|