supervisely 6.73.325__py3-none-any.whl → 6.73.327__py3-none-any.whl

supervisely/nn/inference/inference.py

@@ -75,6 +75,7 @@ from supervisely.nn.utils import (
     ModelPrecision,
     ModelSource,
     RuntimeType,
+    _get_model_name,
 )
 from supervisely.project import ProjectType
 from supervisely.project.download import download_to_cache, read_from_cached_project
@@ -173,9 +174,7 @@ class Inference:
             self._use_gui = False
             deploy_params, need_download = self._get_deploy_params_from_args()
             if need_download:
-                local_model_files = self._download_model_files(
-                    deploy_params["model_source"], deploy_params["model_files"], False
-                )
+                local_model_files = self._download_model_files(deploy_params, False)
                 deploy_params["model_files"] = local_model_files
             self._load_model_headless(**deploy_params)
 
@@ -210,14 +209,12 @@ class Inference:
                     self.initialize_gui()
 
             def on_serve_callback(
-                gui: Union[GUI.InferenceGUI, GUI.ServingGUI, GUI.ServingGUITemplate]
+                gui: Union[GUI.InferenceGUI, GUI.ServingGUI, GUI.ServingGUITemplate],
             ):
                 Progress("Deploying model ...", 1)
                 if isinstance(self.gui, GUI.ServingGUITemplate):
                     deploy_params = self.get_params_from_gui()
-                    model_files = self._download_model_files(
-                        deploy_params["model_source"], deploy_params["model_files"]
-                    )
+                    model_files = self._download_model_files(deploy_params)
                     deploy_params["model_files"] = model_files
                     self._load_model_headless(**deploy_params)
                 elif isinstance(self.gui, GUI.ServingGUI):
@@ -230,7 +227,7 @@ class Inference:
                     gui.show_deployed_model_info(self)
 
             def on_change_model_callback(
-                gui: Union[GUI.InferenceGUI, GUI.ServingGUI, GUI.ServingGUITemplate]
+                gui: Union[GUI.InferenceGUI, GUI.ServingGUI, GUI.ServingGUITemplate],
             ):
                 self.shutdown_model()
                 if isinstance(self.gui, (GUI.ServingGUI, GUI.ServingGUITemplate)):
@@ -567,13 +564,23 @@ class Inference:
     def _checkpoints_cache_dir(self):
         return os.path.join(os.path.expanduser("~"), ".cache", "supervisely", "checkpoints")
 
