supervisely 6.73.420__py3-none-any.whl → 6.73.422__py3-none-any.whl

supervisely/nn/training/gui/train_val_splits_selector.py

@@ -3,6 +3,7 @@ from typing import List
 from supervisely import Api, Project
 from supervisely.app.widgets import Button, Card, Container, Text, TrainValSplits
 from supervisely.api.module_api import ApiField
+from supervisely.api.entities_collection_api import EntitiesCollectionInfo
 
 class TrainValSplitsSelector:
     title = "Train / Val Splits"
@@ -18,6 +19,13 @@ class TrainValSplitsSelector:
         self.card = None
         # -------------------------------- #
 
+        # Automated Splits
+        self._all_train_collections = []
+        self._all_val_collections = []
+        self._latest_train_collection = None
+        self._latest_val_collection = None
+        # -------------------------------- #
+
         self.display_widgets = []
         self.app_options = app_options
         self.api = api
@@ -32,75 +40,9 @@ class TrainValSplitsSelector:
         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, 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
-
-        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)
-
-                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_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:
-                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("")
-                self.validator_text.hide()
+        self.train_val_splits = TrainValSplits(project_id, None, random_split, tag_split, ds_split, collections_splits=coll_split)
 
+        self._detect_splits(coll_split, ds_split)
         self.button = Button("Select")
         self.display_widgets.extend([self.train_val_splits, self.validator_text, self.button])
         # -------------------------------- #
@@ -115,6 +57,22 @@ class TrainValSplitsSelector:
         )
         self.card.lock()
 
+    @property
+    def all_train_collections(self) -> List[EntitiesCollectionInfo]:
+        return self._all_train_collections
+
+    @property
+    def all_val_collections(self) -> List[EntitiesCollectionInfo]:
+        return self._all_val_collections
+
+    @property
+    def latest_train_collection(self) -> EntitiesCollectionInfo:
+        return self._latest_train_collection
+
+    @property
+    def latest_val_collection(self) -> EntitiesCollectionInfo:
+        return self._latest_val_collection
+
     @property
     def widgets_to_disable(self) -> list:
         return [self.train_val_splits]
@@ -313,10 +271,11 @@ class TrainValSplitsSelector:
                 return False
             if len(empty_train_collections) > 0 or len(empty_val_collections) > 0:
                 empty_collections_text = "Selected collections are empty. "
+                # @TODO: Use collection names instead of ids
                 if len(empty_train_collections) > 0:
-                    empty_collections_text += f"train: {', '.join(empty_train_collections)}. "
+                    empty_collections_text += f"train: {', '.join([str(collection_id) for collection_id in empty_train_collections])}. "
                 if len(empty_val_collections) > 0:
-                    empty_collections_text += f"val: {', '.join(empty_val_collections)}. "
+                    empty_collections_text += f"val: {', '.join([str(collection_id) for collection_id in empty_val_collections])}. "
                 empty_collections_text += f"{ensure_text}"
                 self.validator_text.set(
                     text=empty_collections_text,
@@ -372,3 +331,96 @@ class TrainValSplitsSelector:
 
     def set_val_collection_ids(self, collection_ids: List[int]) -> None:
         self.train_val_splits._val_collections_select.set_selected_ids(collection_ids)
+
+    def _detect_splits(self, collections_split: bool, datasets_split: bool) -> bool:
+        """Detect splits based on the selected method"""
+        splits_found = False
+        if collections_split:
+            splits_found = self._detect_collections()
+        if not splits_found and datasets_split:
+            splits_found = self._detect_datasets()
+        return splits_found
+
+    def _detect_collections(self) -> bool:
+        """Find collections with train and val prefixes and set them to train_val_splits"""
+        def _get_latest_collection(collections: List[EntitiesCollectionInfo]) -> EntitiesCollectionInfo:
+            curr_collection = None
+            curr_idx = 0
+            for collection in collections:
+                collection_idx = int(collection.name.rsplit('_', 1)[-1])
+                if collection_idx > curr_idx:
+                    curr_idx = collection_idx
+                    curr_collection = collection
+            return curr_collection
+
+        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)
+            elif collection.name.lower().startswith("val_"):
+                val_collections.append(collection)
+
+        train_collection = _get_latest_collection(train_collections)
+        val_collection = _get_latest_collection(val_collections)
+        if train_collection is not None and val_collection is not None:
+            self.train_val_splits.set_collections_splits([train_collection.id], [val_collection.id])
+            self.validator_text = Text("Train and val collections are detected", status="info")
+            self.validator_text.show()
+            collections_found = True
+            self._all_train_collections = train_collections
+            self._all_val_collections = val_collections
+            self._latest_train_collection = train_collection
+            self._latest_val_collection = val_collection
+        else:
+            self.validator_text = Text("")
+            self.validator_text.hide()
+            collections_found = False
+        return collections_found
+        
+    def _detect_datasets(self) -> bool:
+        """Find datasets with train and val prefixes and set them to train_val_splits"""
+        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
+
+        datasets_found = False
+        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)
+
+            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_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"])
+            datasets_found = True
+
+        if train_count > 0 and val_count > 0:
+            if train_count == val_count == 1:
+                message = "train and val datasets are detected"
+            else:
+                message = "Multiple train and val datasets are detected. Check manually if selection is correct"
+            
+            self.validator_text = Text(message, status="info")
+            self.validator_text.show()
+            datasets_found = True
+        else:
+            self.validator_text = Text("")
+            self.validator_text.hide()
+            datasets_found = False
+        return datasets_found

