Options to save multiple meshes from segmentations

docs/Data_Preparation.md

@@ -58,9 +58,16 @@ This command will generate mesh representations of the membrane segmentations an
 - `--tomo-path` (TEXT, required): Path to the tomogram to be projected. [default: None] 
 - `--out-folder` (TEXT): Path to the folder where mesh projections should be stored. [default: ./mesh_data] 
 - `--step-numbers` (INTEGER): Step numbers for the normal vectors. [default: (-10, 10)] 
-- `--step-size` (FLOAT): Step size for the normal vectors. [default: 2.5] 
-- `--mesh-smoothing` (INTEGER): Smoothing factor for the mesh. [default: 1000] 
-- `--barycentric-area` (FLOAT): Barycentric area for the mesh. [default: 400.0] 
+- `--step-size` (FLOAT): Step size for the normal vectors. [default: 2.5]
+- `--mesh-smoothing` (INTEGER): Smoothing factor for the mesh. [default: 1000]
+- `--barycentric-area` (FLOAT): Barycentric area for the mesh. [default: 400.0]
+- `--min-segmentation-size` (INTEGER): Export connected components whose voxel count
+  is greater than or equal to this threshold. [default: 10000]
+- `--max-segmentations` (INTEGER): Maximum number of connected segmentations to convert per file. Use `None` to export all components. [default: None]
+
+When both `--min-segmentation-size` and `--max-segmentations` are set, MemBrain Pick
+exports the largest components that exceed the size threshold, capped at the
+requested maximum.
 
 Note: options may change faster than these docs. You can check the available options by running `membrain_pick convert_mb_folder`.
 
@@ -76,7 +83,10 @@ This command will generate a mesh representation of the membrane segmentation an
 
 #### More Options
 
-Other options are similar to the `convert_mb_folder` command. You can check the available options by running `membrain-pick convert_file`.
+Other options mirror the `convert_mb_folder` command, including `--min-segmentation-size`
+to discard small connected components and `--max-segmentations` to control how many
+meshes are exported per membrane file. You can check the available options by running
+`membrain-pick convert_file`.
 
 
 ### Outputs

src/membrain_pick/cli/convert_meshes_cli.py

@@ -1,5 +1,5 @@
 import os
-from typing import List
+from typing import List, Optional
 
 from typer import Option
 from .cli import OPTION_PROMPT_KWARGS as PKWARGS
