279 lines
9.8 KiB
Python
279 lines
9.8 KiB
Python
"""
|
|
Streaming frame writer using ffmpeg pipe for zero-copy output
|
|
"""
|
|
|
|
import subprocess
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from typing import Optional, Dict, Any
|
|
import signal
|
|
import atexit
|
|
import warnings
|
|
|
|
|
|
class StreamingFrameWriter:
|
|
"""Write frames directly to ffmpeg via pipe for memory-efficient output"""
|
|
|
|
def __init__(self,
|
|
output_path: str,
|
|
width: int,
|
|
height: int,
|
|
fps: float,
|
|
audio_source: Optional[str] = None,
|
|
video_codec: str = 'h264_nvenc',
|
|
quality_preset: str = 'p4', # NVENC preset
|
|
crf: int = 18,
|
|
pixel_format: str = 'bgr24'):
|
|
|
|
self.output_path = Path(output_path)
|
|
self.output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
self.width = width
|
|
self.height = height
|
|
self.fps = fps
|
|
self.audio_source = audio_source
|
|
self.pixel_format = pixel_format
|
|
self.frames_written = 0
|
|
self.ffmpeg_process = None
|
|
|
|
# Build ffmpeg command
|
|
self.ffmpeg_cmd = self._build_ffmpeg_command(
|
|
video_codec, quality_preset, crf
|
|
)
|
|
|
|
# Start ffmpeg process
|
|
self._start_ffmpeg()
|
|
|
|
# Register cleanup
|
|
atexit.register(self.close)
|
|
|
|
print(f"📼 Streaming writer initialized:")
|
|
print(f" Output: {self.output_path}")
|
|
print(f" Resolution: {width}x{height} @ {fps}fps")
|
|
print(f" Codec: {video_codec}")
|
|
print(f" Audio: {'Yes' if audio_source else 'No'}")
|
|
|
|
def _build_ffmpeg_command(self, video_codec: str, preset: str, crf: int) -> list:
|
|
"""Build ffmpeg command with optimal settings"""
|
|
|
|
cmd = ['ffmpeg', '-y'] # Overwrite output
|
|
|
|
# Video input from pipe
|
|
cmd.extend([
|
|
'-f', 'rawvideo',
|
|
'-pix_fmt', self.pixel_format,
|
|
'-s', f'{self.width}x{self.height}',
|
|
'-r', str(self.fps),
|
|
'-i', 'pipe:0' # Read from stdin
|
|
])
|
|
|
|
# Audio input if provided
|
|
if self.audio_source and Path(self.audio_source).exists():
|
|
cmd.extend(['-i', str(self.audio_source)])
|
|
|
|
# Try GPU encoding first, fallback to CPU
|
|
if video_codec == 'h264_nvenc':
|
|
# NVIDIA GPU encoding
|
|
cmd.extend([
|
|
'-c:v', 'h264_nvenc',
|
|
'-preset', preset, # p1-p7, higher = better quality
|
|
'-rc', 'vbr', # Variable bitrate
|
|
'-cq', str(crf), # Quality level (0-51, lower = better)
|
|
'-b:v', '0', # Let VBR decide bitrate
|
|
'-maxrate', '50M', # Max bitrate for 8K
|
|
'-bufsize', '100M' # Buffer size
|
|
])
|
|
elif video_codec == 'hevc_nvenc':
|
|
# NVIDIA HEVC/H.265 encoding (better for 8K)
|
|
cmd.extend([
|
|
'-c:v', 'hevc_nvenc',
|
|
'-preset', preset,
|
|
'-rc', 'vbr',
|
|
'-cq', str(crf),
|
|
'-b:v', '0',
|
|
'-maxrate', '40M', # HEVC is more efficient
|
|
'-bufsize', '80M'
|
|
])
|
|
else:
|
|
# CPU fallback (libx264)
|
|
cmd.extend([
|
|
'-c:v', 'libx264',
|
|
'-preset', 'medium',
|
|
'-crf', str(crf),
|
|
'-pix_fmt', 'yuv420p'
|
|
])
|
|
|
|
# Audio settings
|
|
if self.audio_source:
|
|
cmd.extend([
|
|
'-c:a', 'copy', # Copy audio without re-encoding
|
|
'-map', '0:v:0', # Map video from pipe
|
|
'-map', '1:a:0', # Map audio from file
|
|
'-shortest' # Match shortest stream
|
|
])
|
|
else:
|
|
cmd.extend(['-map', '0:v:0']) # Video only
|
|
|
|
# Output file
|
|
cmd.append(str(self.output_path))
|
|
|
|
return cmd
|
|
|
|
def _start_ffmpeg(self) -> None:
|
|
"""Start ffmpeg subprocess"""
|
|
try:
|
|
print(f"🎬 Starting ffmpeg: {' '.join(self.ffmpeg_cmd[:10])}...")
|
|
|
|
self.ffmpeg_process = subprocess.Popen(
|
|
self.ffmpeg_cmd,
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
bufsize=10**8 # Large buffer for performance
|
|
)
|
|
|
|
# Set process to ignore SIGINT (Ctrl+C) - we'll handle it
|
|
if hasattr(signal, 'pthread_sigmask'):
|
|
signal.pthread_sigmask(signal.SIG_BLOCK, [signal.SIGINT])
|
|
|
|
except Exception as e:
|
|
# Try CPU fallback if GPU encoding fails
|
|
if 'nvenc' in self.ffmpeg_cmd:
|
|
print(f"⚠️ GPU encoding failed, trying CPU fallback...")
|
|
self.ffmpeg_cmd = self._build_ffmpeg_command('libx264', 'medium', 18)
|
|
self._start_ffmpeg()
|
|
else:
|
|
raise RuntimeError(f"Failed to start ffmpeg: {e}")
|
|
|
|
def write_frame(self, frame: np.ndarray) -> bool:
|
|
"""
|
|
Write a single frame to the video
|
|
|
|
Args:
|
|
frame: Frame to write (BGR format)
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
if self.ffmpeg_process is None or self.ffmpeg_process.poll() is not None:
|
|
raise RuntimeError("FFmpeg process is not running")
|
|
|
|
try:
|
|
# Ensure correct shape
|
|
if frame.shape != (self.height, self.width, 3):
|
|
raise ValueError(
|
|
f"Frame shape {frame.shape} doesn't match expected "
|
|
f"({self.height}, {self.width}, 3)"
|
|
)
|
|
|
|
# Ensure correct dtype
|
|
if frame.dtype != np.uint8:
|
|
frame = frame.astype(np.uint8)
|
|
|
|
# Write raw frame data to pipe
|
|
self.ffmpeg_process.stdin.write(frame.tobytes())
|
|
self.ffmpeg_process.stdin.flush()
|
|
|
|
self.frames_written += 1
|
|
|
|
# Periodic progress update
|
|
if self.frames_written % 100 == 0:
|
|
print(f" Written {self.frames_written} frames...", end='\r')
|
|
|
|
return True
|
|
|
|
except BrokenPipeError:
|
|
# Check if ffmpeg failed
|
|
if self.ffmpeg_process.poll() is not None:
|
|
stderr = self.ffmpeg_process.stderr.read().decode()
|
|
raise RuntimeError(f"FFmpeg process died: {stderr}")
|
|
raise
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to write frame: {e}")
|
|
|
|
def write_frame_alpha(self, frame: np.ndarray, alpha: np.ndarray) -> bool:
|
|
"""
|
|
Write frame with alpha channel (converts to green screen)
|
|
|
|
Args:
|
|
frame: RGB frame
|
|
alpha: Alpha mask (0-255)
|
|
|
|
Returns:
|
|
True if successful
|
|
"""
|
|
# Create green screen composite
|
|
green_bg = np.full_like(frame, [0, 255, 0], dtype=np.uint8)
|
|
|
|
# Normalize alpha to 0-1
|
|
if alpha.dtype == np.uint8:
|
|
alpha_float = alpha.astype(np.float32) / 255.0
|
|
else:
|
|
alpha_float = alpha
|
|
|
|
# Expand alpha to 3 channels if needed
|
|
if alpha_float.ndim == 2:
|
|
alpha_float = np.expand_dims(alpha_float, axis=2)
|
|
alpha_float = np.repeat(alpha_float, 3, axis=2)
|
|
|
|
# Composite
|
|
composite = (frame * alpha_float + green_bg * (1 - alpha_float)).astype(np.uint8)
|
|
|
|
return self.write_frame(composite)
|
|
|
|
def get_progress(self) -> Dict[str, Any]:
|
|
"""Get writing progress"""
|
|
return {
|
|
'frames_written': self.frames_written,
|
|
'duration_seconds': self.frames_written / self.fps if self.fps > 0 else 0,
|
|
'output_path': str(self.output_path),
|
|
'process_alive': self.ffmpeg_process is not None and self.ffmpeg_process.poll() is None
|
|
}
|
|
|
|
def close(self) -> None:
|
|
"""Close ffmpeg process and finalize video"""
|
|
if self.ffmpeg_process is not None:
|
|
try:
|
|
# Close stdin to signal end of input
|
|
if self.ffmpeg_process.stdin:
|
|
self.ffmpeg_process.stdin.close()
|
|
|
|
# Wait for ffmpeg to finish (with timeout)
|
|
print(f"\n🎬 Finalizing video with {self.frames_written} frames...")
|
|
self.ffmpeg_process.wait(timeout=30)
|
|
|
|
# Check return code
|
|
if self.ffmpeg_process.returncode != 0:
|
|
stderr = self.ffmpeg_process.stderr.read().decode()
|
|
warnings.warn(f"FFmpeg exited with code {self.ffmpeg_process.returncode}: {stderr}")
|
|
else:
|
|
print(f"✅ Video saved: {self.output_path}")
|
|
|
|
except subprocess.TimeoutExpired:
|
|
print("⚠️ FFmpeg taking too long, terminating...")
|
|
self.ffmpeg_process.terminate()
|
|
self.ffmpeg_process.wait(timeout=5)
|
|
|
|
except Exception as e:
|
|
warnings.warn(f"Error closing ffmpeg: {e}")
|
|
if self.ffmpeg_process.poll() is None:
|
|
self.ffmpeg_process.kill()
|
|
|
|
finally:
|
|
self.ffmpeg_process = None
|
|
|
|
def __enter__(self):
|
|
"""Context manager support"""
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Context manager cleanup"""
|
|
self.close()
|
|
|
|
def __del__(self):
|
|
"""Ensure cleanup on deletion"""
|
|
try:
|
|
self.close()
|
|
except:
|
|
pass |