supervisely 6.73.373__py3-none-any.whl → 6.73.375__py3-none-any.whl

supervisely/app/widgets/train_val_splits/train_val_splits.py

@@ -1,4 +1,5 @@
 import os
+from collections import defaultdict
 from typing import Dict, List, Literal, Optional, Tuple, Union
 
 import supervisely as sly
@@ -17,6 +18,7 @@ from supervisely.app.widgets import (
 from supervisely.app.widgets.random_splits_table.random_splits_table import (
     RandomSplitsTable,
 )
+from supervisely.app.widgets.select_collection.select_collection import SelectCollection
 from supervisely.app.widgets.select_dataset_tree.select_dataset_tree import (
     SelectDatasetTree,
 )
@@ -41,28 +43,27 @@ class TrainValSplits(Widget):
         tags_splits: Optional[bool] = True,
         datasets_splits: Optional[bool] = True,
         widget_id: Optional[int] = None,
+        collections_splits: Optional[bool] = False,
     ):
         self._project_id = project_id
         self._project_fs = project_fs
 
-        if project_fs is not None and project_id is not None:
-            raise ValueError(
-                "You can not provide both project_id and project_fs parameters to TrainValSplits widget."
-            )
-        if project_fs is None and project_id is None:
-            raise ValueError(
-                "You should provide at least one of: project_id or project_fs parameters to TrainValSplits widget."
-            )
-
         self._project_info = None
+        self._project_type = None
+        self._project_class = None
+        self._api = None
         if project_id is not None:
             self._api = Api()
             self._project_info = self._api.project.get_info_by_id(
                 self._project_id, raise_error=True
             )
 
-        self._project_type = project_fs.type if project_id is None else self._project_info.type
-        self._project_class = get_project_class(self._project_type)
+        if project_fs is not None:
+            self._project_type = project_fs.type
+        elif self._project_info is not None:
+            self._project_type = self._project_info.type
+        if self._project_type is not None:
+            self._project_class = get_project_class(self._project_type)
 
         self._random_splits_table: RandomSplitsTable = None
         self._train_tag_select: SelectTagMeta = None
@@ -70,6 +71,8 @@ class TrainValSplits(Widget):
         self._untagged_select: SelectString = None
         self._train_ds_select: Union[SelectDatasetTree, SelectString] = None
         self._val_ds_select: Union[SelectDatasetTree, SelectString] = None
+        self._train_collections_select: SelectCollection = None
+        self._val_collections_select: SelectCollection = None
         self._split_methods = []
 
         contents = []
@@ -80,12 +83,18 @@ class TrainValSplits(Widget):
             contents.append(self._get_random_content())
         if tags_splits:
             self._split_methods.append("Based on item tags")
-            tabs_descriptions.append(f"{self._project_type.capitalize()} should have assigned train or val tag")
+            tabs_descriptions.append(
+                f"{self._project_type.capitalize()} should have assigned train or val tag"
+            )
             contents.append(self._get_tags_content())
         if datasets_splits:
             self._split_methods.append("Based on datasets")
             tabs_descriptions.append("Select one or several datasets for every split")
             contents.append(self._get_datasets_content())
+        if collections_splits:
+            self._split_methods.append("Based on collections")
+            tabs_descriptions.append("Select one or several collections for every split")
+            contents.append(self._get_collections_content())
         if not self._split_methods:
             raise ValueError(
                 "Any of split methods [random_splits, tags_splits, datasets_splits] must be specified in TrainValSplits."
@@ -216,6 +225,32 @@ class TrainValSplits(Widget):
             widgets=[notification_box, train_field, val_field], direction="vertical", gap=5
         )
 
