feat(Python): Support Image IO in pipelines

src/python/itkwasm/itkwasm/image.py

@@ -45,5 +45,4 @@ def __post_init__(self):
             self.direction = np.eye(dimension).astype(np.float32).ravel()
 
         if len(self.size) == 0:
-            self.size += [1,] * dimension
-
+            self.size += [1,] * dimension

src/python/itkwasm/itkwasm/pipeline.py

@@ -1,17 +1,48 @@
 import json
 from pathlib import Path
+from dataclasses import asdict
 from typing import List, Union, Dict, Tuple
+
+import numpy as np
+
 from .interface_types import InterfaceTypes
 from .pipeline_input import PipelineInput
 from .pipeline_output import PipelineOutput
 from .text_stream import TextStream
 from .binary_stream import BinaryStream
 from .text_file import TextFile
 from .binary_file import BinaryFile
+from .image import Image, ImageType
+from .int_types import IntTypes
+from .float_types import FloatTypes
 
 from wasmer import engine, wasi, Store, Module, ImportObject, Instance
 from wasmer_compiler_cranelift import Compiler
 
+def _memoryview_to_numpy_array(component_type, buf):
+    if component_type == IntTypes.UInt8:
+        return np.frombuffer(buf, dtype=np.uint8)
+    elif component_type == IntTypes.Int8:
+        return np.frombuffer(buf, dtype=np.int8)
+    elif component_type == IntTypes.UInt16:
+        return np.frombuffer(buf, dtype=np.uint16)
+    elif component_type == IntTypes.Int16:
+        return np.frombuffer(buf, dtype=np.int16)
+    elif component_type == IntTypes.UInt32:
+        return np.frombuffer(buf, dtype=np.uint32)
+    elif component_type == IntTypes.Int32:
+        return np.frombuffer(buf, dtype=np.int32)
+    elif component_type == IntTypes.UInt64:
+        return np.frombuffer(buf, dtype=np.uint64)
+    elif component_type == IntTypes.Int64:
+        return np.frombuffer(buf, dtype=np.int64)
+    elif component_type == FloatTypes.Float32:
+        return np.frombuffer(buf, dtype=np.float32)
+    elif component_type == FloatTypes.Float64:
+        return np.frombuffer(buf, dtype=np.float64)
+    else:
+        raise ValueError('Unsupported component type')
+
 
 class Pipeline:
     """Run an itk-wasm WASI pipeline."""
