Fixes annotation filtering (#246)

NKI-AI · Aug 1, 2024 · 1b9a2fc · 1b9a2fc
1 parent 1f05317
commit 1b9a2fc
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/dlup/annotations.py b/dlup/annotations.py
@@ -650,13 +650,14 @@ def filter(self, labels: str | list[str] | tuple[str]) -> None:
         None
         """
         self._available_classes = []
-        self._layers = []
+        self._new_layers = []
         _labels = [labels] if isinstance(labels, str) else labels
         for layer in self._layers:
             if layer.label in _labels:
                 self._available_classes += [layer.annotation_class]
-                self._layers += [layer]
+                self._new_layers += [layer]
 
+        self._layers = self._new_layers
         self._str_tree = STRtree(self._layers)
 
     @property

diff --git a/tests/test_annotations.py b/tests/test_annotations.py
@@ -205,3 +205,12 @@ def test_point_pickling(self):
             pickled_point_file.seek(0)
             loaded_point = pickle.load(pickled_point_file)
         assert dlup_point == loaded_point
+
+    def test_annotation_filter(self):
+        annotations = self.asap_annotations.copy()
+        annotations.filter(["healthy glands"])
+        assert len(annotations._layers) == 1
+        assert annotations.available_classes[0].label == "healthy glands"
+
+        annotations.filter(["non-existing"])
+        assert len(annotations._layers) == 0