2d imod export

synapse_net/imod/to_imod.py

@@ -176,7 +176,7 @@ def write_points_to_imod(
     """Write point annotations to a .mod file for IMOD.
 
     Args:
-        coordinates: Array with the point coordinates.
+        coordinates: Array with the point coordinates. #2D or 3D
         radii: Array with the point radii.
         shape: Shape of the volume to be exported.
         min_radius: Minimum radius for export.
@@ -193,19 +193,32 @@ def _pad(inp, n=3):
         pw = plen * " "
         return f"{pw}{inp}.00"
 
+    is_3d = coordinates.shape[1] == 3
+    if not is_3d and coordinates.shape[1] != 2:
+        raise ValueError("Coordinates must have shape (N, 2) for 2D or (N, 3) for 3D.")
+
     with tempfile.NamedTemporaryFile() as tmp_file:
         fname = tmp_file.name
         with open(fname, "w") as f:
             for coord, radius in zip(coordinates, radii):
                 if radius < min_radius:
                     continue
-                # IMOD needs peculiar whitespace padding
-                x = _pad(coord[2])
-                y = _pad(shape[1] - coord[1])
-                z = _pad(coord[0])
-                f.write(f"{x}{y}{z}{_pad(radius, 2)}\n")
 
+                if is_3d:
+                    # (z, y, x) indexing
+                    x = _pad(coord[2])
+                    y = _pad(shape[1] - coord[1])
+                    z = _pad(coord[0])
+
+                else:
+                    # (y, x) indexing, single z-plane
+                    x = _pad(coord[1])
+                    y = _pad(shape[0] - coord[0])
+                    z=_pad(1)
+
+                f.write(f"{x}{y}{z}{_pad(radius, 2)}\n")
         cmd = [cmd, "-si", "-scat", fname, output_path]
+
         if color is not None:
             assert len(color) == 3
             r, g, b = [str(co) for co in color]
@@ -214,6 +227,7 @@ def _pad(inp, n=3):
         run(cmd)
 
 
+
 def write_segmentation_to_imod_as_points(
     mrc_path: str,
     segmentation: Union[str, np.ndarray],