fix generic channel_names in ChannelAxis

Author: FynnBe

bioimageio/spec/model/v0_5.py

@@ -21,16 +21,16 @@
 from bioimageio.spec._internal.types import DeprecatedLicenseId as DeprecatedLicenseId
 from bioimageio.spec._internal.types import FileSource as FileSource
 from bioimageio.spec._internal.types import Identifier as Identifier
-from bioimageio.spec._internal.types import IdentifierStr, LowerCaseIdentifierStr
 from bioimageio.spec._internal.types import LicenseId as LicenseId
+from bioimageio.spec._internal.types import LowerCaseIdentifierStr
 from bioimageio.spec._internal.types import NotEmpty as NotEmpty
 from bioimageio.spec._internal.types import RdfContent as RdfContent
 from bioimageio.spec._internal.types import RelativeFilePath as RelativeFilePath
 from bioimageio.spec._internal.types import Sha256 as Sha256
 from bioimageio.spec._internal.types import Unit as Unit
 from bioimageio.spec._internal.types import Version as Version
 from bioimageio.spec._internal.types.field_validation import AfterValidator
-from bioimageio.spec._internal.validation_context import InternalValidationContext
+from bioimageio.spec._internal.validation_context import InternalValidationContext, get_internal_validation_context
 from bioimageio.spec.dataset.v0_3 import Dataset as Dataset
 from bioimageio.spec.dataset.v0_3 import LinkedDataset as LinkedDataset
 from bioimageio.spec.generic.v0_3 import Attachment as Attachment
@@ -185,6 +185,20 @@ class AxisBase(NodeWithExplicitlySetFields, frozen=True):
     id: AxisId
     """An axis id unique across all axes of one tensor."""
 
+    @model_validator(mode="before")
+    @classmethod
+    def convert_name_to_id(cls, data: Dict[str, Any], info: ValidationInfo):
+        context = get_internal_validation_context(info.context)
+        if (
+            "name" in data
+            and "id" not in data
+            and "original_format" in context
+            and context["original_format"].release[:2] == (0, 4)
+        ):
+            data["id"] = data.pop("name")
+
+        return data
+
     description: Annotated[str, MaxLen(128)] = ""
 
     __hash__ = NodeWithExplicitlySetFields.__hash__
@@ -205,20 +219,23 @@ class BatchAxis(AxisBase, frozen=True):
     otherwise (the default) it may be chosen arbitrarily depending on available memory"""
 
 
-ChannelName = Annotated[IdentifierStr, StringConstraints(min_length=3, max_length=16, pattern=r"^.*\{i\}.*$")]
+GenericChannelName = Annotated[str, StringConstraints(min_length=3, max_length=16, pattern=r"^.*\{i\}.*$")]
 
 
 class ChannelAxis(AxisBase, frozen=True):
     type: Literal["channel"] = "channel"
     id: AxisId = AxisId("channel")
-    channel_names: Union[Tuple[ChannelName, ...], ChannelName] = "channel{i}"
-    size: Union[Annotated[int, Gt(0)], SizeReference, Literal["#channel_names"]] = "#channel_names"
+    channel_names: Union[Tuple[Identifier, ...], Identifier, GenericChannelName] = "channel{i}"
+    size: Union[Annotated[int, Gt(0)], SizeReference] = "#channel_names"  # type: ignore
 
     @model_validator(mode="before")
     @classmethod
     def set_size_or_channel_names(cls, data: Dict[str, Any]):
         channel_names: Union[Any, Sequence[Any]] = data.get("channel_names", "channel{i}")
         size = data.get("size", "#channel_names")
+        if size == "#channel_names" and channel_names == "channel{i}":
+            raise ValueError("Channel dimension has unknown size. Please specify `size` or `channel_names`.")
+
         if (
             size == "#channel_names"
             and not isinstance(channel_names, str)
@@ -231,13 +248,6 @@ def set_size_or_channel_names(cls, data: Dict[str, Any]):
 
         return data
 
-    @model_validator(mode="after")
-    def validate_size_is_known(self):
-        if self.size == "#channel_names":
-            raise ValueError("Channel dimension has unknown size. Please specify `size` or `channel_names`.")
-
-        return self
-
 
 class IndexTimeSpaceAxisBase(AxisBase, frozen=True):
     size: Annotated[
@@ -948,14 +958,6 @@ def check_entries(self, info: ValidationInfo) -> Self:
         return self
 
 
-# def get_default_partial_inputs():
-#     return (
-#         InputTensor(axes=(BatchAxis(),), test_tensor=HttpUrl("https://example.com/test.npy")).model_dump(
-#             exclude_unset=False, exclude={"axes", "test_tensor"}
-#         ),
-#     )
-
-
 class Model(GenericBaseNoSource, frozen=True, title="bioimage.io model specification"):
     """Specification of the fields used in a bioimage.io-compliant RDF to describe AI models with pretrained weights.
     These fields are typically stored in a YAML file which we call a model resource description file (model RDF).
