python-detector-worker/core/utils/hardware_encoder.py

"""
Hardware-accelerated image encoding using NVIDIA NVENC or Intel QuickSync
"""

import cv2
import numpy as np
import logging
from typing import Optional, Tuple
import os

logger = logging.getLogger("detector_worker")


class HardwareEncoder:
    """Hardware-accelerated JPEG encoder using GPU."""

    def __init__(self):
        """Initialize hardware encoder."""
        self.nvenc_available = False
        self.vaapi_available = False
        self.turbojpeg_available = False

        # Check for TurboJPEG (fastest CPU-based option)
        try:
            from turbojpeg import TurboJPEG
            self.turbojpeg = TurboJPEG()
            self.turbojpeg_available = True
            logger.info("TurboJPEG accelerated encoding available")
        except ImportError:
            logger.debug("TurboJPEG not available")

        # Check for NVIDIA NVENC support
        try:
            # Test if we can create an NVENC encoder
            test_frame = np.zeros((720, 1280, 3), dtype=np.uint8)
            fourcc = cv2.VideoWriter_fourcc(*'H264')
            test_writer = cv2.VideoWriter(
                "test.mp4",
                fourcc,
                30,
                (1280, 720),
                [cv2.CAP_PROP_HW_ACCELERATION, cv2.VIDEO_ACCELERATION_ANY]
            )
            if test_writer.isOpened():
                self.nvenc_available = True
                logger.info("NVENC hardware encoding available")
            test_writer.release()
            if os.path.exists("test.mp4"):
                os.remove("test.mp4")
        except Exception as e:
            logger.debug(f"NVENC not available: {e}")

    def encode_jpeg(self, frame: np.ndarray, quality: int = 85) -> Optional[bytes]:
        """
        Encode frame to JPEG using the fastest available method.

        Args:
            frame: BGR image frame
            quality: JPEG quality (1-100)

        Returns:
            Encoded JPEG bytes or None on failure
        """
        try:
            # Method 1: TurboJPEG (3-5x faster than cv2.imencode)
            if self.turbojpeg_available:
                # Convert BGR to RGB for TurboJPEG
                rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                encoded = self.turbojpeg.encode(rgb_frame, quality=quality)
                return encoded

            # Method 2: Hardware-accelerated encoding via GStreamer (if available)
            if self.nvenc_available:
                return self._encode_with_nvenc(frame, quality)

            # Fallback: Standard OpenCV encoding
            encode_params = [cv2.IMWRITE_JPEG_QUALITY, quality]
            success, encoded = cv2.imencode('.jpg', frame, encode_params)
            if success:
                return encoded.tobytes()

            return None

        except Exception as e:
            logger.error(f"Failed to encode frame: {e}")
            return None

    def _encode_with_nvenc(self, frame: np.ndarray, quality: int) -> Optional[bytes]:
        """
        Encode using NVIDIA NVENC hardware encoder.

        This is complex to implement directly, so we'll use a GStreamer pipeline
        if available.
        """
        try:
            # Create a GStreamer pipeline for hardware encoding
            height, width = frame.shape[:2]
            gst_pipeline = (
                f"appsrc ! "
                f"video/x-raw,format=BGR,width={width},height={height},framerate=30/1 ! "
                f"videoconvert ! "
                f"nvvideoconvert ! "  # GPU color conversion
                f"nvjpegenc quality={quality} ! "  # Hardware JPEG encoder
                f"appsink"
            )

            # This would require GStreamer Python bindings
            # For now, fall back to TurboJPEG or standard encoding
            logger.debug("NVENC JPEG encoding not fully implemented, using fallback")
            encode_params = [cv2.IMWRITE_JPEG_QUALITY, quality]
            success, encoded = cv2.imencode('.jpg', frame, encode_params)
            if success:
                return encoded.tobytes()

            return None

        except Exception as e:
            logger.error(f"NVENC encoding failed: {e}")
            return None

    def encode_batch(self, frames: list, quality: int = 85) -> list:
        """
        Batch encode multiple frames for better GPU utilization.

        Args:
            frames: List of BGR frames
            quality: JPEG quality

        Returns:
            List of encoded JPEG bytes
        """
        encoded_frames = []

        if self.turbojpeg_available:
            # TurboJPEG can handle batch encoding efficiently
            for frame in frames:
                rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                encoded = self.turbojpeg.encode(rgb_frame, quality=quality)
                encoded_frames.append(encoded)
        else:
            # Fallback to sequential encoding
            for frame in frames:
                encoded = self.encode_jpeg(frame, quality)
                encoded_frames.append(encoded)

        return encoded_frames


# Global encoder instance
_hardware_encoder = None


def get_hardware_encoder() -> HardwareEncoder:
    """Get or create the global hardware encoder instance."""
    global _hardware_encoder
    if _hardware_encoder is None:
        _hardware_encoder = HardwareEncoder()
    return _hardware_encoder


def encode_frame_hardware(frame: np.ndarray, quality: int = 85) -> Optional[bytes]:
    """
    Convenience function to encode a frame using hardware acceleration.

    Args:
        frame: BGR image frame
        quality: JPEG quality (1-100)

    Returns:
        Encoded JPEG bytes or None on failure
    """
    encoder = get_hardware_encoder()
    return encoder.encode_jpeg(frame, quality)