improve tests

Author: Jiaqi-Lv

tests/engines/test_nucleus_detection_engine.py

@@ -4,9 +4,13 @@
 import shutil
 from collections.abc import Callable
 
+import pandas as pd
+import pytest
+
 from tiatoolbox.annotation.storage import SQLiteStore
 from tiatoolbox.models.engine.nucleus_detector import NucleusDetector
 from tiatoolbox.utils import env_detection as toolbox_env
+from tiatoolbox.wsicore.wsireader import WSIReader
 
 device = "cuda" if toolbox_env.has_gpu() else "cpu"
 
@@ -19,8 +23,91 @@ def _rm_dir(path: pathlib.Path) -> None:
 
 def check_output(path: pathlib.Path) -> None:
     """Check NucleusDetector output."""
-    store = SQLiteStore.open(path)
-    assert len(store.values()) == 281
+
+
+def test_nucleus_detection_nms_empty_dataframe() -> None:
+    """nucleus_detection_nms should return a copy for empty inputs."""
+    df = pd.DataFrame(columns=["x", "y", "type", "prob"])
+
+    result = NucleusDetector.nucleus_detection_nms(df, radius=3)
+
+    assert result.empty
+    assert result is not df
+    assert list(result.columns) == ["x", "y", "type", "prob"]
+
+
+def test_nucleus_detection_nms_invalid_radius() -> None:
+    """Radius must be strictly positive."""
+    df = pd.DataFrame({"x": [0], "y": [0], "type": [1], "prob": [0.9]})
+
+    with pytest.raises(ValueError, match="radius must be > 0"):
+        NucleusDetector.nucleus_detection_nms(df, radius=0)
+
+
+def test_nucleus_detection_nms_invalid_overlap_threshold() -> None:
+    """overlap_threshold must lie in (0, 1]."""
+    df = pd.DataFrame({"x": [0], "y": [0], "type": [1], "prob": [0.9]})
+
+    message = r"overlap_threshold must be in \(0\.0, 1\.0\], got 0"
+    with pytest.raises(ValueError, match=message):
+        NucleusDetector.nucleus_detection_nms(df, radius=1, overlap_threshold=0)
+
+
+def test_nucleus_detection_nms_suppresses_overlapping_detections() -> None:
+    """Lower-probability overlapping detections are removed."""
+    df = pd.DataFrame(
+        {
+            "x": [2, 0, 20],
+            "y": [1, 0, 20],
+            "type": [1, 1, 2],
+            "prob": [0.6, 0.9, 0.7],
+        }
+    )
+
+    result = NucleusDetector.nucleus_detection_nms(df, radius=5)
+
+    expected = pd.DataFrame(
+        {"x": [0, 20], "y": [0, 20], "type": [1, 2], "prob": [0.9, 0.7]}
+    )
+    pd.testing.assert_frame_equal(result.reset_index(drop=True), expected)
+
+
+def test_nucleus_detection_nms_suppresses_across_types() -> None:
+    """Overlapping detections of different types are also suppressed."""
+    df = pd.DataFrame(
+        {
+            "x": [0, 0, 20],
+            "y": [0, 0, 0],
+            "type": [1, 2, 1],
+            "prob": [0.6, 0.95, 0.4],
+        }
+    )
+
+    result = NucleusDetector.nucleus_detection_nms(df, radius=5)
+
+    expected = pd.DataFrame(
+        {"x": [0, 20], "y": [0, 0], "type": [2, 1], "prob": [0.95, 0.4]}
+    )
+    pd.testing.assert_frame_equal(result.reset_index(drop=True), expected)
+
+
+def test_nucleus_detection_nms_retains_non_overlapping_candidates() -> None:
+    """Detections with IoU below the threshold are preserved."""
+    df = pd.DataFrame(
+        {
+            "x": [0, 10],
+            "y": [0, 0],
+            "type": [1, 1],
+            "prob": [0.8, 0.5],
+        }
+    )
+
+    result = NucleusDetector.nucleus_detection_nms(df, radius=5, overlap_threshold=0.5)
+
+    expected = pd.DataFrame(
+        {"x": [0, 10], "y": [0, 0], "type": [1, 1], "prob": [0.8, 0.5]}
+    )
+    pd.testing.assert_frame_equal(result.reset_index(drop=True), expected)
 
 
 def test_nucleus_detector_wsi(remote_sample: Callable, tmp_path: pathlib.Path) -> None:
