supervisely 6.73.249__py3-none-any.whl → 6.73.251__py3-none-any.whl

supervisely/nn/inference/inference.py

@@ -90,7 +90,10 @@ class Inference:
     """Path to file with list of models"""
     APP_OPTIONS: str = None
     """Path to file with app options"""
-    DEFAULT_BATCH_SIZE = 16
+    DEFAULT_BATCH_SIZE: str = 16
+    """Default batch size for inference"""
+    INFERENCE_SETTINGS: str = None
+    """Path to file with custom inference settings"""
 
     def __init__(
         self,
@@ -125,8 +128,12 @@ class Inference:
         self._autostart_delay_time = 5 * 60  # 5 min
         self._tracker = None
         self._hardware: str = None
+        self.pretrained_models = self._load_models_json(self.MODELS) if self.MODELS else None
         if custom_inference_settings is None:
-            custom_inference_settings = {}
+            if self.INFERENCE_SETTINGS is not None:
+                custom_inference_settings = self.INFERENCE_SETTINGS
+            else:
+                custom_inference_settings = {}
         if isinstance(custom_inference_settings, str):
             if fs.file_exists(custom_inference_settings):
                 with open(custom_inference_settings, "r") as f:
@@ -153,7 +160,7 @@ class Inference:
                 if self.FRAMEWORK_NAME is None:
                     raise ValueError("FRAMEWORK_NAME is not defined")
                 self._gui = GUI.ServingGUITemplate(
-                    self.FRAMEWORK_NAME, self.MODELS, self.APP_OPTIONS
+                    self.FRAMEWORK_NAME, self.pretrained_models, self.APP_OPTIONS
                 )
                 self._user_layout = self._gui.widgets
                 self._user_layout_card = self._gui.card
@@ -239,6 +246,38 @@ class Inference:
                 )
                 device = "cpu"
 
+    def _load_models_json(self, models: str) -> List[Dict[str, Any]]:
+        """
+        Loads models from the provided file or list of model configurations.
+        """
+        if isinstance(models, str):
+            if sly_fs.file_exists(models) and sly_fs.get_file_ext(models) == ".json":
+                models = sly_json.load_json_file(models)
+            else:
+                raise ValueError("File not found or invalid file format.")
+        else:
+            raise ValueError(
+                "Invalid models file. Please provide a valid '.json' file with list of model configurations."
+            )
+
+        if not isinstance(models, list):
+            raise ValueError("models parameters must be a list of dicts")
+        for item in models:
+            if not isinstance(item, dict):
+                raise ValueError(f"Each item in models must be a dict.")
+            model_meta = item.get("meta")
+            if model_meta is None:
+                raise ValueError(
+                    "Model metadata not found. Please update provided models parameter to include key 'meta'."
+                )
+            model_files = model_meta.get("model_files")
+            if model_files is None:
+                raise ValueError(
+                    "Model files not found in model metadata. "
+                    "Please update provided models oarameter to include key 'model_files' in 'meta' key."
+                )
+        return models
+
     def get_ui(self) -> Widget:
         if not self._use_gui:
             return None
@@ -487,6 +526,9 @@ class Inference:
     def load_model_meta(self, model_tab: str, local_weights_path: str):
         raise NotImplementedError("Have to be implemented in child class after inheritance")
 
+    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]) -> dict:
         if model_source == ModelSource.PRETRAINED:
             return self._download_pretrained_model(model_files)