@@ -1070,8 +1072,21 @@ def _validate_axis(
                     f"Self-referencing not allowed for {field_name}[{i}].axes[{a}].size.reference: "
                     f"{axis.size.reference}"
                 )
-            if axis.type == "channel" and valid_independent_refs[axis.size.reference][1].type != "channel":
-                raise ValueError("A channel axis' size may only reference another fixed size channel axis.")
+            if axis.type == "channel":
+                if valid_independent_refs[axis.size.reference][1].type != "channel":
+                    raise ValueError("A channel axis' size may only reference another fixed size channel axis.")
+                if isinstance(axis.channel_names, str) and "{i}" in axis.channel_names:
+                    ref_size = valid_independent_refs[axis.size.reference][2]
+                    assert isinstance(
+                        ref_size, int
+                    ), "channel axis ref (another channel axis) has to specify fixed size"
+                    generated_channel_names = tuple(
+                        Identifier(axis.channel_names.format(i=i)) for i in range(1, ref_size + 1)
+                    )
+                    axis.model_config["frozen"] = False
+                    axis.channel_names = generated_channel_names  # type: ignore
+                    axis.model_config["frozen"] = True
+
             if (ax_unit := getattr(axis, "unit", None)) != (
                 ref_unit := getattr(valid_independent_refs[axis.size.reference][1], "unit", None)
             ):
@@ -1090,8 +1105,20 @@ def _validate_axis(
                 raise ValueError(f"Invalid tensor axis reference at {field_name}[{i}].axes[{a}].size: {axis.size}.")
             if axis.size in (axis.id, f"{tensor_id}.{axis.id}"):
                 raise ValueError(f"Self-referencing not allowed for {field_name}[{i}].axes[{a}].size: {axis.size}.")
-            if axis.type == "channel" and valid_independent_refs[axis.size][1].type != "channel":
-                raise ValueError("A channel axis' size may only reference another fixed size channel axis.")
+            if axis.type == "channel":
+                if valid_independent_refs[axis.size][1].type != "channel":
+                    raise ValueError("A channel axis' size may only reference another fixed size channel axis.")
+                if isinstance(axis.channel_names, str) and "{i}" in axis.channel_names:
+                    ref_size = valid_independent_refs[axis.size][2]
+                    assert isinstance(
+                        ref_size, int
+                    ), "channel axis ref (another channel axis) has to specify fixed size"
+                    generated_channel_names = tuple(
+                        Identifier(axis.channel_names.format(i=i)) for i in range(1, ref_size + 1)
+                    )
+                    axis.model_config["frozen"] = False
+                    axis.channel_names = generated_channel_names  # type: ignore
+                    axis.model_config["frozen"] = True
 
     license: Annotated[
         Union[LicenseId, DeprecatedLicenseId],