@@ -36,13 +123,50 @@ def test_nucleus_detector_wsi(remote_sample: Callable, tmp_path: pathlib.Path) -
         patch_mode=False,
         device=device,
         output_type="annotationstore",
-        auto_get_mask=True,
         memory_threshold=50,
         images=[mini_wsi_svs],
         save_dir=save_dir,
         overwrite=True,
     )
 
-    check_output(save_dir / "wsi4_512_512.db")
+    store = SQLiteStore.open(save_dir / "wsi4_512_512.db")
+    assert len(store.values()) == 281
+    store.close()
+
+    _rm_dir(save_dir)
+
+
+def test_nucleus_detector_patch(
+    remote_sample: Callable, tmp_path: pathlib.Path
+) -> None:
+    """Test for nucleus detection engine in patch mode."""
+    mini_wsi_svs = pathlib.Path(remote_sample("wsi4_512_512_svs"))
+
+    wsi_reader = WSIReader.open(mini_wsi_svs)
+    patch_1 = wsi_reader.read_rect((0, 0), (252, 252), resolution=0.5, units="mpp")
+    patch_2 = wsi_reader.read_rect((252, 252), (252, 252), resolution=0.5, units="mpp")
+
+    pretrained_model = "mapde-conic"
+
+    save_dir = tmp_path
+
+    nucleus_detector = NucleusDetector(model=pretrained_model)
+    _ = nucleus_detector.run(
+        patch_mode=True,
+        device=device,
+        output_type="annotationstore",
+        memory_threshold=50,
+        images=[patch_1, patch_2],
+        save_dir=save_dir,
+        overwrite=True,
+    )
+
+    store_1 = SQLiteStore.open(save_dir / "0.db")
+    assert len(store_1.values()) == 270
+    store_1.close()
+
+    store_2 = SQLiteStore.open(save_dir / "1.db")
+    assert len(store_2.values()) == 52
+    store_2.close()
 
     _rm_dir(save_dir)

tests/models/test_arch_mapde.py

@@ -53,6 +53,31 @@ def test_functionality(remote_sample: Callable) -> None:
 
     np.testing.assert_array_equal(xs[0:2], np.array([242, 192]))
     np.testing.assert_array_equal(ys[0:2], np.array([10, 13]))
+
+    patch = reader.read_bounds(
+        (0, 0, 252, 252),
+        resolution=0.50,
+        units="mpp",
+        coord_space="resolution",
+    )
+
+    model, weights_path = _load_mapde(name="mapde-conic")
+    patch = model.preproc(patch)
+    batch = torch.from_numpy(patch)[None]
+    output = model.infer_batch(model, batch, device=select_device(on_gpu=ON_GPU))
+    block_info = {
+        0: {
+            "array-location": [
+                [0, 1],
+                [0, 1],
+            ],  # dummy block to test no valid detections
+        }
+    }
+    output = model.postproc(output[0], block_info=block_info)
+    xs, ys, _, _ = NucleusDetector._centroid_maps_to_detection_records(output, None)
+    np.testing.assert_array_equal(xs, np.array([]))
+    np.testing.assert_array_equal(ys, np.array([]))
+
     Path(weights_path).unlink()
 
 

tests/models/test_arch_sccnn.py

@@ -60,7 +60,31 @@ def test_functionality(remote_sample: Callable) -> None:
         batch,
         device=select_device(on_gpu=env_detection.has_gpu()),
     )
-    output = model.postproc(output[0])
+    block_info = {
+        0: {
+            "array-location": [[0, 31], [0, 31]],
+        }
+    }
+    output = model.postproc(output[0], block_info=block_info)
     xs, ys, _, _ = NucleusDetector._centroid_maps_to_detection_records(output, None)
     np.testing.assert_array_equal(xs, np.array([7]))
     np.testing.assert_array_equal(ys, np.array([8]))
+
+    model = _load_sccnn(name="sccnn-conic")
+    output = model.infer_batch(
+        model,
+        batch,
+        device=select_device(on_gpu=env_detection.has_gpu()),
+    )
+    block_info = {
+        0: {
+            "array-location": [
+                [0, 1],
+                [0, 1],
+            ],  # dummy block to test no valid detections
+        }
+    }
+    output = model.postproc(output[0], block_info=block_info)
+    xs, ys, _, _ = NucleusDetector._centroid_maps_to_detection_records(output, None)
+    np.testing.assert_array_equal(xs, np.array([]))
+    np.testing.assert_array_equal(ys, np.array([]))