@@ -498,17 +540,28 @@ class Inference:
         Downloads the pretrained model data.
         """
         local_model_files = {}
+        cache_dir = self._checkpoints_cache_dir()
 
         for file in model_files:
             file_url = model_files[file]
-            file_path = os.path.join(self.model_dir, file)
+            file_name = sly_fs.get_file_name_with_ext(file_url)
             if file_url.startswith("http"):
                 with urlopen(file_url) as f:
                     file_size = f.length
                     file_name = get_filename_from_headers(file_url)
-                    file_path = os.path.join(self.model_dir, file_name)
                     if file_name is None:
                         file_name = file
+                    file_path = os.path.join(self.model_dir, file_name)
+                    cached_path = os.path.join(cache_dir, file_name)
+                    if os.path.exists(cached_path):
+                        local_model_files[file] = cached_path
+                        logger.debug(f"Model: '{file_name}' was found in checkpoint cache")
+                        continue
+                    if os.path.exists(file_path):
+                        local_model_files[file] = file_path
+                        logger.debug(f"Model: '{file_name}' was found in model dir")
+                        continue
+
                     with self.gui.download_progress(
                         message=f"Downloading: '{file_name}'",
                         total=file_size,
@@ -614,13 +667,23 @@ class Inference:
         model_files = deploy_params.get("model_files", {})
         if model_info:
             checkpoint_name = os.path.basename(model_files.get("checkpoint"))
+            checkpoint_file_path = os.path.join(
+                model_info.get("artifacts_dir"), "checkpoints", checkpoint_name
+            )
+            checkpoint_file_info = self.api.file.get_info_by_path(
+                env.team_id(), checkpoint_file_path
+            )
+            if checkpoint_file_info is None:
+                checkpoint_url = None
+            else:
+                checkpoint_url = self.api.file.get_url(checkpoint_file_info.id)
+
             self.checkpoint_info = CheckpointInfo(
                 checkpoint_name=checkpoint_name,
                 model_name=model_info.get("model_name"),
                 architecture=model_info.get("framework_name"),
-                custom_checkpoint_path=os.path.join(
-                    model_info.get("artifacts_dir"), checkpoint_name
-                ),
+                checkpoint_url=checkpoint_url,
+                custom_checkpoint_path=checkpoint_file_path,
                 model_source=ModelSource.CUSTOM,
             )
 

supervisely/nn/training/__init__.py

@@ -0,0 +1,2 @@
+from supervisely.nn.training.train_app import TrainApp
+from supervisely.nn.training.loggers import train_logger

supervisely/nn/training/gui/classes_selector.py

@@ -11,36 +11,36 @@ class ClassesSelector:
     lock_message = "Select training and validation splits to unlock"
 
     def __init__(self, project_id: int, classes: list, app_options: dict = {}):
-        self.classes_table = ClassesTable(project_id=project_id)  # use dataset_ids
-        if len(classes) > 0:
-            self.classes_table.select_classes(classes)  # from app options
-        else:
-            self.classes_table.select_all()
+        self.display_widgets = []
 
+        # 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(
             text=f"<i class='zmdi zmdi-chart-donut' style='color: #7f858e'></i> <a href='{qa_stats_link}' target='_blank'> <b> QA & Stats </b></a>"
         )
 
+        self.classes_table = ClassesTable(project_id=project_id)
+        if len(classes) > 0:
+            self.classes_table.select_classes(classes)
+        else:
+            self.classes_table.select_all()
+
         self.validator_text = Text("")
         self.validator_text.hide()
         self.button = Button("Select")
-        container = Container(
-            [
-                qa_stats_text,
-                self.classes_table,
-                self.validator_text,
-                self.button,
-            ]
+        self.display_widgets.extend(
+            [qa_stats_text, self.classes_table, self.validator_text, self.button]
         )
+        # -------------------------------- #
+
+        self.container = Container(self.display_widgets)
         self.card = Card(
             title=self.title,
             description=self.description,
-            content=container,
+            content=self.container,
             lock_message=self.lock_message,
             collapsable=app_options.get("collapsable", False),
         )

supervisely/nn/training/gui/gui.py

@@ -6,7 +6,7 @@ training workflows in Supervisely.
 """
 
 import supervisely.io.env as sly_env
-from supervisely import Api
+from supervisely import Api, ProjectMeta
 from supervisely._utils import is_production
 from supervisely.app.widgets import Stepper, Widget
 from supervisely.nn.training.gui.classes_selector import ClassesSelector
