| # 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() | |
| buf_ptr = None | |
| 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() | |
| buf_ptr = None | |
| 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") |