diff --git a/rerun_py/rerun_sdk/rerun/color_conversion.py b/rerun_py/rerun_sdk/rerun/color_conversion.py index 01b79adaf005..acebb0560b7d 100644 --- a/rerun_py/rerun_sdk/rerun/color_conversion.py +++ b/rerun_py/rerun_sdk/rerun/color_conversion.py @@ -1,30 +1,36 @@ """Color conversion utilities.""" -from typing import Sequence, Union +from typing import Union import numpy as np import numpy.typing as npt -def u8_array_to_rgba(arr: Sequence[int]) -> np.uint32: +def u8_array_to_rgba(arr: npt.NDArray[np.uint8]) -> npt.NDArray[np.uint32]: """ - Convert an array[4] of uint8 values into a uint32. + Convert an array with inner dimension [R,G,B,A] into packed uint32 values. Parameters ---------- - arr : Sequence[int] - The array of uint8 values to convert in RGBA order. + arr : + Nx3 or Nx4 `[[r,g,b,a], ... ]` of uint8 values Returns ------- - int - The uint32 value as 0xRRGGBBAA. + npt.NDArray[np.uint32] + Array of uint32 value as 0xRRGGBBAA. """ - red = arr[0] - green = arr[1] - blue = arr[2] - alpha = arr[3] if len(arr) == 4 else 0xFF - return np.uint32((red << 24) + (green << 16) + (blue << 8) + alpha) + r = arr[:, 0] + g = arr[:, 1] + b = arr[:, 2] + a = arr[:, 3] if arr.shape[1] == 4 else np.repeat(0xFF, len(arr)) + # Reverse the byte order because this is how we encode into uint32 + arr = np.vstack([a, b, g, r]).T + # Make contiguous and then reinterpret + arr = np.ascontiguousarray(arr, dtype=np.uint8) + arr = arr.view(np.uint32) + arr = np.squeeze(arr, axis=1) + return arr # type: ignore[return-value] def linear_to_gamma_u8_value(linear: npt.NDArray[Union[np.float32, np.float64]]) -> npt.NDArray[np.uint8]: diff --git a/rerun_py/rerun_sdk/rerun/components/color.py b/rerun_py/rerun_sdk/rerun/components/color.py index 12862c8c71c8..d2ff97cd9965 100644 --- a/rerun_py/rerun_sdk/rerun/components/color.py +++ b/rerun_py/rerun_sdk/rerun/components/color.py @@ -16,7 +16,7 @@ class ColorRGBAArray(pa.ExtensionArray): # type: ignore[misc] def from_numpy(array: npt.NDArray[np.uint8]) -> ColorRGBAArray: """Build a `ColorRGBAArray` from an numpy array.""" - storage = pa.array([u8_array_to_rgba(c) for c in array], type=ColorRGBAType.storage_type) + storage = pa.array(u8_array_to_rgba(array), type=ColorRGBAType.storage_type) # TODO(john) enable extension type wrapper # return cast(ColorRGBAArray, pa.ExtensionArray.from_storage(ColorRGBAType(), storage)) return storage # type: ignore[no-any-return]