@@ -17,7 +17,7 @@ from supervisely.nn.training.gui.train_val_splits_selector import TrainValSplits
 from supervisely.nn.training.gui.training_logs import TrainingLogs
 from supervisely.nn.training.gui.training_process import TrainingProcess
 from supervisely.nn.training.gui.utils import set_stepper_step, wrap_button_click
-from supervisely.nn.utils import ModelSource
+from supervisely.nn.utils import ModelSource, RuntimeType
 
 
 class TrainGUI:
@@ -62,6 +62,7 @@ class TrainGUI:
         self.workspace_id = sly_env.workspace_id()
         self.project_id = sly_env.project_id()  # from app options?
         self.project_info = self._api.project.get_info_by_id(self.project_id)
+        self.project_meta = ProjectMeta.from_json(self._api.project.get_meta(self.project_id))
 
         # 1. Project selection + Train/val split
         self.input_selector = InputSelector(self.project_info, self.app_options)
@@ -399,7 +400,7 @@ class TrainGUI:
 
             app_state = {
                 "input": {"project_id": 43192},
-                "train_val_splits": {
+                "train_val_split": {
                     "method": "random",
                     "split": "train",
                     "percent": 90
@@ -415,13 +416,18 @@ class TrainGUI:
                         "enable": True,
                         "speed_test": True
                     },
-                    "cache_project": True
+                    "cache_project": True,
+                "export": {
+                    "enable": True,
+                    "ONNXRuntime": True,
+                    "TensorRT": True
+                    },
                 }
             }
         """
         app_state = self.validate_app_state(app_state)
 
-        options = app_state["options"]
+        options = app_state.get("options", {})
         input_settings = app_state["input"]
         train_val_splits_settings = app_state["train_val_split"]
         classes_settings = app_state["classes"]
@@ -444,7 +450,7 @@ class TrainGUI:
         :type options: dict
         """
         # Set Input
-        self.input_selector.set_cache(options["cache_project"])
+        self.input_selector.set_cache(options.get("cache_project", True))
         self.input_selector_cb()
         # ----------------------------------------- #
 
