bithuman 1.0.2__py3-none-any.whl

bithuman/lib/pth2h5.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Union
+
+import cv2
+import h5py
+import numpy as np
+
+
+def encode_mask_as_jpg(mask: np.ndarray, quality: int = 85) -> bytes:
+    """Encode a mask array as JPG bytes.
+
+    Args:
+        mask: Numpy array of mask (single channel)
+        quality: JPG compression quality (1-100)
+
+    Returns:
+        JPG encoded bytes
+    """
+    # Convert to BGR for OpenCV
+    mask = (mask * 255).astype(np.uint8)
+    mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
+
+    # Encode as JPG
+    _, encoded = cv2.imencode(".jpg", mask, [int(cv2.IMWRITE_JPEG_QUALITY), quality])
+    return encoded.tobytes()
+
+
+def convert_pth_to_h5(
+    pth_path: Union[str, Path], h5_path: Union[str, Path] = None
+) -> str:
+    """Convert a PyTorch model file (.pth) to a HDF5 file (.h5).
+
+    Args:
+        pth_path: Path to the PyTorch model file
+        h5_path: Path to save the HDF5 file. If None, saves to same location as pth_path
+    """
+    try:
+        import torch
+    except ImportError:
+        raise ImportError(
+            "PyTorch is not installed. Please install it using 'pip install torch'."
+        )
+
+    pth_path = Path(pth_path)
+    if h5_path is None:
+        h5_path = pth_path.with_suffix(".h5")
+
+    # Load the PyTorch model
+    data = torch.load(str(pth_path))
+
+    # Extract data
+    face_masks: list[bytes] = []
+    face_coords: list[np.ndarray] = []
+    frame_wh = data[0]["frame_wh"].numpy().astype(np.int32)
+
+    for item in data:
+        padded_crop_coords = item["padded_crop_coords"].numpy()
+        face_xyxy = item["face_coords"].numpy()
+        face_mask = item["face_mask"]
+
+        # Encode face mask if needed
+        if not isinstance(face_mask, bytes):
+            face_mask = encode_mask_as_jpg(face_mask.numpy())
+
+        # Adjust coordinates
+        shift_x, shift_y = padded_crop_coords[:2].astype(np.int32)
+        face_xyxy = face_xyxy[:4].astype(np.int32)
+        face_xyxy[0::2] += shift_x
+        face_xyxy[1::2] += shift_y
+
+        face_masks.append(face_mask)
+        face_coords.append(face_xyxy)
+
+    # Save to H5 file
+    with h5py.File(h5_path, "w") as f:
+        f.create_dataset("face_coords", data=face_coords)
+
+        dt = h5py.special_dtype(vlen=np.dtype("uint8"))
+        masks_dataset = f.create_dataset(
+            "face_masks", shape=(len(face_masks),), dtype=dt
+        )
+        for i, mask in enumerate(face_masks):
+            masks_dataset[i] = np.frombuffer(mask, dtype=np.uint8)
+        f.attrs["frame_wh"] = frame_wh
+
+    return str(h5_path)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert PyTorch .pth files to HDF5 format"
+    )
+    parser.add_argument("pth_path", type=str, help="Path to input .pth file")
+    parser.add_argument(
+        "--output", "-o", type=str, help="Path to output .h5 file (optional)"
+    )
+
+    args = parser.parse_args()
+    convert_pth_to_h5(args.pth_path, args.output)
+
+
+if __name__ == "__main__":
+    main()

