diff --git a/static_encoder.py b/static_encoder.py index 38ab423..9ffe68f 100644 --- a/static_encoder.py +++ b/static_encoder.py @@ -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 + + 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, pbar_pos, total_frames): +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,59 +70,85 @@ 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 ] - 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') - - if ffmpeg_process.stdout: - ffmpeg_process.stdout.close() - - progress_regex = re.compile(r"Encoding frame\s+(\d+)/") - - 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: + 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)) - 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() 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) + + 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)] - # --- 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() + 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 - 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)) - - # Move cursor below the progress bars - sys.stdout.write(f'\033[{args.workers + 1};0H') - sys.stdout.write('\n') + 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__":