feat: downsample oversized textures (#260)

* feat: downsample oversided 3d views

* typing

* one threed check

Author: tlambert03

src/ndv/models/_data_wrapper.py

@@ -209,7 +209,7 @@ def create(cls, data: ArrayT) -> DataWrapper[ArrayT]:
         for subclass in sorted(_recurse_subclasses(cls), key=lambda x: x.PRIORITY):
             try:
                 if subclass.supports(data):
-                    logger.debug(f"Using {subclass.__name__} to wrap {type(data)}")
+                    logger.debug("Using %s to wrap %s", subclass.__name__, type(data))
                     return subclass(data)
             except Exception as e:
                 warnings.warn(

src/ndv/views/_pygfx/_array_canvas.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from contextlib import suppress
-from typing import TYPE_CHECKING, Any, Literal, cast
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast
 from weakref import ReferenceType, WeakValueDictionary, ref
 
 import cmap as _cmap
@@ -18,6 +19,7 @@
 )
 from ndv.models._viewer_model import ArrayViewerModel, InteractionMode
 from ndv.views._app import filter_mouse_events
+from ndv.views._util import downsample_data
 from ndv.views.bases import ArrayCanvas, CanvasElement, ImageHandle
 from ndv.views.bases._graphics._canvas_elements import RectangularROIHandle, ROIMoveMode
 
@@ -63,22 +65,33 @@ def __init__(self, image: pygfx.Image | pygfx.Volume, render: Callable) -> None:
         self._render = render
         self._grid = cast("Texture", image.geometry.grid)
         self._material = cast("ImageBasicMaterial", image.material)
+        # per-axis downsample strides applied to fit GPU texture limits
+        self._downsample_factors: tuple[int, ...] = ()
 
     def data(self) -> np.ndarray:
         return self._grid.data  # type: ignore [no-any-return]
 
     def set_data(self, data: np.ndarray) -> None:
+        is_three_d = isinstance(self._image, pygfx.Volume)
+        data, self._downsample_factors = _downcast_and_downsample(
+            data,
+            three_d=is_three_d,
+            warn=False,
+            copy=False,
+        )
         # If dimensions are unchanged, reuse the buffer
         if data.shape == self._grid.data.shape:
             self._grid.data[:] = data  # pyright: ignore[reportOptionalSubscript]
             self._grid.update_range((0, 0, 0), self._grid.size)
         # Otherwise, the size (and maybe number of dimensions) changed
         # - we need a new buffer
         else:
-            self._grid = pygfx.Texture(data, dim=2)
+            dim = 3 if is_three_d else 2
+            self._grid = pygfx.Texture(data, dim=dim)
             self._image.geometry = pygfx.Geometry(grid=self._grid)
             # RGB images (i.e. 3D datasets) cannot have a colormap
-            self._material.map = None if self._is_rgb() else self._cmap.to_pygfx()
+            if not is_three_d:
+                self._material.map = None if self._is_rgb() else self._cmap.to_pygfx()
 
     def visible(self) -> bool:
         return bool(self._image.visible)
@@ -465,11 +478,7 @@ def set_ndim(self, ndim: Literal[2, 3]) -> None:
 
     def add_image(self, data: np.ndarray | None = None) -> PyGFXImageHandle:
         """Add a new Image node to the scene."""
-        if data is not None:
-            # pygfx uses a view of the data without copy, so if we don't
-            # copy it here, the original data will be modified when the
-            # texture changes.
-            data = data.copy()
+        data, downsample_factors = _downcast_and_downsample(data, three_d=False)
         tex = pygfx.Texture(data, dim=2)
         image = pygfx.Image(
             pygfx.Geometry(grid=tex),
@@ -485,15 +494,12 @@ def add_image(self, data: np.ndarray | None = None) -> PyGFXImageHandle:
         # FIXME: I suspect there are more performant ways to refresh the canvas
         # look into it.
         handle = PyGFXImageHandle(image, self.refresh)
+        handle._downsample_factors = downsample_factors
         self._elements[image] = handle
         return handle
 
     def add_volume(self, data: np.ndarray | None = None) -> PyGFXImageHandle:
-        if data is not None:
-            # pygfx uses a view of the data without copy, so if we don't
-            # copy it here, the original data will be modified when the
-            # texture changes.
-            data = data.copy()
+        data, downsample_factors = _downcast_and_downsample(data, three_d=True)
         tex = pygfx.Texture(data, dim=3)
         vol = pygfx.Volume(
             pygfx.Geometry(grid=tex),
@@ -512,6 +518,7 @@ def add_volume(self, data: np.ndarray | None = None) -> PyGFXImageHandle:
         # FIXME: I suspect there are more performant ways to refresh the canvas
         # look into it.
         handle = PyGFXImageHandle(vol, self.refresh)
+        handle._downsample_factors = downsample_factors
         self._elements[vol] = handle
         return handle
 
@@ -539,10 +546,23 @@ def set_scales(self, scales: tuple[float, ...]) -> None:
             gfx_scales.append(1.0)
         sx, sy, sz = gfx_scales[0], gfx_scales[1], gfx_scales[2]
         has_visuals = False
-        for child in self._scene.children:
-            if isinstance(child, (pygfx.Image, pygfx.Volume)):
-                child.local.scale = (sx, sy, sz)
-                has_visuals = True
+        for handle in self._elements.values():
+            if not isinstance(handle, PyGFXImageHandle):
+                continue
+            child = handle._image
+            if not isinstance(child, (pygfx.Image, pygfx.Volume)):
+                continue
+            _sx, _sy, _sz = sx, sy, sz
+            # compensate for downsampling so coordinates stay correct
+            # factors are in data order; pygfx order is (x, y, z) = reversed
+            factors = handle._downsample_factors
+            if factors and any(f > 1 for f in factors):
+                rev = list(reversed(factors))
+                _sx *= rev[0]
+                _sy *= rev[1] if len(rev) > 1 else 1
+                _sz *= rev[2] if len(rev) > 2 else 1
+            child.local.scale = (_sx, _sy, _sz)
+            has_visuals = True
         if has_visuals:
             self.set_range()
 
@@ -710,3 +730,37 @@ def get_cursor(self, event: MouseMoveEvent) -> CursorType:
             if cursor := vis.get_cursor(event):
                 return cursor
         return CursorType.DEFAULT
+
+
+T = TypeVar("T", bound=np.ndarray | None)
+
+
+@lru_cache(maxsize=1)
+def _get_max_texture_sizes() -> tuple[int | None, int | None]:
+    """Return (max_2d, max_3d) texture dimensions from the wgpu adapter."""
+    try:
+        import wgpu
+
+        adapter = wgpu.gpu.request_adapter_sync()
+        limits = adapter.limits
+        max_2d = limits.get("max-texture-dimension-2d")
+        max_3d = limits.get("max-texture-dimension-3d")
+        return max_2d, max_3d
+    except Exception:
+        return None, None
+
+
+def _downcast_and_downsample(
+    data: T, three_d: bool, *, warn: bool = True, copy: bool = True
+) -> tuple[T, tuple[int, ...]]:
+    downsample_factors: tuple[int, ...] = ()
+    if data is not None:
+        if copy:
+            # pygfx uses a view of the data without copy, so if we don't
+            # copy it here, the original data will be modified when the
+            # texture changes.
+            data = data.copy()
+        maxd = _get_max_texture_sizes()[1 if three_d else 0]
+        if maxd is not None:
+            data, downsample_factors = downsample_data(data, maxd, warn=warn)  # type: ignore[assignment]
+    return data, downsample_factors  # pyright: ignore[reportReturnType]

src/ndv/views/_util.py

@@ -0,0 +1,33 @@
+"""Shared utilities for canvas backends."""
+
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+
+logger = logging.getLogger("ndv")
+
+
+def downsample_data(
+    data: np.ndarray, max_size: int, *, warn: bool = True
+) -> tuple[np.ndarray, tuple[int, ...]]:
+    """Downsample data so no axis exceeds max_size.
+
+    Returns the (possibly downsampled view) data and the per-axis stride factors.
+    """
+    factors = tuple(
+        int(np.ceil(s / max_size)) if s > max_size else 1 for s in data.shape
+    )
+    if any(f > 1 for f in factors):
+        if warn:
+            logger.warning(
+                "Data shape %s exceeds max texture dimension (%d) and will be "
+                "downsampled for rendering (strides: %s).",
+                data.shape,
+                max_size,
+                factors,
+            )
+        slices = tuple(slice(None, None, f) for f in factors)
+        data = data[slices]
+    return data, factors

src/ndv/views/_vispy/_array_canvas.py

@@ -3,7 +3,7 @@
 
 import warnings
 from contextlib import suppress
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast
 from weakref import ReferenceType, WeakValueDictionary
 
 import cmap as _cmap
@@ -23,6 +23,8 @@
 )
 from ndv.models._viewer_model import ArrayViewerModel, InteractionMode
 from ndv.views._app import filter_mouse_events
+from ndv.views._util import downsample_data
+from ndv.views._vispy._util import get_max_texture_sizes
 from ndv.views.bases import ArrayCanvas
 from ndv.views.bases._graphics._canvas_elements import (
     CanvasElement,
@@ -43,6 +45,8 @@ class VispyImageHandle(ImageHandle):
     def __init__(self, visual: visuals.ImageVisual | visuals.VolumeVisual) -> None:
         self._visual = visual
         self._allowed_dims = {2, 3} if isinstance(visual, visuals.ImageVisual) else {3}
+        # per-axis downsample strides applied to fit GPU texture limits
+        self._downsample_factors: tuple[int, ...] = ()
 
     def data(self) -> np.ndarray:
         try:
@@ -58,6 +62,13 @@ def set_data(self, data: np.ndarray) -> None:
                 stacklevel=2,
             )
             return
+
+        data, downsample_factors = _downcast_and_downsample(
+            data,
+            three_d=isinstance(self._visual, visuals.VolumeVisual),
+            warn=False,
+        )
+        self._downsample_factors = downsample_factors
         self._visual.set_data(data)
 
     def visible(self) -> bool:
@@ -358,7 +369,7 @@ def refresh(self) -> None:
 
     def add_image(self, data: np.ndarray | None = None) -> VispyImageHandle:
         """Add a new Image node to the scene."""
-        data = _downcast(data)
+        data, downsample_factors = _downcast_and_downsample(data, three_d=False)
         try:
             img = scene.visuals.Image(
                 data, parent=self._view.scene, texture_format="auto"
@@ -370,13 +381,14 @@ def add_image(self, data: np.ndarray | None = None) -> VispyImageHandle:
         img.set_gl_state("additive", depth_test=False)
         img.interactive = True
         handle = VispyImageHandle(img)
+        handle._downsample_factors = downsample_factors
         self._elements[img] = handle
         if data is not None:
             self.set_range()
         return handle
 
     def add_volume(self, data: np.ndarray | None = None) -> VispyImageHandle:
-        data = _downcast(data)
+        data, downsample_factors = _downcast_and_downsample(data, three_d=True)
         try:
             vol = scene.visuals.Volume(
                 data,
@@ -393,6 +405,7 @@ def add_volume(self, data: np.ndarray | None = None) -> VispyImageHandle:
         vol.set_gl_state("additive", depth_test=False)
         vol.interactive = True
         handle = VispyImageHandle(vol)
+        handle._downsample_factors = downsample_factors
         self._elements[vol] = handle
         if data is not None:
             self.set_range()
@@ -418,11 +431,24 @@ def set_scales(self, scales: tuple[float, ...]) -> None:
         while len(vis_scales) < 3:
             vis_scales.append(1.0)
         sx, sy, sz = vis_scales[0], vis_scales[1], vis_scales[2]
-        for child in self._view.scene.children:
-            if isinstance(child, (visuals.ImageVisual, visuals.VolumeVisual)):
-                child.transform = vispy.visuals.transforms.STTransform(
-                    scale=(sx, sy, sz)
-                )
+        for handle in self._elements.values():
+            if not isinstance(handle, VispyImageHandle):
+                continue
+            child = handle._visual
+            if not isinstance(child, (visuals.ImageVisual, visuals.VolumeVisual)):
+                continue
+            _sx, _sy, _sz = sx, sy, sz
+            # compensate for downsampling so coordinates stay correct
+            # factors are in data order; scene order is (x, y, z) = reversed
+            factors = handle._downsample_factors
+            if factors and any(f > 1 for f in factors):
+                rev = list(reversed(factors))
+                _sx *= rev[0]
+                _sy *= rev[1] if len(rev) > 1 else 1
+                _sz *= rev[2] if len(rev) > 2 else 1
+            child.transform = vispy.visuals.transforms.STTransform(
+                scale=(_sx, _sy, _sz)
+            )
         self.set_range()
 
     def set_range(
@@ -561,13 +587,29 @@ def get_cursor(self, event: MouseMoveEvent) -> CursorType:
         return CursorType.DEFAULT
 
 
-def _downcast(data: np.ndarray | None) -> np.ndarray | None:
+T = TypeVar("T", bound="np.ndarray | None")
+
+
+def _downcast(data: T) -> T:
     """Downcast >32bit data to 32bit."""
     # downcast to 32bit, preserving int/float
     if data is not None:
         if np.issubdtype(data.dtype, np.integer) and data.dtype.itemsize > 2:
             warnings.warn("Downcasting integer data to uint16.", stacklevel=2)
-            data = data.astype(np.uint16)
+            data = data.astype(np.uint16)  # type: ignore[assignment]
         elif np.issubdtype(data.dtype, np.floating) and data.dtype.itemsize > 4:
-            data = data.astype(np.float32)
+            data = data.astype(np.float32)  # type: ignore[assignment]
     return data
+
+
+def _downcast_and_downsample(
+    data: T, three_d: bool, warn: bool = True
+) -> tuple[T, tuple[int, ...]]:
+    """Downcast >32bit data to 32bit, and downsample GPU texture limits are exceeded."""
+    data = _downcast(data)
+    downsample_factors: tuple[int, ...] = ()
+    if data is not None:
+        maxd = get_max_texture_sizes()[1 if three_d else 0]
+        if maxd is not None:
+            data, downsample_factors = downsample_data(data, maxd, warn=warn)  # type: ignore[assignment]
+    return data, downsample_factors

src/ndv/views/_vispy/_util.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+from vispy.app import Canvas
+from vispy.gloo import gl
+from vispy.gloo.context import get_current_canvas
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+
+@contextmanager
+def _opengl_context() -> Generator[None, None, None]:
+    """Assure we are running with a valid OpenGL context.
+
+    Only create a Canvas is one doesn't exist. Creating and closing a
+    Canvas causes vispy to process Qt events which can cause problems.
+    """
+    canvas = Canvas(show=False) if get_current_canvas() is None else None
+    try:
+        yield
+    finally:
+        if canvas is not None:
+            canvas.close()
+
+
+@lru_cache
+def get_max_texture_sizes() -> tuple[int | None, int | None]:
+    """Return the maximum texture sizes for 2D and 3D rendering.
+
+    Returns
+    -------
+    Tuple[int | None, int | None]
+        The max textures sizes for (2d, 3d) rendering.
+    """
+    with _opengl_context():
+        max_size_2d = gl.glGetParameter(gl.GL_MAX_TEXTURE_SIZE)
+
+    if not max_size_2d:
+        max_size_2d = None
+
+    # vispy/gloo doesn't provide the GL_MAX_3D_TEXTURE_SIZE location,
+    # but it can be found in this list of constants
+    # http://pyopengl.sourceforge.net/documentation/pydoc/OpenGL.GL.html
+    with _opengl_context():
+        GL_MAX_3D_TEXTURE_SIZE = 32883
+        max_size_3d = gl.glGetParameter(GL_MAX_3D_TEXTURE_SIZE)
+
+    if not max_size_3d:
+        max_size_3d = None
+
+    return max_size_2d, max_size_3d