basis_universal/python/basisu_py/codec.py

# basisu_py/codec.py

import importlib
import numpy as np
from PIL import Image
import ctypes

from .constants import BasisTexFormat, BasisQuality, BasisEffort, BasisFlags
from pathlib import Path

class EncoderBackend:
    NATIVE = "native"
    WASM = "wasm"
    AUTO = "auto"

class Encoder:

    def __init__(self, backend=EncoderBackend.AUTO):
        self.backend = backend
        self._native = None
        self._wasm = None
        self.backend_name = None

        # ------------------------------------------------------------------
        # Try native first (AUTO or NATIVE modes)
        # ------------------------------------------------------------------
        if backend in (EncoderBackend.AUTO, EncoderBackend.NATIVE):
            try:
                import basisu_py.basisu_python as native_encoder
                native_encoder.init()

                self._native = native_encoder
                self._wasm = None
                self.backend_name = "NATIVE"

                print("[Encoder] Using native backend")
                return

            except Exception as e:
                if backend == EncoderBackend.NATIVE:
                    raise RuntimeError(
                        f"[Encoder] Native backend requested but unavailable: {e}"
                    )
                print("[Encoder] Native unavailable; falling back to WASM:", e)

        # ------------------------------------------------------------------
        # Fallback to WASM (AUTO or explicitly WASM)
        # ------------------------------------------------------------------
        try:
            from basisu_py.wasm.wasm_encoder import BasisuWasmEncoder
        except Exception as e:
            raise RuntimeError(
                f"[Encoder] WASM backend cannot be imported: {e}\n"
                "Make sure wasmtime is installed and basisu_py/wasm/*.wasm exist."
            )

        wasm_path = Path(__file__).parent / "wasm" / "basisu_module_st.wasm"
        self._wasm = BasisuWasmEncoder(str(wasm_path))
        self._wasm.load()
        self._native = None
        self.backend_name = "WASM"

        print("[Encoder] Using WASM backend")           
            

    # ------------------------------------------------------
    # Public API
    # ------------------------------------------------------
    def compress(self,
                 image,
                 format=-1,
                 quality=BasisQuality.MAX,
                 effort=BasisEffort.DEFAULT,
                 flags=BasisFlags.KTX2_OUTPUT | BasisFlags.SRGB | BasisFlags.THREADED | BasisFlags.XUASTC_LDR_FULL_ZSTD):
                
        rgba_bytes, w, h, is_hdr = self._convert_input_to_rgba_bytes(image)
        
        # Auto-select format if user passed -1
        if format == -1:
            if is_hdr:
                format = BasisTexFormat.cUASTC_HDR_6x6
            else:
                format = BasisTexFormat.cXUASTC_LDR_6x6

        if self._native:
            return self._compress_native(rgba_bytes, w, h, format, quality, effort, flags, is_hdr)
        else:
            return self._compress_wasm(rgba_bytes, w, h, format, quality, effort, flags, is_hdr)
            
    def compress_float32(self, arr, **kwargs):
        if not isinstance(arr, np.ndarray) or arr.dtype != np.float32:
            raise ValueError("compress_float32 requires float32 NumPy HxWx4 array")

        return self.compress(arr, **kwargs)

    # ------------------------------------------------------
    # Native backend
    # ------------------------------------------------------
    def _compress_native(self, bytes_data, w, h, fmt, quality, effort, flags, is_hdr=False):
        enc = self._native

        params = enc.new_params()

        try:
            buf_ptr = enc.alloc(len(bytes_data))

            # Write raw bytes (uint8 or float32)
            ctypes.memmove(buf_ptr, bytes_data, len(bytes_data))

            if is_hdr:
                ok = enc.set_image_float_rgba(params, 0, buf_ptr, w, h, w * 16)  # 4 floats = 16 bytes per pixel
            else:
                ok = enc.set_image_rgba32(params, 0, buf_ptr, w, h, w * 4)

            if not ok:
                raise RuntimeError("Native encoder: set_image failed (HDR or LDR)")

            ok = enc.compress(params, fmt, quality, effort, flags, 0.0)
            if not ok:
                raise RuntimeError("Native encoder: compress() failed")

            size = enc.get_comp_data_size(params)
            ofs = enc.get_comp_data_ofs(params)
            blob = enc.read_memory(ofs, size)
            return blob

        finally:
            enc.delete_params(params)
            if buf_ptr:
                enc.free(buf_ptr)

    # ------------------------------------------------------
    # WASM backend
    # ------------------------------------------------------
    def _compress_wasm(self, bytes_data, w, h, fmt, quality, effort, flags, is_hdr=False):
        enc = self._wasm

        params = enc.new_params()

        try:
            buf_ptr = enc.alloc(len(bytes_data))
            enc.write_bytes(buf_ptr, bytes_data)

            if is_hdr:
                ok = enc.set_image_float_rgba(params, 0, buf_ptr, w, h, w * 16)
            else:
                ok = enc.set_image_rgba32(params, 0, buf_ptr, w, h, w * 4)

            if not ok:
                raise RuntimeError("WASM encoder: set_image failed (HDR or LDR)")

            ok = enc.compress(params, fmt, quality, effort, flags, 0.0)
            if not ok:
                raise RuntimeError("WASM encoder: compress() failed")

            size = enc.get_comp_data_size(params)
            ofs = enc.get_comp_data_ofs(params)
            blob = enc.read_bytes(ofs, size)
            return blob

        finally:
            enc.delete_params(params)
            if buf_ptr:
                enc.free(buf_ptr)

    # ------------------------------------------------------
    # Image conversion
    # ------------------------------------------------------
    def _convert_input_to_rgba_bytes(self, image):
        """
        Accept:
          - Pillow Image (LDR) -> returns uint8 bytes
          - NumPy uint8 LDR -> returns uint8 bytes
          - NumPy float32 HDR -> returns float32 bytes
        Returns (bytes, width, height, is_hdr)
        """

        # Pillow image -> LDR
        if isinstance(image, Image.Image):
            image = image.convert("RGBA")
            arr = np.array(image, dtype=np.uint8)
            h, w = arr.shape[:2]
            return arr.tobytes(), w, h, False

        # NumPy array
        elif isinstance(image, np.ndarray):

            # HDR float32 image
            if image.dtype == np.float32:
                if image.ndim != 3 or image.shape[2] not in (3,4):
                    raise ValueError("HDR NumPy image must be HxWx3 or HxWx4 float32")

                h, w, c = image.shape

                # Expand RGB -> RGBA if needed
                if c == 3:
                    alpha = np.ones((h, w, 1), dtype=np.float32)
                    arr = np.concatenate([image, alpha], axis=2)
                else:
                    arr = image

                return arr.tobytes(), w, h, True

            # LDR uint8 image
            if image.dtype == np.uint8:
                if image.ndim != 3 or image.shape[2] not in (3,4):
                    raise ValueError("LDR NumPy image must be HxWx3 or HxWx4 uint8")

                h, w, c = image.shape

                if c == 3:
                    alpha = np.full((h, w, 1), 255, dtype=np.uint8)
                    arr = np.concatenate([image, alpha], axis=2)
                else:
                    arr = image

                return arr.tobytes(), w, h, False

            raise ValueError("NumPy image must be uint8 (LDR) or float32 (HDR)")

        else:
            raise TypeError("compress() expects Pillow Image or NumPy array")