""" H.264 Encoder for WebSocket Streaming Encodes frames as H.264 NAL units for efficient delta-based streaming. Uses PyAV (ffmpeg) for hardware-accelerated encoding. """ import av import numpy as np from fractions import Fraction from typing import Generator, Tuple import time class H264Encoder: """ H.264 encoder that produces NAL units for WebSocket streaming. Benefits over JPEG: - I-frames (keyframes): Full frame, sent periodically - P-frames (delta): Only differences, ~10x smaller than JPEG - Hardware acceleration with NVENC on GPU """ def __init__( self, width: int = 256, height: int = 256, fps: int = 25, bitrate: int = 500_000, # 500kbps for 256x256 keyframe_interval: int = 30, # I-frame every 30 frames preset: str = "ultrafast", tune: str = "zerolatency" ): self.width = width self.height = height self.fps = fps self.bitrate = bitrate self.keyframe_interval = keyframe_interval self.preset = preset self.tune = tune self.encoder = None self.frame_count = 0 self.codec_config = None # SPS/PPS for decoder init self._init_encoder() def _init_encoder(self): """Initialize H.264 encoder with optimal settings for low latency.""" # Try NVENC first (GPU), fall back to libx264 (CPU) codecs_to_try = ['h264_nvenc', 'libx264'] for codec_name in codecs_to_try: try: codec = av.codec.Codec(codec_name, 'w') self.encoder = codec.create() self.encoder.width = self.width self.encoder.height = self.height self.encoder.pix_fmt = 'yuv420p' self.encoder.time_base = Fraction(1, self.fps) self.encoder.framerate = Fraction(self.fps, 1) self.encoder.bit_rate = self.bitrate self.encoder.gop_size = self.keyframe_interval # Codec-specific options for low latency if codec_name == 'h264_nvenc': self.encoder.options = { 'preset': 'p1', # Fastest NVENC preset 'tune': 'ull', # Ultra low latency 'zerolatency': '1', 'rc': 'cbr', # Constant bitrate } else: # libx264 self.encoder.options = { 'preset': self.preset, 'tune': self.tune, 'profile': 'baseline', # Best compatibility 'level': '3.1', } self.encoder.open() print(f"[H264] Initialized {codec_name} encoder: {self.width}x{self.height} @ {self.fps}fps") return except Exception as e: print(f"[H264] {codec_name} not available: {e}") continue raise RuntimeError("No H.264 encoder available (tried h264_nvenc, libx264)") def encode_frame(self, frame: np.ndarray, force_keyframe: bool = False) -> Tuple[bytes, bool]: """ Encode a single frame to H.264 NAL units. Args: frame: BGR or RGB numpy array (H, W, 3) force_keyframe: Force this frame to be an I-frame Returns: Tuple of (encoded_data, is_keyframe) """ import cv2 # Ensure correct dimensions if frame.shape[0] != self.height or frame.shape[1] != self.width: frame = cv2.resize(frame, (self.width, self.height)) # Create AV frame from numpy array directly # PyAV can convert from RGB/BGR automatically if len(frame.shape) == 3 and frame.shape[2] == 3: # Assume BGR from OpenCV, convert to RGB for PyAV frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) else: frame_rgb = frame # Create VideoFrame from numpy array av_frame = av.VideoFrame.from_ndarray(frame_rgb, format='rgb24') # Convert to YUV420P (required for H.264) av_frame = av_frame.reformat(format='yuv420p') av_frame.pts = self.frame_count av_frame.time_base = self.encoder.time_base # Determine if this should be a keyframe is_keyframe = force_keyframe or (self.frame_count % self.keyframe_interval == 0) if is_keyframe: av_frame.pict_type = av.video.frame.PictureType.I self.frame_count += 1 # Encode packets = self.encoder.encode(av_frame) # Collect all packet data data = b'' for packet in packets: data += bytes(packet) return data, is_keyframe def get_codec_config(self) -> bytes: """ Get SPS/PPS NAL units needed to initialize the decoder. Must be sent before any frames. """ if self.encoder.extradata: return bytes(self.encoder.extradata) return b'' def flush(self) -> bytes: """Flush any remaining frames from encoder.""" packets = self.encoder.encode(None) data = b'' for packet in packets: data += bytes(packet) return data def close(self): """Close the encoder.""" if self.encoder: try: # Flush remaining packets self.encoder.encode(None) except: pass self.encoder = None def __del__(self): self.close() class H264StreamEncoder: """ Streaming H.264 encoder optimized for real-time avatar generation. Manages frame buffering and produces a continuous stream. """ def __init__(self, width: int = 256, height: int = 256, fps: int = 25): self.encoder = H264Encoder( width=width, height=height, fps=fps, bitrate=self._calculate_bitrate(width, height), keyframe_interval=fps, # I-frame every second ) self.started = False self.start_time = None def _calculate_bitrate(self, width: int, height: int) -> int: """Calculate appropriate bitrate based on resolution.""" pixels = width * height # Rough estimate: 2 bits per pixel at 25fps return int(pixels * 2 * 25) def start(self) -> bytes: """ Start the stream, returns codec config (SPS/PPS). Client must receive this before any frames. """ self.started = True self.start_time = time.time() return self.encoder.get_codec_config() def encode(self, frame: np.ndarray, force_keyframe: bool = False) -> Tuple[bytes, dict]: """ Encode a frame and return data with metadata. Returns: Tuple of (h264_data, metadata_dict) """ if not self.started: self.start() data, is_keyframe = self.encoder.encode_frame(frame, force_keyframe) metadata = { 'frame': self.encoder.frame_count - 1, 'keyframe': is_keyframe, 'size': len(data), 'timestamp': time.time() - self.start_time, } return data, metadata def stop(self) -> bytes: """Stop the stream, returns any remaining data.""" self.started = False return self.encoder.flush() # Test the encoder if __name__ == '__main__': import cv2 # Create test frames encoder = H264StreamEncoder(width=256, height=256, fps=25) # Get codec config config = encoder.start() print(f"Codec config size: {len(config)} bytes") # Encode some test frames total_jpeg = 0 total_h264 = 0 for i in range(30): # Create a test frame with some variation frame = np.zeros((256, 256, 3), dtype=np.uint8) cv2.putText(frame, f"Frame {i}", (50, 128), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2) cv2.circle(frame, (128 + i*2, 128), 30, (0, 255, 0), -1) # H.264 encode h264_data, meta = encoder.encode(frame) # JPEG encode for comparison _, jpeg_data = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80]) total_jpeg += len(jpeg_data) total_h264 += len(h264_data) frame_type = "I" if meta['keyframe'] else "P" print(f"Frame {i} ({frame_type}): H264={len(h264_data):,}b, JPEG={len(jpeg_data):,}b, ratio={len(jpeg_data)/max(1,len(h264_data)):.1f}x") print(f"\nTotal: H264={total_h264:,}b, JPEG={total_jpeg:,}b") print(f"H.264 is {total_jpeg/total_h264:.1f}x more efficient than JPEG") encoder.stop()