supervisely/nn/training/gui/training_artifacts.py

@@ -65,11 +65,6 @@ class TrainingArtifacts:
 
         # Outputs
         need_generate_report = self.app_options.get("generate_report", False)
-
-        # @TODO: temporary code to generate report for dev only
-        is_dev = "dev.internal" in api.server_address
-        if not is_dev:
-            need_generate_report = False
         # ------------------------------------------------------------ #
 
         if need_generate_report:

supervisely/nn/training/train_app.py

@@ -56,7 +56,7 @@ from supervisely.nn.benchmark import (
     SemanticSegmentationEvaluator,
 )
 from supervisely.nn.inference import RuntimeType, SessionJSON
-from supervisely.nn.inference.inference import Inference
+from supervisely.nn.inference.inference import Inference, torch_load_safe
 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
@@ -72,6 +72,7 @@ from supervisely.project.download import (
     is_cached,
 )
 from supervisely.template.experiment.experiment_generator import ExperimentGenerator
+from supervisely.api.entities_collection_api import EntitiesCollectionInfo
 
 
 class TrainApp:
@@ -159,8 +160,14 @@ class TrainApp:
         self.sly_project = None
         # -------------------------- #
 
-        self._train_split = None
-        self._val_split = None
+        # Train Val Splits
+        self._train_split = []
+        self._train_split_item_ids = set()
+        self._train_collection_id = None
+
+        self._val_split = []
+        self._val_split_item_ids = set()
+        self._val_collection_id = None
         # -------------------------- #
 
         # Input
@@ -232,19 +239,33 @@ class TrainApp:
                 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}
+        @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 = getattr(self.progress_bar_main, "total", None)
+                    current_epoch = getattr(self.progress_bar_main, "current", None)
+                    if total_epochs is not None and current_epoch is not None:
+                        status += f" (Epoch {current_epoch}/{total_epochs})"
+                except Exception:
+                    pass
+            return {"status": status}
+
+        # Read GUI State when launched from experiment modal
+        state = self.gui._extract_state_from_env()
+        logger.debug(f"State: {state}")
+        gui_state_raw = state.get("guiState")
+        if gui_state_raw is not None:
+            logger.info("Loading GUI from state")
+            logger.debug(f"GUI State: {gui_state_raw}")
+            try:
+                self.gui.load_from_app_state(gui_state_raw)
+                logger.info("Successfully loaded GUI from state")
+            except Exception as e:
+                raise e
+        # ----------------------------------------- #
 
     def _register_routes(self):
         """
@@ -399,14 +420,14 @@ class TrainApp:
         :rtype: str
         """
         return self.gui.training_process.get_device()