-    def _download_model_files(
-        self, model_source: str, model_files: List[str], log_progress: bool = True
-    ) -> dict:
-        if model_source == ModelSource.PRETRAINED:
-            return self._download_pretrained_model(model_files, log_progress)
-        elif model_source == ModelSource.CUSTOM:
-            return self._download_custom_model(model_files, log_progress)
+    def _download_model_files(self, deploy_params: dict, log_progress: bool = True) -> dict:
+        if deploy_params["runtime"] != RuntimeType.PYTORCH:
+            export = deploy_params["model_info"].get("export", {})
+            export_model = export.get(deploy_params["runtime"], None)
+            if export_model is not None:
+                if sly_fs.get_file_name(export_model) == sly_fs.get_file_name(
+                    deploy_params["model_files"]["checkpoint"]
+                ):
+                    deploy_params["model_files"]["checkpoint"] = (
+                        deploy_params["model_info"]["artifacts_dir"] + export_model
+                    )
+                    logger.info(f"Found model checkpoint for '{deploy_params['runtime']}'")
+
+        if deploy_params["model_source"] == ModelSource.PRETRAINED:
+            return self._download_pretrained_model(deploy_params["model_files"], log_progress)
+        elif deploy_params["model_source"] == ModelSource.CUSTOM:
+            return self._download_custom_model(deploy_params["model_files"], log_progress)
 
     def _download_pretrained_model(self, model_files: dict, log_progress: bool = True):
         """
@@ -2929,9 +2936,7 @@ class Inference:
                 state = request.state.state
                 deploy_params = state["deploy_params"]
                 if isinstance(self.gui, GUI.ServingGUITemplate):
-                    model_files = self._download_model_files(
-                        deploy_params["model_source"], deploy_params["model_files"]
-                    )
+                    model_files = self._download_model_files(deploy_params)
                     deploy_params["model_files"] = model_files
                     self._load_model_headless(**deploy_params)
                 elif isinstance(self.gui, GUI.ServingGUI):
@@ -3061,7 +3066,7 @@ class Inference:
                 raise ValueError("No pretrained models found.")
 
             model = self.pretrained_models[0]
-            model_name = model.get("meta", {}).get("model_name", None)
+            model_name = _get_model_name(model)
             if model_name is None:
                 raise ValueError("No model name found in the first pretrained model.")
 
@@ -3126,7 +3131,7 @@ class Inference:
             meta = m.get("meta", None)
             if meta is None:
                 continue
-            model_name = meta.get("model_name", None)
+            model_name = _get_model_name(m)
             if model_name is None:
                 continue
             m_files = meta.get("model_files", None)
@@ -3135,7 +3140,7 @@ class Inference:
             checkpoint = m_files.get("checkpoint", None)
             if checkpoint is None:
                 continue
-            if model == m["meta"]["model_name"]:
+            if model.lower() == model_name.lower():
                 model_info = m
                 model_source = ModelSource.PRETRAINED
                 model_files = {"checkpoint": checkpoint}
@@ -3153,8 +3158,6 @@ class Inference:
             model_meta_path = os.path.join(artifacts_dir, "model_meta.json")
             model_info["model_meta"] = self._load_json_file(model_meta_path)
             original_model_files = model_info.get("model_files")
-            if not original_model_files:
-                raise ValueError("Invalid 'experiment_info.json'. Missing 'model_files' key.")
             return model_info, original_model_files
 
         def _prepare_local_model_files(artifacts_dir, checkpoint_path, original_model_files):
@@ -3201,6 +3204,7 @@ class Inference:
             model_files = _prepare_local_model_files(
                 artifacts_dir, checkpoint_path, original_model_files
             )
+
         else:
             local_artifacts_dir = os.path.join(
                 self.model_dir, "local_deploy", os.path.basename(artifacts_dir)
@@ -3298,7 +3302,11 @@ class Inference:
             if draw:
                 raise ValueError("Draw visualization is not supported for project inference")
 
-            state = {"projectId": project_id, "dataset_ids": dataset_ids, "settings": settings}
+            state = {
+                "projectId": project_id,
+                "dataset_ids": dataset_ids,
+                "settings": settings,
+            }
             if upload:
                 source_project = api.project.get_info_by_id(project_id)
                 workspace_id = source_project.workspace_id
@@ -3472,7 +3480,7 @@ class Inference:
     def _add_workflow_input(self, model_source: str, model_files: dict, model_info: dict):
         if model_source == ModelSource.PRETRAINED:
             checkpoint_url = model_info["meta"]["model_files"]["checkpoint"]
-            checkpoint_name = model_info["meta"]["model_name"]
+            checkpoint_name = _get_model_name(model_info)
         else:
             checkpoint_name = sly_fs.get_file_name_with_ext(model_files["checkpoint"])
             checkpoint_url = os.path.join(

supervisely/nn/training/gui/classes_selector.py

@@ -4,21 +4,28 @@ from supervisely.app.widgets import Button, Card, ClassesTable, Container, Text
 
 class ClassesSelector:
     title = "Classes Selector"
-    description = (
-        "Select classes that will be used for training. "
-        "Supported shapes are Bitmap, Polygon, Rectangle."
-    )
+    description = "Select classes that will be used for training"
     lock_message = "Select training and validation splits to unlock"
 
     def __init__(self, project_id: int, classes: list, app_options: dict = {}):
+        # Init widgets
+        self.qa_stats_text = None
+        self.classes_table = None
+        self.validator_text = None
+        self.button = None
+        self.container = None
+        self.card = None
+        # -------------------------------- #
+
         self.display_widgets = []
+        self.app_options = app_options
 
         # GUI Components
         if is_development() or is_debug_with_sly_net():
             qa_stats_link = abs_url(f"projects/{project_id}/stats/datasets")
         else:
             qa_stats_link = f"/projects/{project_id}/stats/datasets"
-        qa_stats_text = Text(
+        self.qa_stats_text = Text(
             text=f"<i class='zmdi zmdi-chart-donut' style='color: #7f858e'></i> <a href='{qa_stats_link}' target='_blank'> <b> QA & Stats </b></a>"
         )
 
@@ -32,7 +39,7 @@ class ClassesSelector:
         self.validator_text.hide()
         self.button = Button("Select")
         self.display_widgets.extend(
-            [qa_stats_text, self.classes_table, self.validator_text, self.button]
+            [self.qa_stats_text, self.classes_table, self.validator_text, self.button]
         )
         # -------------------------------- #
 
@@ -42,7 +49,7 @@ class ClassesSelector:
             description=self.description,
             content=self.container,
             lock_message=self.lock_message,
-            collapsable=app_options.get("collapsable", False),
+            collapsable=self.app_options.get("collapsable", False),
         )
         self.card.lock()
 
@@ -62,14 +69,14 @@ class ClassesSelector:
     def validate_step(self) -> bool:
         self.validator_text.hide()
 
-        if len(self.classes_table.project_meta.obj_classes) == 0:
+        project_classes = self.classes_table.project_meta.obj_classes
+        if len(project_classes) == 0:
             self.validator_text.set(text="Project has no classes", status="error")
             self.validator_text.show()
             return False
 
         selected_classes = self.classes_table.get_selected_classes()
         table_data = self.classes_table._table_data
-
         empty_classes = [
             row[0]["data"]
             for row in table_data
@@ -78,23 +85,21 @@ class ClassesSelector:
 
         n_classes = len(selected_classes)
         if n_classes == 0:
-            self.validator_text.set(text="Please select at least one class", status="error")
+            message = "Please select at least one class"
+            status = "error"
         else:
-            warning_text = ""
+            class_text = "class" if n_classes == 1 else "classes"
+            message = f"Selected {n_classes} {class_text}"
             status = "success"
             if empty_classes:
                 intersections = set(selected_classes).intersection(empty_classes)
                 if intersections:
-                    warning_text = (
-                        f". Selected class has no annotations: {', '.join(intersections)}"
-                        if len(intersections) == 1
-                        else f". Selected classes have no annotations: {', '.join(intersections)}"
+                    class_text = "class" if len(intersections) == 1 else "classes"
+                    message += (
+                        f". Selected {class_text} have no annotations: {', '.join(intersections)}"
                     )
                     status = "warning"
 
-            class_text = "class" if n_classes == 1 else "classes"
-            self.validator_text.set(
-                text=f"Selected {n_classes} {class_text}{warning_text}", status=status
-            )
+        self.validator_text.set(text=message, status=status)
         self.validator_text.show()
         return n_classes > 0

supervisely/nn/training/gui/gui.py

@@ -6,11 +6,15 @@ training workflows in Supervisely.
 """
 
 from os import environ
