Update radius extraction for IMOD

Author: constantinpape

scripts/cooper/export_vesicles_to_imod.py

@@ -5,7 +5,9 @@
 
 
 def export_vesicles_to_imod(args):
-    export_function = partial(write_segmentation_to_imod_as_points, min_radius=args.min_radius, radius_factor=args.increase_radius)
+    export_function = partial(
+        write_segmentation_to_imod_as_points, min_radius=args.min_radius, radius_factor=args.increase_radius
+    )
     export_helper(args.input_path, args.segmentation_path, args.output_path, export_function, force=args.force)
 
 
@@ -32,7 +34,7 @@ def main():
         help="Whether to over-write already present export results."
     )
     parser.add_argument(
-        "--increase_radius", type=float, default=1.5,
+        "--increase_radius", type=float, default=1.3,
         help="The factor to increase the radius of the exported vesicles.",
     )
     args = parser.parse_args()

synaptic_reconstruction/imod/to_imod.py

@@ -22,6 +22,7 @@ def convert_segmentation_to_spheres(
     num_workers: Optional[int] = None,
     resolution: Optional[Tuple[float, float, float]] = None,
     radius_factor: float = 1.0,
+    estimate_radius_2d: bool = True,
 ):
     """Extract spheres parameterized by center and radius from a segmentation.
 
@@ -31,6 +32,9 @@ def convert_segmentation_to_spheres(
         num_workers: Number of workers to use for parallelization.
         resolution: The physical resolution of the data.
         radius_factor: Factor for increasing the radius to account for too small exported spheres.
+        estimate_radius_2d: If true the distance to boundary for determining the centroid and computing
+            the radius will be computed only in 2d rather than in 3d. This can lead to better results
+            in case of deformation across the depth axis.
 
     Returns:
         np.array: the center coordinates
@@ -45,10 +49,14 @@ def coords_and_rads(prop):
         bbox = prop.bbox
         bb = np.s_[bbox[0]:bbox[3], bbox[1]:bbox[4], bbox[2]:bbox[5]]
         mask = segmentation[bb] == seg_id
-        seg_radii = distance_transform_edt(mask, sampling=resolution)
 
-        max_coord = np.unravel_index(np.argmax(seg_radii), mask.shape)
-        radius = seg_radii[max_coord] * radius_factor
+        if estimate_radius_2d:
+            dists = np.array([distance_transform_edt(ma, sampling=resolution[1:]) for ma in mask])
+        else:
+            dists = distance_transform_edt(mask, sampling=resolution)
+
+        max_coord = np.unravel_index(np.argmax(dists), mask.shape)
+        radius = dists[max_coord] * radius_factor
 
         offset = np.array(bbox[:3])
         coord = np.array(max_coord) + offset
@@ -120,6 +128,7 @@ def write_segmentation_to_imod_as_points(
     output_path: str,
     min_radius: Union[int, float],
     radius_factor: float = 1.0,
+    estimate_radius_2d: bool = True,
 ):
     """Write segmentation results to .mod file with imod point annotations.
 
@@ -131,6 +140,9 @@ def write_segmentation_to_imod_as_points(
         output_path: Where to save the .mod file.
         min_radius: Minimum radius for export.
         radius_factor: Factor for increasing the radius to account for too small exported spheres.
+        estimate_radius_2d: If true the distance to boundary for determining the centroid and computing
+            the radius will be computed only in 2d rather than in 3d. This can lead to better results
+            in case of deformation across the depth axis.
     """
 
     # Read the resolution information from the mrcfile.
@@ -142,7 +154,9 @@ def write_segmentation_to_imod_as_points(
 
     # Extract the center coordinates and radii from the segmentation.
     segmentation = imageio.imread(segmentation_path)
-    coordinates, radii = convert_segmentation_to_spheres(segmentation, resolution=resolution, radius_factor=radius_factor)
+    coordinates, radii = convert_segmentation_to_spheres(
+        segmentation, resolution=resolution, radius_factor=radius_factor, estimate_radius_2d=estimate_radius_2d
+    )
 
     # Write the point annotations to imod.
     write_points_to_imod(coordinates, radii, segmentation.shape, min_radius, output_path)