feat: inference subsystem and optimization to decoder

services/stream_decoder.py

@@ -0,0 +1,481 @@
+import threading
+from typing import Optional
+from collections import deque
+from enum import Enum
+import torch
+import PyNvVideoCodec as nvc
+import av
+import numpy as np
+from cuda.bindings import driver as cuda_driver
+from .jpeg_encoder import encode_frame_to_jpeg
+
+
+def nv12_to_rgb_gpu(nv12_tensor: torch.Tensor, height: int, width: int) -> torch.Tensor:
+    """
+    Convert NV12 format to RGB on GPU using PyTorch operations.
+
+    NV12 format:
+    - Y plane: height x width (luminance)
+    - UV plane: (height/2) x width (interleaved U and V, subsampled by 2)
+
+    Total tensor size: (height * 3/2) x width
+
+    Args:
+        nv12_tensor: Input tensor in NV12 format, shape (H*3/2, W)
+        height: Original frame height
+        width: Original frame width
+
+    Returns:
+        RGB tensor, shape (3, H, W) in range [0, 255]
+    """
+    device = nv12_tensor.device
+
+    # Split Y and UV planes
+    y_plane = nv12_tensor[:height, :].float()  # (H, W)
+    uv_plane = nv12_tensor[height:, :].float()  # (H/2, W)
+
+    # Reshape UV plane to separate U and V channels
+    # UV is interleaved: U0V0U1V1... we need to deinterleave
+    uv_plane = uv_plane.reshape(height // 2, width // 2, 2)  # (H/2, W/2, 2)
+    u_plane = uv_plane[:, :, 0]  # (H/2, W/2)
+    v_plane = uv_plane[:, :, 1]  # (H/2, W/2)
+
+    # Upsample U and V to full resolution using bilinear interpolation
+    u_upsampled = torch.nn.functional.interpolate(
+        u_plane.unsqueeze(0).unsqueeze(0),  # (1, 1, H/2, W/2)
+        size=(height, width),
+        mode='bilinear',
+        align_corners=False
+    ).squeeze(0).squeeze(0)  # (H, W)
+
+    v_upsampled = torch.nn.functional.interpolate(
+        v_plane.unsqueeze(0).unsqueeze(0),  # (1, 1, H/2, W/2)
+        size=(height, width),
+        mode='bilinear',
+        align_corners=False
+    ).squeeze(0).squeeze(0)  # (H, W)
+
+    # YUV to RGB conversion using BT.601 standard
+    # R = Y + 1.402 * (V - 128)
+    # G = Y - 0.344136 * (U - 128) - 0.714136 * (V - 128)
+    # B = Y + 1.772 * (U - 128)
+
+    y = y_plane
+    u = u_upsampled - 128.0
+    v = v_upsampled - 128.0
+
+    r = y + 1.402 * v
+    g = y - 0.344136 * u - 0.714136 * v
+    b = y + 1.772 * u
+
+    # Clamp to [0, 255] and convert to uint8
+    r = torch.clamp(r, 0, 255).to(torch.uint8)
+    g = torch.clamp(g, 0, 255).to(torch.uint8)
+    b = torch.clamp(b, 0, 255).to(torch.uint8)
+
+    # Stack to (3, H, W)
+    rgb = torch.stack([r, g, b], dim=0)
+
+    return rgb
+
+
+class ConnectionStatus(Enum):
+    DISCONNECTED = "disconnected"
+    CONNECTING = "connecting"
+    CONNECTED = "connected"
+    ERROR = "error"
+    RECONNECTING = "reconnecting"
+
+
+class StreamDecoderFactory:
+    """
+    Factory for creating StreamDecoder instances with shared CUDA context.
+    This minimizes VRAM overhead by sharing the CUDA context across all decoders.
+    """
+
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls, gpu_id: int = 0):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super(StreamDecoderFactory, cls).__new__(cls)
+                    cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self, gpu_id: int = 0):
+        if self._initialized:
+            return
+
+        self.gpu_id = gpu_id
+
+        # Initialize CUDA and get device
+        err, = cuda_driver.cuInit(0)
+        if err != cuda_driver.CUresult.CUDA_SUCCESS:
+            raise RuntimeError(f"Failed to initialize CUDA: {err}")
+
+        # Get CUDA device
+        err, self.cuda_device = cuda_driver.cuDeviceGet(gpu_id)
+        if err != cuda_driver.CUresult.CUDA_SUCCESS:
+            raise RuntimeError(f"Failed to get CUDA device {gpu_id}: {err}")
+
+        # Retain primary context (shared across all decoders)
+        err, self.cuda_context = cuda_driver.cuDevicePrimaryCtxRetain(self.cuda_device)
+        if err != cuda_driver.CUresult.CUDA_SUCCESS:
+            raise RuntimeError(f"Failed to retain CUDA primary context: {err}")
+
+        self._initialized = True
+        print(f"StreamDecoderFactory initialized with shared CUDA context on GPU {gpu_id}")
+
+    def create_decoder(self, rtsp_url: str, buffer_size: int = 30,
+                      codec: str = "h264") -> 'StreamDecoder':
+        """
+        Create a new StreamDecoder instance with shared CUDA context.
+
+        Args:
+            rtsp_url: RTSP stream URL
+            buffer_size: Number of frames to buffer in VRAM
+            codec: Video codec (h264, hevc, etc.)
+
+        Returns:
+            StreamDecoder instance
+        """
+        return StreamDecoder(
+            rtsp_url=rtsp_url,
+            cuda_context=self.cuda_context,
+            gpu_id=self.gpu_id,
+            buffer_size=buffer_size,
+            codec=codec
+        )
+
+    def __del__(self):
+        """Cleanup shared CUDA context on factory destruction"""
+        if hasattr(self, 'cuda_device') and hasattr(self, 'gpu_id'):
+            cuda_driver.cuDevicePrimaryCtxRelease(self.cuda_device)
+
+
+class StreamDecoder:
+    """
+    Decodes RTSP stream using NVDEC and maintains a ring buffer of frames in GPU VRAM.
+    Thread-safe for concurrent read/write operations.
+    """
+
+    def __init__(self, rtsp_url: str, cuda_context, gpu_id: int,
+                 buffer_size: int = 30, codec: str = "h264"):
+        """
+        Initialize StreamDecoder.
+
+        Args:
+            rtsp_url: RTSP stream URL
+            cuda_context: Shared CUDA context handle
+            gpu_id: GPU device ID
+            buffer_size: Number of frames to keep in ring buffer
+            codec: Video codec type
+        """
+        self.rtsp_url = rtsp_url
+        self.cuda_context = cuda_context
+        self.gpu_id = gpu_id
+        self.buffer_size = buffer_size
+        self.codec = codec
+
+        # Connection status
+        self.status = ConnectionStatus.DISCONNECTED
+        self._status_lock = threading.Lock()
+
+        # Frame buffer (ring buffer) - stores CUDA device pointers
+        self.frame_buffer = deque(maxlen=buffer_size)
+        self._buffer_lock = threading.RLock()
+
+        # Decoder and container instances
+        self.decoder = None
+        self.container = None
+
+        # Decode thread
+        self._decode_thread: Optional[threading.Thread] = None
+        self._stop_flag = threading.Event()
+
+        # Frame metadata
+        self.frame_width: Optional[int] = None
+        self.frame_height: Optional[int] = None
+        self.frame_count: int = 0
+
+    def start(self):
+        """Start the RTSP stream decoding in background thread"""
+        if self._decode_thread is not None and self._decode_thread.is_alive():
+            print(f"Decoder already running for {self.rtsp_url}")
+            return
+
+        self._stop_flag.clear()
+        self._decode_thread = threading.Thread(target=self._decode_loop, daemon=True)
+        self._decode_thread.start()
+        print(f"Started decoder thread for {self.rtsp_url}")
+
+    def stop(self):
+        """Stop the decoding thread and cleanup resources"""
+        self._stop_flag.set()
+        if self._decode_thread is not None:
+            self._decode_thread.join(timeout=5.0)
+        self._cleanup()
+        print(f"Stopped decoder for {self.rtsp_url}")
+
+    def _set_status(self, status: ConnectionStatus):
+        """Thread-safe status update"""
+        with self._status_lock:
+            self.status = status
+
+    def get_status(self) -> ConnectionStatus:
+        """Get current connection status"""
+        with self._status_lock:
+            return self.status
+
+    def _init_rtsp_connection(self) -> bool:
+        """Initialize RTSP connection using PyAV + PyNvVideoCodec"""
+        try:
+            self._set_status(ConnectionStatus.CONNECTING)
+
+            # Open RTSP stream with PyAV
+            options = {
+                'rtsp_transport': 'tcp',
+                'max_delay': '500000',  # 500ms
+                'rtsp_flags': 'prefer_tcp',
+                'timeout': '5000000',  # 5 seconds
+            }
+
+            self.container = av.open(self.rtsp_url, options=options)
+
+            # Get video stream
+            video_stream = self.container.streams.video[0]
+            self.frame_width = video_stream.width
+            self.frame_height = video_stream.height
+
+            print(f"RTSP connected: {self.frame_width}x{self.frame_height}")
+
+            # Map codec name to PyNvVideoCodec codec enum
+            codec_map = {
+                'h264': nvc.cudaVideoCodec.H264,
+                'hevc': nvc.cudaVideoCodec.HEVC,
+                'h265': nvc.cudaVideoCodec.HEVC,
+            }
+
+            codec_id = codec_map.get(self.codec.lower(), nvc.cudaVideoCodec.H264)
+
+            # Initialize NVDEC decoder with shared CUDA context
+            self.decoder = nvc.CreateDecoder(
+                gpuid=self.gpu_id,
+                codec=codec_id,
+                cudacontext=self.cuda_context,
+                usedevicememory=True
+            )
+
+            self._set_status(ConnectionStatus.CONNECTED)
+            return True
+
+        except Exception as e:
+            print(f"Failed to connect to RTSP stream {self.rtsp_url}: {e}")
+            self._set_status(ConnectionStatus.ERROR)
+            return False
+
+    def _decode_loop(self):
+        """Main decode loop running in background thread"""
+        retry_count = 0
+        max_retries = 5
+
+        while not self._stop_flag.is_set():
+            # Initialize connection
+            if not self._init_rtsp_connection():
+                retry_count += 1
+                if retry_count >= max_retries:
+                    print(f"Max retries reached for {self.rtsp_url}")
+                    self._set_status(ConnectionStatus.ERROR)
+                    break
+
+                self._set_status(ConnectionStatus.RECONNECTING)
+                self._stop_flag.wait(timeout=2.0)
+                continue
+
+            retry_count = 0  # Reset on successful connection
+
+            try:
+                # Decode loop - iterate through packets from PyAV
+                for packet in self.container.demux(video=0):
+                    if self._stop_flag.is_set():
+                        break
+
+                    if packet.dts is None:
+                        continue
+
+                    # Convert packet to numpy array
+                    packet_data = np.frombuffer(bytes(packet), dtype=np.uint8)
+
+                    # Create PacketData and pass numpy array pointer
+                    pkt = nvc.PacketData()
+                    pkt.bsl_data = packet_data.ctypes.data
+                    pkt.bsl = len(packet_data)
+
+                    # Decode using NVDEC
+                    decoded_frames = self.decoder.Decode(pkt)
+
+                    if not decoded_frames:
+                        continue
+
+                    # Add frames to ring buffer (thread-safe)
+                    with self._buffer_lock:
+                        for frame in decoded_frames:
+                            self.frame_buffer.append(frame)
+                            self.frame_count += 1
+
+            except Exception as e:
+                print(f"Error in decode loop for {self.rtsp_url}: {e}")
+                self._set_status(ConnectionStatus.RECONNECTING)
+                self._cleanup()
+                self._stop_flag.wait(timeout=2.0)
+
+    def _cleanup(self):
+        """Cleanup resources"""
+        if self.container:
+            try:
+                self.container.close()
+            except:
+                pass
+            self.container = None
+
+        self.decoder = None
+
+        with self._buffer_lock:
+            self.frame_buffer.clear()
+
+    def get_frame(self, index: int = -1, rgb: bool = True) -> Optional[torch.Tensor]:
+        """
+        Get a frame from the buffer as a CUDA tensor (in VRAM).
+
+        Args:
+            index: Frame index in buffer (-1 for latest, -2 for second latest, etc.)
+            rgb: If True, convert NV12 to RGB. If False, return raw NV12 format.
+
+        Returns:
+            torch.Tensor in CUDA memory (device tensor) or None if buffer empty
+            - If rgb=True: Shape (3, H, W) in RGB format, dtype uint8
+            - If rgb=False: Shape (H*3/2, W) in NV12 format, dtype uint8
+        """
+        with self._buffer_lock:
+            if len(self.frame_buffer) == 0:
+                return None
+
+            try:
+                decoded_frame = self.frame_buffer[index]
+
+                # Convert DecodedFrame to PyTorch tensor using DLPack (zero-copy)
+                # This keeps the data in GPU memory
+                nv12_tensor = torch.from_dlpack(decoded_frame)
+
+                if not rgb:
+                    # Return raw NV12 format
+                    return nv12_tensor
+
+                # Convert NV12 to RGB on GPU
+                if self.frame_height is None or self.frame_width is None:
+                    print("Frame dimensions not available")
+                    return None
+
+                rgb_tensor = nv12_to_rgb_gpu(nv12_tensor, self.frame_height, self.frame_width)
+                return rgb_tensor
+
+            except (IndexError, Exception) as e:
+                print(f"Error getting frame: {e}")
+                return None
+
+    def get_latest_frame(self, rgb: bool = True) -> Optional[torch.Tensor]:
+        """
+        Get the most recent decoded frame as CUDA tensor.
+
+        Args:
+            rgb: If True, convert to RGB. If False, return raw NV12.
+
+        Returns:
+            torch.Tensor on GPU in RGB (3, H, W) or NV12 (H*3/2, W) format
+        """
+        return self.get_frame(-1, rgb=rgb)
+
+    def get_frame_cpu(self, index: int = -1, rgb: bool = True) -> Optional[np.ndarray]:
+        """
+        Get a frame from the buffer and copy it to CPU memory as numpy array.
+
+        Args:
+            index: Frame index in buffer (-1 for latest, -2 for second latest, etc.)
+            rgb: If True, convert NV12 to RGB. If False, return raw NV12 format.
+
+        Returns:
+            numpy.ndarray in CPU memory or None if buffer empty
+            - If rgb=True: Shape (H, W, 3) in RGB format, dtype uint8 (HWC format for easy display)
+            - If rgb=False: Shape (H*3/2, W) in NV12 format, dtype uint8
+        """
+        # Get frame on GPU
+        gpu_frame = self.get_frame(index=index, rgb=rgb)
+
+        if gpu_frame is None:
+            return None
+
+        # Transfer from GPU to CPU
+        cpu_tensor = gpu_frame.cpu()
+
+        # Convert to numpy array
+        if rgb:
+            # Convert from (3, H, W) to (H, W, 3) for standard image format
+            cpu_array = cpu_tensor.permute(1, 2, 0).numpy()
+        else:
+            # Keep NV12 format as-is
+            cpu_array = cpu_tensor.numpy()
+
+        return cpu_array
+
+    def get_latest_frame_cpu(self, rgb: bool = True) -> Optional[np.ndarray]:
+        """
+        Get the most recent decoded frame as CPU numpy array.
+
+        Args:
+            rgb: If True, convert to RGB. If False, return raw NV12.
+
+        Returns:
+            numpy.ndarray in CPU memory
+            - If rgb=True: Shape (H, W, 3) in RGB format, dtype uint8
+            - If rgb=False: Shape (H*3/2, W) in NV12 format, dtype uint8
+        """
+        return self.get_frame_cpu(-1, rgb=rgb)
+
+    def get_buffer_size(self) -> int:
+        """Get current number of frames in buffer"""
+        with self._buffer_lock:
+            return len(self.frame_buffer)
+
+    def is_connected(self) -> bool:
+        """Check if stream is actively connected"""
+        return self.get_status() == ConnectionStatus.CONNECTED
+
+    def get_frame_as_jpeg(self, index: int = -1, quality: int = 95) -> Optional[bytes]:
+        """
+        Get a frame from the buffer and encode to JPEG.
+
+        This method:
+        1. Gets RGB frame from buffer (stays on GPU)
+        2. Encodes to JPEG using nvJPEG (GPU operation via shared encoder)
+        3. Transfers JPEG bytes to CPU
+        4. Returns bytes for saving to disk
+
+        Args:
+            index: Frame index in buffer (-1 for latest)
+            quality: JPEG quality (0-100, default 95)
+
+        Returns:
+            JPEG encoded bytes or None if frame unavailable
+        """
+        # Get RGB frame (on GPU)
+        rgb_frame = self.get_frame(index=index, rgb=True)
+
+        # Use the shared JPEG encoder from jpeg_encoder module
+        return encode_frame_to_jpeg(rgb_frame, quality=quality)
+
+    def __repr__(self):
+        return (f"StreamDecoder(url={self.rtsp_url}, status={self.status.value}, "
+                f"buffer={self.get_buffer_size()}/{self.buffer_size}, "
+                f"frames_decoded={self.frame_count})")