🎨 Improve structure of the test

Author: shaneahmed

tests/engines/test_nucleus_instance_segmentor.py

@@ -1,8 +1,8 @@
 """Test tiatoolbox.models.engine.nucleus_instance_segmentor."""
 
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from pathlib import Path
-from typing import Final
+from typing import Any, Final
 
 import numpy as np
 import torch
@@ -13,9 +13,55 @@
 from tiatoolbox.wsicore import WSIReader
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
+OutputType = dict[str, Any] | Any
 
 
-def test_functionality_patch_mode(  # noqa: PLR0915
+def assert_output_lengths(output: OutputType, expected_counts: Sequence[int]) -> None:
+    """Assert lengths of output dict fields against expected counts."""
+    for field in ["box", "centroid", "contour", "prob", "type"]:
+        for i, expected in enumerate(expected_counts):
+            assert len(output[field][i]) == expected, f"{field}[{i}] mismatch"
+
+
+def assert_output_equal(
+    output_a: OutputType,
+    output_b: OutputType,
+    fields: Sequence[str],
+    indices_a: Sequence[int],
+    indices_b: Sequence[int],
+) -> None:
+    """Assert equality of arrays across outputs for given fields/indices."""
+    for field in fields:
+        for i_a, i_b in zip(indices_a, indices_b, strict=False):
+            left = output_a[field][i_a]
+            right = output_b[field][i_b]
+            assert all(
+                np.array_equal(a, b) for a, b in zip(left, right, strict=False)
+            ), f"{field}[{i_a}] vs {field}[{i_b}] mismatch"
+
+
+def assert_predictions_and_boxes(
+    output: OutputType, expected_counts: Sequence[int], *, is_zarr: bool = False
+) -> None:
+    """Assert predictions maxima and box lengths against expected counts."""
+    # predictions maxima
+    for idx, expected in enumerate(expected_counts):
+        if is_zarr and idx == 2:
+            # zarr output doesn't store predictions for patch 2
+            continue
+        assert np.max(output["predictions"][idx][:]) == expected, (
+            f"predictions[{idx}] mismatch"
+        )
+
+    # box lengths
+    for idx, expected in enumerate(expected_counts):
+        if is_zarr and idx < 2:
+            # for zarr, compare boxes only for patches 0 and 1
+            continue
+        assert len(output["box"][idx]) == expected, f"box[{idx}] mismatch"
+
+
+def test_functionality_patch_mode(
     remote_sample: Callable, track_tmp_path: Path
 ) -> None:
     """Patch mode functionality test for nuclei instance segmentor."""
@@ -24,178 +70,105 @@ def test_functionality_patch_mode(  # noqa: PLR0915
     size = (256, 256)
     resolution = 0.25
     units: Final = "mpp"
+
     patch1 = mini_wsi.read_rect(
-        location=(0, 0),
-        size=size,
-        resolution=resolution,
-        units=units,
+        location=(0, 0), size=size, resolution=resolution, units=units
     )
     patch2 = mini_wsi.read_rect(
-        location=(512, 512),
-        size=size,
-        resolution=resolution,
-        units=units,
+        location=(512, 512), size=size, resolution=resolution, units=units
     )
-
-    # Test dummy input, should result in no output segmentation
     patch3 = np.zeros_like(patch1)
-
-    patches = np.stack(arrays=[patch1, patch2, patch3], axis=0)
+    patches = np.stack([patch1, patch2, patch3], axis=0)
 
     inst_segmentor = NucleusInstanceSegmentor(
-        batch_size=1,
-        num_workers=0,
-        model="hovernet_fast-pannuke",
+        batch_size=1, num_workers=0, model="hovernet_fast-pannuke"
     )
-    output = inst_segmentor.run(
-        images=patches,
-        patch_mode=True,
-        device=device,
-        output_type="dict",
+    output_dict = inst_segmentor.run(
+        images=patches, patch_mode=True, device=device, output_type="dict"
     )
 
-    assert np.max(output["predictions"][0][:]) == 41
-    assert np.max(output["predictions"][1][:]) == 17
-    assert np.max(output["predictions"][2][:]) == 0
-
-    assert len(output["box"][0]) == 41
-    assert len(output["box"][1]) == 17
-    assert len(output["box"][2]) == 0
-
-    assert len(output["centroid"][0]) == 41
-    assert len(output["centroid"][1]) == 17
-    assert len(output["centroid"][2]) == 0
-
-    assert len(output["contour"][0]) == 41
-    assert len(output["contour"][1]) == 17
-    assert len(output["contour"][2]) == 0
-
-    assert len(output["prob"][0]) == 41
-    assert len(output["prob"][1]) == 17
-    assert len(output["prob"][2]) == 0
+    expected_counts = [41, 17, 0]
 
-    assert len(output["type"][0]) == 41
-    assert len(output["type"][1]) == 17
-    assert len(output["type"][2]) == 0
+    assert_predictions_and_boxes(output_dict, expected_counts, is_zarr=False)
+    assert_output_lengths(output_dict, expected_counts)
 
-    output_ = output
-
-    output = inst_segmentor.run(
+    # Zarr output comparison
+    output_zarr = inst_segmentor.run(
         images=patches,
         patch_mode=True,
         device=device,
         output_type="zarr",
         save_dir=track_tmp_path / "patch_output_zarr",
     )
+    output_zarr = zarr.open(output_zarr, mode="r")
+    assert_predictions_and_boxes(output_zarr, expected_counts, is_zarr=True)
 
-    output = zarr.open(output, mode="r")
-
-    assert np.max(output["predictions"][0][:]) == 41
-    assert np.max(output["predictions"][1][:]) == 17
-
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(output["box"][0], output_["box"][0], strict=False)
+    assert_output_equal(
+        output_zarr,
+        output_dict,
+        fields=["box", "centroid", "contour", "prob", "type"],
+        indices_a=[0, 1, 2],
+        indices_b=[0, 1, 2],
     )
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(output["box"][1], output_["box"][1], strict=False)
-    )
-    assert len(output["box"][2]) == 0
 
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(output["centroid"][0], output_["centroid"][0], strict=False)
-    )
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(output["centroid"][1], output_["centroid"][1], strict=False)
-    )
-
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(output["contour"][0], output_["contour"][0], strict=False)
-    )
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(output["contour"][1], output_["contour"][1], strict=False)
-    )
-
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(output["prob"][0], output_["prob"][0], strict=False)
-    )
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(output["prob"][1], output_["prob"][1], strict=False)
-    )
-
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(output["type"][0], output_["type"][0], strict=False)
-    )
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(output["type"][1], output_["type"][1], strict=False)
-    )
-
-    inst_segmentor = NucleusInstanceSegmentor(
-        batch_size=1,
-        num_workers=0,
-        model="hovernet_fast-pannuke",
-    )
-    output = inst_segmentor.run(
+    # AnnotationStore output comparison
+    output_ann = inst_segmentor.run(
         images=patches,
         patch_mode=True,
         device=device,
         output_type="annotationstore",
         save_dir=track_tmp_path / "patch_output_annotationstore",
     )
+    assert len(output_ann) == 3
+    assert output_ann[0] == track_tmp_path / "patch_output_annotationstore" / "0.db"
 
-    assert output[0] == track_tmp_path / "patch_output_annotationstore" / "0.db"
-    assert len(output) == 3
-    store_ = SQLiteStore.open(output[0])
-    annotations_ = store_.values()
-    annotations_geometry_type = [
-        str(annotation_.geometry_type) for annotation_ in annotations_
-    ]
-    assert "Polygon" in annotations_geometry_type
-
-    annotations_list = list(annotations_)
-    ann_properties = [ann.properties for ann in annotations_list]
-
-    result = {}
-    for d in ann_properties:
-        for key, value in d.items():
-            result.setdefault(key, []).append(value)
-
-    polygons = [ann.geometry for ann in annotations_list]
-    result["contour"] = [list(poly.exterior.coords) for poly in polygons]
-
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(result["box"], output_["box"][0], strict=False)
-    )
-
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(result["centroid"], output_["centroid"][0], strict=False)
-    )
-
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(result["prob"], output_["prob"][0], strict=False)
-    )
-
-    assert all(
-        np.array_equal(a, b)
-        for a, b in zip(result["type"], output_["type"][0], strict=False)
-    )
-
-    assert all(
-        np.array_equal(
-            np.array(a[:-1], dtype=int),  # discard last point
-            np.array(b, dtype=int),
+    for patch_idx, db_path in enumerate(output_ann):
+        assert (
+            db_path
+            == track_tmp_path / "patch_output_annotationstore" / f"{patch_idx}.db"
         )
-        for a, b in zip(result["contour"], output_["contour"][0], strict=False)
-    )
+        store_ = SQLiteStore.open(db_path)
+        annotations_ = store_.values()
+        annotations_geometry_type = [
+            str(annotation_.geometry_type) for annotation_ in annotations_
+        ]
+        annotations_list = list(annotations_)
+        if expected_counts[patch_idx] > 0:
+            assert "Polygon" in annotations_geometry_type
+
+            # Build result dict from annotation properties
+            result = {}
+            for ann in annotations_list:
+                for key, value in ann.properties.items():
+                    result.setdefault(key, []).append(value)
+            result["contour"] = [
+                list(poly.exterior.coords)
+                for poly in (a.geometry for a in annotations_list)
+            ]
+
+            # wrap it to make it compatible to assert_output_lengths
+            result_ = {
+                field: [result[field]]
+                for field in ["box", "centroid", "contour", "prob", "type"]
+            }
+
+            # Lengths and equality checks for this patch
+            assert_output_lengths(result_, [expected_counts[patch_idx]])
+            assert_output_equal(
+                result_,
+                output_dict,
+                fields=["box", "centroid", "prob", "type"],
+                indices_a=[0],
+                indices_b=[patch_idx],
+            )
+
+            # Contour check (discard last point)
+            assert all(
+                np.array_equal(np.array(a[:-1], dtype=int), np.array(b, dtype=int))
+                for a, b in zip(
+                    result["contour"], output_dict["contour"][patch_idx], strict=False
+                )
+            )
+        else:
+            assert annotations_geometry_type == []
+            assert annotations_list == []