bithuman 1.0.2__py3-none-any.whl

bithuman/engine/compression.py

@@ -0,0 +1,257 @@
+"""Frame compression/decompression — replaces image.cpp encode/decode."""
+
+from __future__ import annotations
+
+import os
+import struct
+import tempfile
+from typing import Optional
+import numpy as np
+from .enums import CompressionType
+
+IMAGE_TYPE_BGR8 = 16  # Matches CV_8UC3
+
+# ---------------------------------------------------------------------------
+# TurboJPEG (preferred) with cv2 fallback
+# ---------------------------------------------------------------------------
+try:
+    from turbojpeg import TurboJPEG, TJPF_BGR, TJFLAG_FASTDCT, TJSAMP_420
+
+    _TURBOJPEG_LIB_PATHS = [
+        None,  # auto-detect first
+        "/usr/lib/x86_64-linux-gnu/libturbojpeg.so",
+        "/usr/lib/libturbojpeg.so",
+        "/opt/libjpeg-turbo/lib64/libturbojpeg.so",
+        "/usr/lib/aarch64-linux-gnu/libturbojpeg.so",
+    ]
+    _tj = None
+    for _lib_path in _TURBOJPEG_LIB_PATHS:
+        try:
+            _tj = TurboJPEG(_lib_path) if _lib_path else TurboJPEG()
+            break
+        except (RuntimeError, OSError):
+            continue
+    _HAS_TURBOJPEG = _tj is not None
+except ImportError:
+    _HAS_TURBOJPEG = False
+
+# ---------------------------------------------------------------------------
+# LZ4
+# ---------------------------------------------------------------------------
+try:
+    import lz4.block
+
+    _HAS_LZ4 = True
+except ImportError:
+    _HAS_LZ4 = False
+
+
+# ---------------------------------------------------------------------------
+# JPEG
+# ---------------------------------------------------------------------------
+def encode_jpeg(image: np.ndarray, quality: int = 95) -> bytes:
+    """Encode BGR uint8 image to JPEG bytes.
+
+    Matches image.cpp encodeJPEG: TurboJPEG with TJPF_BGR, TJSAMP_420.
+    """
+    if _HAS_TURBOJPEG:
+        return _tj.encode(
+            image, quality=quality, pixel_format=TJPF_BGR, jpeg_subsample=TJSAMP_420
+        )
+    import cv2
+
+    _, buf = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, quality])
+    return buf.tobytes()
+
+
+def decode_jpeg(data: bytes) -> np.ndarray:
+    """Decode JPEG bytes to BGR uint8 numpy array (H, W, 3).
+
+    Matches image.cpp decodeJPEG: TurboJPEG with TJPF_BGR.
+    """
+    if not data:
+        raise RuntimeError("Empty JPEG data")
+    if _HAS_TURBOJPEG:
+        return _tj.decode(data, pixel_format=TJPF_BGR)
+    import cv2
+
+    arr = np.frombuffer(data, dtype=np.uint8)
+    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+    if img is None:
+        raise RuntimeError("Failed to decompress JPEG")
+    return img
+
+
+# ---------------------------------------------------------------------------
+# LZ4
+# ---------------------------------------------------------------------------
+def encode_lz4(image: np.ndarray) -> bytes:
+    """Encode BGR image with 16-byte header + LZ4 compressed data.
+
+    Header format matches image.cpp encodeLZ4:
+    [uint32 width, uint32 height, uint32 type(16), uint32 original_size]
+    """
+    if not _HAS_LZ4:
+        raise RuntimeError("lz4 package not installed")
+    h, w = image.shape[:2]
+    data_size = w * h * 3
+    header = struct.pack("<IIII", w, h, IMAGE_TYPE_BGR8, data_size)
+    # lz4.block.compress accepts buffer-protocol objects; pass raw memory
+    # buffer directly when array is already C-contiguous to avoid .tobytes() copy.
+    if image.flags.c_contiguous:
+        compressed = lz4.block.compress(image.data, store_size=False)
+    else:
+        compressed = lz4.block.compress(
+            np.ascontiguousarray(image).tobytes(), store_size=False
+        )
+    return header + compressed
+
+
+def decode_lz4(data: bytes) -> np.ndarray:
+    """Decode 16-byte header + LZ4 compressed data to BGR image.
+
+    Matches image.cpp decodeLZ4.
+    """
+    if not _HAS_LZ4:
+        raise RuntimeError("lz4 package not installed")
+    if len(data) < 16:
+        raise RuntimeError("Invalid LZ4 data: header too small")
+    w, h, img_type, orig_size = struct.unpack("<IIII", data[:16])
+    if img_type != IMAGE_TYPE_BGR8:
+        raise RuntimeError(f"Invalid image type in LZ4 data: {img_type}")
+    expected = w * h * 3
+    if orig_size != expected:
+        raise RuntimeError(
+            f"Invalid data size in LZ4 header: {orig_size} != {expected}"
+        )
+    decompressed = lz4.block.decompress(data[16:], uncompressed_size=orig_size)
+    return np.frombuffer(decompressed, dtype=np.uint8).reshape(h, w, 3).copy()
+
+
+# ---------------------------------------------------------------------------
+# NONE (raw with header)
+# ---------------------------------------------------------------------------
+def encode_none(image: np.ndarray) -> bytes:
+    """Encode with 16-byte header + raw BGR pixel data.
+
+    Matches image.cpp encodeNONE.
+    """
+    h, w = image.shape[:2]
+    data_size = w * h * 3
+    header = struct.pack("<IIII", w, h, IMAGE_TYPE_BGR8, data_size)
+    return header + np.ascontiguousarray(image).tobytes()
+
+
+def decode_none(data: bytes) -> np.ndarray:
+    """Decode 16-byte header + raw BGR pixel data.
+
+    Matches image.cpp decodeNONE.
+    """
+    if len(data) < 16:
+        raise RuntimeError("Invalid raw data: header too small")
+    w, h, img_type, data_size = struct.unpack("<IIII", data[:16])
+    if img_type != IMAGE_TYPE_BGR8:
+        raise RuntimeError(f"Invalid image type in raw data: {img_type}")
+    expected = w * h * 3
+    if data_size != expected or len(data) != expected + 16:
+        raise RuntimeError("Invalid data size in raw header")
+    return np.frombuffer(data[16:], dtype=np.uint8).reshape(h, w, 3).copy()
+
+
+# ---------------------------------------------------------------------------
+# TEMP_FILE
+# ---------------------------------------------------------------------------
+_temp_counter = 0
+
+
+def encode_temp_file(image: np.ndarray, temp_dir: str, quality: int = 95) -> bytes:
+    """Encode image to JPEG, write to temp file, return path as bytes.
+
+    Matches image.cpp encodeTempFile.
+    """
+    global _temp_counter
+    if not temp_dir:
+        raise RuntimeError("Temp dir is not set")
+    _temp_counter += 1
+    path = os.path.join(temp_dir, f"frame_{_temp_counter}.bin")
+    jpeg_data = encode_jpeg(image, quality)
+    with open(path, "wb") as f:
+        f.write(jpeg_data)
+    return path.encode("utf-8")
+
+
+def decode_temp_file(path_data: bytes) -> np.ndarray:
+    """Read JPEG data from temp file, decode to BGR image.
+
+    Matches image.cpp decodeTempFile.
+    """
+    if not path_data:
+        raise RuntimeError("Empty path data")
+    path = path_data.decode("utf-8")
+    with open(path, "rb") as f:
+        jpeg_data = f.read()
+    return decode_jpeg(jpeg_data)
+
+
+# ---------------------------------------------------------------------------
+# Temp dir management
+# ---------------------------------------------------------------------------
+def create_temp_dir() -> str:
+    """Create a unique temp directory for frame storage.
+
+    Matches image.cpp createTempDir.
+    """
+    root = os.path.join(tempfile.gettempdir(), "bithuman_tmp")
+    os.makedirs(root, exist_ok=True)
+    return tempfile.mkdtemp(prefix="video_", dir=root)
+
+
+def cleanup_temp_dir(path: str) -> None:
+    """Remove temp directory and all contents.
+
+    Matches image.cpp cleanupTempDir.
+    """
+    import shutil
+
+    if path and os.path.isdir(path):
+        shutil.rmtree(path, ignore_errors=True)
+
+
+# ---------------------------------------------------------------------------
+# Unified encode/decode API
+# ---------------------------------------------------------------------------
+def encode_image(
+    image: np.ndarray,
+    compression: CompressionType,
+    quality: int = 95,
+    temp_dir: str = "",
+) -> bytes:
+    """Encode image with specified compression type.
+
+    Matches image.cpp encodeImage.
+    """
+    if compression == CompressionType.JPEG:
+        return encode_jpeg(image, quality)
+    elif compression == CompressionType.LZ4:
+        return encode_lz4(image)
+    elif compression == CompressionType.NONE:
+        return encode_none(image)
+    elif compression == CompressionType.TEMP_FILE:
+        return encode_temp_file(image, temp_dir, quality)
+    raise ValueError(f"Unknown compression type: {compression}")
+
+
+def decode_image(data: bytes, compression: CompressionType) -> np.ndarray:
+    """Decode image with specified compression type.
+
+    Matches image.cpp decodeImage.
+    """
+    if compression == CompressionType.JPEG:
+        return decode_jpeg(data)
+    elif compression == CompressionType.LZ4:
+        return decode_lz4(data)
+    elif compression == CompressionType.NONE:
+        return decode_none(data)
+    elif compression == CompressionType.TEMP_FILE:
+        return decode_temp_file(data)
+    raise ValueError(f"Unknown compression type: {compression}")

