add all the IJ dimensions to the image

Author: carlosuc3m

bioimageio/core/build_spec/build_model.py

@@ -420,24 +420,32 @@ def write_im(path, im, axes, pixel_size=None):
         # convert the image to expects (Z)CYX axis order
         if im.ndim == 4:
             assert set(axes) == {"b", "x", "y", "c"}, f"{axes}"
-            axes_ij = "bcyx"
+            resolution_axes_ij = "cyxb"
         else:
             assert set(axes) == {"b", "x", "y", "z", "c"}, f"{axes}"
-            axes_ij = "bzcyx"
-
-        axis_permutation = tuple(axes.index(ax) for ax in axes_ij)
+            resolution_axes_ij = "bzcyx"
+        
+        def addMissingAxes(im_axes):
+            needed_axes = ["b", "c", "x", "y", "z", "s"]
+            for ax in needed_axes:
+                if not ax in im_axes:
+                im_axes += ax
+            return im_axes
+        
+        axes_ij = "bzcyxs"
+        # Expand the image to ImageJ dimensions
+        im = np.expand_dims(im, axis=tuple(range(len(axes),len(axes_ij))))
+
+
+        axis_permutation = tuple(addMissingAxes(axes).index(ax) for ax in axes_ij)
         im = im.transpose(axis_permutation)
-        # 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)
+            spatial_axes = list(set(resolution_axes_ij) - set("bc"))
+            resolution = tuple(1.0 / pixel_size[ax] for ax in resolution_axes_ij if ax in spatial_axes)
         # does not work for double
         if np.dtype(im.dtype) == np.dtype("float64"):
             im = im.astype("float32")

tests/build_spec/test_build_spec.py

@@ -94,12 +94,6 @@ def _test_build_spec(
         input_names=[inp.name for inp in model_spec.inputs],
         output_names=[out.name for out in model_spec.outputs],
     )
-    print("**********************************************")
-    print("**********************************************")
-    print("**********************************************")
-    print(model_spec.name)
-    print(model_spec.root_path)
-    print(model_spec.rdf_source)
     if architecture is not None:
         kwargs["architecture"] = architecture
     if model_kwargs is not None: