Bugfix: Model storage named after template instead of manifest (#1058)

Author: maxxgx

interactive_ai/libs/iai_core_py/iai_core/utils/project_builder.py

@@ -64,24 +64,23 @@ class ProjectBuilder:
     """
 
     @staticmethod
-    def get_default_model_template_by_task_type(task_type: TaskType) -> ModelTemplate:
+    def get_default_model_template_by_task_type(
+        task_type: TaskType, default_models_per_task: dict[str, str]
+    ) -> ModelTemplate:
         """
         Get the default ModelTemplate which belongs to the task type.
 
         :param task_type: the task type which requires a model template
+        :param default_models_per_task: a dictionary mapping task type names to default model template IDs
         :return: the ModelTemplate associated with the task type
         """
-        if task_type in [TaskType.DATASET, TaskType.CROP]:
-            default_model_template = ModelTemplateList().get_by_id(task_type.name.lower())
-        else:
-            default_model_template = next(
-                (
-                    model_template
-                    for model_template in ModelTemplateList().get_all()
-                    if model_template.task_type == task_type and model_template.is_default_for_task
-                ),
-                NullModelTemplate(),
-            )
+        task_type_name = task_type.name.lower()
+        model_template_id = (
+            task_type_name
+            if task_type in [TaskType.DATASET, TaskType.CROP]
+            else default_models_per_task[task_type_name]
+        )
+        default_model_template = ModelTemplateList().get_by_id(model_template_id)
 
         if isinstance(default_model_template, NullModelTemplate):
             raise ModelTemplateError("A NullModelTemplate was created.")
@@ -565,6 +564,7 @@ def build_full_project(
         creator_id: str,
         parser_class: type[ProjectParser],
         parser_kwargs: dict,
+        default_models_per_task: dict[str, str],
     ) -> tuple[Project, LabelSchema, dict[str, LabelSchemaView]]:
         """
         This does NOT save anything to the database.
@@ -583,6 +583,7 @@ def build_full_project(
         :param creator_id: the ID of who created the project
         :param parser_class: a parser which will be used to get relevant information from the REST response
         :param parser_kwargs: arguments to pass to the parser for initialization
+        :param default_models_per_task: a dictionary mapping task type names to default model template IDs
         :return: the project, the label schema, and a mapping of task to label schema view
         """
         is_keypoint_detection_enabled = FeatureFlagProvider.is_enabled(FEATURE_FLAG_KEYPOINT_DETECTION)
@@ -608,7 +609,9 @@ def build_full_project(
             custom_labels_names = parser.get_custom_labels_names_by_task(
                 task_name=task_name,
             )
-            model_template = cls.get_default_model_template_by_task_type(task_type=task_type)
+            model_template = cls.get_default_model_template_by_task_type(
+                task_type=task_type, default_models_per_task=default_models_per_task
+            )
             task_node = cls._build_task_node(
                 project_id=project_id,
                 task_name=task_name,
@@ -1060,7 +1063,11 @@ class PersistedProjectBuilder(ProjectBuilder):
 
     @classmethod
     def build_full_project(
-        cls, creator_id: str, parser_class: type[ProjectParser], parser_kwargs: dict
+        cls,
+        creator_id: str,
+        parser_class: type[ProjectParser],
+        parser_kwargs: dict,
+        default_models_per_task: dict[str, str],
     ) -> tuple[Project, LabelSchema, dict[str, LabelSchemaView]]:
         """
         Overrides the method from the super ProjectBuilder
@@ -1070,13 +1077,19 @@ def build_full_project(
         :param creator_id: the ID of who created the project
         :param parser_class: a parser which will be used to get relevant information from the REST response
         :param parser_kwargs: arguments to pass to the parser for initialization
+        :param default_models_per_task: a dictionary mapping task type names to default model template IDs
         :return: the project, the label schema, and a mapping of task to label schema view
         """
         (
             project,
             label_schema,
             task_to_label_schema_view,
-        ) = super().build_full_project(creator_id=creator_id, parser_class=parser_class, parser_kwargs=parser_kwargs)
+        ) = super().build_full_project(
+            creator_id=creator_id,
+            parser_class=parser_class,
+            parser_kwargs=parser_kwargs,
+            default_models_per_task=default_models_per_task,
+        )
         return project, label_schema, task_to_label_schema_view
 
     @classmethod

interactive_ai/libs/iai_core_py/tests/utils/test_project_builder.py

@@ -99,6 +99,7 @@ def test_detection_label_schema_from_pipeline(
                 creator_id="Geti",
                 parser_class=CustomTestProjectParser,
                 parser_kwargs=fxt_detection_classification_project_data,
+                default_models_per_task={},  # not necessary as get_default_model_template_by_task_type is mocked
             )
         detection_label_group = label_schema.get_label_group_by_name("Default Detection")
         classification_label_group = label_schema.get_label_group_by_name("Default Classification")
@@ -127,6 +128,7 @@ def test_anomaly_label_schema_from_pipeline(self, fxt_anomaly_project_data):
                 creator_id="Geti",
                 parser_class=CustomTestProjectParser,
                 parser_kwargs=fxt_anomaly_project_data,
+                default_models_per_task={},  # not necessary as get_default_model_template_by_task_type is mocked
             )
 
         anomaly_label_group = label_schema.get_label_group_by_name("default - anomaly")
@@ -176,6 +178,7 @@ def test_label_schema_with_hierarchy_task_chain_from_pipeline(
                 creator_id="Geti",
                 parser_class=CustomTestProjectParser,
                 parser_kwargs=fxt_label_hierarchy_task_chain_project_data,
+                default_models_per_task={},  # not necessary as get_default_model_template_by_task_type is mocked
             )
 
         expected_group_names = {
@@ -213,6 +216,7 @@ def test_label_schema_with_hierarchy_from_pipeline(
                 creator_id="Geti",
                 parser_class=CustomTestProjectParser,
                 parser_kwargs=fxt_hierarchy_classification_project_data_2,
+                default_models_per_task={},  # not necessary as get_default_model_template_by_task_type is mocked
             )
 
         expected_group_names = {
@@ -268,6 +272,7 @@ def test_label_schema_for_binary_pipeline(
                 creator_id="Geti",
                 parser_class=CustomTestProjectParser,
                 parser_kwargs=fxt_binary_classification_project_data,
+                default_models_per_task={},  # not necessary as get_default_model_template_by_task_type is mocked
             )
 
         label_groups = label_schema.get_groups(include_empty=True)
@@ -291,6 +296,7 @@ def test_label_schema_for_multiclass_pipeline(
                 creator_id="Geti",
                 parser_class=CustomTestProjectParser,
                 parser_kwargs=fxt_multiclass_classification_project_data,
+                default_models_per_task={},  # not necessary as get_default_model_template_by_task_type is mocked
             )
         label_groups = label_schema.get_groups(include_empty=True)
         labels = label_schema.get_labels(include_empty=True)
@@ -317,6 +323,7 @@ def test_label_schema_for_multilabel_pipeline(
                 creator_id="Geti",
                 parser_class=CustomTestProjectParser,
                 parser_kwargs=fxt_multilabel_classification_project_data,
+                default_models_per_task={},  # not necessary as get_default_model_template_by_task_type is mocked
             )
 
         label_groups = label_schema.get_groups(include_empty=True)
@@ -339,7 +346,10 @@ def test_label_schema_for_segmentation_pipeline(
             ],
         ):
             project, label_schema, task_schema = ProjectBuilder.build_full_project(
-                creator_id="Geti", parser_class=CustomTestProjectParser, parser_kwargs=fxt_segmentation_project_data
+                creator_id="Geti",
+                parser_class=CustomTestProjectParser,
+                parser_kwargs=fxt_segmentation_project_data,
+                default_models_per_task={},  # not necessary as get_default_model_template_by_task_type is mocked
             )
 
         label_groups = label_schema.get_groups(include_empty=True)
@@ -391,6 +401,7 @@ def test_detection_to_segmentation_labels_from_pipeline(
                 creator_id="Geti",
                 parser_class=CustomTestProjectParser,
                 parser_kwargs=fxt_detection_to_segmentation_project_data,
+                default_models_per_task={},  # not necessary as get_default_model_template_by_task_type is mocked
             )
         expected_group_names = {
             "Default Detection",

interactive_ai/services/director/tests/fixtures/database.py

@@ -18,6 +18,7 @@
 import jsonschema
 import numpy as np
 import pytest
+from geti_supported_models.default_models import DefaultModels
 
 from coordination.configuration_manager.task_node_config import TaskNodeConfig
 
@@ -721,6 +722,7 @@ def create_empty_project(
                 creator_id="",
                 parser_class=RestProjectParser,
                 parser_kwargs=parser_kwargs,
+                default_models_per_task=DefaultModels.get_default_models_per_task(),
             )
 
         return self._project

interactive_ai/services/resource/Dockerfile

@@ -21,6 +21,7 @@ ENV UV_PYTHON_DOWNLOADS=0
 COPY --link --from=libs . libs
 COPY --link --from=iai_core . interactive_ai/libs/iai_core_py
 COPY --link --from=media_utils . interactive_ai/libs/media_utils
+COPY --link --from=supported_models . interactive_ai/supported_models
 
 WORKDIR /interactive_ai/services/resource
 

interactive_ai/services/resource/Makefile

@@ -5,3 +5,4 @@ DOCKER_BUILD_CONTEXT := --build-context schemas=../api \
 						--build-context libs=../../../libs \
 						--build-context iai_core=../../../interactive_ai/libs/iai_core_py \
 						--build-context media_utils=../../../interactive_ai/libs/media_utils \
+						--build-context supported_models=../../../interactive_ai/supported_models \

interactive_ai/services/resource/app/communication/rest_views/model_rest_views.py

@@ -6,6 +6,8 @@
 from enum import Enum, auto
 from typing import TYPE_CHECKING, Any, cast
 
+from geti_supported_models import SupportedModels
+
 from communication.rest_views.label_rest_views import LabelRESTViews
 from communication.rest_views.performance_rest_views import PerformanceRESTViews
 
@@ -125,12 +127,13 @@ def model_storage_to_rest(
         :return: dictionary with REST representation of the model storage
         """
         models_rest_list: list[dict] = []
+        model_architecture_name = ModelRESTViews._get_model_architecture_name(model_storage=model_storage)
         for model_info in per_model_info:
             model = model_info.model
             models_rest_list.append(
                 {
                     ID_: str(model.id_),
-                    NAME: model.model_storage.name,
+                    NAME: model_architecture_name,
                     CREATION_DATE: model.creation_date.isoformat(),
                     SCORE_UP_TO_DATE: True,  # this field is deprecated, value is always 'True' and should not be used
                     ACTIVE_MODEL: model == active_model,
@@ -151,7 +154,7 @@ def model_storage_to_rest(
 
         rest_view = {
             ID_: str(model_storage.id_),
-            NAME: model_storage.name,
+            NAME: model_architecture_name,
             MODEL_TEMPLATE_ID: model_storage.model_template.model_template_id,
             TASK_ID: str(task_node_id),
             MODELS: models_rest_list,
@@ -186,7 +189,7 @@ def model_to_rest(
         """
         return {
             ID_: str(model.id_),
-            NAME: model.model_storage.name,
+            NAME: ModelRESTViews._get_model_architecture_name(model_storage=model.model_storage),
             CREATION_DATE: model.creation_date.isoformat(),
             SCORE_UP_TO_DATE: True,  # Note: this field is deprecated, value is always 'True' and should not be used
             ACTIVE_MODEL: model == active_model,
@@ -233,10 +236,11 @@ def model_info_to_rest(
         """
         label_schema = model.configuration.get_label_schema()
         labels = label_schema.get_labels(include_empty=True)
+        model_architecture_name = ModelRESTViews._get_model_architecture_name(model_storage=model.model_storage)
         result = {
             ID_: model.id_,
-            NAME: model.model_storage.name,
-            ARCHITECTURE: model.model_storage.model_template.name,
+            NAME: model_architecture_name,
+            ARCHITECTURE: model_architecture_name,
             VERSION: model.version,
             CREATION_DATE: model.creation_date.isoformat(),
             SIZE: model.size,
@@ -313,7 +317,7 @@ def optimized_model_to_rest(optimized_model: Model, performance: Performance) ->
             filter(
                 None,
                 [
-                    optimized_model.model_storage.name,
+                    ModelRESTViews._get_model_architecture_name(model_storage=optimized_model.model_storage),
                     model_format,
                     precision_text,
                     with_xai,
@@ -364,3 +368,14 @@ def optimized_model_to_rest(optimized_model: Model, performance: Performance) ->
             result[OPTIMIZATION_METHODS] = [method.name for method in optimized_model.optimization_methods]
 
         return result
+
+    @staticmethod
+    def _get_model_architecture_name(model_storage: ModelStorage) -> str:
+        """
+        Get the model architecture name from the model storage, based on the model manifests.
+
+        :param model_storage: ModelStorage to get the architecture name from
+        :return: The name of the model architecture
+        """
+        model_manifest = SupportedModels.get_model_manifest_by_id(model_manifest_id=model_storage.model_manifest_id)
+        return model_manifest.name

interactive_ai/services/resource/app/managers/project_manager.py

@@ -6,6 +6,7 @@
 from typing import cast
 
 from geti_spicedb_tools import Permissions, SpiceDB
+from geti_supported_models.default_models import DefaultModels
 
 from communication.exceptions import DatasetStorageNotInProjectException, LabelNotFoundException, ProjectLockedException
 from managers.annotation_manager import AnnotationManager
@@ -86,6 +87,7 @@ def create_project(
             creator_id=creator_id,
             parser_class=project_parser,
             parser_kwargs=parser_kwargs,
+            default_models_per_task=DefaultModels.get_default_models_per_task(),
         )
 
         # TODO: CVS-89772 call spicedb before storing a project in database

interactive_ai/services/resource/pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
     "iai_core",
     "media_utils",
     "grpc_interfaces",
+    "geti-supported-models",
 ]
 
 [tool.uv.sources]
@@ -32,6 +33,7 @@ geti_spicedb_tools = { path = "../../../libs/spicedb_tools", editable = true }
 iai_core = { path = "../../libs/iai_core_py", editable = true }
 media_utils = { path = "../../libs/media_utils", editable = true }
 grpc_interfaces = { path = "../../../libs/grpc_interfaces", editable = true }
+geti-supported-models = { path = "../../supported_models", editable = true }
 
 
 [dependency-groups]

interactive_ai/services/resource/tests/fixtures/database.py

@@ -13,6 +13,7 @@
 import cv2
 import numpy as np
 import pytest
+from geti_supported_models.default_models import DefaultModels
 
 from communication.rest_parsers import RestProjectParser
 from service.label_schema_service import LabelSchemaService
@@ -446,6 +447,7 @@ def create_empty_project(
                 creator_id="",
                 parser_class=RestProjectParser,
                 parser_kwargs=parser_kwargs,
+                default_models_per_task=DefaultModels.get_default_models_per_task(),
             )
 
         return self._project

interactive_ai/services/resource/tests/fixtures/model.py

@@ -431,13 +431,13 @@ def fxt_model_empty_dataset(
 @pytest.fixture
 def fxt_model_group_rest(fxt_model, fxt_mongo_id):
     yield {
-        "id": fxt_mongo_id(0),
-        "name": "Sample Detection Template",
-        "task_id": fxt_mongo_id(1),
-        "model_template_id": "test_template_detection",
+        "id": str(fxt_mongo_id(0)),
+        "name": fxt_model.model_storage.name,
+        "task_id": str(fxt_mongo_id(1)),
+        "model_template_id": fxt_model.model_storage.model_template_id,
         "models": [
             {
-                "id": fxt_model.id_,
+                "id": str(fxt_model.id_),
                 "name": fxt_model.model_storage.name,
                 "creation_date": "2020-01-01T00:00:00+00:00",
                 "performance": {"score": get_performance_score(fxt_model.performance)},