+from typing import Union
 
 import supervisely.io.env as sly_env
+import supervisely.io.json as sly_json
 from supervisely import Api, ProjectMeta
 from supervisely._utils import is_production
 from supervisely.app.widgets import Stepper, Widget
+from supervisely.geometry.bitmap import Bitmap
+from supervisely.geometry.graph import GraphNodes
 from supervisely.geometry.polygon import Polygon
 from supervisely.geometry.rectangle import Rectangle
 from supervisely.nn.task_type import TaskType
@@ -63,7 +67,7 @@ class TrainGUI:
         self.models = models
         self.hyperparameters = hyperparameters
         self.app_options = app_options
-        self.collapsable = app_options.get("collapsable", False)
+        self.collapsable = self.app_options.get("collapsable", False)
         self.need_convert_shapes_for_bm = False
 
         self.team_id = sly_env.team_id(raise_not_found=False)
@@ -142,33 +146,73 @@ class TrainGUI:
             self.training_process.set_experiment_name(experiment_name)
 
         def need_convert_class_shapes() -> bool:
-            if not self.hyperparameters_selector.run_model_benchmark_checkbox.is_checked():
-                self.hyperparameters_selector.model_benchmark_auto_convert_warning.hide()
-                self.need_convert_shapes_for_bm = False
-            else:
-                task_type = self.model_selector.get_selected_task_type()
-
-                def _need_convert(shape):
-                    if task_type == TaskType.OBJECT_DETECTION:
-                        return shape != Rectangle.geometry_name()
-                    elif task_type in [
-                        TaskType.INSTANCE_SEGMENTATION,
-                        TaskType.SEMANTIC_SEGMENTATION,
-                    ]:
-                        return shape == Polygon.geometry_name()
-                    return
-
-                data = self.classes_selector.classes_table._table_data
-                selected_classes = set(self.classes_selector.classes_table.get_selected_classes())
-                empty = set(r[0]["data"] for r in data if r[2]["data"] == 0 and r[3]["data"] == 0)
-                need_convert = set(r[0]["data"] for r in data if _need_convert(r[1]["data"]))
-
-                if need_convert.intersection(selected_classes - empty):
-                    self.hyperparameters_selector.model_benchmark_auto_convert_warning.show()
-                    self.need_convert_shapes_for_bm = True
-                else:
+            if self.hyperparameters_selector.run_model_benchmark_checkbox is not None:
+                if not self.hyperparameters_selector.run_model_benchmark_checkbox.is_checked():
                     self.hyperparameters_selector.model_benchmark_auto_convert_warning.hide()
                     self.need_convert_shapes_for_bm = False
