ENH: c_index argument to ngff_image_to_itk_image

Only extract the given channel index.

Author: thewtex

ngff_zarr/ngff_image_to_itk_image.py

@@ -38,6 +38,7 @@ def ngff_image_to_itk_image(
     ngff_image: NgffImage,
     wasm: bool = True,
     t_index: Optional[int] = None,
+    c_index: Optional[int] = None,
 ):
     """Convert a NgffImage to an ITK image."""
     from itkwasm import IntTypes, PixelTypes
@@ -65,6 +66,29 @@ def ngff_image_to_itk_image(
             axes_units=new_axes_units,
         )
 
+    if c_index is not None and "c" in ngff_image.dims:
+        c_dim_index = ngff_image.dims.index("c")
+        new_dims = list(ngff_image.dims)
+        new_dims.remove("c")
+        new_dims = tuple(new_dims)
+        new_scale = {dim: ngff_image.scale[dim] for dim in new_dims}
+        new_translation = {dim: ngff_image.translation[dim] for dim in new_dims}
+        new_axes_units = {dim: ngff_image.axes_units[dim] for dim in new_dims}
+        if isinstance(ngff_image.data, DaskArray):
+            from dask.array import take
+
+            new_data = take(ngff_image.data, c_index, axis=c_dim_index)
+        else:
+            new_data = ngff_image.data.take(c_index, axis=c_dim_index)
+        ngff_image = NgffImage(
+            data=new_data,
+            dims=new_dims,
+            name=ngff_image.name,
+            scale=new_scale,
+            translation=new_translation,
+            axes_units=new_axes_units,
+        )
+
     ngff_image = _channel_dim_last(ngff_image)
 
     dims = ngff_image.dims

pixi.lock

@@ -89,3 +89,40 @@ def test_t_index(input_images):  # noqa: ARG001
     assert len(itk_image.origin) == 3
     assert len(itk_image.direction) == 3
     assert itk_image.data.shape == (12, 223, 198, 6)
+
+
+def test_c_index(input_images):  # noqa: ARG001
+    dataset_name = "13457537"
+    store_path = test_data_dir / "input" / f"{dataset_name}.zarr"
+    multiscales = from_ngff_zarr(store_path)
+    ngff_image = multiscales.images[0]
+
+    itk_image = ngff_image_to_itk_image(ngff_image)
+
+    assert itk_image.imageType.dimension == 4
+    assert itk_image.imageType.components == 6
+    assert len(itk_image.size) == 4
+    assert len(itk_image.spacing) == 4
+    assert len(itk_image.origin) == 4
+    assert len(itk_image.direction) == 4
+    assert itk_image.data.shape == (18, 12, 223, 198, 6)
+
+    itk_image = ngff_image_to_itk_image(ngff_image, c_index=0)
+
+    assert itk_image.imageType.dimension == 4
+    assert itk_image.imageType.components == 1
+    assert len(itk_image.size) == 4
+    assert len(itk_image.spacing) == 4
+    assert len(itk_image.origin) == 4
+    assert len(itk_image.direction) == 4
+    assert itk_image.data.shape == (18, 12, 223, 198)
+
+    itk_image = ngff_image_to_itk_image(ngff_image, t_index=0, c_index=0)
+
+    assert itk_image.imageType.dimension == 3
+    assert itk_image.imageType.components == 1
+    assert len(itk_image.size) == 3
+    assert len(itk_image.spacing) == 3
+    assert len(itk_image.origin) == 3
+    assert len(itk_image.direction) == 3
+    assert itk_image.data.shape == (12, 223, 198)