@@ -527,13 +533,22 @@ class TrainGUI:
         """
         self.hyperparameters_selector.set_hyperparameters(hyperparameters_settings)
 
-        model_benchmark_settings = options["model_benchmark"]
-        self.hyperparameters_selector.set_model_benchmark_checkbox_value(
-            model_benchmark_settings["enable"]
-        )
-        self.hyperparameters_selector.set_speedtest_checkbox_value(
-            model_benchmark_settings["speed_test"]
-        )
+        model_benchmark_settings = options.get("model_benchmark", None)
+        if model_benchmark_settings is not None:
+            self.hyperparameters_selector.set_model_benchmark_checkbox_value(
+                model_benchmark_settings["enable"]
+            )
+            self.hyperparameters_selector.set_speedtest_checkbox_value(
+                model_benchmark_settings["speed_test"]
+            )
+        export_weights_settings = options.get("export", None)
+        if export_weights_settings is not None:
+            self.hyperparameters_selector.set_export_onnx_checkbox_value(
+                export_weights_settings.get(RuntimeType.ONNXRUNTIME, False)
+            )
+            self.hyperparameters_selector.set_export_tensorrt_checkbox_value(
+                export_weights_settings.get(RuntimeType.TENSORRT, False)
+            )
         self.hyperparameters_selector_cb()
 
     # ----------------------------------------- #

supervisely/nn/training/gui/hyperparameters_selector.py

@@ -9,6 +9,7 @@ from supervisely.app.widgets import (
     Field,
     Text,
 )
+from supervisely.nn.utils import RuntimeType
 
 
 class HyperparametersSelector:
@@ -17,55 +18,75 @@ class HyperparametersSelector:
     lock_message = "Select model to unlock"
 
     def __init__(self, hyperparameters: dict, app_options: dict = {}):
+        self.display_widgets = []
         self.app_options = app_options
+
+        # GUI Components
         self.editor = Editor(
             hyperparameters, height_lines=50, language_mode="yaml", auto_format=True
         )
+        self.display_widgets.extend([self.editor])
 
-        # Model Benchmark
-        self.run_model_benchmark_checkbox = Checkbox(
-            content="Run Model Benchmark evaluation", checked=True
-        )
-        self.run_speedtest_checkbox = Checkbox(content="Run speed test", checked=True)
-
-        self.model_benchmark_field = Field(
-            Container(
-                widgets=[
-                    self.run_model_benchmark_checkbox,
-                    self.run_speedtest_checkbox,
-                ]
-            ),
-            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.",
-        )
-        docs_link = '<a href="https://docs.supervisely.com/neural-networks/model-evaluation-benchmark/" target="_blank">documentation</a>'
-        self.model_benchmark_learn_more = Text(
-            f"Learn more about Model Benchmark in the {docs_link}.", status="info"
-        )
-
+        # Optional Model Benchmark
         if app_options.get("model_benchmark", True):
-            self.model_benchmark_field.show()
-            self.model_benchmark_learn_more.show()
-        else:
-            self.model_benchmark_field.hide()
-            self.model_benchmark_learn_more.hide()
+            # Model Benchmark
+            self.run_model_benchmark_checkbox = Checkbox(
+                content="Run Model Benchmark evaluation", checked=True
+            )
+            self.run_speedtest_checkbox = Checkbox(content="Run speed test", checked=True)
+
+            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.",
+                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>'
+            self.model_benchmark_learn_more = Text(
+                f"Learn more about Model Benchmark in the {docs_link}.", status="info"
+            )
+            self.display_widgets.extend(
+                [self.model_benchmark_field, self.model_benchmark_learn_more]
+            )
+        # -------------------------------- #
+
+        # Optional Export Weights
+        export_onnx_supported = app_options.get("export_onnx_supported", False)
+        export_tensorrt_supported = app_options.get("export_tensorrt_supported", False)
+
+        onnx_name = "ONNX"
+        tensorrt_name = "TensorRT engine"
+        export_runtimes = []
+        export_runtime_names = []
+        if export_onnx_supported:
+            self.export_onnx_checkbox = Checkbox(content=f"Export to {onnx_name}")
+            export_runtimes.append(self.export_onnx_checkbox)
+            export_runtime_names.append(onnx_name)
+        if export_tensorrt_supported:
+            self.export_tensorrt_checkbox = Checkbox(content=f"Export to {tensorrt_name}")
+            export_runtimes.append(self.export_tensorrt_checkbox)
+            export_runtime_names.append(tensorrt_name)
+        if export_onnx_supported or export_tensorrt_supported:
+            export_field_description = ", ".join(export_runtime_names)
+            runtime_container = Container(export_runtimes)
+            self.export_field = Field(
+                title="Export model",
+                description=f"Export best checkpoint to the following formats: {export_field_description}.",
+                content=runtime_container,
+            )
+            self.display_widgets.extend([self.export_field])
+        # -------------------------------- #
 
         self.validator_text = Text("")
         self.validator_text.hide()
         self.button = Button("Select")
-        container = Container(
-            [
-                self.editor,
-                self.model_benchmark_field,
-                self.model_benchmark_learn_more,
-                self.validator_text,
-                self.button,
-            ]
-        )
+        self.display_widgets.extend([self.validator_text, self.button])
+        # -------------------------------- #
+
+        self.container = Container(self.display_widgets)
         self.card = Card(
             title=self.title,
             description=self.description,
-            content=container,
+            content=self.container,
             lock_message=self.lock_message,
             collapsable=app_options.get("collapsable", False),
         )
@@ -73,11 +94,14 @@ class HyperparametersSelector:
 
     @property
     def widgets_to_disable(self) -> list:
-        return [
-            self.editor,
-            self.run_model_benchmark_checkbox,
-            self.run_speedtest_checkbox,
-        ]
+        widgets = [self.editor]
+        if self.app_options.get("model_benchmark", True):
+            widgets.extend([self.run_model_benchmark_checkbox, self.run_speedtest_checkbox])
+        if self.app_options.get("export_onnx_supported", False):
+            widgets.append(self.export_onnx_checkbox)
+        if self.app_options.get("export_tensorrt_supported", False):
+            widgets.append(self.export_tensorrt_checkbox)
+        return widgets
 
     def set_hyperparameters(self, hyperparameters: Union[str, dict]) -> None:
         self.editor.set_text(hyperparameters)
@@ -113,5 +137,30 @@ class HyperparametersSelector:
         else:
             self.run_speedtest_checkbox.hide()
 
+    def get_export_onnx_checkbox_value(self) -> bool:
+        if self.app_options.get("export_onnx_supported", False):
+            return self.export_onnx_checkbox.is_checked()
+        return False
+
+    def set_export_onnx_checkbox_value(self, value: bool) -> None:
+        if value:
+            self.export_onnx_checkbox.check()
+        else:
+            self.export_onnx_checkbox.uncheck()
+
+    def get_export_tensorrt_checkbox_value(self) -> bool:
+        if self.app_options.get("export_tensorrt_supported", False):
+            return self.export_tensorrt_checkbox.is_checked()
+        return False
+
+    def set_export_tensorrt_checkbox_value(self, value: bool) -> None:
+        if value:
+            self.export_tensorrt_checkbox.check()
+        else:
+            self.export_tensorrt_checkbox.uncheck()
+
+    def is_export_required(self) -> bool:
+        return self.get_export_onnx_checkbox_value() or self.get_export_tensorrt_checkbox_value()
+
     def validate_step(self) -> bool:
         return True

supervisely/nn/training/gui/input_selector.py

@@ -17,9 +17,11 @@ class InputSelector:
     lock_message = None
 
     def __init__(self, project_info: ProjectInfo, app_options: dict = {}):
+        self.display_widgets = []
         self.project_id = project_info.id
         self.project_info = project_info
 
+        # GUI Components
         self.project_thumbnail = ProjectThumbnail(self.project_info)
 
         if is_cached(self.project_id):
@@ -32,27 +34,27 @@ class InputSelector:
         self.validator_text = Text("")
         self.validator_text.hide()
         self.button = Button("Select")
-        container = Container(
-            widgets=[
+        self.display_widgets.extend(
+            [
                 self.project_thumbnail,
                 self.use_cache_checkbox,
                 self.validator_text,
                 self.button,
             ]
         )
+        # -------------------------------- #
 
+        self.container = Container(self.display_widgets)
         self.card = Card(
             title=self.title,
             description=self.description,
-            content=container,
+            content=self.container,
             collapsable=app_options.get("collapsable", False),
         )
 
     @property
     def widgets_to_disable(self) -> list:
-        return [
-            self.use_cache_checkbox,
-        ]
+        return [self.use_cache_checkbox]
 
     def get_project_id(self) -> int:
         return self.project_id

supervisely/nn/training/gui/model_selector.py

@@ -21,15 +21,14 @@ class ModelSelector:
     lock_message = "Select classes to unlock"
 
     def __init__(self, api: Api, framework: str, models: list, app_options: dict = {}):
+        self.display_widgets = []
         self.team_id = sly_env.team_id()  # get from project id
         self.models = models
 
-        # Pretrained models
+        # GUI Components
         self.pretrained_models_table = PretrainedModelsSelector(self.models)
-
         experiment_infos = get_experiment_infos(api, self.team_id, framework)
         self.experiment_selector = ExperimentSelector(self.team_id, experiment_infos)
-        # Model source tabs
         self.model_source_tabs = RadioTabs(
             titles=[ModelSource.PRETRAINED, ModelSource.CUSTOM],
             descriptions=[
@@ -42,11 +41,14 @@ class ModelSelector:
         self.validator_text = Text("")
         self.validator_text.hide()
         self.button = Button("Select")
-        container = Container([self.model_source_tabs, self.validator_text, self.button])
+        self.display_widgets.extend([self.model_source_tabs, self.validator_text, self.button])
+        # -------------------------------- #
+
+        self.container = Container(self.display_widgets)
         self.card = Card(
             title=self.title,
             description=self.description,
-            content=container,
+            content=self.container,
             lock_message=self.lock_message,
             collapsable=app_options.get("collapsable", False),
         )

supervisely/nn/training/gui/train_val_splits_selector.py

@@ -10,10 +10,12 @@ class TrainValSplitsSelector:
     lock_message = "Select input options to unlock"
 
     def __init__(self, api: Api, project_id: int, app_options: dict = {}):
+        self.display_widgets = []
         self.api = api
         self.project_id = project_id
-        self.train_val_splits = TrainValSplits(project_id)
 
+        # GUI Components
+        self.train_val_splits = TrainValSplits(project_id)
         train_val_dataset_ids = {"train": [], "val": []}
         for _, dataset in api.dataset.tree(project_id):
             if dataset.name.lower() == "train" or dataset.name.lower() == "training":
@@ -39,17 +41,14 @@ class TrainValSplitsSelector:
             self.validator_text.hide()
 
         self.button = Button("Select")
-        container = Container(
-            [
-                self.train_val_splits,
-                self.validator_text,
-                self.button,
-            ]
-        )
+        self.display_widgets.extend([self.train_val_splits, self.validator_text, self.button])
+        # -------------------------------- #
+
+        self.container = Container(self.display_widgets)
         self.card = Card(
             title=self.title,
             description=self.description,
-            content=container,
+            content=self.container,
             lock_message=self.lock_message,
             collapsable=app_options.get("collapsable", False),
         )

supervisely/nn/training/gui/training_logs.py

@@ -12,14 +12,13 @@ class TrainingLogs:
     lock_message = "Start training to unlock"
 
     def __init__(self, app_options: Dict[str, Any]):
+        self.display_widgets = []
         api = Api.from_env()
         self.app_options = app_options
 
-        self.progress_bar_main = Progress(hide_on_finish=False)
-        self.progress_bar_main.hide()
-
-        self.progress_bar_secondary = Progress(hide_on_finish=False)
-        self.progress_bar_secondary.hide()
+        # GUI Components
+        self.validator_text = Text("")
+        self.validator_text.hide()
 
         if is_production():
             task_id = get_task_id(raise_not_found=False)
@@ -43,16 +42,9 @@ class TrainingLogs:
         )
         self.tensorboard_button.disable()
 
-        self.validator_text = Text("")
-        self.validator_text.hide()
-
-        container_widgets = [
-            self.validator_text,
-            self.tensorboard_button,
-            self.progress_bar_main,
-            self.progress_bar_secondary,
-        ]
+        self.display_widgets.extend([self.validator_text, self.tensorboard_button])
 
+        # Optional Show logs button
         if app_options.get("show_logs_in_gui", False):
             self.logs_button = Button(
                 text="Show logs",
@@ -63,14 +55,22 @@ class TrainingLogs:
             self.task_logs = TaskLogs(task_id)
             self.task_logs.hide()
             logs_container = Container([self.logs_button, self.task_logs])
-            container_widgets.insert(2, logs_container)
+            self.display_widgets.extend([logs_container])
+        # -------------------------------- #
 
-        container = Container(container_widgets)
+        # Progress bars
+        self.progress_bar_main = Progress(hide_on_finish=False)
+        self.progress_bar_main.hide()
+        self.progress_bar_secondary = Progress(hide_on_finish=False)
+        self.progress_bar_secondary.hide()
+        self.display_widgets.extend([self.progress_bar_main, self.progress_bar_secondary])
+        # -------------------------------- #
 
+        self.container = Container(self.display_widgets)
         self.card = Card(
             title=self.title,
             description=self.description,
-            content=container,
+            content=self.container,
             lock_message=self.lock_message,
         )
         self.card.lock()