+                else:
+                    task_type = self.model_selector.get_selected_task_type()
+
+                    def _need_convert(shape):
+                        if task_type == TaskType.OBJECT_DETECTION:
+                            return shape != Rectangle.geometry_name()
+                        elif task_type in [
+                            TaskType.INSTANCE_SEGMENTATION,
+                            TaskType.SEMANTIC_SEGMENTATION,
+                        ]:
+                            return shape == Polygon.geometry_name()
+                        return
+
+                    data = self.classes_selector.classes_table._table_data
+                    selected_classes = set(
+                        self.classes_selector.classes_table.get_selected_classes()
+                    )
+                    empty = set(
+                        r[0]["data"] for r in data if r[2]["data"] == 0 and r[3]["data"] == 0
+                    )
+                    need_convert = set(r[0]["data"] for r in data if _need_convert(r[1]["data"]))
+
+                    if need_convert.intersection(selected_classes - empty):
+                        self.hyperparameters_selector.model_benchmark_auto_convert_warning.show()
+                        self.need_convert_shapes_for_bm = True
+                    else:
+                        self.hyperparameters_selector.model_benchmark_auto_convert_warning.hide()
+                        self.need_convert_shapes_for_bm = False
+            else:
+                self.need_convert_shapes_for_bm = False
+
+        def validate_class_shape_for_model_task():
+            task_type = self.model_selector.get_selected_task_type()
+            classes = self.classes_selector.get_selected_classes()
+
+            required_geometries = {
+                TaskType.INSTANCE_SEGMENTATION: {Polygon, Bitmap},
+                TaskType.SEMANTIC_SEGMENTATION: {Polygon, Bitmap},
+                TaskType.POSE_ESTIMATION: {GraphNodes},
+            }
+            task_specific_texts = {
+                TaskType.INSTANCE_SEGMENTATION: "Only polygon and bitmap shapes are supported for segmentation task",
+                TaskType.SEMANTIC_SEGMENTATION: "Only polygon and bitmap shapes are supported for segmentation task",
+                TaskType.POSE_ESTIMATION: "Only keypoint (graph) shape is supported for pose estimation task",
+            }
+
+            if task_type not in required_geometries:
+                return
+
+            wrong_shape_classes = [
+                class_name
+                for class_name in classes
+                if self.project_meta.get_obj_class(class_name).geometry_type
+                not in required_geometries[task_type]
+            ]
+
+            if wrong_shape_classes:
+                specific_text = task_specific_texts[task_type]
+                message_text = f"Model task type is {task_type}. {specific_text}. Selected classes have wrong shapes for the model task: {', '.join(wrong_shape_classes)}"
+                self.model_selector.validator_text.set(
+                    text=message_text,
+                    status="warning",
+                )
 
         # ------------------------------------------------- #
 