@@ -78,6 +109,23 @@ def run(self, args: List[str], outputs: List[PipelineOutput]=[], inputs: List[Pi
                 pass
             elif input_.type == InterfaceTypes.BinaryFile:
                 pass
+            elif input_.type == InterfaceTypes.Image:
+                image = input_.data
+                mv = bytes(image.data.data)
+                # self.memory.grow(15)
+                data_ptr = self._set_input_array(mv, index, 0)
+                dv = bytes(image.direction.data)
+                direction_ptr = self._set_input_array(dv, index, 1)
+                image_json = {
+                    "imageType": asdict(image.imageType),
+                    "name": image.name,
+                    "origin": image.origin,
+                    "spacing": image.spacing,
+                    "direction": f"data:application/vnd.itk.address,0:{direction_ptr}",
+                    "size": image.size,
+                    "data": f"data:application/vnd.itk.address,0:{data_ptr}"
+                }
+                self._set_input_json(image_json, index)
             else:
                 raise ValueError(f'Unexpected/not yet supported input.type {input_.type}')
 
@@ -102,6 +150,25 @@ def run(self, args: List[str], outputs: List[PipelineOutput]=[], inputs: List[Pi
                     output_data = PipelineOutput(InterfaceTypes.TextFile, TextFile(output.data.path))
                 elif output.type == InterfaceTypes.BinaryFile:
                     output_data = PipelineOutput(InterfaceTypes.BinaryFile, BinaryFile(output.data.path))
+                elif output.type == InterfaceTypes.Image:
+                    image_json = self._get_output_json(index)
+
+                    image = Image(**image_json)
+                    image.name = 'aoeu'
+
+                    data_ptr = self.output_array_address(0, index, 0)
+                    data_size = self.output_array_size(0, index, 0)
+                    data_array = _memoryview_to_numpy_array(image.imageType.componentType, memoryview(self.memory.buffer)[data_ptr:data_ptr+data_size])
+                    image.data = data_array
+
+                    direction_ptr = self.output_array_address(0, index, 1)
+                    direction_size = self.output_array_size(0, index, 1)
+                    direction_array = _memoryview_to_numpy_array(FloatTypes.Float64, memoryview(self.memory.buffer)[direction_ptr:direction_ptr+direction_size])
+                    dimension = image.imageType.dimension
+                    direction_array.shape = (dimension, dimension)
+                    image.direction = direction_array
+
+                    output_data = PipelineOutput(InterfaceTypes.Image, image)
                 populated_outputs.append(output_data)
 
         delayed_exit = instance.exports.itk_wasm_delayed_exit
@@ -110,7 +177,7 @@ def run(self, args: List[str], outputs: List[PipelineOutput]=[], inputs: List[Pi
         # Should we be returning the return_code?
         return tuple(populated_outputs)
 
-    def _set_input_array(self, data_array: bytes, input_index: int, sub_index: int) -> int:
+    def _set_input_array(self, data_array: Union[bytes, bytearray], input_index: int, sub_index: int) -> int:
         data_ptr = 0
         if data_array != None:
             data_ptr = self.input_array_alloc(0, input_index, sub_index, len(data_array))
@@ -122,4 +189,11 @@ def _set_input_json(self, data_object: Dict, input_index: int) -> None:
         data_json = json.dumps(data_object).encode()
         json_ptr = self.input_json_alloc(0, input_index, len(data_json))
         buf = memoryview(self.memory.buffer)
-        buf[json_ptr:json_ptr+len(data_json)] = data_json
+        buf[json_ptr:json_ptr+len(data_json)] = data_json
+
+    def _get_output_json(self, output_index: int) -> Dict:
+        json_ptr = self.output_json_address(0, output_index)
+        json_len = self.output_json_size(0, output_index)
+        json_str = bytes(memoryview(self.memory.buffer)[json_ptr:json_ptr+json_len]).decode()
+        json_result = json.loads(json_str)
+        return json_result

src/python/itkwasm/test/test_pipeline.py

@@ -1,9 +1,14 @@
 from pathlib import Path, PurePosixPath
+import tempfile
+from dataclasses import asdict
+
+import itk
+import numpy as np
 
-from itkwasm import InterfaceTypes, TextStream, BinaryStream, PipelineInput, PipelineOutput, Pipeline, TextFile, BinaryFile
+from itkwasm import InterfaceTypes, TextStream, BinaryStream, PipelineInput, PipelineOutput, Pipeline, TextFile, BinaryFile, Image
 
 test_input_dir = Path(__file__).resolve().parent / 'input'
-import tempfile
+test_baseline_dir = Path(__file__).resolve().parent / 'baseline'
 
 
 def test_stdout_stderr():
@@ -87,4 +92,36 @@ def test_pipeline_input_output_files():
             assert content[0] == 222
             assert content[1] == 173
             assert content[2] == 190
-            assert content[3] == 239
+            assert content[3] == 239
+
+def test_pipeline_write_read_image():
+    pipeline = Pipeline(test_input_dir / 'median-filter-test.wasi.wasm')
+
+    data = test_input_dir / "cthead1.png"
+    itk_image = itk.imread(data, itk.UC)
+    itk_image_dict = itk.dict_from_image(itk_image)
+    itkwasm_image = Image(**itk_image_dict)
+
+    pipeline_inputs = [
+        PipelineInput(InterfaceTypes.Image, itkwasm_image),
+    ]
+
+    pipeline_outputs = [
+        PipelineOutput(InterfaceTypes.Image),
+    ]
+
+    args = [
+        '0',
+        '0',
+        '--radius', '2', '--memory-io',]
+
+    outputs = pipeline.run(args, pipeline_outputs, pipeline_inputs)
+
+    out_image = itk.image_from_dict(asdict(outputs[0].data))
+    # To be addresses in itk-5.3.1
+    out_image.SetRegions([256,256])
+
+    baseline = itk.imread(test_baseline_dir / "test_pipeline_write_read_image.png")
+
+    difference = np.sum(itk.comparison_image_filter(out_image, baseline))
+    assert difference == 0.0

test/node/pipeline/runPipelineNodeTest.js

@@ -112,10 +112,10 @@ test('runPipelineNode writes and reads an itk.Image via memory io', (t) => {
   return readImageLocalFile(testInputFilePath)
     .then(function (image) {
       const pipelinePath = path.resolve('test', 'pipelines', 'median-filter-pipeline', 'emscripten-build', 'median-filter-test')
-      const args = ['--memory-io',
+      const args = [
         '0',
         '0',
-        '--radius', '4']
+        '--radius', '4', '--memory-io']
       const desiredOutputs = [
         { type: InterfaceTypes.Image }
       ]

test/pipelines/median-filter-pipeline/median-filter-test.cxx

@@ -77,12 +77,12 @@ MedianFilter(itk::wasm::Pipeline & pipeline, const TImage * inputImage)
     splitter->GetSplit( split, numberOfSplits, requestedRegion );
     roiFilter->SetInput( smoother->GetOutput() );
     roiFilter->SetExtractionRegion( requestedRegion );
-    roiFilter->Update();
+    roiFilter->UpdateLargestPossibleRegion();
     outputImage.Set( roiFilter->GetOutput() );
   }
   else
   {
-    smoother->Update();
+    smoother->UpdateLargestPossibleRegion();
     outputImage.Set( smoother->GetOutput() );
   }
 
@@ -103,7 +103,8 @@ class PipelineFunctor
 
     ITK_WASM_PRE_PARSE(pipeline);
 
-    return MedianFilter<ImageType>(pipeline, inputImage.Get());
+    typename ImageType::ConstPointer image = inputImage.Get();
+    return MedianFilter<ImageType>(pipeline, image);
   }
 };