220 lines
8.3 KiB
Python
220 lines
8.3 KiB
Python
#!/usr/bin/env python3
|
|
import subprocess
|
|
import os
|
|
import sys
|
|
import argparse
|
|
import shutil
|
|
import multiprocessing
|
|
import re
|
|
import time
|
|
|
|
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):
|
|
return None, None
|
|
|
|
def get_frame_count(video_path):
|
|
"""Gets the total number of frames in a video file using ffprobe."""
|
|
command = [
|
|
'ffprobe', '-v', 'error', '-select_streams', 'v:0',
|
|
'-count_frames', '-show_entries', 'stream=nb_read_frames',
|
|
'-of', 'default=nokey=1:noprint_wrappers=1', video_path
|
|
]
|
|
try:
|
|
result = subprocess.run(command, capture_output=True, text=True, check=True)
|
|
return int(result.stdout.strip())
|
|
except (subprocess.CalledProcessError, ValueError):
|
|
return 0
|
|
|
|
def encode_segment_worker(task_args):
|
|
"""
|
|
Wrapper for the multiprocessing pool.
|
|
Reports progress to shared memory objects.
|
|
"""
|
|
segment_path, crf, worker_id, progress_dict, total_processed_frames, lock = task_args
|
|
|
|
progress_dict[worker_id] = {'fps': 0.0, 'status': 'Starting'}
|
|
|
|
success = encode_segment(segment_path, crf, worker_id, progress_dict, total_processed_frames, lock)
|
|
|
|
status = 'Finished' if success else 'FAILED'
|
|
progress_dict[worker_id] = {'fps': 0.0, 'status': status}
|
|
|
|
return success
|
|
|
|
def encode_segment(segment_path, crf, worker_id, progress_dict, total_processed_frames, lock):
|
|
"""
|
|
Encodes a single segment, reporting progress via shared objects.
|
|
Calculates FPS manually based on frame processing time.
|
|
"""
|
|
output_dir = "segments"
|
|
base_name = os.path.basename(segment_path)
|
|
final_output_path = os.path.join(output_dir, base_name)
|
|
|
|
width, height = get_video_resolution(segment_path)
|
|
if not width or not height:
|
|
return False
|
|
|
|
ffmpeg_command = [
|
|
'ffmpeg', '-hide_banner', '-loglevel', 'error', '-i', segment_path,
|
|
'-pix_fmt', 'yuv420p10le', '-f', 'yuv4mpegpipe', '-strict', '-1', '-'
|
|
]
|
|
svt_command = [
|
|
'SvtAv1EncApp.exe', '-i', 'stdin', '--width', str(width), '--height', str(height),
|
|
'--progress', '2', '--preset', '2', '--input-depth', '10',
|
|
'--crf', str(crf), '--film-grain', '8', '--tune', '2', '--keyint', '-1',
|
|
'--color-primaries', '1', '--transfer-characteristics', '1', '--matrix-coefficients', '1',
|
|
'-b', final_output_path
|
|
]
|
|
|
|
try:
|
|
ffmpeg_process = subprocess.Popen(ffmpeg_command, stdout=subprocess.PIPE)
|
|
svt_process = subprocess.Popen(svt_command, stdin=ffmpeg_process.stdout, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, encoding='utf-8')
|
|
|
|
if ffmpeg_process.stdout:
|
|
ffmpeg_process.stdout.close()
|
|
|
|
progress_regex = re.compile(r"Encoding frame\s+(\d+)")
|
|
last_frame = 0
|
|
frames_processed_this_segment = 0
|
|
last_update_time = time.time()
|
|
frames_at_last_update = 0
|
|
|
|
for line in iter(svt_process.stderr.readline, ''):
|
|
line = line.strip()
|
|
match = progress_regex.search(line)
|
|
if match:
|
|
current_frame = int(match.group(1))
|
|
|
|
current_time = time.time()
|
|
time_since_last_update = current_time - last_update_time
|
|
|
|
fps = 0.0
|
|
if time_since_last_update > 1.0:
|
|
frames_since_last_update = current_frame - frames_at_last_update
|
|
if frames_since_last_update > 0 and time_since_last_update > 0:
|
|
fps = frames_since_last_update / time_since_last_update
|
|
|
|
last_update_time = current_time
|
|
frames_at_last_update = current_frame
|
|
|
|
delta = current_frame - last_frame
|
|
if delta > 0:
|
|
with lock:
|
|
total_processed_frames.value += delta
|
|
frames_processed_this_segment += delta
|
|
last_frame = current_frame
|
|
|
|
if fps > 0.0:
|
|
progress_dict[worker_id] = {'fps': fps, 'status': 'Encoding'}
|
|
|
|
svt_process.communicate()
|
|
ffmpeg_process.wait()
|
|
|
|
if svt_process.returncode != 0:
|
|
raise subprocess.CalledProcessError(svt_process.returncode, svt_command)
|
|
|
|
total_segment_frames = get_frame_count(segment_path)
|
|
remaining_frames = total_segment_frames - frames_processed_this_segment
|
|
if remaining_frames > 0:
|
|
with lock:
|
|
total_processed_frames.value += remaining_frames
|
|
|
|
return True
|
|
except (subprocess.CalledProcessError, FileNotFoundError):
|
|
if os.path.exists(final_output_path):
|
|
os.remove(final_output_path)
|
|
return False
|
|
|
|
def draw_global_progress(processed_frames, total_frames, progress_dict):
|
|
"""Draws a single line global progress bar."""
|
|
bar_width = 50
|
|
percentage = processed_frames / total_frames if total_frames > 0 else 0
|
|
filled_len = int(round(bar_width * percentage))
|
|
bar = '█' * filled_len + '-' * (bar_width - filled_len)
|
|
|
|
total_fps = sum(worker.get('fps', 0.0) for worker in progress_dict.values())
|
|
|
|
status_str = f"Progress: |{bar}| {processed_frames}/{total_frames} ({percentage:.1%}) @ {total_fps:.2f} FPS"
|
|
|
|
sys.stdout.write('\r' + status_str)
|
|
sys.stdout.flush()
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Encodes video segments from the 'cuts' directory to AV1 using a static CRF value.",
|
|
formatter_class=argparse.RawTextHelpFormatter
|
|
)
|
|
parser.add_argument("--crf", type=int, default=27, help="The static CRF value to use for all segments. Default: 27")
|
|
parser.add_argument(
|
|
"--workers",
|
|
type=int,
|
|
default=4,
|
|
help="Number of segments to encode in parallel. Default: 4"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
for exe in ["SvtAv1EncApp.exe", "ffprobe"]:
|
|
if not shutil.which(exe):
|
|
print(f"Error: '{exe}' not found. Please ensure it is in your system's PATH.")
|
|
sys.exit(1)
|
|
|
|
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)
|
|
|
|
segments = sorted([os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith('.mkv')])
|
|
total_segments = len(segments)
|
|
if total_segments == 0:
|
|
print(f"No segments found in '{input_dir}'.")
|
|
sys.exit(0)
|
|
|
|
print(f"Found {total_segments} segments to process from '{input_dir}'.")
|
|
print(f"Encoding with static CRF {args.crf} using {args.workers} parallel worker(s).")
|
|
|
|
print("\nGathering segment information...")
|
|
grand_total_frames = sum(get_frame_count(s) for s in segments)
|
|
|
|
manager = multiprocessing.Manager()
|
|
progress_dict = manager.dict()
|
|
total_processed_frames = manager.Value('i', 0)
|
|
lock = manager.Lock()
|
|
|
|
tasks = [(segments[i], args.crf, i, progress_dict, total_processed_frames, lock) for i in range(total_segments)]
|
|
|
|
pool = multiprocessing.Pool(processes=args.workers)
|
|
future_results = pool.imap_unordered(encode_segment_worker, tasks)
|
|
|
|
while total_processed_frames.value < grand_total_frames:
|
|
draw_global_progress(total_processed_frames.value, grand_total_frames, progress_dict)
|
|
time.sleep(0.1)
|
|
if all(p.get('status') in ['Finished', 'FAILED'] for p in progress_dict.values()) and len(progress_dict) == total_segments:
|
|
break
|
|
|
|
pool.close()
|
|
pool.join()
|
|
|
|
draw_global_progress(grand_total_frames, grand_total_frames, progress_dict)
|
|
print()
|
|
|
|
results = list(future_results)
|
|
successful_encodes = sum(1 for r in results if r)
|
|
print(f"\n--- All segments processed. ---")
|
|
print(f"Successfully encoded {successful_encodes}/{total_segments} segments.")
|
|
|
|
if __name__ == "__main__":
|
|
main() |