🆕 Add Cerberus

tests/engines/test_multi_task_segmentor.py

@@ -28,7 +28,10 @@
     DaskDelayedJSONStore,
     MultiTaskSegmentor,
     _clear_zarr,
+    _crop_halo_post_process_output,
+    _get_postproc_tile_read_bounds,
     _get_sel_indices_margin_lines,
+    _normalise_postproc_halo,
     _post_save_json_store,
     _process_instance_predictions,
     _save_multitask_vertical_to_cache,
@@ -889,9 +892,10 @@ class FakeVM:
     )
 
     # --- Call function ---
-    new_zarr, new_da = _save_multitask_vertical_to_cache(
+    new_zarr, new_da, zarr_group = _save_multitask_vertical_to_cache(
         probabilities_zarr=probabilities_zarr,
         probabilities_da=probabilities_da,
+        zarr_group=None,
         probabilities=probabilities,
         idx=idx,
         tqdm_loop=tqdm_loop,
@@ -905,11 +909,44 @@ class FakeVM:
 
     # new_zarr must be a real zarr array
     assert isinstance(new_zarr[idx], zarr.Array)
+    assert zarr_group is not None
 
     # Data was written correctly
     assert np.array_equal(new_zarr[idx][:], np.array([[1, 2, 3]]))
 
 
+def test_multitask_vertical_merge_continues_after_zarr_spill(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Test multitask vertical merge appends all chunks after spilling to Zarr."""
+
+    class FakeVM:
+        """Fake psutil.virtual_memory() with extremely low available memory."""
+
+        available = 1
+
+    monkeypatch.setattr(psutil, "virtual_memory", FakeVM)
+
+    values = np.arange(8 * 3, dtype=np.float32).reshape(8, 3, 1)
+    canvas = [da.from_array(values, chunks=(2, 3, 1))]
+    count = [da.from_array(np.ones_like(values), chunks=(2, 3, 1))]
+    output_locs_y = np.array([[0, 2], [2, 4], [4, 6], [6, 8]])
+
+    result = merge_multitask_vertical_chunkwise(
+        canvas=canvas,
+        count=count,
+        output_locs_y_=output_locs_y,
+        zarr_group=None,
+        save_path=tmp_path / "vertical.zarr",
+        memory_threshold=0,
+        output_shape=(8, 3),
+        verbose=False,
+    )
+
+    assert result[0].shape == values.shape
+    assert np.array_equal(result[0].compute(), values)
+
+
 def test_qupath_feature_class_dict_lookup_fails() -> None:
     """Test qupath_feature_class_dict lookup fails."""
     qupath_json = DaskDelayedJSONStore.__new__(DaskDelayedJSONStore)
@@ -1079,6 +1116,215 @@ def test_get_tile_info_small_image_triggers_early_return(
     assert np.all(flag == 0)
 
 
+def test_postproc_halo_bounds_and_output_crop() -> None:
+    """Test halo-expanded tile output is cropped and shifted to core space."""
+    halo_xy = _normalise_postproc_halo((3, 2))
+    assert np.array_equal(halo_xy, np.array([2, 3]))
+
+    read_bounds = _get_postproc_tile_read_bounds(
+        tile_bounds=(4, 5, 10, 11),
+        postproc_halo_xy=halo_xy,
+        image_shape=(12, 13),
+    )
+    assert read_bounds == (2, 2, 12, 13)
+
+    predictions = np.arange(11 * 10).reshape(11, 10)
+    info_dict = {
+        "box": np.array(
+            [
+                [2, 3, 4, 5],
+                [5, 6, 7, 8],
+                [9, 6, 11, 8],
+            ],
+            dtype=np.int32,
+        ),
+        "centroid": np.array(
+            [
+                [3, 4],
+                [6, 7],
+                [10, 7],
+            ],
+            dtype=np.float32,
+        ),
+        "contours": np.array(
+            [
+                [[2, 3], [4, 3], [4, 5], [2, 5]],
+                [[5, 6], [7, 6], [7, 8], [5, 8]],
+                [[9, 6], [11, 6], [11, 8], [9, 8]],
+            ],
+            dtype=np.int32,
+        ),
+        "type": np.array([1, 2, 3], dtype=np.int32),
+    }
+
+    cropped = _crop_halo_post_process_output(
+        post_process_output=(
+            {
+                "task_type": "gland",
+                "seg_type": "instance",
+                "predictions": predictions,
+                "info_dict": info_dict,
+            },
+        ),
+        tile_bounds=(4, 5, 10, 11),
+        tile_read_bounds=read_bounds,
+    )[0]
+
+    assert np.array_equal(cropped["predictions"], predictions[3:9, 2:8])
+    assert np.array_equal(cropped["info_dict"]["type"], np.array([1, 2]))
+    assert np.array_equal(
+        cropped["info_dict"]["box"],
+        np.array([[0, 0, 2, 2], [3, 3, 5, 5]], dtype=np.int32),
+    )
+    assert np.array_equal(
+        cropped["info_dict"]["centroid"],
+        np.array([[1, 1], [4, 4]], dtype=np.float32),
+    )
+    assert np.array_equal(
+        cropped["info_dict"]["contours"][0],
+        np.array([[0, 0], [2, 0], [2, 2], [0, 2]], dtype=np.int32),
+    )
+
+
+def test_postproc_halo_ownership_without_centroids() -> None:
+    """Test halo ownership falls back to boxes and padded contours."""
+    read_bounds = (2, 2, 12, 13)
+    predictions = np.arange(11 * 10).reshape(11, 10)
+
+    box_cropped = _crop_halo_post_process_output(
+        post_process_output=(
+            {
+                "task_type": "gland",
+                "seg_type": "instance",
+                "predictions": predictions,
+                "info_dict": {
+                    "box": np.array(
+                        [
+                            [2, 3, 4, 5],
+                            [9, 6, 11, 8],
+                        ],
+                        dtype=np.int32,
+                    ),
+                    "type": np.array([1, 2], dtype=np.int32),
+                },
+            },
+        ),
+        tile_bounds=(4, 5, 10, 11),
+        tile_read_bounds=read_bounds,
+    )[0]
+    assert np.array_equal(
+        box_cropped["info_dict"]["box"],
+        np.array([[0, 0, 2, 2]], dtype=np.int32),
+    )
+    assert np.array_equal(box_cropped["info_dict"]["type"], np.array([1]))
+
+    pad_value = np.iinfo(np.int32).min
+    contour_cropped = _crop_halo_post_process_output(
+        post_process_output=(
+            {
+                "task_type": "gland",
+                "seg_type": "instance",
+                "predictions": predictions,
+                "info_dict": {
+                    "contours": np.array(
+                        [
+                            [[2, 3], [4, 3], [4, 5], [2, 5]],
+                            [[9, 6], [11, 6], [11, 8], [9, 8]],
+                            [[pad_value, pad_value]] * 4,
+                        ],
+                        dtype=np.int32,
+                    ),
+                    "type": np.array([1, 2, 3], dtype=np.int32),
+                },
+            },
+        ),
+        tile_bounds=(4, 5, 10, 11),
+        tile_read_bounds=read_bounds,
+    )[0]
+    assert np.array_equal(contour_cropped["info_dict"]["type"], np.array([1]))
+    assert np.array_equal(
+        contour_cropped["info_dict"]["contours"][0],
+        np.array([[0, 0], [2, 0], [2, 2], [0, 2]], dtype=np.int32),
+    )
+
+
+def test_process_tile_mode_uses_postproc_halo(
+    track_tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test tile mode expands reads and crops outputs when halo is set."""
+    seg = MultiTaskSegmentor.__new__(MultiTaskSegmentor)
+    seg.verbose = False
+    seg.num_workers = 1
+    seg.mask_bounds = (0, 0, 10, 10)
+    seg.mask_padding = (0, 0, 0, 0)
+    seg.dataloader = SimpleNamespace(
+        dataset=SimpleNamespace(
+            reader=SimpleNamespace(slide_dimensions=lambda **_: (10, 10)),
+        ),
+    )
+    seg._ioconfig = SimpleNamespace(
+        highest_input_resolution={},
+        tile_shape=(4, 4),
+        to_baseline=lambda: SimpleNamespace(margin=0),
+    )
+
+    tile_info_sets = [
+        [
+            np.array([[2, 2, 6, 6]], dtype=np.int32),
+            np.array([[1, 1, 1, 1]], dtype=np.int32),
+        ],
+        [
+            np.array([[6, 6, 10, 10]], dtype=np.int32),
+            np.array([[1, 1, 1, 1]], dtype=np.int32),
+        ],
+    ]
+    seg._get_tile_info = lambda **_: tile_info_sets
+    recorded_bounds = []
+    expanded_predictions = np.arange(64, dtype=np.uint8).reshape(8, 8)
+
+    def _compute_tile(tile_bounds: tuple[int, int, int, int]) -> tuple[dict]:
+        """Return one halo-expanded post-processing output."""
+        recorded_bounds.append(tile_bounds)
+        return (
+            {
+                "task_type": "instance",
+                "seg_type": "instance",
+                "predictions": expanded_predictions,
+                "info_dict": {
+                    "box": np.empty((0, 4), dtype=np.int32),
+                    "centroid": np.empty((0, 2), dtype=np.float32),
+                    "contours": np.empty((0, 0, 2), dtype=np.int32),
+                    "prob": np.empty((0,), dtype=np.float32),
+                    "type": np.empty((0,), dtype=np.int32),
+                },
+            },
+        )
+
+    seg._compute_tile = _compute_tile
+    monkeypatch.setattr(
+        "tiatoolbox.models.engine.multi_task_segmentor.tqdm_dask_progress_bar",
+        lambda **kwargs: kwargs["write_tasks"],
+    )
+
+    output = seg._process_tile_mode(
+        probabilities=[da.zeros((10, 10, 1), chunks=(10, 10, 1))],
+        save_path=track_tmp_path / "halo.zarr",
+        memory_threshold=100,
+        return_predictions=(True,),
+        postproc_halo=2,
+    )
+
+    assert recorded_bounds == [(0, 0, 8, 8)]
+    assert len(output) == 1
+    predictions = output[0]["predictions"]
+    assert np.array_equal(predictions[2:6, 2:6], expanded_predictions[2:6, 2:6])
+    assert np.count_nonzero(predictions[:2, :]) == 0
+    assert np.count_nonzero(predictions[:, :2]) == 0
+    assert np.count_nonzero(predictions[6:, :]) == 0
+    assert np.count_nonzero(predictions[:, 6:]) == 0
+
+
 class FakeSeg(MultiTaskSegmentor):
     """Minimal subclass that allows us to override internals cleanly."""
 
@@ -1166,6 +1412,7 @@ def fake_store_probabilities(
         *_: Any,  # noqa: ANN401
         **__: Any,  # noqa: ANN401
     ) -> tuple[zarr.Array | None, da.Array | None]:
+        """Record unexpected probability-store calls during merge tests."""
         nonlocal called_store
         called_store = True
         return None, None
@@ -1612,26 +1859,31 @@ def test_post_save_json_store_deletes_empty_store(
     # ---- Proxy object that LOOKS like a zarr.Group ----
     class GroupProxy:
         def __init__(self: GroupProxy, group: zarr.Group, path: Path | str) -> None:
+            """Wrap a Zarr group with a path used by cleanup code."""
             self._group = group
             self.path = path
             self.store = group.store
 
         # Make isinstance(proxy, zarr.Group) return True
         @property
         def __class__(self: GroupProxy) -> type[zarr.Group]:
+            """Expose the wrapped object as a Zarr group for isinstance."""
             return zarr.Group
 
         # Delegate attribute access
         def __getattr__(
             self: GroupProxy, item: str
         ) -> zarr.Group | zarr.Array | str | int | float | Iterable[str]:
+            """Delegate unknown attributes to the wrapped Zarr group."""
             return getattr(self._group, item)
 
         # Delegate mapping behavior
         def keys(self: GroupProxy) -> Iterable[str]:
+            """Return keys from the wrapped Zarr group."""
             return self._group.keys()
 
         def __getitem__(self: GroupProxy, item: str) -> zarr.Group | zarr.Array:
+            """Return an item from the wrapped Zarr group."""
             return self._group[item]
 
     processed_predictions = GroupProxy(root, "dummy")
@@ -1640,6 +1892,7 @@ def __getitem__(self: GroupProxy, item: str) -> zarr.Group | zarr.Array:
     called = {"flag": False}
 
     def fake_rmtree(path: Path | str, *, ignore_errors: bool) -> None:  # noqa: ARG001
+        """Record that cleanup attempted to remove an empty Zarr store."""
         called["flag"] = True
 
     monkeypatch.setattr(shutil, "rmtree", fake_rmtree)
@@ -1723,6 +1976,7 @@ def fake_save_qupath_json(
         save_path: Path | None,  # noqa: ARG001
         qupath_json: dict[str, Any],
     ) -> dict[str, Any]:
+        """Return generated QuPath JSON instead of writing it to disk."""
         return qupath_json
 
     monkeypatch.setattr(
@@ -1759,6 +2013,7 @@ def _build_single_qupath_feature(
         scale_factor: tuple[float, float],
         class_colors: dict[int, Any],
     ) -> dict[str, Any]:
+        """Delegate feature construction to the production JSON store."""
         return DaskDelayedJSONStore._build_single_qupath_feature(
             self, i, class_dict, origin, scale_factor, class_colors
         )