+    def _get_collections_content(self):
+        notification_box = NotificationBox(
+            title="Notice: How to make equal splits",
+            description="Choose the same collection(s) for train/validation to make splits equal. Can be used for debug and for tiny projects",
+            box_type="info",
+        )
+
+        self._train_collections_select = SelectCollection(multiselect=True, compact=True)
+        self._val_collections_select = SelectCollection(multiselect=True, compact=True)
+        if self._project_id is not None:
+            self._train_collections_select.set_project_id(self._project_id)
+            self._val_collections_select.set_project_id(self._project_id)
+        train_field = Field(
+            self._train_collections_select,
+            title="Train collection(s)",
+            description="all images in selected collection(s) are considered as training set",
+        )
+        val_field = Field(
+            self._val_collections_select,
+            title="Validation collection(s)",
+            description="all images in selected collection(s) are considered as validation set",
+        )
+        return Container(
+            widgets=[notification_box, train_field, val_field], direction="vertical", gap=5
+        )
+
     def get_json_data(self):
         return {}
 
@@ -223,6 +258,8 @@ class TrainValSplits(Widget):
         return {}
 
     def get_splits(self) -> Tuple[List[ItemInfo], List[ItemInfo]]:
+        if self._project_id is None and self._project_fs is None:
+            raise ValueError("Both project_id and project_fs are None.")
         split_method = self._content.get_active_tab()
         tmp_project_dir = None
         train_set, val_set = [], []
@@ -276,18 +313,35 @@ class TrainValSplits(Widget):
             train_set, val_set = self._project_class.get_train_val_splits_by_dataset(
                 project_dir, train_ds_names, val_ds_names
             )
+        elif split_method == "Based on collections":
+            if self._project_id is None:
+                raise ValueError(
+                    "You can not use collections_splits parameter without project_id parameter."
+                )
+            train_collections = self._train_collections_select.get_selected_ids()
+            val_collections = self._val_collections_select.get_selected_ids()
+
+            train_set, val_set = self._project_class.get_train_val_splits_by_collections(
+                project_dir,
+                train_collections,
+                val_collections,
+                self._project_id,
+                self._api,
+            )
 
         if tmp_project_dir is not None:
             remove_dir(tmp_project_dir)
         return train_set, val_set
 
-    def set_split_method(self, split_method: Literal["random", "tags", "datasets"]):
+    def set_split_method(self, split_method: Literal["random", "tags", "datasets", "collections"]):
         if split_method == "random":
             split_method = "Random"
         elif split_method == "tags":
             split_method = "Based on item tags"
         elif split_method == "datasets":
             split_method = "Based on datasets"
+        elif split_method == "collections":
+            split_method = "Based on collections"
         self._content.set_active_tab(split_method)
         StateJson().send_changes()
         DataJson().send_changes()
@@ -337,6 +391,42 @@ class TrainValSplits(Widget):
     def get_val_dataset_ids(self) -> List[int]:
         return self._val_ds_select.get_selected_ids()
 
+    def set_project_id_for_collections(self, project_id: int):
+        if not isinstance(project_id, int):
+            raise ValueError("Project ID must be an integer.")
+        self._project_id = project_id
+        self._project_type = None
+        if self._api is None:
+            self._api = Api()
+        self._project_info = self._api.project.get_info_by_id(self._project_id, raise_error=True)
+        self._project_type = self._project_info.type
+        self._project_class = get_project_class(self._project_type)
+        if not self._train_collections_select or not self._val_collections_select:
+            raise ValueError("Collections select widgets are not initialized.")
+        self._train_collections_select.set_project_id(project_id)
+        self._val_collections_select.set_project_id(project_id)
+
+    def get_train_collections_ids(self) -> List[int]:
+        return self._train_collections_select.get_selected_ids() or []
+
+    def get_val_collections_ids(self) -> List[int]:
+        return self._val_collections_select.get_selected_ids() or []
+
+    def set_collections_splits(self, train_collections: List[int], val_collections: List[int]):
+        self._content.set_active_tab("Based on collections")
+        self.set_collections_splits_by_ids("train", train_collections)
+        self.set_collections_splits_by_ids("val", val_collections)
+
+    def set_collections_splits_by_ids(
+        self, split: Literal["train", "val"], collection_ids: List[int]
+    ):
+        if split == "train":
+            self._train_collections_select.set_collections(collection_ids)
+        elif split == "val":
+            self._val_collections_select.set_collections(collection_ids)
+        else:
+            raise ValueError("Split value must be 'train' or 'val'")
+
     def get_untagged_action(self) -> str:
         return self._untagged_select.get_value()
 