-    
+
     @property
     def base_checkpoint(self) -> str:
         """
         Returns the name of the base checkpoint.
         """
         return self.gui.model_selector.get_checkpoint_name()
-    
+
     @property
     def base_checkpoint_link(self) -> str:
         """
@@ -596,13 +617,18 @@ class TrainApp:
             if self.gui.classes_selector is not None:
                 if self.gui.classes_selector.is_convert_class_shapes_enabled():
                     self._convert_project_to_model_task()
+
         # Step 4. Split Project
         self._split_project()
         # Step 5. Remove classes except selected
         if self.sly_project.type == ProjectType.IMAGES.value:
             self.sly_project.remove_classes_except(self.project_dir, self.classes, True)
             self._read_project()
-        # Step 6. Download Model files
+
+        # Step 6. Create collections
+        self._create_collection_splits()
+
+        # Step 7. Download Model files
         self._download_model()
 
     def _finalize(self, experiment_info: dict) -> None:
@@ -847,7 +873,7 @@ class TrainApp:
                         "TensorRT": True
                     },
                 },
-                "experiment_name": "my_experiment",
+                "experiment_name": "My Experiment",
             }
         """
         self.gui.load_from_app_state(app_state)
@@ -1162,13 +1188,15 @@ class TrainApp:
             self._train_val_split_file = None
             self._train_split = []
             self._val_split = []
+            self._train_split_item_ids = set()
+            self._val_split_item_ids = set()
             return
 
         # Load splits
         self.gui.train_val_splits_selector.set_sly_project(self.sly_project)
-        self._train_split, self._val_split = (
-            self.gui.train_val_splits_selector.train_val_splits.get_splits()
-        )
+        self._train_split, self._val_split = self.gui.train_val_splits_selector.train_val_splits.get_splits()
+        self._train_split_ids, self._val_split_ids = [], []
+        self._train_split_item_ids, self._val_split_item_ids = set(), set()
 
         # Prepare paths
         project_split_path = join(self.work_dir, "splits")
@@ -1177,11 +1205,13 @@ class TrainApp:
                 "split_path": join(project_split_path, "train"),
                 "img_dir": join(project_split_path, "train", "img"),
                 "ann_dir": join(project_split_path, "train", "ann"),
+                "img_info_dir": join(project_split_path, "train", "img_info"),
             },
             "val": {
                 "split_path": join(project_split_path, "val"),
                 "img_dir": join(project_split_path, "val", "img"),
                 "ann_dir": join(project_split_path, "val", "ann"),
+                "img_info_dir": join(project_split_path, "val", "img_info"),
             },
         }
 
@@ -1199,7 +1229,7 @@ class TrainApp:
         }
 
         # Utility function to move files
-        def move_files(split, paths, img_name_format, pbar):
+        def move_files(split, split_name, paths, img_name_format, pbar):
             """
             Move files to the appropriate directories.
             """
@@ -1208,6 +1238,19 @@ class TrainApp:
                 ann_name = f"{item_name}.json"
                 shutil.copy(item.img_path, join(paths["img_dir"], item_name))
                 shutil.copy(item.ann_path, join(paths["ann_dir"], ann_name))
+
+                # Move img_info
+                img_info_name = f"{sly_fs.get_file_name_with_ext(item.img_path)}.json"
+                img_info_path = join(dirname(dirname(item.img_path)), "img_info", img_info_name)
+                # shutil.copy(img_info_path, join(paths["img_info_dir"], ann_name))
+
+                # Write split ids to img_info
+                img_info = sly_json.load_json_file(img_info_path)
+                if split_name == "train":
+                    self._train_split_item_ids.add(img_info["id"])
+                else:
+                    self._val_split_item_ids.add(img_info["id"])
+
                 pbar.update(1)
 
         # Main split processing
@@ -1216,12 +1259,16 @@ class TrainApp:
         ) as main_pbar:
             self.progress_bar_main.show()
             for dataset in ["train", "val"]:
-                split = self._train_split if dataset == "train" else self._val_split
+                split_name = dataset
+                if split_name == "train":
+                    split = self._train_split
+                else:
+                    split = self._val_split
                 with self.progress_bar_secondary(
                     message=f"Preparing '{dataset}'", total=len(split)
                 ) as second_pbar:
                     self.progress_bar_secondary.show()
