add transforms (#39)

* start chunk selection

* fix circular import

* add transform tests

* remove old transform implementation

* move chunked image utils

* add init

* rename utils

* rename utils

* update chunk

* add transform to visual

* add transform setting

* implement set_transform

Author: kevinyamauchi

examples/oblique_slicing.py

@@ -8,7 +8,7 @@
 from qtpy.QtWidgets import QApplication, QGridLayout, QLabel, QWidget
 from wgpu.gui.qt import WgpuCanvas
 
-from cellier.slicer.transforms import AffineTransform
+from cellier.transform import AffineTransform
 
 
 def make_plane_mesh(

examples/tiled_rendering/octree_pyfx.py

@@ -7,8 +7,8 @@
 from skimage.transform import resize
 from wgpu.gui.qt import WgpuCanvas
 
-from cellier.util.chunk import ChunkedArray3D, MultiScaleChunkedArray3D
-from cellier.util.geometry import (
+from cellier.utils.chunk import ChunkedArray3D, MultiScaleChunkedArray3D
+from cellier.utils.geometry import (
     frustum_planes_from_corners,
 )
 

examples/tiled_rendering/tiled_render_zarr.py

@@ -13,12 +13,12 @@
 from wgpu.gui.qt import WgpuCanvas
 
 from cellier.models.data_stores.image import MultiScaleImageZarrStore
-from cellier.util.chunk import (
+from cellier.utils.chunk import (
     ChunkedArray3D,
     ImageDataStoreChunk,
     MultiScaleChunkedArray3D,
 )
-from cellier.util.geometry import (
+from cellier.utils.geometry import (
     frustum_planes_from_corners,
 )
 

src/cellier/models/data_stores/image.py

@@ -24,7 +24,7 @@
     TilingMethod,
     VisualId,
 )
-from cellier.util.geometry import generate_2d_grid
+from cellier.utils.geometry import generate_2d_grid
 
 
 class BaseImageDataStore(BaseDataStore):

src/cellier/models/viewer.py

@@ -7,6 +7,8 @@
 
 from cellier.models.data_manager import DataManager
 from cellier.models.scene import DimsManager, DimsState, Scene
+from cellier.models.visuals.base import BaseVisual
+from cellier.types import VisualId
 
 
 class SceneManager(EventedModel):
@@ -52,6 +54,31 @@ class ViewerModel(EventedModel):
     data: DataManager
     scenes: SceneManager
 
+    def get_visual_by_id(self, visual_id: VisualId) -> tuple[Scene, BaseVisual]:
+        """Get a visual model by its id.
+
+        This iterates through all scenes. If you know the scene the visual
+        is in, it is likely more efficient to get the scene first and then
+        get the visual from the scene.
+
+        Parameters
+        ----------
+        visual_id : VisualId
+            The ID of the visual model to retrieve.
+
+        Returns
+        -------
+        Scene
+            The scene containing the visual model.
+        BaseVisual
+            The visual model with the specified ID.
+        """
+        for scene in self.scenes.scenes.values():
+            visual_model = scene.get_visual_by_id(visual_id)
+            if visual_model is not None:
+                return scene, visual_model
+        raise ValueError(f"Visual model with id {visual_id} not found.")
+
     def to_json_file(self, file_path: str, indent: int = 2) -> None:
         """Save the viewer state as a JSON file."""
         with open(file_path, "w") as f:

src/cellier/models/visuals/base.py

@@ -3,9 +3,17 @@
 from typing import Any
 from uuid import uuid4
 
+import numpy as np
 from psygnal import EmissionInfo, EventedModel
 from pydantic import Field
 
+from cellier.transform import AffineTransform
+
+
+def default_transform() -> AffineTransform:
+    """Create a default identity transform."""
+    return AffineTransform(matrix=np.eye(4))
+
 
 class BaseAppearance(EventedModel):
     """Base model for all materials.
@@ -31,13 +39,17 @@ class BaseVisual(EventedModel):
         The appearance of the visual.
         This should be overridden with the visual-specific
         implementation in the subclasses.
+    transform : AffineTransform
+        The transform of the visual from data space to world space.
+        Default is the identity.
     pick_write : bool
         If True, the visual can be picked.
         Default value is True.
     """
 
     name: str
     appearance: BaseAppearance
+    transform: AffineTransform = Field(default_factory=default_transform)
     pick_write: bool = True
 
     # store a UUID to identify this specific scene.

src/cellier/render/_render_manager.py

@@ -23,6 +23,7 @@
 )
 from cellier.render.cameras import construct_pygfx_camera_from_model
 from cellier.render.utils import construct_pygfx_object
+from cellier.transform import BaseTransform
 from cellier.types import (
     CameraControlsUpdateEvent,
     CameraId,
@@ -575,14 +576,42 @@ def _on_visual_model_update(self, new_state: dict):
 
         visual = self._visuals.get(visual_id)
         if visual is None:
-            logger.warning(f"Visual with id {visual_id} not found.")
+            logger.warning(
+                f"RenderManager attempted to update visual with id {visual_id},"
+                " but it wasn't found."
+            )
             return
 
         # update the visual with the new state
         if "appearance" in new_state:
             visual.update_appearance(new_state["appearance"])
 
+        if "transform" in new_state:
+            transform = new_state["transform"]
+            visual.set_transform(transform)
+
         # emit a redraw request for the canvas
         scene_id = new_state.get("scene_id")
         if scene_id:
             self.events.redraw_canvas.emit(CanvasRedrawRequest(scene_id=scene_id))
+
+    def _on_visual_transform_update(
+        self,
+        scene_id: SceneId,
+        visual_id: VisualId,
+        transform: BaseTransform,
+        redraw_canvas: bool = True,
+    ):
+        """Update the visual transform based on a change to the model."""
+        visual = self._visuals.get(visual_id)
+        if visual is None:
+            logger.warning(
+                f"RenderManager attempted to update visual with id {visual_id},"
+                "but it wasn't found."
+            )
+            return
+
+        visual.set_transform(transform)
+
+        if redraw_canvas:
+            self.events.redraw_canvas.emit(CanvasRedrawRequest(scene_id=scene_id))

src/cellier/render/image.py

@@ -7,6 +7,7 @@
 import wgpu
 
 from cellier.models.visuals import ImageAppearance, MultiscaleImageVisual
+from cellier.transform import BaseTransform
 from cellier.types import ImageDataResponse
 
 
@@ -193,3 +194,14 @@ def update_appearance(self, new_state: dict):
         if "visible" in new_state:
             # update the visibility
             self.node.visible = new_state["visible"]
+
+    def set_transform(self, transform: BaseTransform):
+        """Set the local transform of the node.
+
+        Parameters
+        ----------
+        transform : np.ndarray
+            The 4x4 affine transform matrix to set.
+        """
+        # set the local transform
+        self.node.local.matrix = transform.matrix

src/cellier/render/labels.py

@@ -8,6 +8,7 @@
 
 from cellier.models.visuals import LabelsAppearance, MultiscaleLabelsVisual
 from cellier.render.shaders import LabelImageMaterial, LabelIsoMaterial
+from cellier.transform import BaseTransform
 from cellier.types import ImageDataResponse
 
 
@@ -194,3 +195,14 @@ def update_appearance(self, new_state: dict):
         if "visible" in new_state:
             # update the visibility
             self.node.visible = new_state["visible"]
+
+    def set_transform(self, transform: BaseTransform):
+        """Set the local transform of the node.
+
+        Parameters
+        ----------
+        transform : np.ndarray
+            The 4x4 affine transform matrix to set.
+        """
+        # set the local transform
+        self.node.local.matrix = transform.matrix

src/cellier/render/lines.py

@@ -11,6 +11,7 @@
     LinesVisual,
 )
 from cellier.render.constants import cellier_to_gfx_coordinate_space
+from cellier.transform import BaseTransform
 from cellier.types import LinesDataResponse
 
 
@@ -155,3 +156,14 @@ def update_appearance(self, new_state: dict):
         if "visible" in new_state:
             # update the visibility
             self.node.visible = new_state["visible"]
+
+    def set_transform(self, transform: BaseTransform):
+        """Set the local transform of the node.
+
+        Parameters
+        ----------
+        transform : np.ndarray
+            The 4x4 affine transform matrix to set.
+        """
+        # set the local transform
+        self.node.local.matrix = transform.matrix