@@ -355,6 +445,10 @@ class TrainValSplits(Widget):
         if self._val_ds_select is not None:
             self._val_ds_select.disable()
         self._disabled = True
+        if self._train_collections_select is not None:
+            self._train_collections_select.disable()
+        if self._val_collections_select is not None:
+            self._val_collections_select.disable()
         DataJson()[self.widget_id]["disabled"] = self._disabled
         DataJson().send_changes()
 
@@ -373,5 +467,9 @@ class TrainValSplits(Widget):
         if self._val_ds_select is not None:
             self._val_ds_select.enable()
         self._disabled = False
+        if self._train_collections_select is not None:
+            self._train_collections_select.enable()
+        if self._val_collections_select is not None:
+            self._val_collections_select.enable()
         DataJson()[self.widget_id]["disabled"] = self._disabled
         DataJson().send_changes()

supervisely/nn/training/gui/gui.py

@@ -815,6 +815,20 @@ class TrainGUI:
                 raise ValueError("split must be 'train' or 'val'")
             if not isinstance(percent, int) or not 0 < percent < 100:
                 raise ValueError("percent must be an integer in range 1 to 99")
+        elif train_val_splits_settings.get("method") == "collections":
+            train_collections = train_val_splits_settings.get("train_collections", [])
+            val_collections = train_val_splits_settings.get("val_collections", [])
+            collection_ids = set()
+            for collection in self._api.entities_collection.get_list(self.project_id):
+                collection_ids.add(collection.id)
+            missing_collections_ids = set(train_collections + val_collections) - collection_ids
+            if missing_collections_ids:
+                missing_collections_text = ", ".join(
+                    [str(collection_id) for collection_id in missing_collections_ids]
+                )
+                raise ValueError(
+                    f"Collections with ids: {missing_collections_text} not found in the project"
+                )
         return app_state
 
     def load_from_app_state(self, app_state: Union[str, dict]) -> None:
@@ -849,7 +863,8 @@ class TrainGUI:
                     "ONNXRuntime": True,
                     "TensorRT": True
                     },
-                }
+                },
+                "experiment_name": "my_experiment",
             }
         """
         if isinstance(app_state, str):
@@ -863,6 +878,7 @@ class TrainGUI:
         tags_settings = app_state.get("tags", [])
         model_settings = app_state["model"]
         hyperparameters_settings = app_state["hyperparameters"]
+        experiment_name = app_state.get("experiment_name", None)
 
         self._init_input(input_settings, options)
         self._init_train_val_splits(train_val_splits_settings, options)
@@ -870,6 +886,8 @@ class TrainGUI:
         self._init_tags(tags_settings, options)
         self._init_model(model_settings, options)
         self._init_hyperparameters(hyperparameters_settings, options)
+        if experiment_name is not None:
+            self.training_process.set_experiment_name(experiment_name)
 
     def _init_input(self, input_settings: Union[dict, None], options: dict) -> None:
         """
@@ -938,6 +956,15 @@ class TrainGUI:
             self.train_val_splits_selector.train_val_splits.set_datasets_splits(
                 train_datasets, val_datasets
             )
+        elif split_method == "collections":
+            train_collections = train_val_splits_settings["train_collections"]
+            val_collections = train_val_splits_settings["val_collections"]
+            self.train_val_splits_selector.train_val_splits.set_project_id_for_collections(
+                self.project_id
+            )
+            self.train_val_splits_selector.train_val_splits.set_collections_splits(
+                train_collections, val_collections
+            )
         self.train_val_splits_selector_cb()
 
     def _init_classes(self, classes_settings: list, options: dict) -> None:

supervisely/nn/training/gui/train_val_splits_selector.py

@@ -26,53 +26,80 @@ class TrainValSplitsSelector:
         # GUI Components
         split_methods = self.app_options.get("train_val_split_methods", [])
         if len(split_methods) == 0:
-            split_methods = ["Random", "Based on tags", "Based on datasets"]
+            split_methods = ["Random", "Based on tags", "Based on datasets", "Based on collections"]
         random_split = "Random" in split_methods
         tag_split = "Based on tags" in split_methods
         ds_split = "Based on datasets" in split_methods
+        coll_split = "Based on collections" in split_methods
 
-        self.train_val_splits = TrainValSplits(project_id, None, random_split, tag_split, ds_split)
+        self.train_val_splits = TrainValSplits(
+            project_id, None, random_split, tag_split, ds_split, collections_splits=coll_split
+        )
+
+        # check for collections with "train" and "val" prefixes
+        all_collections = self.api.entities_collection.get_list(self.project_id)
+        train_collections = []
+        val_collections = []
+        collections_found = False
+        for collection in all_collections:
+            if collection.name.lower().startswith("train"):
+                train_collections.append(collection.id)
+            elif collection.name.lower().startswith("val"):
+                val_collections.append(collection.id)
+        
+        if len(train_collections) > 0 and len(val_collections) > 0:
+            self.train_val_splits.set_collections_splits(train_collections, val_collections)
+            self.validator_text = Text(
+                "Train and val collections are detected", status="info"
+            )
+            self.validator_text.show()
+            collections_found = True
+        else:
+            self.validator_text = Text("")
+            self.validator_text.hide()
+        
 
         def _extend_with_nested(root_ds):
             nested = self.api.dataset.get_nested(self.project_id, root_ds.id)
             nested_ids = [ds.id for ds in nested]
             return [root_ds.id] + nested_ids
 
-        train_val_dataset_ids = {"train": set(), "val": set()}
-        for _, dataset in self.api.dataset.tree(self.project_id):
-            ds_name = dataset.name.lower()
+        if not collections_found:
+            train_val_dataset_ids = {"train": set(), "val": set()}
+            for _, dataset in self.api.dataset.tree(self.project_id):
+                ds_name = dataset.name.lower()
 
-            if ds_name in {"train", "training"}:
-                for _id in _extend_with_nested(dataset):
-                    train_val_dataset_ids["train"].add(_id)
+                if ds_name in {"train", "training"}:
+                    for _id in _extend_with_nested(dataset):
+                        train_val_dataset_ids["train"].add(_id)
 
-            elif ds_name in {"val", "validation", "test", "testing"}:
-                for _id in _extend_with_nested(dataset):
-                    train_val_dataset_ids["val"].add(_id)
+                elif ds_name in {"val", "validation", "test", "testing"}:
+                    for _id in _extend_with_nested(dataset):
+                        train_val_dataset_ids["val"].add(_id)
 
-        train_val_dataset_ids["train"] = list(train_val_dataset_ids["train"])
-        train_val_dataset_ids["val"] = list(train_val_dataset_ids["val"])
+            train_val_dataset_ids["train"] = list(train_val_dataset_ids["train"])
+            train_val_dataset_ids["val"] = list(train_val_dataset_ids["val"])
 
-        train_count = len(train_val_dataset_ids["train"])
-        val_count = len(train_val_dataset_ids["val"])
+            train_count = len(train_val_dataset_ids["train"])
+            val_count = len(train_val_dataset_ids["val"])
 
