Move relabeling into segment_ROI function to fix #785

Author: jluethi

fractal_tasks_core/tasks/cellpose_segmentation.py

@@ -73,20 +73,26 @@
 
 def segment_ROI(
     x: np.ndarray,
+    num_labels_tot: dict[str, int],
     model: models.CellposeModel = None,
     do_3D: bool = True,
     channels: list[int] = [0, 0],
     diameter: float = 30.0,
     normalize: CellposeCustomNormalizer = CellposeCustomNormalizer(),
     normalize2: Optional[CellposeCustomNormalizer] = None,
     label_dtype: Optional[np.dtype] = None,
+    relabeling: bool = True,
     advanced_cellpose_model_params: CellposeModelParams = CellposeModelParams(),  # noqa: E501
 ) -> np.ndarray:
     """
     Internal function that runs Cellpose segmentation for a single ROI.
 
     Args:
         x: 4D numpy array.
+        num_labels_tot: Number of labels already in total image. Used for
+            relabeling purposes. Using a dict to have a mutable object that
+            can be edited from within the function without having to be passed
+            back through the masked_loading_wrapper.
         model: An instance of `models.CellposeModel`.
         do_3D: If `True`, cellpose runs in 3D mode: runs on xy, xz & yz planes,
             then averages the flows.
@@ -107,6 +113,7 @@ def segment_ROI(
             normalized with default settings, both channels need to be
             normalized with default settings.
         label_dtype: Label images are cast into this `np.dtype`.
+        relabeling: Whether relabeling based on num_labels_tot is performed.
         advanced_cellpose_model_params: Advanced Cellpose model parameters
             that are passed to the Cellpose `model.eval` method.
     """
@@ -163,6 +170,23 @@ def segment_ROI(
         f" {advanced_cellpose_model_params.flow_threshold=}"
     )
 
+    # Shift labels and update relabeling counters
+    if relabeling:
+        num_labels_roi = np.max(mask)
+        mask[mask > 0] += num_labels_tot["num_labels_tot"]
+        num_labels_tot["num_labels_tot"] += num_labels_roi
+
+        # Write some logs
+        logger.info(f"ROI had {num_labels_roi=}, {num_labels_tot=}")
+
+        # Check that total number of labels is under control
+        if num_labels_tot["num_labels_tot"] > np.iinfo(label_dtype).max:
+            raise ValueError(
+                "ERROR in re-labeling:"
+                f"Reached {num_labels_tot} labels, "
+                f"but dtype={label_dtype}"
+            )
+
     return mask.astype(label_dtype)
 
 
@@ -438,8 +462,7 @@ def cellpose_segmentation(
         logger.info(f"{data_zyx_c2.chunks}")
 
     # Counters for relabeling
-    if relabeling:
-        num_labels_tot = 0
+    num_labels_tot = {"num_labels_tot": 0}
 
     # Iterate over ROIs
     num_ROIs = len(list_indices)
@@ -485,13 +508,15 @@ def cellpose_segmentation(
 
         # Prepare keyword arguments for segment_ROI function
         kwargs_segment_ROI = dict(
+            num_labels_tot=num_labels_tot,
             model=model,
             channels=channels,
             do_3D=do_3D,
             label_dtype=label_dtype,
             diameter=diameter_level0 / coarsening_xy**level,
             normalize=channel.normalize,
             normalize2=channel2.normalize,
+            relabeling=relabeling,
             advanced_cellpose_model_params=advanced_cellpose_model_params,
         )
 
@@ -515,23 +540,6 @@ def cellpose_segmentation(
             preprocessing_kwargs=preprocessing_kwargs,
         )
 
-        # Shift labels and update relabeling counters
-        if relabeling:
-            num_labels_roi = np.max(new_label_img)
-            new_label_img[new_label_img > 0] += num_labels_tot
-            num_labels_tot += num_labels_roi
-
-            # Write some logs
-            logger.info(f"ROI {indices}, {num_labels_roi=}, {num_labels_tot=}")
-
-            # Check that total number of labels is under control
-            if num_labels_tot > np.iinfo(label_dtype).max:
-                raise ValueError(
-                    "ERROR in re-labeling:"
-                    f"Reached {num_labels_tot} labels, "
-                    f"but dtype={label_dtype}"
-                )
-
         if output_ROI_table:
             bbox_df = array_to_bounding_box_table(
                 new_label_img,

tests/tasks/test_workflows_cellpose_segmentation.py

@@ -653,10 +653,6 @@ def test_cellpose_within_masked_bb_with_overlap(
         patched_cellpose_eval,
     )
 
-    # Setup caplog fixture, see
-    # https://docs.pytest.org/en/stable/how-to/logging.html#caplog-fixture
-    caplog.set_level(logging.WARNING)
-
     # Use pre-made 3D zarr
     zarr_dir = tmp_path / "tmp_out/"
     zarr_urls = prepare_3D_zarr(str(zarr_dir), zenodo_zarr)
@@ -717,6 +713,16 @@ def test_cellpose_within_masked_bb_with_overlap(
     # mock testing ignores the input
     # assert np.max(secondary_segmentation) == 4
 
+    # Ensure that labels stay in correct proportions => relabeling doesn't do
+    # reassignment
+    label1 = np.sum(secondary_segmentation == 1)
+    label3 = np.sum(secondary_segmentation == 3)
+    label5 = np.sum(secondary_segmentation == 5)
+    label7 = np.sum(secondary_segmentation == 7)
+    assert label1 == label3
+    assert label1 == label5
+    assert label1 == label7
+
 
 def test_workflow_with_per_FOV_labeling_via_script(
     tmp_path: Path,