add from_array method to annotation array attrs

src/cellmap_schemas/annotation.py

@@ -14,13 +14,13 @@
 from typing import (
     Any,
     Generic,
-    List,
     Literal,
     Mapping,
     Optional,
     TypeVar,
-    Union,
 )
+import numpy as np
+from typing_extensions import Self
 from pydantic_zarr.v2 import GroupSpec, ArraySpec
 from pydantic import BaseModel, model_validator, field_serializer
 import zarr
@@ -111,7 +111,7 @@ class Label(BaseModel, extra="forbid"):
 
 
 class LabelList(BaseModel, extra="forbid"):
-    labels: List[Label]
+    labels: list[Label]
     annotation_type: AnnotationType = "semantic"
 
 
@@ -191,7 +191,7 @@ class SemanticSegmentation(BaseModel, extra="forbid"):
     """
 
     type: Literal["semantic_segmentation"] = "semantic_segmentation"
-    encoding: dict[Union[Possibility, Literal["present"]], int]
+    encoding: dict[Literal[Possibility, Literal["present"]], int]
 
 
 class InstanceSegmentation(BaseModel, extra="forbid"):
@@ -221,7 +221,7 @@ class InstanceSegmentation(BaseModel, extra="forbid"):
     encoding: dict[Possibility, int]
 
 
-AnnotationType = Union[SemanticSegmentation, InstanceSegmentation]
+AnnotationType = SemanticSegmentation | InstanceSegmentation
 
 TName = TypeVar("TName", bound=str)
 
@@ -247,16 +247,52 @@ class AnnotationArrayAttrs(BaseModel, Generic[TName]):
 
     class_name: TName
     # a mapping from values to frequencies
-    complement_counts: Optional[dict[Possibility, int]]
+    complement_counts: dict[Possibility, int] | None
     # a mapping from class names to values
     # this is array metadata because labels might disappear during downsampling
     annotation_type: AnnotationType
 
     @model_validator(mode="after")
     def check_encoding(self: "AnnotationArrayAttrs"):
-        assert set(self.annotation_type.encoding.keys()).issuperset((self.complement_counts.keys()))
+        if (
+            isinstance(self.annotation_type, SemanticSegmentation)
+            and self.complement_counts is not None
+        ):
+            assert set(self.annotation_type.encoding.keys()).issuperset(
+                (self.complement_counts.keys())
+            )
         return self
 
+    @classmethod
+    def from_array(
+        cls,
+        data: np.ndarray,
+        class_name: TName,
+        annotation_type: AnnotationType,
+        complement_counts: None | dict[Possibility, int] | Literal["auto"],
+    ) -> Self:
+        if complement_counts == "auto":
+            num_unknown = sum(data == annotation_type.encoding["unknown"])
+            num_absent = sum(data == annotation_type.encoding["absent"])
+            num_present = data.size - (num_unknown + num_absent)
+
+            if isinstance(annotation_type, SemanticSegmentation):
+                complement_counts_parsed = {"unknown": num_unknown, "absent": num_absent}
+            elif isinstance(annotation_type, InstanceSegmentation):
+                complement_counts_parsed = {
+                    "unknown": num_unknown,
+                    "absent": num_absent,
+                    "present": num_present,
+                }
+        else:
+            complement_counts_parsed = complement_counts
+
+        return cls(
+            class_name=class_name,
+            annotation_type=annotation_type,
+            complement_counts=complement_counts_parsed,
+        )
+
 
 class AnnotationGroupAttrs(BaseModel, Generic[TName]):
     """
@@ -327,11 +363,15 @@ class CropGroupAttrs(BaseModel, Generic[TName], validate_assignment=True):
     class_names: list[TName]
 
     @field_serializer("start_date")
-    def ser_end_date(value: date):
+    def ser_end_date(value: date) -> None | str:
+        if value is None:
+            return None
         return serialize_date(value)
 
     @field_serializer("end_date")
-    def ser_start_date(value: date):
+    def ser_start_date(value: date) -> None | str:
+        if value is None:
+            return None
         return serialize_date(value)