From 7d0133712423b8fa3053db267f0949e1dbd5fdab Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Mon, 23 Mar 2026 18:39:57 +0100 Subject: [PATCH 1/6] Add first draft for uSegment3d backbone --- scripts/usegment_3d/segment_3d.py | 68 +++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 scripts/usegment_3d/segment_3d.py diff --git a/scripts/usegment_3d/segment_3d.py b/scripts/usegment_3d/segment_3d.py new file mode 100644 index 000000000..4d5f618ad --- /dev/null +++ b/scripts/usegment_3d/segment_3d.py @@ -0,0 +1,68 @@ +from tqdm import tqdm +from functools import partial + +import numpy as np + +import segment3D.usegment3d as uSegment3D +import segment3D.parameters as uSegment3D_params + + +def run_usegment3d_with_microsam(volume): + # Run MicroSAM on 3d volume along all three directions. + from micro_sam.automatic_segmentation import get_predictor_and_segmenter, automatic_instance_segmentation + + # Load the MicroSAM model. + predictor, segmenter = get_predictor_and_segmenter(model_type="vit_b_lm", segmentation_mode="ais") + + # Run AIS on our volume in a per-slice fashion. + seg_runner = partial( + automatic_instance_segmentation, predictor=predictor, segmenter=segmenter, ndim=2, verbose=False + ) + + instances_xy = np.stack( + [seg_runner(input_path=curr_slice) for curr_slice in tqdm(volume, desc="xy")] + ) + instances_xz = np.stack( + [seg_runner(input_path=curr_slice) for curr_slice in tqdm(volume.transpose(1, 0, 2), desc="xz")] + ) + # skip 'yz' for now + instances_yz = [] + + # Get the default parameters first. + params = uSegment3D_params.get_2D_to_3D_aggregation_params() + + # The available choices are "cellpose_improve", "fmm", "cellpose_skel", "fmm_skel", "edt". + params["indirect_method"]["dtform_method"] = "cellpose_improve" + + # And a few other parameters. + params["indirect_method"]["dtform_method"] = "cellpose_improve" + params["gradient_descent"]["gradient_decay"] = None + params["indirect_method"]["smooth_skel_sigma"] = None + params["indirect_method"]["edt_fixed_point_percentile"] = None + + # Run the uSegment3d's 'indirect' method for the most amount of flexibility. + segmentation_3d, _ = uSegment3D.aggregate_2D_to_3D_segmentation_indirect_method( + segmentations=[instances_xy, instances_xz, instances_yz], + img_xy_shape=instances_xy.shape, # full 3D shape, not one slice + precomputed_binary=None, # Seems like binary segmentations in xy direction. + params=params, + savefolder=None, + basename="usegment3d_indirect_test", + ) + + # Store segmentations + import h5py + with h5py.File("test.h5", "w") as f: + f.create_dataset("raw", data=volume, compression="gzip") + f.create_dataset("segmentation/usegment3d-microsam", data=segmentation_3d, compression="gzip") + + +def main(): + # Let's work with the 'cell3d' example data in scikit-image. + from skimage.data import cells3d + volume = cells3d()[:, 0] # input has shape of (60, 256, 256). + run_usegment3d_with_microsam(volume) + + +if __name__ == "__main__": + main() From 9d0bab8a2b3d2ac8b0e98318055fa9f01fa2764f Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Mon, 23 Mar 2026 18:46:58 +0100 Subject: [PATCH 2/6] Remove unnecessary arguments --- scripts/usegment_3d/segment_3d.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/scripts/usegment_3d/segment_3d.py b/scripts/usegment_3d/segment_3d.py index 4d5f618ad..8ebacb0d1 100644 --- a/scripts/usegment_3d/segment_3d.py +++ b/scripts/usegment_3d/segment_3d.py @@ -34,12 +34,6 @@ def run_usegment3d_with_microsam(volume): # The available choices are "cellpose_improve", "fmm", "cellpose_skel", "fmm_skel", "edt". params["indirect_method"]["dtform_method"] = "cellpose_improve" - # And a few other parameters. - params["indirect_method"]["dtform_method"] = "cellpose_improve" - params["gradient_descent"]["gradient_decay"] = None - params["indirect_method"]["smooth_skel_sigma"] = None - params["indirect_method"]["edt_fixed_point_percentile"] = None - # Run the uSegment3d's 'indirect' method for the most amount of flexibility. segmentation_3d, _ = uSegment3D.aggregate_2D_to_3D_segmentation_indirect_method( segmentations=[instances_xy, instances_xz, instances_yz], From b59efd3554972e5eee22a7d28020b158c3cfc8e3 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Mon, 23 Mar 2026 19:03:55 +0100 Subject: [PATCH 3/6] Remove unnecessary dask warnings --- scripts/usegment_3d/segment_3d.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/scripts/usegment_3d/segment_3d.py b/scripts/usegment_3d/segment_3d.py index 8ebacb0d1..acb7cb351 100644 --- a/scripts/usegment_3d/segment_3d.py +++ b/scripts/usegment_3d/segment_3d.py @@ -1,3 +1,4 @@ +import logging from tqdm import tqdm from functools import partial @@ -7,6 +8,13 @@ import segment3D.parameters as uSegment3D_params +# Suppress noisy Dask distributed worker/nanny startup logs. +for _logger in ( + "distributed", "distributed.nanny", "distributed.worker", "distributed.scheduler", "distributed.core" +): + logging.getLogger(_logger).setLevel(logging.WARNING) + + def run_usegment3d_with_microsam(volume): # Run MicroSAM on 3d volume along all three directions. from micro_sam.automatic_segmentation import get_predictor_and_segmenter, automatic_instance_segmentation @@ -25,14 +33,16 @@ def run_usegment3d_with_microsam(volume): instances_xz = np.stack( [seg_runner(input_path=curr_slice) for curr_slice in tqdm(volume.transpose(1, 0, 2), desc="xz")] ) - # skip 'yz' for now - instances_yz = [] + instances_yz = np.stack( + [seg_runner(input_path=curr_slice) for curr_slice in tqdm(volume.transpose(1, 2, 0), desc="yz")] + ) # Get the default parameters first. params = uSegment3D_params.get_2D_to_3D_aggregation_params() # The available choices are "cellpose_improve", "fmm", "cellpose_skel", "fmm_skel", "edt". params["indirect_method"]["dtform_method"] = "cellpose_improve" + params["indirect_method"]["n_cpu"] = 4 # default spawns (cpu_count-1)//2 dask workers # Run the uSegment3d's 'indirect' method for the most amount of flexibility. segmentation_3d, _ = uSegment3D.aggregate_2D_to_3D_segmentation_indirect_method( @@ -54,7 +64,7 @@ def run_usegment3d_with_microsam(volume): def main(): # Let's work with the 'cell3d' example data in scikit-image. from skimage.data import cells3d - volume = cells3d()[:, 0] # input has shape of (60, 256, 256). + volume = cells3d()[:, 1] # input has shape of (60, 256, 256). run_usegment3d_with_microsam(volume) From 63a0968bd590ff34a27e707f1006a5d8c929489b Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Mon, 23 Mar 2026 19:13:52 +0100 Subject: [PATCH 4/6] Move warning bumper inside function and fix axis alignment --- scripts/usegment_3d/segment_3d.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/scripts/usegment_3d/segment_3d.py b/scripts/usegment_3d/segment_3d.py index acb7cb351..dd88670da 100644 --- a/scripts/usegment_3d/segment_3d.py +++ b/scripts/usegment_3d/segment_3d.py @@ -8,14 +8,13 @@ import segment3D.parameters as uSegment3D_params -# Suppress noisy Dask distributed worker/nanny startup logs. -for _logger in ( - "distributed", "distributed.nanny", "distributed.worker", "distributed.scheduler", "distributed.core" -): - logging.getLogger(_logger).setLevel(logging.WARNING) - - def run_usegment3d_with_microsam(volume): + # Suppress noisy Dask distributed worker/nanny startup logs. + for _logger in ( + "distributed", "distributed.nanny", "distributed.worker", "distributed.scheduler", "distributed.core" + ): + logging.getLogger(_logger).setLevel(logging.WARNING) + # Run MicroSAM on 3d volume along all three directions. from micro_sam.automatic_segmentation import get_predictor_and_segmenter, automatic_instance_segmentation @@ -28,13 +27,13 @@ def run_usegment3d_with_microsam(volume): ) instances_xy = np.stack( - [seg_runner(input_path=curr_slice) for curr_slice in tqdm(volume, desc="xy")] + [seg_runner(input_path=curr_slice) for curr_slice in tqdm(volume, desc="Segment XY")] ) instances_xz = np.stack( - [seg_runner(input_path=curr_slice) for curr_slice in tqdm(volume.transpose(1, 0, 2), desc="xz")] + [seg_runner(input_path=curr_slice) for curr_slice in tqdm(volume.transpose(1, 0, 2), desc="Segment XZ")] ) instances_yz = np.stack( - [seg_runner(input_path=curr_slice) for curr_slice in tqdm(volume.transpose(1, 2, 0), desc="yz")] + [seg_runner(input_path=curr_slice) for curr_slice in tqdm(volume.transpose(2, 0, 1), desc="Segment YZ")] ) # Get the default parameters first. @@ -42,7 +41,7 @@ def run_usegment3d_with_microsam(volume): # The available choices are "cellpose_improve", "fmm", "cellpose_skel", "fmm_skel", "edt". params["indirect_method"]["dtform_method"] = "cellpose_improve" - params["indirect_method"]["n_cpu"] = 4 # default spawns (cpu_count-1)//2 dask workers + params["indirect_method"]["n_cpu"] = 4 # default spawns '(cpu_count - 1) // 2 dask workers' # Run the uSegment3d's 'indirect' method for the most amount of flexibility. segmentation_3d, _ = uSegment3D.aggregate_2D_to_3D_segmentation_indirect_method( From 1da8f336d5ee0f1baf575e4bb9eb5950bd031ac8 Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Mon, 23 Mar 2026 19:41:23 +0100 Subject: [PATCH 5/6] Add evaluation pipeline --- scripts/usegment_3d/segment_3d.py | 80 ++++++++++++++++++++++++++----- 1 file changed, 67 insertions(+), 13 deletions(-) diff --git a/scripts/usegment_3d/segment_3d.py b/scripts/usegment_3d/segment_3d.py index dd88670da..82e2dd0eb 100644 --- a/scripts/usegment_3d/segment_3d.py +++ b/scripts/usegment_3d/segment_3d.py @@ -2,22 +2,44 @@ from tqdm import tqdm from functools import partial +import h5py import numpy as np +import imageio.v3 as imageio -import segment3D.usegment3d as uSegment3D -import segment3D.parameters as uSegment3D_params +from torch_em.data.datasets.light_microscopy.embedseg_data import get_embedseg_paths +from micro_sam.automatic_segmentation import get_predictor_and_segmenter, automatic_instance_segmentation + + +def run_microsam3d(volume, labels=None, save_path="test.h5"): + # Load the MicroSAM model. + predictor, segmenter = get_predictor_and_segmenter(model_type="vit_b_lm", segmentation_mode="ais") + + # Run AIS on our entire volume. + instances = automatic_instance_segmentation( + predictor=predictor, segmenter=segmenter, ndim=3, input_path=volume, + ) + + # Store segmentations + with h5py.File(save_path, "a") as f: + if "raw" not in f: + f.create_dataset("raw", data=volume, compression="gzip") + if "labels" not in f and labels is not None: + f.create_dataset("labels", data=labels, compression="gzip") + f.create_dataset("segmentation/microsam3d", data=instances, compression="gzip") + + +def run_usegment3d_with_microsam(volume, labels=None, save_path="test.h5"): + # Run MicroSAM on 3d volume along all three directions. + import segment3D.usegment3d as uSegment3D + import segment3D.parameters as uSegment3D_params -def run_usegment3d_with_microsam(volume): # Suppress noisy Dask distributed worker/nanny startup logs. for _logger in ( "distributed", "distributed.nanny", "distributed.worker", "distributed.scheduler", "distributed.core" ): logging.getLogger(_logger).setLevel(logging.WARNING) - # Run MicroSAM on 3d volume along all three directions. - from micro_sam.automatic_segmentation import get_predictor_and_segmenter, automatic_instance_segmentation - # Load the MicroSAM model. predictor, segmenter = get_predictor_and_segmenter(model_type="vit_b_lm", segmentation_mode="ais") @@ -38,10 +60,10 @@ def run_usegment3d_with_microsam(volume): # Get the default parameters first. params = uSegment3D_params.get_2D_to_3D_aggregation_params() + params["indirect_method"]["n_cpu"] = 4 # default spawns '(cpu_count - 1) // 2 dask workers' # The available choices are "cellpose_improve", "fmm", "cellpose_skel", "fmm_skel", "edt". params["indirect_method"]["dtform_method"] = "cellpose_improve" - params["indirect_method"]["n_cpu"] = 4 # default spawns '(cpu_count - 1) // 2 dask workers' # Run the uSegment3d's 'indirect' method for the most amount of flexibility. segmentation_3d, _ = uSegment3D.aggregate_2D_to_3D_segmentation_indirect_method( @@ -54,17 +76,49 @@ def run_usegment3d_with_microsam(volume): ) # Store segmentations - import h5py - with h5py.File("test.h5", "w") as f: - f.create_dataset("raw", data=volume, compression="gzip") + with h5py.File(save_path, "a") as f: + if "raw" not in f: + f.create_dataset("raw", data=volume, compression="gzip") + if "labels" not in f and labels is not None: + f.create_dataset("labels", data=labels, compression="gzip") f.create_dataset("segmentation/usegment3d-microsam", data=segmentation_3d, compression="gzip") +def evaluate_results(save_path): + with h5py.File(save_path, "r") as f: + labels = f["labels"][:] + seg_microsam3d = f["segmentation/microsam3d"][:] + seg_microsam_usegment3d = f["segmentation/usegment3d-microsam"][:] + + # Let's evaluate and see how the results are + from elf.evaluation import mean_segmentation_accuracy + msa_microsam3d = mean_segmentation_accuracy(labels, seg_microsam3d) + msa_microsam_usegment3d = mean_segmentation_accuracy(labels, seg_microsam_usegment3d) + print(msa_microsam3d, msa_microsam_usegment3d) + + def main(): # Let's work with the 'cell3d' example data in scikit-image. - from skimage.data import cells3d - volume = cells3d()[:, 1] # input has shape of (60, 256, 256). - run_usegment3d_with_microsam(volume) + # from skimage.data import cells3d + # volume = cells3d()[:, 1] # input has shape of (60, 256, 256). + + # Let's work with labeled data (so that we can evaluate) + data_dir = "/mnt/vast-nhr/projects/cidas/cca/data/embedseg" + raw_paths, label_paths = get_embedseg_paths( + path=data_dir, name="Mouse-Skull-Nuclei-CBG", split="test", + ) + volume, labels = imageio.imread(raw_paths[0]), imageio.imread(label_paths[0]) + + # Run / evaluate segmentation models. + save_path = "embedseg_mouse_skull_nuclei.h5" + + # run_usegment3d_with_microsam(volume, labels, save_path) + # run_microsam3d(volume, labels, save_path) + + # Observations on Mouse-Skull-Nuclei-CBG data: + # MicroSAM3d: 0.359 | MicroSAM2d + uSegment3d: 0.479 + + evaluate_results(save_path) if __name__ == "__main__": From 95f7fbf37c5b4d2b160d89a2eb3d406e2ca75f7c Mon Sep 17 00:00:00 2001 From: Anwai Archit Date: Mon, 23 Mar 2026 19:50:22 +0100 Subject: [PATCH 6/6] Add installation instructions --- .gitignore | 1 + scripts/usegment_3d/segment_3d.py | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 445e35ef7..39b637aff 100644 --- a/.gitignore +++ b/.gitignore @@ -202,3 +202,4 @@ clf-test-data data/ embeddings/ *.zarr +*.h5 diff --git a/scripts/usegment_3d/segment_3d.py b/scripts/usegment_3d/segment_3d.py index 82e2dd0eb..9671dfa05 100644 --- a/scripts/usegment_3d/segment_3d.py +++ b/scripts/usegment_3d/segment_3d.py @@ -31,6 +31,7 @@ def run_microsam3d(volume, labels=None, save_path="test.h5"): def run_usegment3d_with_microsam(volume, labels=None, save_path="test.h5"): # Run MicroSAM on 3d volume along all three directions. + # NOTE: Install uSegment3D using `pip install u-Segment3D`. import segment3D.usegment3d as uSegment3D import segment3D.parameters as uSegment3D_params @@ -112,8 +113,8 @@ def main(): # Run / evaluate segmentation models. save_path = "embedseg_mouse_skull_nuclei.h5" - # run_usegment3d_with_microsam(volume, labels, save_path) - # run_microsam3d(volume, labels, save_path) + run_usegment3d_with_microsam(volume, labels, save_path) + run_microsam3d(volume, labels, save_path) # Observations on Mouse-Skull-Nuclei-CBG data: # MicroSAM3d: 0.359 | MicroSAM2d + uSegment3d: 0.479