complete rewrite of the logic of the static_encoder

This commit is contained in:
2025-07-20 17:42:27 +02:00
parent 7a8a85d953
commit 6af7ac5c25

View File

@@ -6,15 +6,7 @@ import argparse
import shutil
import multiprocessing
import re
import math
# --- Global lock for synchronized screen writing ---
LOCK = None
def init_worker(lock):
"""Initializer for each worker process to share the lock."""
global LOCK
LOCK = lock
import time
def get_video_resolution(video_path):
"""Gets video resolution using ffprobe."""
@@ -27,8 +19,7 @@ def get_video_resolution(video_path):
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 for '{video_path}': {e}")
except (subprocess.CalledProcessError, ValueError):
return None, None
def get_frame_count(video_path):
@@ -42,19 +33,28 @@ def get_frame_count(video_path):
result = subprocess.run(command, capture_output=True, text=True, check=True)
return int(result.stdout.strip())
except (subprocess.CalledProcessError, ValueError):
return 0 # Return 0 if we can't get the frame count
return 0
def encode_segment_worker(task_args):
"""
Wrapper function for multiprocessing pool.
Unpacks arguments and calls the main encoding function.
Wrapper for the multiprocessing pool.
Reports progress to shared memory objects.
"""
segment_path, crf, pbar_pos, total_frames = task_args
return encode_segment(segment_path, crf, pbar_pos, total_frames)
segment_path, crf, worker_id, progress_dict, total_processed_frames, lock = task_args
def encode_segment(segment_path, crf, pbar_pos, total_frames):
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, drawing a manual progress bar to the console.
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)
@@ -70,45 +70,52 @@ def encode_segment(segment_path, crf, pbar_pos, total_frames):
]
svt_command = [
'SvtAv1EncApp.exe', '-i', 'stdin', '--width', str(width), '--height', str(height),
'--progress', '2', '--preset', '2', '--lp', '1', '--input-depth', '10',
'--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.PIPE, stderr=subprocess.PIPE, text=True, encoding='utf-8')
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+)/")
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
def draw_progress(frame, total):
if total == 0: return
bar_width = 40
percentage = frame / total
filled_len = int(round(bar_width * percentage))
bar = '' * filled_len + '-' * (bar_width - filled_len)
# Truncate filename for display
display_name = (base_name[:23] + '..') if len(base_name) > 25 else base_name
progress_str = f"{display_name:<25} |{bar}| {frame}/{total} ({percentage:.1%})"
with LOCK:
sys.stdout.write(f'\033[{pbar_pos};0H') # Move cursor to line `pbar_pos`
sys.stdout.write('\033[K') # Clear the line
sys.stdout.write(progress_str)
sys.stdout.flush()
try:
for line in iter(svt_process.stderr.readline, ''):
line = line.strip()
match = progress_regex.search(line)
if match:
current_frame = int(match.group(1))
draw_progress(current_frame, total_frames)
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()
@@ -116,13 +123,32 @@ def encode_segment(segment_path, crf, pbar_pos, total_frames):
if svt_process.returncode != 0:
raise subprocess.CalledProcessError(svt_process.returncode, svt_command)
draw_progress(total_frames, total_frames) # Final update to 100%
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 Exception:
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.",
@@ -132,19 +158,17 @@ def main():
parser.add_argument(
"--workers",
type=int,
default=max(1, multiprocessing.cpu_count() // 2),
help="Number of segments to encode in parallel. Default: half of available CPU cores."
default=4,
help="Number of segments to encode in parallel. Default: 4"
)
args = parser.parse_args()
# --- Pre-flight check for executables ---
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)
# --- Setup Directories ---
input_dir = "cuts"
output_dir = "segments"
if not os.path.isdir(input_dir):
@@ -153,7 +177,6 @@ def main():
os.makedirs(output_dir, exist_ok=True)
# --- Process Segments ---
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:
@@ -162,30 +185,35 @@ def main():
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(f"Final files will be saved in '{output_dir}'.")
print("\nGathering segment information...")
frame_counts = [get_frame_count(s) for s in segments]
tasks = [(segments[i], args.crf, i + 1, frame_counts[i]) for i in range(total_segments)]
grand_total_frames = sum(get_frame_count(s) for s in segments)
# --- Run Encoding ---
# Clear screen and reserve space for progress bars
os.system('cls' if os.name == 'nt' else 'clear')
print(f"--- Starting {args.workers} encoding worker(s) ---")
sys.stdout.write('\n' * args.workers)
sys.stdout.flush()
manager = multiprocessing.Manager()
progress_dict = manager.dict()
total_processed_frames = manager.Value('i', 0)
lock = manager.Lock()
lock = multiprocessing.Lock()
with multiprocessing.Pool(processes=args.workers, initializer=init_worker, initargs=(lock,)) as pool:
# Use imap to maintain order, so pbar_pos is consistent
results = list(pool.imap(encode_segment_worker, tasks))
tasks = [(segments[i], args.crf, i, progress_dict, total_processed_frames, lock) for i in range(total_segments)]
# Move cursor below the progress bars
sys.stdout.write(f'\033[{args.workers + 1};0H')
sys.stdout.write('\n')
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"--- All segments processed. ---")
print(f"\n--- All segments processed. ---")
print(f"Successfully encoded {successful_encodes}/{total_segments} segments.")
if __name__ == "__main__":