-                    move_files(split, paths[dataset], image_name_formats[dataset], second_pbar)
+                    move_files(split, split_name, paths[dataset], image_name_formats[dataset], second_pbar)
                     main_pbar.update(1)
                 self.progress_bar_secondary.hide()
             self.progress_bar_main.hide()
@@ -1241,10 +1288,7 @@ class TrainApp:
         with self.progress_bar_main(message="Processing splits", total=2) as pbar:
             self.progress_bar_main.show()
             for dataset in ["train", "val"]:
-                shutil.move(
-                    paths[dataset]["split_path"],
-                    train_ds_path if dataset == "train" else val_ds_path,
-                )
+                shutil.move(paths[dataset]["split_path"], train_ds_path if dataset == "train" else val_ds_path)
                 pbar.update(1)
             self.progress_bar_main.hide()
 
@@ -1667,7 +1711,8 @@ class TrainApp:
             try:
                 # pylint: disable=import-error
                 import torch
-                state_dict = torch.load(new_checkpoint_path)
+
+                state_dict = torch_load_safe(new_checkpoint_path)
                 state_dict["model_info"] = {
                     "task_id": self.task_id,
                     "model_name": experiment_info["model_name"],
@@ -1679,9 +1724,7 @@ class TrainApp:
                 state_dict["model_files"] = ckpt_files
                 torch.save(state_dict, new_checkpoint_path)
             except Exception as e:
-                logger.warning(
-                    f"Error writing info to checkpoint: '{checkpoint_name}'. Error:{e}"
-                )
+                logger.warning(f"Error writing info to checkpoint: '{checkpoint_name}'. Error:{e}")
                 continue
 
             new_checkpoint_paths.append(new_checkpoint_path)
@@ -1817,7 +1860,6 @@ class TrainApp:
         :type export_weights: dict
         """
         logger.debug("Updating experiment info")
-
         experiment_info = {
             "experiment_name": self.gui.training_process.get_experiment_name(),
             "framework_name": self.framework_name,
@@ -1826,6 +1868,7 @@ class TrainApp:
             "base_checkpoint_link": self.base_checkpoint_link,
             "task_type": experiment_info["task_type"],
             "project_id": self.project_info.id,
+            "project_version": self.project_info.version,
             "task_id": self.task_id,
             "model_files": experiment_info["model_files"],
             "checkpoints": experiment_info["checkpoints"],
@@ -1843,6 +1886,8 @@ class TrainApp:
             "logs": {"type": "tensorboard", "link": f"{remote_dir}logs/"},
             "device": self.gui.training_process.get_device_name(),
             "training_duration": self._training_duration,
+            "train_collection_id": self._train_collection_id,
+            "val_collection_id": self._val_collection_id,
         }
 
         if self._has_splits_selector:
@@ -1997,13 +2042,7 @@ class TrainApp:
         :rtype: tuple
         """
         need_generate_report = self._app_options.get("generate_report", False)
-        # @TODO: temporary code to generate report for dev only
-        is_dev = "dev.internal" in self._api.server_address
-        if not is_dev:
-            need_generate_report = False
-        # ------------------------------------------------------------ #
-
-        if need_generate_report: # link to experiment page
+        if need_generate_report:  # link to experiment page
             try:
                 output_file_info = self._generate_experiment_report(experiment_info, model_meta)
                 experiment_info["has_report"] = True
@@ -2011,7 +2050,7 @@ class TrainApp:
                 logger.error(f"Error generating experiment report: {e}")
                 output_file_info = session_link_file_info
                 experiment_info["has_report"] = False
-        else: # link to artifacts directory
+        else:  # link to artifacts directory
             output_file_info = session_link_file_info
             experiment_info["has_report"] = False
         return output_file_info, experiment_info
@@ -3089,4 +3128,82 @@ class TrainApp:
         # 4. Match splits with original project
         gt_split_data = self._postprocess_splits(gt_project_info.id)
         return gt_project_info.id, gt_split_data
-        return gt_project_info.id, gt_split_data
+
+    def _create_collection_splits(self):
+        def _check_match(current_selected_collection_ids: List[int], all_split_collections: List[EntitiesCollectionInfo]):
+            if len(current_selected_collection_ids) > 0:
+                if len(current_selected_collection_ids) == 1:
+                    current_selected_collection_id = current_selected_collection_ids[0]
+                    for collection in all_split_collections:
+                        if collection.id == current_selected_collection_id:
+                            return True
+            return False
+
+        # Case 1: Use existing collections for training. No need to create new collections
+        split_method = self.gui.train_val_splits_selector.get_split_method()
+        all_train_collections = self.gui.train_val_splits_selector.all_train_collections
+        all_val_collections = self.gui.train_val_splits_selector.all_val_collections
+        if split_method == "Based on collections":
+            current_selected_train_collection_ids = self.gui.train_val_splits_selector.train_val_splits.get_train_collections_ids()
+            train_match = _check_match(current_selected_train_collection_ids, all_train_collections)
+            if train_match:
+                current_selected_val_collection_ids = self.gui.train_val_splits_selector.train_val_splits.get_val_collections_ids()
+                val_match = _check_match(current_selected_val_collection_ids, all_val_collections)
+                if val_match:
+                    self._train_collection_id = current_selected_train_collection_ids[0]
+                    self._val_collection_id = current_selected_val_collection_ids[0]
+                    self._update_project_custom_data(self._train_collection_id, self._val_collection_id)
+                    return
+        # ------------------------------------------------------------ #
+
+        # Case 2: Create new collections for selected train val splits. Need to create new collections
+        item_type = self.project_info.type
+        experiment_name = self.gui.training_process.get_experiment_name()
+
+        train_collection_idx = 1
+        val_collection_idx = 1        
+
+        # Get train collection with max idx
+        if len(all_train_collections) > 0:
+            train_collection_idx = max([int(collection.name.split("_")[1]) for collection in all_train_collections])
+            train_collection_idx += 1
+        # Get val collection with max idx
+        if len(all_val_collections) > 0:
+            val_collection_idx = max([int(collection.name.split("_")[1]) for collection in all_val_collections])
+            val_collection_idx += 1
+        # -------------------------------- #
+
+        # Create Train Collection
+        train_img_ids = list(self._train_split_item_ids)
+        train_collection_description = f"Collection with train {item_type} for experiment: {experiment_name}"
+        train_collection = self._api.entities_collection.create(self.project_id, f"train_{train_collection_idx}", train_collection_description)
+        train_collection_id = getattr(train_collection, "id", None)
+        if train_collection_id is None:
+            raise AttributeError("Train EntitiesCollectionInfo object does not have 'id' attribute")
+        self._api.entities_collection.add_items(train_collection_id, train_img_ids)
+        self._train_collection_id = train_collection_id
+
+        # Create Val Collection
+        val_img_ids = list(self._val_split_item_ids)
+        val_collection_description = f"Collection with val {item_type} for experiment: {experiment_name}"
+        val_collection = self._api.entities_collection.create(self.project_id, f"val_{val_collection_idx}", val_collection_description)
+        val_collection_id = getattr(val_collection, "id", None)
+        if val_collection_id is None:
+            raise AttributeError("Val EntitiesCollectionInfo object does not have 'id' attribute")
+        self._api.entities_collection.add_items(val_collection_id, val_img_ids)
+        self._val_collection_id = val_collection_id
+
+        # Update Project Custom Data
+        self._update_project_custom_data(train_collection_id, val_collection_id)
+
+    def _update_project_custom_data(self, train_collection_id: int, val_collection_id: int):
+        train_info = {
+            "task_id": self.task_id,
+            "framework_name": self.framework_name,
+            "splits": {"train_collection": train_collection_id, "val_collection": val_collection_id}
+        }
+        custom_data = self._api.project.get_info_by_id(self.project_id).custom_data
+        train_info_list = custom_data.get("train_info", [])
+        train_info_list.append(train_info)
+        custom_data.update({"train_info": train_info_list})
+        self._api.project.update_custom_data(self.project_id, custom_data)