ENH: Add support for label maps

Author: bnmajor

itkwidgets/integrations/__init__.py

@@ -11,25 +11,34 @@
 
 _image_count = 1
 
-async def _set_viewer_image(itk_viewer, image, name=None):
+async def _set_viewer_image(itk_viewer, image, name=None, is_label=False):
     global _image_count
     if isinstance(image, itkwasm.Image):
         if not name:
             name = image.name
             if not name:
                 name = f"image {_image_count}"
                 _image_count += 1
-        await itk_viewer.setImage(image, name)
+        if is_label:
+            await itk_viewer.setLabelImage(image)
+        else:
+            await itk_viewer.setImage(image, name)
     elif isinstance(image, np.ndarray):
         if not name:
             name = f"image {_image_count}"
             _image_count += 1
-        await itk_viewer.setImage(image, name)
+        if is_label:
+            await itk_viewer.setLabelImage(image)
+        else:
+            await itk_viewer.setImage(image, name)
     elif isinstance(image, zarr.Group):
         if not name:
             name = f"image {_image_count}"
             _image_count += 1
-        await itk_viewer.setImage(image, name)
+        if is_label:
+            await itk_viewer.setLabelImage(image)
+        else:
+            await itk_viewer.setImage(image, name)
     elif HAVE_ITK:
         import itk
         if isinstance(image, itk.Image):
@@ -38,15 +47,21 @@ async def _set_viewer_image(itk_viewer, image, name=None):
             if not name:
                 name = f"image {_image_count}"
                 _image_count += 1
-            await itk_viewer.setImage(wasm_image, name)
+            if is_label:
+                await itk_viewer.setLabelImage(wasm_image)
+            else:
+                await itk_viewer.setImage(wasm_image, name)
     if HAVE_VTK:
         import vtk
         if isinstance(image, vtk.vtkImageData):
             ndarray = vtk_image_to_ndarray(image)
             if not name:
                 name = f"image {_image_count}"
                 _image_count += 1
-            await itk_viewer.setImage(ndarray, name)
+            if is_label:
+                await itk_viewer.setLabelImage(ndarray)
+            else:
+                await itk_viewer.setImage(ndarray, name)
     if HAVE_DASK:
         import dask
         if isinstance(image, dask.array.core.Array):
@@ -55,14 +70,20 @@ async def _set_viewer_image(itk_viewer, image, name=None):
             if not name:
                 name = f"image {_image_count}"
                 _image_count += 1
-            await itk_viewer.setImage(ndarray, name)
+            if is_label:
+                await itk_viewer.setLabelImage(ndarray)
+            else:
+                await itk_viewer.setImage(ndarray, name)
     if HAVE_TORCH:
         import torch
         if isinstance(image, torch.Tensor):
             if not name:
                 name = f"image {_image_count}"
                 _image_count += 1
-            await itk_viewer.setImage(image.numpy(), name)
+            if is_label:
+                await itk_viewer.setLabelImage(image.numpy())
+            else:
+                await itk_viewer.setImage(image.numpy(), name)
     if HAVE_XARRAY:
         import xarray
         if isinstance(image, xarray.DataArray):
@@ -71,13 +92,19 @@ async def _set_viewer_image(itk_viewer, image, name=None):
             if not name:
                 name = f"image {_image_count}"
                 _image_count += 1
-            await itk_viewer.setImage(ndarray, name)
+            if is_label:
+                await itk_viewer.setLabelImage(ndarray)
+            else:
+                await itk_viewer.setImage(ndarray, name)
         if isinstance(image, xarray.Dataset):
             ndarray = xarray_data_set_to_numpy(image)
             if not name:
                 name = f"image {_image_count}"
                 _image_count += 1
-            await itk_viewer.setImage(ndarray, name)
+            if is_label:
+                await itk_viewer.setLabelImage(ndarray)
+            else:
+                await itk_viewer.setImage(ndarray, name)
 
 
 async def _set_viewer_point_sets(itk_viewer, point_sets):

itkwidgets/viewer.py

@@ -22,12 +22,13 @@ def __init__(self, ui_collapsed=True, rotate=False, **add_data_kwargs):
         self._init_viewer_kwargs.update(**add_data_kwargs)
 
     def _get_input_data(self):
-        input_options = ['data', 'image', 'point_sets']
+        input_options = ['data', 'image', 'point_sets', 'label_image']
+        inputs = []
         for option in input_options:
             data = self._init_viewer_kwargs.get(option, None)
             if data is not None:
-                break
-        return data, option
+                inputs.append((option, data))
+        return inputs
 
     async def setup(self):
         pass
@@ -43,13 +44,15 @@ async def run(self, ctx):
         )
         _viewer_count += 1
 
-        data, input_type = self._get_input_data()
-        if data is not None:
-            render_type = _detect_render_type(data, input_type)
-            if render_type is RenderType.IMAGE:
-                await _set_viewer_image(itk_viewer, data)
-            elif render_type is RenderType.POINT_SET:
-                await _set_viewer_point_sets(itk_viewer, data)
+        inputs = self._get_input_data()
+        if inputs:
+            for (input_type, data) in inputs:
+                render_type = _detect_render_type(data, input_type)
+                if render_type is RenderType.IMAGE:
+                    is_label = input_type == 'label_image'
+                    await _set_viewer_image(itk_viewer, data, is_label)
+                elif render_type is RenderType.POINT_SET:
+                    await _set_viewer_point_sets(itk_viewer, data)
 
             self.set_default_ui_values(itk_viewer)
 
@@ -115,6 +118,13 @@ def set_image_shadow_enabled(self, enabled: bool):
     def set_image_volume_sample_distance(self, distance: float):
         self.viewer_rpc.itk_viewer.setImageVolumeSampleDistance(distance)
 
+    async def set_label_image(self, label_image: Image):
+        render_type = _detect_render_type(label_image, 'image')
+        if render_type is RenderType.IMAGE:
+            await _set_viewer_image(self.viewer_rpc.itk_viewer, label_image, is_label=True)
+        elif render_type is RenderType.POINT_SET:
+            await _set_viewer_point_sets(self.viewer_rpc.itk_viewer, label_image)
+
     def set_label_image_blend(self, blend: float):
         self.viewer_rpc.itk_viewer.setLabelImageBlend(blend)