add more elaborate tests + corresponding logic in from_array

Author: d-v-b

Diff for: src/cellmap_schemas/annotation.py

@@ -181,7 +181,7 @@ class SemanticSegmentation(BaseModel, extra="forbid"):
 
     type: Literal["semantic_segmentation"]
         Must be the string 'semantic_segmentation'.
-    encoding: dict[Union[Possibility, Literal["present"]], int]
+    encoding: dict[SemanticPossibility, int]
         This dict represents the mapping from possibilities to numeric values. The keys
         must be strings in the set `{'unknown', 'absent', 'present'}`, and the values
         must be numeric values contained in the array described by this metadata.
@@ -280,24 +280,36 @@ def from_array(
         if complement_counts == "auto":
             complement_counts_parsed: dict[SemanticPossibility, int] | dict[
                 InstancePossibility, int
-            ]
-            num_unknown = (array == annotation_type.encoding["unknown"]).sum()
-            num_absent = (array == annotation_type.encoding["absent"]).sum()
-            num_present = array.size - (num_unknown + num_absent)
+            ] = {}
+
+            num_unknown: int = 0
+            num_absent: int = 0
+            num_present: int = 0
+
+            if "unknown" in annotation_type.encoding:
+                num_unknown = (array == annotation_type.encoding["unknown"]).sum()
+
+            if "absent" in annotation_type.encoding:
+                num_absent = (array == annotation_type.encoding["absent"]).sum()
+
+            if "present" in annotation_type.encoding:
+                num_present = array.size - (num_unknown + num_absent)
 
             if isinstance(annotation_type, SemanticSegmentation):
-                complement_counts_parsed: dict[SemanticPossibility, int] = {
-                    "unknown": num_unknown,
-                    "absent": num_absent,
-                    "present": num_present,
-                }
+                if "unknown" in annotation_type.encoding:
+                    complement_counts_parsed["unknown"] = num_unknown
+                if "absent" in annotation_type.encoding:
+                    complement_counts_parsed["absent"] = num_absent
+                if "present" in annotation_type.encoding:
+                    complement_counts_parsed["present"] = num_present
+
             elif isinstance(annotation_type, InstanceSegmentation):
-                complement_counts_parsed: dict[InstancePossibility, int] = {
-                    "unknown": num_unknown,
-                    "absent": num_absent,
-                }
-        else:
-            complement_counts_parsed = complement_counts  # type: ignore
+                if "unknown" in annotation_type.encoding:
+                    complement_counts_parsed["unknown"] = num_unknown
+                if "absent" in annotation_type.encoding:
+                    complement_counts_parsed["absent"] = num_absent
+            else:
+                complement_counts_parsed = complement_counts  # type: ignore
 
         return cls(
             class_name=class_name,

Diff for: tests/test_annotation.py

@@ -9,17 +9,33 @@
     AnnotationGroupAttrs,
     CropGroup,
     CropGroupAttrs,
+    InstancePossibility,
+    InstanceSegmentation,
+    SemanticPossibility,
     SemanticSegmentation,
     wrap_attributes,
 )
 import numpy as np
 import zarr
 
 
-def test_cropgroup():
+@pytest.mark.parametrize(
+    "annotation_cls, encoding",
+    [
+        (SemanticSegmentation, {"absent": 0}),
+        (SemanticSegmentation, {"absent": 0, "present": 1}),
+        (SemanticSegmentation, {"absent": 0, "present": 1, "unknown": 3}),
+        (InstanceSegmentation, {"absent": 0}),
+        (InstanceSegmentation, {"absent": 0, "unknown": 3}),
+    ],
+)
+def test_cropgroup(
+    annotation_cls: type[SemanticSegmentation] | type[InstanceSegmentation],
+    encoding: dict[SemanticPossibility, int] | dict[InstancePossibility, int],
+):
     ClassNamesT = Literal["foo", "bar"]
-    class_names = ["foo", "bar"]
-    ann_type = SemanticSegmentation(encoding={"absent": 0})
+    class_names: list[Literal["foo", "bar"]] = ["foo", "bar"]
+    ann_type = annotation_cls(encoding=encoding)  # type: ignore
     arrays = [np.zeros(10) for class_name in class_names]
     crop_group_attrs = CropGroupAttrs[ClassNamesT](
         name="foo",
@@ -79,3 +95,24 @@ def test_cropgroup():
     )
     with pytest.raises(ValueError, match=match):
         CropGroup.from_zarr(stored)
+
+
+@pytest.mark.parametrize(
+    "annotation_cls, encoding",
+    [
+        (SemanticSegmentation, {"absent": 0}),
+        (SemanticSegmentation, {"absent": 0, "present": 1}),
+        (SemanticSegmentation, {"absent": 0, "present": 1, "unknown": 3}),
+        (InstanceSegmentation, {"absent": 0}),
+        (InstanceSegmentation, {"absent": 0, "unknown": 3}),
+    ],
+)
+def test_annotation_attrs_from_array(annotation_cls, encoding) -> None:
+    array = np.array([0, 1, 0, 3])
+    attrs = AnnotationArrayAttrs.from_array(
+        array,
+        class_name="foo",
+        annotation_type=annotation_cls(encoding=encoding),
+        complement_counts="auto",
+    )
+    assert attrs.complement_counts["absent"] == 2