AA functions: add more stringent label/shape type compatibility checking (#9005)

These changes are meant to enforce several conditions:

1. A function should not be able to create shapes of type incompatible
   with the type of the label it declares (for example, if a
   function's label has type `rectangle`, it should not be able to create
   ellipses with that label).

2. A function should not be able to create shapes of type incompatible
   with the type of the label in the task being annotated.

3. If a function's declared label has a type incompatible with the type
   of the corresponding task label, then it should not run at all (since it
   would be impossible for it to output a shape with that label that
   wouldn't violate either condition 1 or 2).

Altogether, these restrictions ensure that we don't create any shapes in
a task that aren't compatible with that task's label types.

In addition, set explicit label types for the predefined functions.

Author: SpecLad

changelog.d/20250128_144218_roman_aa_fun_label_types.md

@@ -0,0 +1,5 @@
+### Added
+
+- \[SDK\] The shapes output by auto-annotation functions are now checked
+  for compatibility with the function's and the task's label specs
+  (<https://github.com/cvat-ai/cvat/pull/9005>)

cvat-sdk/cvat_sdk/auto_annotation/driver.py

@@ -56,12 +56,28 @@ class _AnnotationMapper:
     class _LabelIdMapping:
         id: int
         sublabels: Mapping[int, Optional[_AnnotationMapper._SublabelIdMapping]]
-        expected_num_elements: int = 0
+        expected_num_elements: int
+        expected_type: str
 
     _SpecIdMapping: TypeAlias = Mapping[int, Optional[_LabelIdMapping]]
 
     _spec_id_mapping: _SpecIdMapping
 
+    def _get_expected_function_output_type(self, fun_label, ds_label):
+        fun_output_type = getattr(fun_label, "type", "any")
+        if fun_output_type == "any":
+            return ds_label.type
+
+        if self._conv_mask_to_poly and fun_output_type == "mask":
+            fun_output_type = "polygon"
+
+        if not self._are_label_types_compatible(fun_output_type, ds_label.type):
+            raise BadFunctionError(
+                f"label {fun_label.name!r} has type {fun_output_type!r} in the function,"
+                f" but {ds_label.type!r} in the dataset"
+            )
+        return fun_output_type
+
     def _build_label_id_mapping(
         self,
         fun_label: models.ILabel,
@@ -70,36 +86,37 @@ def _build_label_id_mapping(
         label_nm: _LabelNameMapping,
         allow_unmatched_labels: bool,
     ) -> Optional[_LabelIdMapping]:
-        sl_map = {}
-
-        if getattr(fun_label, "sublabels", []):
-            ds_sublabels_by_name = {ds_sl.name: ds_sl for ds_sl in ds_label.sublabels}
-
-            def sublabel_mapping(fun_sl: models.ILabel) -> Optional[int]:
-                sublabel_nm = label_nm.map_sublabel(fun_sl.name)
-                if sublabel_nm is None:
-                    return None
-
-                ds_sl = ds_sublabels_by_name.get(sublabel_nm.name)
-                if not ds_sl:
-                    if not allow_unmatched_labels:
-                        raise BadFunctionError(
-                            f"sublabel {fun_sl.name!r} of label {fun_label.name!r} is not in dataset"
-                        )
-
-                    self._logger.info(
-                        "sublabel %r of label %r is not in dataset; any annotations using it will be ignored",
-                        fun_sl.name,
-                        fun_label.name,
+        ds_sublabels_by_name = {ds_sl.name: ds_sl for ds_sl in ds_label.sublabels}
+
+        def sublabel_mapping(fun_sl: models.ILabel) -> Optional[int]:
+            sublabel_nm = label_nm.map_sublabel(fun_sl.name)
+            if sublabel_nm is None:
+                return None
+
+            ds_sl = ds_sublabels_by_name.get(sublabel_nm.name)
+            if not ds_sl:
+                if not allow_unmatched_labels:
+                    raise BadFunctionError(
+                        f"sublabel {fun_sl.name!r} of label {fun_label.name!r} is not in dataset"
                     )
-                    return None
 
-                return ds_sl.id
+                self._logger.info(
+                    "sublabel %r of label %r is not in dataset; any annotations using it will be ignored",
+                    fun_sl.name,
+                    fun_label.name,
+                )
+                return None
 
-            sl_map = {fun_sl.id: sublabel_mapping(fun_sl) for fun_sl in fun_label.sublabels}
+            return ds_sl.id
 
         return self._LabelIdMapping(
-            ds_label.id, sublabels=sl_map, expected_num_elements=len(ds_label.sublabels)
+            ds_label.id,
+            sublabels={
+                fun_sl.id: sublabel_mapping(fun_sl)
+                for fun_sl in getattr(fun_label, "sublabels", [])
+            },
+            expected_num_elements=len(ds_label.sublabels),
+            expected_type=self._get_expected_function_output_type(fun_label, ds_label),
         )
 
     def _build_spec_id_mapping(
@@ -254,6 +271,12 @@ def _remap_shape(self, shape: models.LabeledShapeRequest, ds_frame: int) -> bool
 
         shape.label_id = label_id_mapping.id
 
+        if not self._are_label_types_compatible(shape.type.value, label_id_mapping.expected_type):
+            raise BadFunctionError(
+                f"function output shape of type {shape.type.value!r}"
+                f" (expected {label_id_mapping.expected_type!r})"
+            )
+
         if shape.type.value == "mask" and self._conv_mask_to_poly:
             raise BadFunctionError("function output mask shape despite conv_mask_to_poly=True")
 
@@ -269,6 +292,11 @@ def _remap_shape(self, shape: models.LabeledShapeRequest, ds_frame: int) -> bool
     def validate_and_remap(self, shapes: list[models.LabeledShapeRequest], ds_frame: int) -> None:
         shapes[:] = [shape for shape in shapes if self._remap_shape(shape, ds_frame)]
 
+    @staticmethod
+    def _are_label_types_compatible(source_type: str, destination_type: str) -> bool:
+        assert source_type != "any"
+        return destination_type == "any" or destination_type == source_type
+
 
 @attrs.frozen(kw_only=True)
 class _DetectionFunctionContextImpl(DetectionFunctionContext):

cvat-sdk/cvat_sdk/auto_annotation/functions/_torchvision.py

@@ -10,6 +10,8 @@
 
 
 class TorchvisionFunction:
+    _label_type = "any"
+
     def __init__(self, model_name: str, weights_name: str = "DEFAULT", **kwargs) -> None:
         weights_enum = torchvision.models.get_model_weights(model_name)
         self._weights = weights_enum[weights_name]
@@ -21,7 +23,7 @@ def __init__(self, model_name: str, weights_name: str = "DEFAULT", **kwargs) ->
     def spec(self) -> cvataa.DetectionFunctionSpec:
         return cvataa.DetectionFunctionSpec(
             labels=[
-                cvataa.label_spec(cat, i)
+                cvataa.label_spec(cat, i, type=self._label_type)
                 for i, cat in enumerate(self._weights.meta["categories"])
                 if cat != "N/A"
             ]

cvat-sdk/cvat_sdk/auto_annotation/functions/torchvision_detection.py

@@ -11,6 +11,8 @@
 
 
 class _TorchvisionDetectionFunction(TorchvisionFunction):
+    _label_type = "rectangle"
+
     def detect(
         self, context: cvataa.DetectionFunctionContext, image: PIL.Image.Image
     ) -> list[models.LabeledShapeRequest]:

cvat-sdk/cvat_sdk/auto_annotation/functions/torchvision_instance_segmentation.py

@@ -50,6 +50,8 @@ def _generate_shapes(
 
 
 class _TorchvisionInstanceSegmentationFunction(TorchvisionFunction):
+    _label_type = "mask"
+
     def detect(
         self, context: cvataa.DetectionFunctionContext, image: PIL.Image.Image
     ) -> list[models.LabeledShapeRequest]:

site/content/en/docs/api_sdk/sdk/auto-annotation.md

@@ -63,8 +63,9 @@ class TorchvisionDetectionFunction:
         # describe the annotations
         return cvataa.DetectionFunctionSpec(
             labels=[
-                cvataa.label_spec(cat, i)
-                for i, cat in enumerate(self._weights.meta['categories'])
+                cvataa.label_spec(cat, i, type="rectangle")
+                for i, cat in enumerate(self._weights.meta["categories"])
+                if cat != "N/A"
             ]
         )