-        if train_count > 0 and val_count > 0:
-            self.train_val_splits.set_datasets_splits(
-                train_val_dataset_ids["train"], train_val_dataset_ids["val"]
-            )
+            if train_count > 0 and val_count > 0:
+                self.train_val_splits.set_datasets_splits(
+                    train_val_dataset_ids["train"], train_val_dataset_ids["val"]
+                )
 
-        if train_count > 0 and val_count > 0:
-            if train_count == val_count == 1:
-                self.validator_text = Text("train and val datasets are detected", status="info")
+            if train_count > 0 and val_count > 0:
+                if train_count == val_count == 1:
+                    self.validator_text = Text("train and val datasets are detected", status="info")
+                else:
+                    self.validator_text = Text(
+                        "Multiple train and val datasets are detected. Check manually if selection is correct",
+                        status="info",
+                    )
+                self.validator_text.show()
             else:
-                self.validator_text = Text(
-                    "Multiple train and val datasets are detected. Check manually if selection is correct",
-                    status="info",
-                )
-            self.validator_text.show()
-        else:
-            self.validator_text = Text("")
-            self.validator_text.hide()
+                self.validator_text = Text("")
+                self.validator_text.hide()
 
         self.button = Button("Select")
         self.display_widgets.extend([self.train_val_splits, self.validator_text, self.button])
@@ -238,6 +265,68 @@ class TrainValSplitsSelector:
                 self.validator_text.set("Train and val datasets are selected", status="success")
                 return True
 
+        def validate_based_on_collections():
+            train_collection_id = self.train_val_splits.get_train_collections_ids()
+            val_collection_id = self.train_val_splits.get_val_collections_ids()
+            if train_collection_id is None and val_collection_id is None:
+                self.validator_text.set("No collections are selected", status="error")
+                return False
+            if len(train_collection_id) == 0 or len(val_collection_id) == 0:
+                self.validator_text.set("Collections are not selected", status="error")
+                return False
+            if set(train_collection_id) == set(val_collection_id):
+                self.validator_text.set(
+                    text=f"Same collections are selected for both train and val splits. {ensure_text} {warning_text}",
+                    status="warning",
+                )
+                return True
+            from supervisely.api.entities_collection_api import CollectionTypeFilter
+
+            train_items = set()
+            empty_train_collections = []
+            for collection_id in train_collection_id:
+                items = self.api.entities_collection.get_items(
+                    collection_id=collection_id,
+                    project_id=self.project_id,
+                    collection_type=CollectionTypeFilter.DEFAULT,
+                )
+                train_items.update([item.id for item in items])
+                if len(items) == 0:
+                    empty_train_collections.append(collection_id)
+            val_items = set()
+            empty_val_collections = []
+            for collection_id in val_collection_id:
+                items = self.api.entities_collection.get_items(
+                    collection_id=collection_id,
+                    project_id=self.project_id,
+                    collection_type=CollectionTypeFilter.DEFAULT,
+                )
+                val_items.update([item.id for item in items])
+                if len(items) == 0:
+                    empty_val_collections.append(collection_id)
+            if len(train_items) == 0 and len(val_items) == 0:
+                self.validator_text.set(
+                    text="All selected collections are empty. ",
+                    status="error",
+                )
+                return False
+            if len(empty_train_collections) > 0 or len(empty_val_collections) > 0:
+                empty_collections_text = "Selected collections are empty. "
+                if len(empty_train_collections) > 0:
+                    empty_collections_text += f"train: {', '.join(empty_train_collections)}. "
+                if len(empty_val_collections) > 0:
+                    empty_collections_text += f"val: {', '.join(empty_val_collections)}. "
+                empty_collections_text += f"{ensure_text}"
+                self.validator_text.set(
+                    text=empty_collections_text,
+                    status="error",
+                )
+                return True
+
+            else:
+                self.validator_text.set("Train and val collections are selected", status="success")
+                return True
+
         if split_method == "Random":
             is_valid = validate_random_split()
 
@@ -246,6 +335,8 @@ class TrainValSplitsSelector:
 
         elif split_method == "Based on datasets":
             is_valid = validate_based_on_datasets()
