Repository: fractal-analytics-platform/fractal-tasks-core
Issue category: Performance Issue

Issue:
Review/optimize use of `get_overlapping_pairs_3D` in Cellpose task
Branching from #764 (fixed in principle with #778)

After a fix, it'd be useful to have a rough benchmark of `get_overlapping_pairs_3D for` one of those many-labels cases (say 1000 labels for a given `i_ROI`). Since this function does nothing else than printing a warning, I would be very strict in how long a runtime we can accept.

If it turns out that it is slow, we can easily improve it in a trivial way (simply stop after the first overlap is found) or also in more systematic ways (the current approach is close to being the definition of a slow Python function: a quadratic-scaling for loop, which calls a pure-Python function and then even appends to a list).

If the function is not slow, then there's no special need for a refactor.


Hints:
Another thing to consider here:
This (as well as the original implementation) will only check for potential bounding box overlaps within a given ROI that's being processed, right?

A future improvement may be to check for overlaps across ROIs as well, as in theory, they could also overlap. Given that this is just for warnings, not sure how relevant this is though. Our masking should still work for this after all.
> A future improvement may be to check for overlaps across ROIs as well, as in theory, they could also overlap. Given that this is just for warnings, not sure how relevant this is though. Our masking should still work for this after all.

For the record, making this change is quite straightforward, since we should only move the check from within the ROI loop to its end (EDIT: this diff is not working, but the idea remains valid):
```diff
diff --git a/fractal_tasks_core/tasks/cellpose_segmentation.py b/fractal_tasks_core/tasks/cellpose_segmentation.py
index ea78661a..403630a1 100644
--- a/fractal_tasks_core/tasks/cellpose_segmentation.py
+++ b/fractal_tasks_core/tasks/cellpose_segmentation.py
@@ -541,15 +541,6 @@ def cellpose_segmentation(
 
             bbox_dataframe_list.append(bbox_df)
 
-            overlap_list = get_overlapping_pairs_3D(
-                bbox_df, full_res_pxl_sizes_zyx
-            )
-            if len(overlap_list) > 0:
-                logger.warning(
-                    f"ROI {indices} has "
-                    f"{len(overlap_list)} bounding-box pairs overlap"
-                )
-
         # Compute and store 0-th level to disk
         da.array(new_label_img).to_zarr(
             url=mask_zarr,
@@ -582,6 +573,15 @@ def cellpose_segmentation(
             bbox_dataframe_list = [empty_bounding_box_table()]
         # Concatenate all ROI dataframes
         df_well = pd.concat(bbox_dataframe_list, axis=0, ignore_index=True)
+
+        overlap_list = get_overlapping_pairs_3D(
+            bbox_dataframe_list, full_res_pxl_sizes_zyx
+        )
+        if len(overlap_list) > 0:
+            logger.warning(
+                f"{len(overlap_list)} bounding-box pairs overlap"
+            )
+
         df_well.index = df_well.index.astype(str)
         # Extract labels and drop them from df_well
         labels = pd.DataFrame(df_well["label"].astype(str))
```

---

However we should not do this without actual benchmark, or we risk introducing a very bad scaling. If there are `N` organoids with `M` labels each (under the fake assumption that each organoid has the same number of internal objects), we have:
1. Current `main`: `(N*(N-1)/2) * (M*(M-1)/2)` overlap checks
2. After #778: N*M*(M-1)/2 overlap checks
3. After the patch described in this comment: (N*M)*(N*M-1)/2 overlap checks

E.g. for N=138 organoids (ref #764) and 100 labels (random guess), we'd have:
1. 46*10^6 checks
2. 683k checks
3. 95*10^6 checks

Note that:
> If it turns out that it is slow, we can easily improve it in a trivial way (simply stop after the first overlap is found) or also in more systematic ways (the current approach is close to being the definition of a slow Python function: a quadratic-scaling for loop, which calls a pure-Python function and then even appends to a list).
Agreed on not doing this now. More something to test if this gets benchmarked further :)

I'd say we go with the fix now as you implemented it here. And consider most of this benchmarking in the context of a future Cellpose task refactor to use new OME-Zarr reader/writer strategy
I just debugged this with some users who (mistakenly) had output ROI tables set for single cell segmentation where they have tens of thousands of segmented objects. In that case, this performance became super prohibitive, even in the slightly optimized version.

I'd say we should remove this in the next Cellpose refactor. The warning is not very useful, but sometimes highly performance critcial.

