server/h264_encoder.py

"""
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()