+        elif split_method == "Based on collections":
+            is_valid = validate_based_on_collections()
 
         # @TODO: handle button correctly if validation fails. Do not unlock next card until validation passes if returned False
         self.validator_text.show()
@@ -268,3 +359,15 @@ class TrainValSplitsSelector:
 
     def set_val_dataset_ids(self, dataset_ids: List[int]) -> None:
         self.train_val_splits._val_ds_select.set_selected_ids(dataset_ids)
+
+    def get_train_collection_ids(self) -> List[int]:
+        return self.train_val_splits._train_collections_select.get_selected_ids()
+
+    def set_train_collection_ids(self, collection_ids: List[int]) -> None:
+        self.train_val_splits._train_collections_select.set_selected_ids(collection_ids)
+
+    def get_val_collection_ids(self) -> List[int]:
+        return self.train_val_splits._val_collections_select.get_selected_ids()
+
+    def set_val_collection_ids(self, collection_ids: List[int]) -> None:
+        self.train_val_splits._val_collections_select.set_selected_ids(collection_ids)

supervisely/nn/training/gui/training_logs.py

@@ -36,6 +36,7 @@ class TrainingLogs:
         self.display_widgets = []
         self.app_options = app_options
         api = Api.from_env()
+        team_id = sly_env.team_id()
 
         # GUI Components
         self.validator_text = Text("")
@@ -73,6 +74,7 @@ class TrainingLogs:
             module_info = gui_utils.get_module_info_by_name(api, app_name)
             if module_info is not None:
                 self.tensorboard_offline_button = RunAppButton(
+                    team_id=team_id,
                     workspace_id=workspace_id,
                     module_id=module_info["id"],
                     payload={},
@@ -81,9 +83,10 @@ class TrainingLogs:
                     plain=True,
                     icon="zmdi zmdi-chart",
                     available_in_offline=True,
-                    visible_by_vue_field="isStaticVersion",
+                    visible_by_vue_field=None,
                 )
                 self.tensorboard_offline_button.disable()