bithuman/engine/enums.py

@@ -0,0 +1,16 @@
+"""Pure Python enums replacing pybind11-exported C++ enums."""
+
+from enum import IntEnum
+
+
+class CompressionType(IntEnum):
+    NONE = 0
+    JPEG = 1
+    LZ4 = 2
+    TEMP_FILE = 3
+
+
+class LoadingMode(IntEnum):
+    SYNC = 0
+    ASYNC = 1
+    ON_DEMAND = 2

bithuman/engine/image_ops.py

@@ -0,0 +1,192 @@
+"""Image operations — replaces image_ops.cpp (SSE2 blend + bilinear resize).
+
+Provides two blend implementations:
+1. numba JIT (primary) — per-pixel with exact div255 formula, SIMD-comparable
+2. numpy fallback — vectorized uint16 arithmetic, ~3-5ms for typical face ROI
+
+Resize delegates to cv2.resize (IPP/SIMD-optimized internally).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import cv2
+
+try:
+    from numba import njit, prange
+
+    _HAS_NUMBA = True
+except ImportError:
+    _HAS_NUMBA = False
+
+
+# ---------------------------------------------------------------------------
+# Blend: numba JIT (primary path)
+# ---------------------------------------------------------------------------
+if _HAS_NUMBA:
+
+    @njit(cache=True, parallel=True)
+    def _blend_3ch_numba(
+        frame: np.ndarray,
+        lip: np.ndarray,
+        mask: np.ndarray,
+        roi_x: int,
+        roi_y: int,
+        roi_h: int,
+        roi_w: int,
+    ) -> None:
+        """In-place alpha blend with 3-channel mask using exact div255."""
+        for y in prange(roi_h):
+            for x in range(roi_w):
+                for c in range(3):
+                    m = np.uint16(mask[y, x, c])
+                    im = np.uint16(255) - m
+                    val = np.uint16(lip[y, x, c]) * m + np.uint16(
+                        frame[roi_y + y, roi_x + x, c]
+                    ) * im
+                    # div255: exact integer division by 255
+                    val = (val + np.uint16(1) + ((val + np.uint16(1)) >> np.uint16(8))) >> np.uint16(8)
+                    frame[roi_y + y, roi_x + x, c] = np.uint8(val)
+
+    @njit(cache=True, parallel=True)
+    def _blend_1ch_numba(
+        frame: np.ndarray,
+        lip: np.ndarray,
+        mask: np.ndarray,
+        roi_x: int,
+        roi_y: int,
+        roi_h: int,
+        roi_w: int,
+    ) -> None:
+        """In-place alpha blend with 1-channel mask using exact div255."""
+        for y in prange(roi_h):
+            for x in range(roi_w):
+                m = np.uint16(mask[y, x, 0] if mask.ndim == 3 else mask[y, x])
+                im = np.uint16(255) - m
+                for c in range(3):
+                    val = np.uint16(lip[y, x, c]) * m + np.uint16(
+                        frame[roi_y + y, roi_x + x, c]
+                    ) * im
+                    val = (val + np.uint16(1) + ((val + np.uint16(1)) >> np.uint16(8))) >> np.uint16(8)
+                    frame[roi_y + y, roi_x + x, c] = np.uint8(val)
+
+
+# ---------------------------------------------------------------------------
+# Blend: numpy fallback
+# ---------------------------------------------------------------------------
+def _blend_numpy(
+    frame: np.ndarray,
+    lip: np.ndarray,
+    mask: np.ndarray,
+    roi_x: int,
+    roi_y: int,
+    roi_h: int,
+    roi_w: int,
+) -> None:
+    """Vectorized numpy blend with exact div255 formula."""
+    roi = frame[roi_y : roi_y + roi_h, roi_x : roi_x + roi_w]
+
+    # Ensure mask is 3-channel
+    if mask.ndim == 2:
+        m = mask[:, :, np.newaxis].astype(np.uint16)
+    elif mask.ndim == 3 and mask.shape[2] == 1:
+        m = mask.astype(np.uint16)
+    else:
+        m = mask.astype(np.uint16)
+
+    # Broadcast to 3 channels if needed
+    if m.shape[2] == 1:
+        m = np.broadcast_to(m, roi.shape).copy()
+
+    im = np.uint16(255) - m
+    val = lip[:roi_h, :roi_w].astype(np.uint16) * m + roi.astype(np.uint16) * im
+    # div255: (val + 1 + ((val + 1) >> 8)) >> 8
+    val_p1 = val + np.uint16(1)
+    result = ((val_p1 + (val_p1 >> 8)) >> 8).astype(np.uint8)
+    frame[roi_y : roi_y + roi_h, roi_x : roi_x + roi_w] = result
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+def blend_face_region(
+    frame: np.ndarray,
+    roi_x: int,
+    roi_y: int,
+    roi_w: int,
+    roi_h: int,
+    lip: np.ndarray,
+    mask: np.ndarray,
+) -> None:
+    """Alpha-blend lip overlay into face region of frame (in-place).
+
+    Matches image_ops.cpp blend_face_region (lines 429-503):
+    Formula per byte: out = div255(lip * mask + face * (255 - mask))
+    where div255(x) = (x + 1 + ((x + 1) >> 8)) >> 8
+
+    Args:
+        frame: (H, W, 3) uint8 BGR — modified in-place.
+        roi_x, roi_y: top-left corner of face region in frame.
+        roi_w, roi_h: dimensions of face region.
+        lip: (roi_h, roi_w, 3) uint8 BGR lip overlay.
+        mask: (roi_h, roi_w, 1) or (roi_h, roi_w, 3) uint8 blend mask.
+    """
+    is_3ch = mask.ndim == 3 and mask.shape[2] == 3
+
+    if _HAS_NUMBA:
+        if is_3ch:
+            _blend_3ch_numba(frame, lip, mask, roi_x, roi_y, roi_h, roi_w)
+        else:
+            _blend_1ch_numba(frame, lip, mask, roi_x, roi_y, roi_h, roi_w)
+    else:
+        _blend_numpy(frame, lip, mask, roi_x, roi_y, roi_h, roi_w)
+
+
+def resize_image(image: np.ndarray, new_width: int, new_height: int) -> np.ndarray:
+    """Bilinear resize using OpenCV (IPP/SIMD optimized).
+
+    Replaces image_ops.cpp resize_image (lines 347-375).
+    """
+    if image.size == 0:
+        return image
+    h, w = image.shape[:2]
+    if w == new_width and h == new_height:
+        return image
+    return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LINEAR)
+
+
+def resize_image_scale(image: np.ndarray, scale: float) -> np.ndarray:
+    """Resize by scale factor. Replaces resize_image_scale in C++."""
+    if scale == 1.0 or image.size == 0:
+        return image
+    new_w = round(image.shape[1] * scale)
+    new_h = round(image.shape[0] * scale)
+    return resize_image(image, new_w, new_h)
+
+
+# ---------------------------------------------------------------------------
+# Pre-warm numba JIT in background thread (avoids ~200ms spike on first blend)
+# ---------------------------------------------------------------------------
+def _warmup_numba_jit() -> None:
+    """Trigger JIT compilation of blend functions with tiny arrays.
+
+    With cache=True, the compiled code persists to __pycache__ for
+    subsequent process starts.
+    """
+    if not _HAS_NUMBA:
+        return
+    try:
+        _tiny_frame = np.zeros((2, 2, 3), dtype=np.uint8)
+        _tiny_lip = np.zeros((1, 1, 3), dtype=np.uint8)
+        _tiny_mask_3ch = np.zeros((1, 1, 3), dtype=np.uint8)
+        _tiny_mask_1ch = np.zeros((1, 1, 1), dtype=np.uint8)
+        _blend_3ch_numba(_tiny_frame, _tiny_lip, _tiny_mask_3ch, 0, 0, 1, 1)
+        _blend_1ch_numba(_tiny_frame, _tiny_lip, _tiny_mask_1ch, 0, 0, 1, 1)
+    except Exception:
+        pass  # Warmup failure is non-fatal
+
+
+import threading as _threading
+
+_warmup_thread = _threading.Thread(target=_warmup_numba_jit, daemon=True)
+_warmup_thread.start()

bithuman/engine/inference.py

@@ -0,0 +1,108 @@
+"""ONNX audio encoder inference — replaces C++ melChunkToAudioEmbedding."""
+
+from __future__ import annotations
+
+import os
+import threading
+from typing import Dict, Optional
+
+import numpy as np
+import onnxruntime as ort
+
+# Module-level ONNX session cache keyed by absolute model path.
+# ONNX Runtime internally shares model weights across sessions, but this
+# avoids Python-side overhead (~8 KB per duplicate session object).
+_onnx_session_cache: Dict[str, ort.InferenceSession] = {}
+_onnx_session_lock = threading.Lock()
+
+
+class AudioEncoder:
+    """Wraps an ONNX audio encoder session for mel-to-embedding inference.
+
+    Matches the C++ BithumanRuntime ONNX setup:
+    - intra_op_num_threads=1 (avoid nested parallelism)
+    - disable CPU memory arena (container compatibility)
+    - fallback from EXTENDED to BASIC optimization level
+    - pre-allocated input buffer reused across calls
+    """
+
+    def __init__(self, model_path: str = ""):
+        self._session: Optional[ort.InferenceSession] = None
+        self._input_name: Optional[str] = None
+        self._output_name: Optional[str] = None
+        # Pre-allocated input buffer matching C++ onnx_input_data_ (1*1*80*16)
+        self._input_buf: np.ndarray = np.zeros((1, 1, 80, 16), dtype=np.float32)
+        if model_path:
+            self.load(model_path)
+
+    def load(self, model_path: str) -> None:
+        """Load ONNX model, reusing cached session if available.
+
+        Mirrors generator.cpp setAudioEncoder (lines 439-502):
+        - Try EXTENDED first, fall back to BASIC on failure.
+        """
+        abs_path = os.path.abspath(model_path)
+
+        with _onnx_session_lock:
+            if abs_path in _onnx_session_cache:
+                self._session = _onnx_session_cache[abs_path]
+                self._input_name = self._session.get_inputs()[0].name
+                self._output_name = self._session.get_outputs()[0].name
+                return
+
+        # Create new session outside the lock (loading can be slow)
+        opts = ort.SessionOptions()
+        opts.intra_op_num_threads = 1
+        opts.enable_cpu_mem_arena = False
+
+        try:
+            opts.graph_optimization_level = (
+                ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
+            )
+            session = ort.InferenceSession(model_path, opts)
+        except Exception:
+            opts.graph_optimization_level = (
+                ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
+            )
+            try:
+                session = ort.InferenceSession(model_path, opts)
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to load audio encoder model '{model_path}': {e}"
+                ) from e
+
+        with _onnx_session_lock:
+            # Double-check: another thread may have cached it while we were loading
+            if abs_path in _onnx_session_cache:
+                session = _onnx_session_cache[abs_path]
+            else:
+                _onnx_session_cache[abs_path] = session
+
+        self._session = session
+        self._input_name = self._session.get_inputs()[0].name
+        self._output_name = self._session.get_outputs()[0].name
+
+    def encode(self, mel_chunk: np.ndarray) -> np.ndarray:
+        """Convert a mel spectrogram chunk to an audio embedding.
+
+        Args:
+            mel_chunk: float32 array of shape (80, 16).
+
+        Returns:
+            float32 array of shape (embedding_dim,).
+
+        Mirrors generator.cpp melChunkToAudioEmbedding (lines 505-533):
+        - Input shape: (1, 1, 80, 16), row-major (C-contiguous = numpy default).
+        - Output: squeezed to 1-D embedding vector.
+        """
+        if self._session is None:
+            raise RuntimeError("Audio encoder not initialized")
+
+        # Copy into pre-allocated buffer (matches C++ zero-copy pattern)
+        np.copyto(self._input_buf[0, 0], mel_chunk.astype(np.float32))
+
+        result = self._session.run(
+            [self._output_name],
+            {self._input_name: self._input_buf},
+        )
+        return result[0].squeeze()

bithuman/engine/knn.py

@@ -0,0 +1,58 @@
+"""Audio feature KNN search — replaces Eigen squaredNorm + minCoeff."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+class AudioFeatureIndex:
+    """Stores audio feature cluster centers and finds the nearest match.
+
+    Mirrors generator.cpp processAudio (lines 337-342):
+    ``(audio_features_.rowwise() - audio_embed.transpose())
+        .rowwise().squaredNorm().minCoeff(&min_index)``
+    """
+
+    def __init__(self) -> None:
+        self._features: np.ndarray = np.empty((0, 0), dtype=np.float32)
+
+    def set_features(self, features: np.ndarray) -> None:
+        """Set audio feature cluster centers.
+
+        Args:
+            features: float32 array shape (num_clusters, embedding_dim), row-major.
+        """
+        self._features = np.ascontiguousarray(features, dtype=np.float32)
+
+    def load_from_h5(self, h5_path: str) -> None:
+        """Load audio features from an HDF5 file.
+
+        Mirrors generator.cpp setAudioFeature (lines 409-432):
+        Reads dataset ``audio_feature`` as row-major float32 matrix.
+        """
+        import h5py
+
+        with h5py.File(h5_path, "r") as f:
+            data = f["audio_feature"][:]
+        self.set_features(data)
+
+    def find_nearest(self, embedding: np.ndarray) -> int:
+        """Return index of nearest cluster center by squared Euclidean distance.
+
+        Args:
+            embedding: float32 array shape (embedding_dim,).
+
+        Returns:
+            Index of the nearest cluster center.
+        """
+        diff = self._features - embedding
+        sq_dists = np.einsum("ij,ij->i", diff, diff)
+        return int(np.argmin(sq_dists))
+
+    @property
+    def num_clusters(self) -> int:
+        return self._features.shape[0]
+
+    @property
+    def embedding_dim(self) -> int:
+        return self._features.shape[1] if self._features.size > 0 else 0