@@ -201,7 +245,11 @@ class TrainGUI:
             callback=self.hyperparameters_selector_cb,
             validation_text=self.model_selector.validator_text,
             validation_func=self.model_selector.validate_step,
-            on_select_click=[set_experiment_name, need_convert_class_shapes],
+            on_select_click=[
+                set_experiment_name,
+                need_convert_class_shapes,
+                validate_class_shape_for_model_task,
+            ],
             collapse_card=(self.model_selector.card, self.collapsable),
         )
 
@@ -299,17 +347,19 @@ class TrainGUI:
         # ------------------------------------------------- #
 
         # Other Buttons
-        if app_options.get("show_logs_in_gui", False):
+        if self.app_options.get("show_logs_in_gui", False):
 
             @self.training_logs.logs_button.click
             def show_logs():
                 self.training_logs.toggle_logs()
 
         # Other handlers
-        @self.hyperparameters_selector.run_model_benchmark_checkbox.value_changed
-        def show_mb_speedtest(is_checked: bool):
-            self.hyperparameters_selector.toggle_mb_speedtest(is_checked)
-            need_convert_class_shapes()
+        if self.hyperparameters_selector.run_model_benchmark_checkbox is not None:
+
+            @self.hyperparameters_selector.run_model_benchmark_checkbox.value_changed
+            def show_mb_speedtest(is_checked: bool):
+                self.hyperparameters_selector.toggle_mb_speedtest(is_checked)
+                need_convert_class_shapes()
 
         # ------------------------------------------------- #
 
