Merge pull request #212 from bioimage-io/fix-tif

Change how the tiff sample data is written

Author: constantinpape

bioimageio/core/build_spec/build_model.py

@@ -338,8 +338,10 @@ def _get_deepimagej_macro(name, kwargs, export_folder):
     return {"spec": "ij.IJ::runMacroFile", "kwargs": macro}
 
 
-def _get_deepimagej_config(export_folder, sample_inputs, sample_outputs, pixel_sizes, preprocessing, postprocessing):
-    assert len(sample_inputs) == len(sample_outputs) == 1, "deepimagej config only valid for single input/output"
+def _get_deepimagej_config(
+    export_folder, test_inputs, test_outputs, input_axes, output_axes, pixel_sizes, preprocessing, postprocessing
+):
+    assert len(test_inputs) == len(test_outputs) == 1, "deepimagej config only valid for single input/output"
 
     if any(preproc is not None for preproc in preprocessing):
         assert len(preprocessing) == 1
@@ -363,13 +365,21 @@ def _get_deepimagej_config(export_folder, sample_inputs, sample_outputs, pixel_s
     else:
         postprocess_ij = [{"spec": None}]
 
-    def get_size(path):
-        assert tifffile is not None, "need tifffile for writing deepimagej config"
-        with tifffile.TiffFile(export_folder / path) as f:
-            shape = f.asarray().shape
-        # add singleton z axis if we have 2d data
+    def get_size(fname, axes):
+        shape = np.load(export_folder / fname).shape
+        assert len(shape) == len(axes)
+        shape = [sh for sh, ax in zip(shape, axes) if ax != "b"]
+        axes = [ax for ax in axes if ax != "b"]
+        # the shape for deepij is always given as xyzc
+        if len(shape) == 3:
+            axes_ij = "xyc"
+        else:
+            axes_ij = "xyzc"
+        assert set(axes) == set(axes_ij)
+        axis_permutation = [axes_ij.index(ax) for ax in axes]
+        shape = [shape[permut] for permut in axis_permutation]
         if len(shape) == 3:
-            shape = shape[:2] + (1,) + shape[-1:]
+            shape = shape[:2] + [1] + shape[-1:]
         assert len(shape) == 4
         return " x ".join(map(str, shape))
 
@@ -378,10 +388,13 @@ def get_size(path):
 
     test_info = {
         "inputs": [
-            {"name": in_path, "size": get_size(in_path), "pixel_size": pix_size}
-            for in_path, pix_size in zip(sample_inputs, pixel_sizes_)
+            {"name": in_path, "size": get_size(in_path, axes), "pixel_size": pix_size}
+            for in_path, axes, pix_size in zip(test_inputs, input_axes, pixel_sizes_)
+        ],
+        "outputs": [
+            {"name": out_path, "type": "image", "size": get_size(out_path, axes)}
+            for out_path, axes in zip(test_outputs, output_axes)
         ],
-        "outputs": [{"name": out_path, "type": "image", "size": get_size(out_path)} for out_path in sample_outputs],
         "memory_peak": None,
         "runtime": None,
     }
@@ -397,36 +410,49 @@ def get_size(path):
     return {"deepimagej": config}, [Path(a) for a in attachments]
 
 
-def _write_sample_data(input_paths, output_paths, input_axes, output_axes, export_folder: Path):
-    def write_im(path, im, axes):
+def _write_sample_data(input_paths, output_paths, input_axes, output_axes, pixel_sizes, export_folder: Path):
+    def write_im(path, im, axes, pixel_size=None):
         assert tifffile is not None, "need tifffile for writing deepimagej config"
-        assert len(axes) == im.ndim
-        assert im.ndim in (3, 4)
+        assert len(axes) == im.ndim, f"{len(axes), {im.ndim}}"
+        assert im.ndim in (4, 5), f"{im.ndim}"
 
-        # deepimagej expects xyzc axis order
-        if im.ndim == 3:
-            assert set(axes) == {"x", "y", "c"}
-            axes_ij = "xyc"
+        # convert the image to expects (Z)CYX axis order
+        if im.ndim == 4:
+            assert set(axes) == {"b", "x", "y", "c"}, f"{axes}"
+            axes_ij = "cyxb"
         else:
-            assert set(axes) == {"x", "y", "z", "c"}
-            axes_ij = "xyzc"
+            assert set(axes) == {"b", "x", "y", "z", "c"}, f"{axes}"
+            axes_ij = "zcyxb"
 
-        axis_permutation = tuple(axes_ij.index(ax) for ax in axes)
+        axis_permutation = tuple(axes.index(ax) for ax in axes_ij)
         im = im.transpose(axis_permutation)
-
-        with tifffile.TiffWriter(path) as f:
-            f.write(im)
+        # expand to TZCYXS
+        if len(axes_ij) == 4:  # add singleton t and z axis
+            im = im[None, None]
+        else:  # add singeton z axis
+            im = im[None]
+
+        if pixel_size is None:
+            resolution = None
+        else:
+            spatial_axes = list(set(axes_ij) - set("bc"))
+            resolution = tuple(1.0 / pixel_size[ax] for ax in axes_ij if ax in spatial_axes)
+        # does not work for double
+        if np.dtype(im.dtype) == np.dtype("float64"):
+            im = im.astype("float32")
+        tifffile.imsave(path, im, imagej=True, resolution=resolution)
 
     sample_in_paths = []
     for i, (in_path, axes) in enumerate(zip(input_paths, input_axes)):
-        inp = np.load(export_folder / in_path)[0]
+        inp = np.load(export_folder / in_path)
         sample_in_path = export_folder / f"sample_input_{i}.tif"
-        write_im(sample_in_path, inp, axes)
+        pixel_size = None if pixel_sizes is None else pixel_sizes[i]
+        write_im(sample_in_path, inp, axes, pixel_size)
         sample_in_paths.append(sample_in_path)
 
     sample_out_paths = []
     for i, (out_path, axes) in enumerate(zip(output_paths, output_axes)):
-        outp = np.load(export_folder / out_path)[0]
+        outp = np.load(export_folder / out_path)
         sample_out_path = export_folder / f"sample_output_{i}.tif"
         write_im(sample_out_path, outp, axes)
         sample_out_paths.append(sample_out_path)
@@ -797,17 +823,15 @@ def build_model(
     # add the deepimagej config if specified
     if add_deepimagej_config:
         if sample_inputs is None:
-            input_axes_ij = [inp.axes[1:] for inp in inputs]
-            output_axes_ij = [out.axes[1:] for out in outputs]
             sample_inputs, sample_outputs = _write_sample_data(
-                test_inputs, test_outputs, input_axes_ij, output_axes_ij, root
+                test_inputs, test_outputs, input_axes, output_axes, pixel_sizes, root
             )
         # deepimagej expect tifs as sample data
         assert all(os.path.splitext(path)[1] in (".tif", ".tiff") for path in sample_inputs)
         assert all(os.path.splitext(path)[1] in (".tif", ".tiff") for path in sample_outputs)
 
         ij_config, ij_attachments = _get_deepimagej_config(
-            root, sample_inputs, sample_outputs, pixel_sizes, preprocessing, postprocessing
+            root, test_inputs, test_outputs, input_axes, output_axes, pixel_sizes, preprocessing, postprocessing
         )
 
         if config is None: