Adding annotation typings (#300)

* Adding annotation typings
This PR takes a first stab at standardizing the labels and annotation columns for our datasets with the hope to merge annotations across datasets easily.

* ruff format

* changing the labels to dicts

* format

* distinction between cell cycle and cell division labels

* ruff format

---------

Co-authored-by: Ziwen Liu <[email protected]>

Author: edyoshikun

viscy/data/cell_classification.py

@@ -11,6 +11,7 @@
 
 from viscy.data.hcs import _read_norm_meta
 from viscy.data.triplet import INDEX_COLUMNS
+from viscy.data.typing import AnnotationColumns
 
 
 class ClassificationDataset(Dataset):
@@ -23,7 +24,7 @@ def __init__(
         transform: Callable | None,
         initial_yx_patch_size: tuple[int, int],
         return_indices: bool = False,
-        label_column: str = "infection_state",
+        label_column: AnnotationColumns = "infection_state",
     ):
         self.plate = plate
         self.z_range = z_range

viscy/data/typing.py

@@ -1,4 +1,4 @@
-from typing import Callable, NamedTuple, Sequence, TypedDict, TypeVar
+from typing import Callable, Literal, NamedTuple, Sequence, TypedDict, TypeVar
 
 from torch import ShortTensor, Tensor
 
@@ -88,3 +88,39 @@ class TripletSample(TypedDict):
     positive: NotRequired[Tensor]
     negative: NotRequired[Tensor]
     index: NotRequired[TrackingIndex]
+
+
+# NOTE: these are the only columns that are allowed for the annotation dataframe.
+AnnotationColumns = Literal[
+    "infection_state",
+    "cell_division_state",
+    "cell_remodeling_state",
+    "cell_cycle_state",
+]
+
+
+# NOTE: The following labels are not mutable. They are used to map the labels to the integer values.
+LABEL_INFECTION_STATE = {"uninfected": 0, "infected": 1, "unknown": -1}
+
+LABEL_CELL_DIVISION_STATE = {
+    "interphase": 0,
+    "mitosis": 1,
+    "unknown": -1,
+}
+
+LABEL_CELL_CYCLE_STATE = {
+    "G1": 0,
+    "S": 1,
+    "G2": 2,
+    "prophase": 3,
+    "metaphase": 4,
+    "anaphase": 5,
+    "telophase": 6,
+    "unknown": -1,
+}
+
+LABEL_CELL_REMODELING_STATE = {
+    "no_remodel": 0,
+    "remodeling": 1,
+    "unknown": -1,
+}