bithuman/plugins/stt.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Union
+
+import aiohttp
+import numpy as np
+from loguru import logger
+
+from bithuman.audio import resample
+
+try:
+    from bithuman_local_voice import BithumanLocalSTT as BithumanSTTImpl
+except ImportError:
+    BithumanSTTImpl = None
+
+try:
+    from livekit.agents import utils
+    from livekit.agents.stt import (
+        STT,
+        SpeechData,
+        SpeechEvent,
+        SpeechEventType,
+        STTCapabilities,
+    )
+    from livekit.agents.types import NOT_GIVEN, APIConnectOptions, NotGivenOr
+except ImportError:
+    raise ImportError(
+        "livekit is required, please install it with `pip install livekit-agents`"
+    )
+
+
+@dataclass
+class _STTOptions:
+    locale: str = "en-US"
+    on_device: bool = True
+    punctuation: bool = True
+    debug: bool = False
+
+
+class BithumanSTTError(Exception):
+    pass
+
+
+class BithumanLocalSTT(STT):
+    _SAMPLE_RATE = 16000
+
+    def __init__(
+        self,
+        *,
+        locale="en-US",
+        server_url=None,
+        on_device=True,
+        punctuation=True,
+        debug=False,
+    ):
+        capabilities = STTCapabilities(streaming=False, interim_results=False)
+        super().__init__(capabilities=capabilities)
+        self._opts = _STTOptions(
+            locale=locale, on_device=on_device, punctuation=punctuation, debug=debug
+        )
+        self._server_url: str | None = None
+        self._session: aiohttp.ClientSession | None = None
+        self._stt_impl = None
+
+        if server_url:
+            self._server_url = server_url
+            self._session = aiohttp.ClientSession()
+        else:
+            if BithumanSTTImpl is None:
+                raise ImportError(
+                    "bithuman_local_voice is required if server_url is not provided, "
+                    "please install it with `pip install bithuman_local_voice`"
+                )
+            self._stt_impl = BithumanSTTImpl(
+                locale=locale,
+                on_device=on_device,
+                punctuation=punctuation,
+                debug=debug,
+            )
+
+    async def _recognize_impl(
+        self,
+        buffer: utils.audio.AudioBuffer,
+        *,
+        language: NotGivenOr[str] = NOT_GIVEN,
+        conn_options: APIConnectOptions,
+    ):
+        if utils.is_given(language) and language != self._opts.locale:
+            try:
+                await self._set_locale(language)
+            except Exception as e:
+                logger.error(f"Failed to set locale: {e}")
+
+        frame = utils.audio.combine_frames(buffer)
+        audio_data = np.frombuffer(frame.data, dtype=np.int16)
+        if frame.sample_rate != self._SAMPLE_RATE:
+            audio_data = resample(audio_data, frame.sample_rate, self._SAMPLE_RATE)
+
+        try:
+            result = await self._recognize_audio(
+                audio_data, sample_rate=self._SAMPLE_RATE
+            )
+        except Exception as e:
+            logger.warning(f"Failed to recognize audio with error: {e}")
+            result = {"text": "", "confidence": 0.0}
+
+        return SpeechEvent(
+            type=SpeechEventType.FINAL_TRANSCRIPT,
+            request_id="",
+            alternatives=[
+                SpeechData(
+                    language=self._opts.locale,
+                    text=result["text"],
+                    confidence=result["confidence"],
+                )
+            ],
+        )
+
+    async def _recognize_audio(
+        self, audio_data: Union[bytes, np.ndarray], sample_rate: int = 16000
+    ) -> Dict:
+        """Recognize speech in the provided audio data.
+
+        Args:
+            audio_data: Audio data as bytes or numpy array
+            sample_rate: Sample rate of the audio data
+
+        Returns:
+            Dict containing transcription and confidence score
+        """
+        if self._stt_impl is not None:
+            return await self._stt_impl.recognize(audio_data, sample_rate)
+
+        # Convert numpy array to bytes if needed
+        if isinstance(audio_data, np.ndarray):
+            # Convert to float in [-1, 1] range if not already
+            if audio_data.dtype != np.int16:
+                audio_data = (audio_data * 32767).astype(np.int16)
+            audio_bytes = audio_data.tobytes()
+        else:
+            audio_bytes = audio_data
+
+        async with self._session.post(
+            f"{self._server_url}/transcribe",
+            data=audio_bytes,
+            headers={"Content-Type": "application/octet-stream"},
+            timeout=5,
+        ) as response:
+            response.raise_for_status()
+            result: dict = await response.json()
+            return {
+                "text": result.get("text", ""),
+                "confidence": result.get("confidence", 0.0),
+            }
+
+    async def _set_locale(self, locale: str):
+        """Set the recognition locale.
+
+        Args:
+            locale: Locale identifier (e.g. en-US, fr-FR)
+        """
+        if self._stt_impl is not None:
+            await self._stt_impl.set_locale(locale)
+            self._opts.locale = locale
+            return
+
+        assert self._server_url is not None
+        async with self._session.post(
+            f"{self._server_url}/setLocale",
+            json={"locale": locale},
+            headers={"Content-Type": "application/json"},
+            timeout=5,
+        ) as response:
+            response.raise_for_status()
+            result: dict = await response.json()
+            if result.get("success"):
+                self._opts.locale = locale
+            raise BithumanSTTError(result.get("message", "Unknown error"))
+
+    async def aclose(self):
+        if self._stt_impl is not None:
+            await self._stt_impl.stop()
+        if self._session is not None:
+            await self._session.close()