For the extreme example, for 1 user, their whole segmentation task ran in under a minute, but then it took almost 9 hours (!!) to do these overlap checks for 28 thousand label objects:
```
2024-10-16 10:28:52,329; INFO; START cellpose_segmentation task
[... Detail logs] 
2024-10-16 10:28:53,449; INFO; [segment_ROI] START | x: <class 'numpy.ndarray'>, (1, 1, 5400, 5120) | do_3D=False | model.diam_mean=17.0 | diameter=10.0 | advanced_cellpose_model_params.flow_threshold=0.4 | normalize.type='default'
2024-10-16 10:29:27,081; INFO; [segment_ROI] END   | Elapsed: 33.627 s | mask.shape=(1, 5400, 5120), mask.dtype=dtype('uint16') (then <class 'numpy.uint32'>), np.max(mask)=28360 | model.diam_mean=17.0 | diameter=10.0 | advanced_cellpose_model_params.flow_threshold=0.4
2024-10-16 10:29:27,134; INFO; ROI had num_labels_roi=28360, num_labels_tot={'num_labels_tot': 28360}
[This is were the cellpose segmentation has ended]
2024-10-16 19:20:10,278; WARNING; ROI [0, 1, 0, 5400, 0, 5120] has 18940 bounding-box pairs overlap
```


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f7aea0d58..35082befe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 **Note**: Numbers like (\#123) point to closed Pull Requests on the fractal-tasks-core repository.
 
+# 1.4.1
+
+* Tasks:
+    * Remove overlap checking for output ROIs in Cellpose task to address performance issues (\#889).
 
 # 1.4.0
 
diff --git a/fractal_tasks_core/tasks/cellpose_segmentation.py b/fractal_tasks_core/tasks/cellpose_segmentation.py
index 9ee23b8e8..499ecc0b7 100644
--- a/fractal_tasks_core/tasks/cellpose_segmentation.py
+++ b/fractal_tasks_core/tasks/cellpose_segmentation.py
@@ -44,7 +44,6 @@
 from fractal_tasks_core.roi import (
     find_overlaps_in_ROI_indices,
 )
-from fractal_tasks_core.roi import get_overlapping_pairs_3D
 from fractal_tasks_core.roi import is_ROI_table_valid
 from fractal_tasks_core.roi import load_region
 from fractal_tasks_core.tables import write_table
@@ -561,15 +560,6 @@ def cellpose_segmentation(
 
             bbox_dataframe_list.append(bbox_df)
 
-            overlap_list = get_overlapping_pairs_3D(
-                bbox_df, full_res_pxl_sizes_zyx
-            )
-            if len(overlap_list) > 0:
-                logger.warning(
-                    f"ROI {indices} has "
-                    f"{len(overlap_list)} bounding-box pairs overlap"
-                )
-
         # Compute and store 0-th level to disk
         da.array(new_label_img).to_zarr(
             url=mask_zarr,


Test patch:
diff --git a/tests/tasks/test_workflows_cellpose_segmentation.py b/tests/tasks/test_workflows_cellpose_segmentation.py
index 2c3c04216..1ce511d30 100644
--- a/tests/tasks/test_workflows_cellpose_segmentation.py
+++ b/tests/tasks/test_workflows_cellpose_segmentation.py
@@ -602,49 +602,6 @@ def test_workflow_bounding_box(
     assert "encoding-version" in table_group.attrs.asdict().keys()
 
 
-def test_workflow_bounding_box_with_overlap(
-    tmp_path: Path,
-    zenodo_zarr: list[str],
-    caplog: pytest.LogCaptureFixture,
-    monkeypatch: MonkeyPatch,
-):
-    monkeypatch.setattr(
-        "fractal_tasks_core.tasks.cellpose_segmentation.cellpose.core.use_gpu",
-        patched_cellpose_core_use_gpu,
-    )
-
-    monkeypatch.setattr(
-        "fractal_tasks_core.tasks.cellpose_segmentation.segment_ROI",
-        patched_segment_ROI_overlapping_organoids,
-    )
-
-    # 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)
-    debug(zarr_dir)
-    debug(zarr_urls)
-
-    # Per-FOV labeling
-    channel = CellposeChannel1InputModel(
-        wavelength_id="A01_C01", normalize=CellposeCustomNormalizer()
-    )
-    for zarr_url in zarr_urls:
-        cellpose_segmentation(
-            zarr_url=zarr_url,
-            channel=channel,
-            level=3,
-            relabeling=True,
-            diameter_level0=80.0,
-            output_ROI_table="bbox_table",
-        )
-        debug(caplog.text)
-        assert "bounding-box pairs overlap" in caplog.text
-
-
 def test_cellpose_within_masked_bb_with_overlap(
     tmp_path: Path,
     zenodo_zarr: list[str],
diff --git a/tests/test_unit_ROIs.py b/tests/test_unit_ROIs.py
index eb8cc3bf4..3806bcdb2 100644
--- a/tests/test_unit_ROIs.py
+++ b/tests/test_unit_ROIs.py
@@ -30,6 +30,7 @@
     find_overlaps_in_ROI_indices,
 )
 from fractal_tasks_core.roi import get_image_grid_ROIs
+from fractal_tasks_core.roi import get_overlapping_pairs_3D
 from fractal_tasks_core.roi import get_single_image_ROI
 from fractal_tasks_core.roi import is_ROI_table_valid
 from fractal_tasks_core.roi import (
@@ -755,3 +756,32 @@ def test_create_roi_table_from_df_list_with_label_repeats():
         ]
     )
     np.testing.assert_allclose(output_array, roi_table.X)
+
+
+def test_get_overlapping_pairs_3D():
+    common_columns = {
+        "y_micrometer": [0.0, 0.0],
+        "z_micrometer": [0.0, 0.0],
+        "len_x_micrometer": [1.0, 1.0],
+        "len_y_micrometer": [1.0, 1.0],
+        "len_z_micrometer": [1.0, 1.0],
+        "label": [1, 2],
+    }
+    df_overlapping = pd.DataFrame(
+        {"x_micrometer": [0.0, 0.5], **common_columns}
+    )
+    full_res_pxl_sizes_zyx = [1.0, 1.0, 1.0]
+    df_non_overlapping = pd.DataFrame(
+        {"x_micrometer": [0.0, 2.0], **common_columns}
+    )
+    res_overlapping = get_overlapping_pairs_3D(
+        df_overlapping,
+        full_res_pxl_sizes_zyx,
+    )
+    assert len(res_overlapping) == 1
+
+    res_non_overlapping = get_overlapping_pairs_3D(
+        df_non_overlapping,
+        full_res_pxl_sizes_zyx,
+    )
+    assert len(res_non_overlapping) == 0