+                self.tensorboard_offline_button.hide()
                 self.display_widgets.extend([self.tensorboard_offline_button])
             else:
                 logger.warning(

supervisely/nn/training/gui/utils.py

@@ -136,3 +136,26 @@ def get_module_info_by_name(api: Api, app_name: str) -> Union[Dict, None]:
         if module["name"] == app_name:
             app_info = api.app.get_info(module["id"])
             return app_info
+
+
+def generate_task_check_function_js(folder: str) -> str:
+    """
+    Returns JavaScript function code for checking existing tasks.
+
+    :param folder: Remote folder to check.
+    :type folder: str
+    :return: JavaScript function code for checking existing tasks.
+    :rtype: str
+    """
+    escaped_folder = folder.replace("'", "\\'")
+    js_code = f"""
+        if (!task || !task.meta || !task.meta.params || !task.meta.params.state) {{
+            return false;
+        }}
+        const taskFolder = task.meta.params.state.slyFolder;
+        if (!taskFolder || typeof taskFolder !== 'string') {{
+            return false;
+        }}
+        return taskFolder === '{escaped_folder}';
+    """
+    return js_code

supervisely/nn/training/train_app.py

@@ -10,6 +10,7 @@ import subprocess
 from datetime import datetime
 from os import getcwd, listdir, walk
 from os.path import basename, dirname, exists, expanduser, isdir, isfile, join
+from time import sleep
 from typing import Any, Dict, List, Literal, Optional, Union
 from urllib.request import urlopen
 
@@ -57,6 +58,7 @@ from supervisely.nn.inference import RuntimeType, SessionJSON
 from supervisely.nn.inference.inference import Inference
 from supervisely.nn.task_type import TaskType
 from supervisely.nn.training.gui.gui import TrainGUI
+from supervisely.nn.training.gui.utils import generate_task_check_function_js
 from supervisely.nn.training.loggers import setup_train_logger, train_logger
 from supervisely.nn.utils import ModelSource, _get_model_name
 from supervisely.output import set_directory
@@ -206,16 +208,41 @@ class TrainApp:
         def _train_from_api(response: Response, request: Request):
             try:
                 state = request.state.state
+                wait = state.get("wait", True)
                 app_state = state["app_state"]
                 self.gui.load_from_app_state(app_state)
 
-                self._wrapped_start_training()
+                if wait:
+                    self._wrapped_start_training()
+                else:
+                    import threading
+
+                    training_thread = threading.Thread(
+                        target=self._wrapped_start_training,
+                        daemon=True,
+                    )
+                    training_thread.start()
+                    return {"result": "model training started"}
 
                 return {"result": "model was successfully trained"}
             except Exception as e:
                 self.gui.training_process.start_button.loading = False
                 raise e
 
+        # # Get training status
+        # @self._server.post("/train_status")
+        # def _train_status(response: Response, request: Request):
+        #     """Returns the current training status."""
+        #     status = self.gui.training_process.validator_text.get_value()
+        #     if status == "Training is in progress...":
+        #         try:
+        #             total_epochs = self.progress_bar_main.total
+        #             current_epoch = self.progress_bar_main.current
+        #             status += f" (Epoch {current_epoch}/{total_epochs})"
+        #         except Exception:
+        #             pass
+        #     return {"status": status}
+
     def _register_routes(self):
         """
         Registers API routes for TensorBoard and training endpoints.
@@ -1868,6 +1895,13 @@ class TrainApp:
                     "val_datasets": self.gui.train_val_splits_selector.train_val_splits.get_val_dataset_ids(),
                 }
             )
+        elif split_method == "Based on collections":
+            train_val_splits.update(
+                {
+                    "train_collections": self.gui.train_val_splits_selector.get_train_collection_ids(),
+                    "val_collections": self.gui.train_val_splits_selector.get_val_collection_ids(),
+                }
+            )
         return train_val_splits
 
     def _get_model_config_for_app_state(self, experiment_info: Dict = None) -> Dict:
@@ -1949,11 +1983,15 @@ class TrainApp:
             self.progress_bar_main.hide()
 
         file_info = self._api.file.get_info_by_path(self.team_id, join(remote_dir, "open_app.lnk"))
+
         # Set offline tensorboard button payload
         if is_production():
-            self.gui.training_logs.tensorboard_offline_button.payload = {
-                "state": {"slyFolder": f"{join(remote_dir, 'logs')}"}
-            }
+            remote_log_dir = join(remote_dir, "logs")
+            tb_btn_payload = {"state": {"slyFolder": remote_log_dir}}
+            self.gui.training_logs.tensorboard_offline_button.payload = tb_btn_payload
+            self.gui.training_logs.tensorboard_offline_button.set_check_existing_task_cb(
+                generate_task_check_function_js(remote_log_dir)
+            )
             self.gui.training_logs.tensorboard_offline_button.enable()
 
         return remote_dir, file_info
@@ -2104,7 +2142,7 @@ class TrainApp:
         ]
         task_type = experiment_info["task_type"]
         if task_type not in supported_task_types:
-            logger.warn(
+            logger.warning(
                 f"Task type: '{task_type}' is not supported for Model Benchmark. "
                 f"Supported tasks: {', '.join(task_type)}"
             )
@@ -2608,6 +2646,10 @@ class TrainApp:
             self.gui.training_process.start_button.loading = False
 
             # Shutdown the app after training is finished
+
+            self.gui.training_logs.tensorboard_button.hide()
+            self.gui.training_logs.tensorboard_offline_button.show()
+            sleep(1)  # wait for the button to be shown
             self.app.shutdown()
         except Exception as e:
             message = f"Error occurred during finalizing and uploading training artifacts. {check_logs_text}"