@@ -361,7 +411,6 @@ class TrainGUI:
             raise ValueError("app_state must be a dictionary")
 
         required_keys = {
-            "input": ["project_id"],
             "train_val_split": ["method"],
             "classes": list,
             "model": ["source"],
@@ -453,7 +502,7 @@ class TrainGUI:
                 raise ValueError("percent must be an integer in range 1 to 99")
         return app_state
 
-    def load_from_app_state(self, app_state: dict) -> None:
+    def load_from_app_state(self, app_state: Union[str, dict]) -> None:
         """
         Load the GUI state from app state dictionary.
 
@@ -463,7 +512,6 @@ class TrainGUI:
         app_state example:
 
             app_state = {
-                "input": {"project_id": 43192},
                 "train_val_split": {
                     "method": "random",
                     "split": "train",
@@ -489,10 +537,13 @@ class TrainGUI:
                 }
             }
         """
+        if isinstance(app_state, str):
+            app_state = sly_json.load_json_file(app_state)
+
         app_state = self.validate_app_state(app_state)
 
         options = app_state.get("options", {})
-        input_settings = app_state["input"]
+        input_settings = app_state.get("input")
         train_val_splits_settings = app_state["train_val_split"]
         classes_settings = app_state["classes"]
         model_settings = app_state["model"]
@@ -504,7 +555,7 @@ class TrainGUI:
         self._init_model(model_settings)
         self._init_hyperparameters(hyperparameters_settings, options)
 
-    def _init_input(self, input_settings: dict, options: dict) -> None:
+    def _init_input(self, input_settings: Union[dict, None], options: dict) -> None:
         """
         Initialize the input selector with the given settings.
 
@@ -604,6 +655,8 @@ class TrainGUI:
             )
             self.hyperparameters_selector.set_speedtest_checkbox_value(
                 model_benchmark_settings["speed_test"]
+                if model_benchmark_settings["enable"]
+                else False
             )
         export_weights_settings = options.get("export", None)
         if export_weights_settings is not None:

supervisely/nn/training/gui/hyperparameters_selector.py

@@ -9,7 +9,6 @@ from supervisely.app.widgets import (
     Field,
     Text,
 )
-from supervisely.nn.utils import RuntimeType
 
 
 class HyperparametersSelector:
@@ -18,6 +17,22 @@ class HyperparametersSelector:
     lock_message = "Select model to unlock"
 
     def __init__(self, hyperparameters: dict, app_options: dict = {}):
+        # Init widgets
+        self.editor = None
+        self.run_model_benchmark_checkbox = None
+        self.run_speedtest_checkbox = None
+        self.model_benchmark_field = None
+        self.model_benchmark_learn_more = None
+        self.model_benchmark_auto_convert_warning = None
+        self.export_onnx_checkbox = None
+        self.export_tensorrt_checkbox = None
+        self.export_field = None
+        self.validator_text = None
+        self.button = None
+        self.container = None
+        self.card = None
+        # -------------------------------- #
+
         self.display_widgets = []
         self.app_options = app_options
 
@@ -28,7 +43,7 @@ class HyperparametersSelector:
         self.display_widgets.extend([self.editor])
 
         # Optional Model Benchmark
-        if app_options.get("model_benchmark", True):
+        if self.app_options.get("model_benchmark", True):
             # Model Benchmark
             self.run_model_benchmark_checkbox = Checkbox(
                 content="Run Model Benchmark evaluation", checked=True
@@ -37,7 +52,7 @@ class HyperparametersSelector:
 
             self.model_benchmark_field = Field(
                 title="Model Evaluation Benchmark",
-                description=f"Generate evalutaion dashboard with visualizations and detailed analysis of the model performance after training. The best checkpoint will be used for evaluation. You can also run speed test to evaluate model inference speed.",
+                description="Generate evaluation dashboard with visualizations and detailed analysis of the model performance after training. The best checkpoint will be used for evaluation. You can also run speed test to evaluate model inference speed.",
                 content=Container([self.run_model_benchmark_checkbox, self.run_speedtest_checkbox]),
             )
             docs_link = '<a href="https://docs.supervisely.com/neural-networks/model-evaluation-benchmark/" target="_blank">documentation</a>'
@@ -60,8 +75,8 @@ class HyperparametersSelector:
         # -------------------------------- #
 
         # Optional Export Weights
-        export_onnx_supported = app_options.get("export_onnx_supported", False)
-        export_tensorrt_supported = app_options.get("export_tensorrt_supported", False)
+        export_onnx_supported = self.app_options.get("export_onnx_supported", False)
+        export_tensorrt_supported = self.app_options.get("export_tensorrt_supported", False)
 
         onnx_name = "ONNX"
         tensorrt_name = "TensorRT engine"
@@ -120,26 +135,28 @@ class HyperparametersSelector:
         return self.editor.get_value()
 
     def get_model_benchmark_checkbox_value(self) -> bool:
-        if self.app_options.get("model_benchmark", True):
+        if self.run_model_benchmark_checkbox is not None:
             return self.run_model_benchmark_checkbox.is_checked()
         return False
 
     def set_model_benchmark_checkbox_value(self, is_checked: bool) -> bool:
-        if is_checked:
-            self.run_model_benchmark_checkbox.check()
-        else:
-            self.run_model_benchmark_checkbox.uncheck()
+        if self.run_model_benchmark_checkbox is not None:
+            if is_checked:
+                self.run_model_benchmark_checkbox.check()
+            else:
+                self.run_model_benchmark_checkbox.uncheck()
 
     def get_speedtest_checkbox_value(self) -> bool:
-        if self.app_options.get("model_benchmark", True):
+        if self.run_speedtest_checkbox is not None:
             return self.run_speedtest_checkbox.is_checked()
         return False
 
     def set_speedtest_checkbox_value(self, is_checked: bool) -> bool:
-        if is_checked:
-            self.run_speedtest_checkbox.check()
-        else:
-            self.run_speedtest_checkbox.uncheck()
+        if self.run_speedtest_checkbox is not None:
+            if is_checked:
+                self.run_speedtest_checkbox.check()
+            else:
+                self.run_speedtest_checkbox.uncheck()
 
     def toggle_mb_speedtest(self, is_checked: bool) -> None:
         if is_checked:

supervisely/nn/training/gui/input_selector.py

@@ -4,7 +4,6 @@ from supervisely.app.widgets import (
     Card,
     Checkbox,
     Container,
-    Field,
     ProjectThumbnail,
     Text,
 )
@@ -17,7 +16,19 @@ class InputSelector:
     lock_message = None
 
     def __init__(self, project_info: ProjectInfo, app_options: dict = {}):
+        # Init widgets
+        self.project_thumbnail = None
+        self.use_cache_text = None
+        self.use_cache_checkbox = None
+        self.validator_text = None
+        self.button = None
+        self.container = None
+        self.card = None
+        # -------------------------------- #
+
         self.display_widgets = []
+        self.app_options = app_options
+
         self.project_id = project_info.id
         self.project_info = project_info
 
@@ -49,7 +60,7 @@ class InputSelector:
             title=self.title,
             description=self.description,
             content=self.container,
-            collapsable=app_options.get("collapsable", False),
+            collapsable=self.app_options.get("collapsable", False),
         )
 
     @property

supervisely/nn/training/gui/model_selector.py

@@ -14,7 +14,7 @@ from supervisely.app.widgets import (
 )
 from supervisely.nn.artifacts.utils import FrameworkMapper
 from supervisely.nn.experiments import get_experiment_infos
-from supervisely.nn.utils import ModelSource
+from supervisely.nn.utils import ModelSource, _get_model_name
 
 
 class ModelSelector:
@@ -23,15 +23,26 @@ class ModelSelector:
     lock_message = "Select classes to unlock"
 
     def __init__(self, api: Api, framework: str, models: list, app_options: dict = {}):
+        # Init widgets
+        self.pretrained_models_table = None
+        self.experiment_selector = None
+        self.model_source_tabs = None
+        self.validator_text = None
+        self.button = None
+        self.container = None
+        self.card = None
+        # -------------------------------- #
+
         self.display_widgets = []
+        self.app_options = app_options
+
         self.team_id = sly_env.team_id()
         self.models = models
 
         # GUI Components
         self.pretrained_models_table = PretrainedModelsSelector(self.models)
-
         experiment_infos = get_experiment_infos(api, self.team_id, framework)
-        if app_options.get("legacy_checkpoints", False):
+        if self.app_options.get("legacy_checkpoints", False):
             try:
                 framework_cls = FrameworkMapper.get_framework_cls(framework, self.team_id)
                 legacy_experiment_infos = framework_cls.get_list_experiment_info()
@@ -61,7 +72,7 @@ class ModelSelector:
             description=self.description,
             content=self.container,
             lock_message=self.lock_message,
-            collapsable=app_options.get("collapsable", False),
+            collapsable=self.app_options.get("collapsable", False),
         )
         self.card.lock()
 
@@ -82,8 +93,7 @@ class ModelSelector:
     def get_model_name(self) -> str:
         if self.get_model_source() == ModelSource.PRETRAINED:
             selected_row = self.pretrained_models_table.get_selected_row()
-            model_meta = selected_row.get("meta", {})
-            model_name = model_meta.get("model_name", None)
+            model_name = _get_model_name(selected_row)
         else:
             selected_row = self.experiment_selector.get_selected_experiment_info()
             model_name = selected_row.get("model_name", None)