🚧 Update script for annotation store

Author: shaneahmed

tests/engines/test_semantic_segmentor.py

@@ -87,6 +87,30 @@ def test_semantic_segmentor_patches(
     assert "probabilities" not in output.keys()  # noqa: SIM118
 
 
+def test_save_annotation_store(sample_patch1: Path, sample_patch2: Path, tmp_path: Path):
+    segmentor = SemanticSegmentor(
+        model="fcn-tissue_mask", batch_size=32, verbose=False, device=device
+    )
+
+    inputs = [Path(sample_patch1), Path(sample_patch2)]
+    output = segmentor.run(
+        images=inputs,
+        return_probabilities=False,
+        return_labels=False,
+        device=device,
+        patch_mode=True,
+        cache_mode=True,
+        save_dir=tmp_path / "output1",
+        output_type="annotationstore",
+    )
+
+    assert output == tmp_path / "output1" / "output.zarr"
+
+    output = zarr.open(output, mode="r")
+    assert 0.24 < np.mean(output["predictions"][:]) < 0.25
+    assert "probabilities" not in output.keys()  # noqa: SIM118
+
+
 def test_hovernet_dat() -> None:
     from tiatoolbox.utils.misc import store_from_dat
     from pathlib import Path

tiatoolbox/models/engine/semantic_segmentor_new.py

@@ -7,13 +7,13 @@
 
 import zarr
 from typing_extensions import Unpack
+from pathlib import Path
 
 from .patch_predictor import PatchPredictor, PredictorRunParams
 from ...utils.misc import dict_to_zarr, dict_to_store_semantic_segmentor
 
 if TYPE_CHECKING:  # pragma: no cover
     import os
-    from pathlib import Path
 
     import numpy as np
 

tiatoolbox/utils/misc.py

@@ -1258,6 +1258,7 @@ def dict_to_store_semantic_segmentor(
 
     """
     preds = patch_output["predictions"]
+    preds = preds[0]
     layer_list = np.unique(preds)
     layer_list = np.delete(layer_list, np.where(layer_list == 0))
     layer_info_dict = {}
@@ -1273,76 +1274,32 @@ def dict_to_store_semantic_segmentor(
             cv2.RETR_TREE,
             cv2.CHAIN_APPROX_NONE,
         )
-        for layer in contours:
-            coords = layer[:, 0, :]
+        for layer_ in contours:
+            coords = layer_[:, 0, :]
             layer_info_dict[count] = {
                 "contours": coords,
-                "type": class_dict[type_class],
+                "type": "mask",
             }
             count += 1
 
             origin = (0, 0)
-
-            annotations.append(
-                Annotation(
-                    geometry=make_valid_poly(
-                        feature2geometry(
+            scaled_coords = np.array([scale_factor * coords])
+            feature_geom = feature2geometry(
                             {
                                 "type": "Polygon",
-                                "coordinates": scale_factor * coords,
+                                "coordinates": scaled_coords,
                             },
-                        ),
+                        )
+            annotations.append(
+                Annotation(
+                    geometry=make_valid_poly(
+                        feature_geom,
                         origin=origin,
                     ),
-                    properties={},
+                    properties={"type": "mask"},
                 )
             )
 
-    # return layer_info_dict
-
-    # if "coordinates" not in patch_output:
-    #     # we cant create annotations without coordinates
-    #     msg = "Patch output must contain coordinates."
-    #     raise ValueError(msg)
-    #
-    # # get relevant keys
-    # class_probs = get_zarr_array(patch_output.get("probabilities", []))
-    # preds = get_zarr_array(patch_output.get("predictions", []))
-    #
-    # patch_coords = np.array(patch_output.get("coordinates", []))
-    # if not np.all(np.array(scale_factor) == 1):
-    #     patch_coords = patch_coords * (np.tile(scale_factor, 2))  # to baseline mpp
-    # patch_coords = patch_coords.astype(float)
-    # labels = patch_output.get("labels", [])
-    # # get classes to consider
-    # if len(class_probs) == 0:
-    #     classes_predicted = np.unique(preds).tolist()
-    # else:
-    #     classes_predicted = range(len(class_probs[0]))
-    #
-    # if class_dict is None:
-    #     # if no class dict create a default one
-    #     if len(class_probs) == 0:
-    #         class_dict = {i: i for i in np.unique(np.append(preds, labels)).tolist()}
-    #     else:
-    #         class_dict = {i: i for i in range(len(class_probs[0]))}
-    #
-    # # find what keys we need to save
-    # keys = ["predictions"]
-    # keys = keys + [key for key in ["probabilities", "labels"] if key in patch_output]
-    #
-    # # put patch predictions into a store
-    # annotations = patch_predictions_as_annotations(
-    #     preds,
-    #     keys,
-    #     class_dict,
-    #     class_probs,
-    #     patch_coords,
-    #     classes_predicted,
-    #     labels,
-    # )
-    #
-    # store = SQLiteStore()
     _ = store.append_many(annotations, [str(i) for i in range(len(annotations))])
 
     # # if a save director is provided, then dump store into a file