add docstrings

Author: measty

tests/engines/test_multi_task_segmentor.py

@@ -1273,6 +1273,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
@@ -1719,26 +1720,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")
@@ -1747,6 +1753,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)
@@ -1830,6 +1837,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(
@@ -1866,6 +1874,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
         )

tests/models/test_arch_cerberus.py

@@ -41,6 +41,7 @@ def _mock_torch_load(
         *_args: object,
         **_kwargs: object,
     ) -> dict[str, dict[str, torch.Tensor]]:
+        """Return a synthetic Cerberus checkpoint for load-weight tests."""
         return checkpoint
 
     monkeypatch.setattr(torch, "load", _mock_torch_load)
@@ -63,6 +64,7 @@ def _mock_torch_load(
         *_args: object,
         **_kwargs: object,
     ) -> dict[str, dict[str, torch.Tensor]]:
+        """Return a synthetic Cerberus checkpoint for registry loading."""
         return checkpoint
 
     monkeypatch.setattr(torch, "load", _mock_torch_load)
@@ -153,6 +155,7 @@ def _mock_post_process(
         tissue_mode: str,
         ds_factor: float,
     ) -> tuple[np.ndarray, np.ndarray | None]:
+        """Return deterministic task maps for Cerberus postproc testing."""
         calls.append((tissue_mode, raw_map.shape, idx_dict, ds_factor))
         inst_map = np.zeros(output_shape, dtype=np.int32)
         type_map = np.ones(output_shape, dtype=np.uint8)
@@ -172,6 +175,7 @@ def _mock_get_instance_info(
         offset: tuple[int, int],
         verbose: object,
     ) -> dict[int, dict]:
+        """Return deterministic instance metadata for Cerberus postproc tests."""
         assert offset == (7, 11)
         assert verbose is False
         type_value = 0 if type_map is None else int(type_map[inst_map > 0][0])

tiatoolbox/models/architecture/cerberus/backbone/resnet.py

@@ -83,6 +83,7 @@ def _make_layer(
         blocks: int,
         stride: int = 1,
     ) -> nn.Sequential:
+        """Build one ResNet stage with optional downsampling."""
         downsample = None
         if stride != 1 or self.inplanes != planes * BasicBlock.expansion:
             downsample = nn.Sequential(

tiatoolbox/models/architecture/cerberus/model.py

@@ -185,6 +185,7 @@ def postproc(
 
 
 def _strip_dataparallel_prefix(state: dict) -> dict:
+    """Remove ``module.`` prefixes from DataParallel checkpoint keys."""
     if all(key.split(".")[0] == "module" for key in state):
         return {".".join(key.split(".")[1:]): value for key, value in state.items()}
     return state
@@ -194,6 +195,7 @@ def _crop_center_tensor(
     tensor: torch.Tensor,
     output_shape: tuple[int, int],
 ) -> torch.Tensor:
+    """Crop a BHWC tensor to the requested center output shape."""
     h, w = tensor.shape[1:3]
     out_h, out_w = output_shape
     top = max((h - out_h) // 2, 0)
@@ -204,6 +206,7 @@ def _crop_center_tensor(
 def _build_tissue_raw_map(
     head_map: dict[str, np.ndarray], tissue_name: str
 ) -> tuple[np.ndarray, dict[str, list[int]]]:
+    """Combine Cerberus heads for one tissue into a raw postproc map."""
     idx_dict = {}
     maps = []
     start = 0
@@ -223,6 +226,7 @@ def _build_tissue_raw_map(
 
 
 def _inst_dict_for_dask_processing(inst_info_dict: dict, *, is_dask: bool) -> dict:
+    """Convert instance metadata into arrays with optional Dask wrapping."""
     if not inst_info_dict:
         output = {
             "box": np.empty((0, 4), dtype=np.int32),

tiatoolbox/models/architecture/cerberus/postproc.py

@@ -26,6 +26,7 @@ class PostProcInstErodedContourMap:
 
     @staticmethod
     def _proc_gland(inst_fg: np.ndarray, ds_factor: float = 1.0) -> np.ndarray:
+        """Extract labelled gland instances from inner and contour maps."""
         ksize = int((11 - 1) * ds_factor)
         k_disk = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ksize, ksize))
 
@@ -44,6 +45,7 @@ def _proc_gland(inst_fg: np.ndarray, ds_factor: float = 1.0) -> np.ndarray:
 
     @staticmethod
     def _proc_lumen(inst_fg: np.ndarray, ds_factor: float = 1.0) -> np.ndarray:
+        """Extract labelled lumen instances from inner and contour maps."""
         ksize = int((3 - 1) * ds_factor)
         k_disk = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ksize, ksize))
 
@@ -62,6 +64,7 @@ def _proc_lumen(inst_fg: np.ndarray, ds_factor: float = 1.0) -> np.ndarray:
 
     @staticmethod
     def _proc_nuclei(inst_fg: np.ndarray, ds_factor: float = 1.0) -> np.ndarray:
+        """Extract labelled nuclei instances from inner and contour maps."""
         _ = ds_factor
         k_disk = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
 

tiatoolbox/models/engine/multi_task_segmentor.py

@@ -4104,6 +4104,7 @@ def _post_save_json_store(
     save_path: Path | None,
     **kwargs: Unpack[MultiTaskSegmentorRunParams],
 ) -> None:
+    """Clean temporary JSON-store data and report unsupported probability saves."""
     for key in keys_to_compute:
         del processed_predictions[key]