@@ -37,6 +37,13 @@ def convert_single_file(
     barycentric_area: float = Option(  # noqa: B008
         400.0, help="Barycentric area for the mesh. Default: 1.0"
     ),
+    min_segmentation_size: int = Option(  # noqa: B008
+        10000,
+        help=(
+            "Export all connected components whose voxel count is greater than or "
+            "equal to this threshold."
+        ),
+    ),
     imod_meshing: bool = Option(  # noqa: B008
         False,
         help="Should the mesh be generated using IMOD? WARNING: This is highly experimental.",
@@ -45,6 +52,20 @@ def convert_single_file(
         False,
         help="Should the mesh be generated using PyMeshLab? WARNING: This is highly experimental.",
     ),
+    max_segmentations: Optional[int] = Option(  # noqa: B008
+        None,
+        help=(
+            "Maximum number of connected segmentations to convert to meshes. "
+            "Leave empty or set to None to export all available segmentations."
+        ),
+    ),
+    merge_outputs: bool = Option(  # noqa: B008
+        False,
+        help=(
+            "Store all exported connected components in a single set of output "
+            "files (one H5/OBJ pair)."
+        ),
+    ),
 ):
     """Convert a single membrane segmentation to a mesh.
 
@@ -67,8 +88,11 @@ def convert_single_file(
         mesh_smoothing=mesh_smoothing,
         barycentric_area=barycentric_area,
         input_pixel_size=input_pixel_size,
+        min_connected_size=min_segmentation_size,
         imod_meshing=imod_meshing,
         pymeshlab_meshing=pymeshlab_meshing,
+        max_segmentations=max_segmentations,
+        merge_outputs=merge_outputs,
     )
 
 
@@ -101,6 +125,13 @@ def convert_mb_folder(
     barycentric_area: float = Option(  # noqa: B008
         400.0, help="Barycentric area for the mesh. Default: 1.0"
     ),
+    min_segmentation_size: int = Option(  # noqa: B008
+        10000,
+        help=(
+            "Export all connected components whose voxel count is greater than or "
+            "equal to this threshold."
+        ),
+    ),
     imod_meshing: bool = Option(  # noqa: B008
         False,
         help="Should the mesh be generated using IMOD? WARNING: This is highly experimental.",
@@ -109,6 +140,20 @@ def convert_mb_folder(
         False,
         help="Should the mesh be generated using PyMeshLab? WARNING: This is highly experimental.",
     ),
+    max_segmentations: Optional[int] = Option(  # noqa: B008
+        None,
+        help=(
+            "Maximum number of connected segmentations to convert to meshes. "
+            "Leave empty or set to None to export all available segmentations."
+        ),
+    ),
+    merge_outputs: bool = Option(  # noqa: B008
+        False,
+        help=(
+            "Store all exported connected components in a single set of output "
+            "files (one H5/OBJ pair) per segmentation."
+        ),
+    ),
 ):
     """Convert a folder of membrane segmentations to meshes.
 
@@ -129,6 +174,9 @@ def convert_mb_folder(
         mesh_smoothing=mesh_smoothing,
         barycentric_area=barycentric_area,
         input_pixel_size=input_pixel_size,
+        min_connected_size=min_segmentation_size,
         imod_meshing=imod_meshing,
         pymeshlab_meshing=pymeshlab_meshing,
+        max_segmentations=max_segmentations,
+        merge_outputs=merge_outputs,
     )

src/membrain_pick/cli/mesh_conversion_cli.py

@@ -3,7 +3,7 @@
 import typer
 from click import Context
 from typer.core import TyperGroup
-from typing import List
+from typing import List, Optional
 from typer import Option
 
 
@@ -82,6 +82,20 @@ def convert_single_file(
         False,
         help="Should the mesh be generated using IMOD? WARNING: This is highly experimental.",
     ),
+    max_segmentations: Optional[int] = Option(  # noqa: B008
+        None,
+        help=(
+            "Maximum number of connected segmentations to convert to meshes. "
+            "Leave empty or set to None to export all available segmentations."
+        ),
+    ),
+    merge_outputs: bool = Option(  # noqa: B008
+        False,
+        help=(
+            "Store all exported connected components in a single set of output "
+            "files (one H5/OBJ pair)."
+        ),
+    ),
 ):
     """Convert a single membrane segmentation to a mesh.
 
@@ -103,6 +117,8 @@ def convert_single_file(
         barycentric_area=barycentric_area,
         input_pixel_size=input_pixel_size,
         imod_meshing=imod_meshing,
+        max_segmentations=max_segmentations,
+        merge_outputs=merge_outputs,
     )
 
 
@@ -139,6 +155,20 @@ def convert_mb_folder(
         False,
         help="Should the mesh be generated using IMOD? WARNING: This is highly experimental.",
     ),
+    max_segmentations: Optional[int] = Option(  # noqa: B008
+        None,
+        help=(
+            "Maximum number of connected segmentations to convert to meshes. "
+            "Leave empty or set to None to export all available segmentations."
+        ),
+    ),
+    merge_outputs: bool = Option(  # noqa: B008
+        False,
+        help=(
+            "Store all exported connected components in a single set of output "
+            "files (one H5/OBJ pair) per segmentation."
+        ),
+    ),
 ):
     """Convert a folder of membrane segmentations to meshes.
 
@@ -158,6 +188,8 @@ def convert_mb_folder(
         barycentric_area=barycentric_area,
         input_pixel_size=input_pixel_size,
         imod_meshing=imod_meshing,
+        max_segmentations=max_segmentations,
+        merge_outputs=merge_outputs,
     )
 
 
@@ -560,7 +592,11 @@ def surforama(
     from surforama.app import QtSurforama
     import numpy as np
     from matplotlib.pyplot import get_cmap
-    from membrain_pick.dataloading.data_utils import load_mesh_from_hdf5, get_csv_data
+    from membrain_pick.dataloading.data_utils import (
+        get_csv_data,
+        iter_mesh_entries,
+        load_mesh_from_hdf5,
+    )
     from membrain_seg.segmentation.dataloading.data_utils import load_tomogram
     from membrain_pick.scalar_selection import ScalarSelectionWidget
 
@@ -587,17 +623,21 @@ def surforama(
     else:
         mesh_files = [h5_path]
 
-    for h5_nr, h5_path in enumerate(mesh_files):
+    mesh_entries = []
+    for h5_path in mesh_files:
         mesh_data = load_mesh_from_hdf5(h5_path)
+        for group_name, group_data in iter_mesh_entries(mesh_data):
+            mesh_entries.append((h5_path, group_name, group_data))
 
-        if h5_nr == 0:
+    for entry_idx, (h5_path, _group_name, mesh_data) in enumerate(mesh_entries):
+        if entry_idx == 0:
             volume_layer = display_tomo(viewer, mesh_data, tomogram_path)
             pixel_size = get_pixel_size(mesh_data, None)
 
         points, faces = get_points_and_faces(mesh_data, pixel_size)
         display_scores(viewer, mesh_data, points, faces)
 
-        if h5_nr == 0:
+        if entry_idx == 0:
             surforama_widget = initialize_surforama_widget(
                 points, faces, volume_layer, viewer, normal_offset=normal_offset
             )
@@ -612,7 +652,7 @@ def surforama(
                 viewer, mesh_data, pixel_size, point_size=point_size
             )
 
-        if h5_nr == 0:
+        if entry_idx == 0:
             display_input_normal_values(viewer, mesh_data, points, faces)
     if return_viewer:
         return viewer

src/membrain_pick/cli/surforama_cli.py

@@ -33,7 +33,10 @@ def surforama(
     """
     import os
     import napari
-    from membrain_pick.dataloading.data_utils import load_mesh_from_hdf5
+    from membrain_pick.dataloading.data_utils import (
+        iter_mesh_entries,
+        load_mesh_from_hdf5,
+    )
     from membrain_pick.napari_utils.surforama_cli_utils import (
         display_tomo,
         get_pixel_size,
@@ -57,17 +60,21 @@ def surforama(
     else:
         mesh_files = [h5_path]
 
-    for h5_nr, h5_path in enumerate(mesh_files):
+    mesh_entries = []
+    for h5_path in mesh_files:
         mesh_data = load_mesh_from_hdf5(h5_path)
+        for group_name, group_data in iter_mesh_entries(mesh_data):
+            mesh_entries.append((h5_path, group_name, group_data))
 
-        if h5_nr == 0:
+    for entry_idx, (h5_path, _group_name, mesh_data) in enumerate(mesh_entries):
+        if entry_idx == 0:
             volume_layer = display_tomo(viewer, mesh_data, tomogram_path)
             pixel_size = get_pixel_size(mesh_data, None)
 
         points, faces = get_points_and_faces(mesh_data, pixel_size)
         display_scores(viewer, mesh_data, points, faces)
 
-        if h5_nr == 0:
+        if entry_idx == 0:
             surforama_widget = initialize_surforama_widget(
                 points, faces, volume_layer, viewer, normal_offset=normal_offset
             )
@@ -82,7 +89,7 @@ def surforama(
                 viewer, mesh_data, pixel_size, point_size=point_size
             )
 
-        if h5_nr == 0:
+        if entry_idx == 0:
             display_input_normal_values(viewer, mesh_data, points, faces)
     if return_viewer:
         return viewer

src/membrain_pick/clustering/mean_shift_inference.py

@@ -1,5 +1,6 @@
 from membrain_pick.dataloading.data_utils import (
     get_csv_data,
+    iter_mesh_entries,
     store_array_in_star,
     load_mesh_from_hdf5,
     store_mesh_in_hdf5,
@@ -58,22 +59,40 @@ def mean_shift_for_h5(
     method: str = "membrain_pick",
     score_threshold: float = 9.0,
 ):
-    mesh_data = load_mesh_from_hdf5(h5_file)
-    verts = mesh_data["points"]
-    scores = mesh_data["scores"]
-    out_pos, _ = mean_shift_for_scores(
-        verts, scores, bandwidth, max_iter, margin, device, score_threshold, method
-    )
-    # add positions to h5 container
-    mesh_data["cluster_centers"] = out_pos
-    out_file = os.path.join(out_dir, os.path.basename(h5_file))
+    mesh_entries = list(iter_mesh_entries(load_mesh_from_hdf5(h5_file)))
     os.makedirs(out_dir, exist_ok=True)
-    store_mesh_in_hdf5(
-        out_file,
-        **mesh_data,
-    )
+    out_file = os.path.join(out_dir, os.path.basename(h5_file))
+    if os.path.exists(out_file):
+        os.remove(out_file)
+
+    for group_name, mesh_data in mesh_entries:
+        verts = mesh_data["points"]
+        scores = mesh_data["scores"]
+        out_pos, _ = mean_shift_for_scores(
+            verts, scores, bandwidth, max_iter, margin, device, score_threshold, method
+        )
+        mesh_data = dict(mesh_data)
+        mesh_data["cluster_centers"] = out_pos
+
+        store_mesh_in_hdf5(
+            out_file,
+            group_name=group_name,
+            **mesh_data,
+        )
 
-    store_clusters(h5_file, out_dir, out_pos, np.zeros((0,)), verts, mesh_data["faces"])
+        group_file = h5_file
+        if group_name is not None:
+            base, ext = os.path.splitext(h5_file)
+            group_file = f"{base}_{group_name}{ext}"
+
+        store_clusters(
+            group_file,
+            out_dir,
+            out_pos,
+            np.zeros((0,)),
+            verts,
+            mesh_data["faces"],
+        )
 
 
 def mean_shift_for_csv(