supervisely 6.73.243__py3-none-any.whl → 6.73.245__py3-none-any.whl

supervisely/app/widgets/experiment_selector/style.css

@@ -0,0 +1,27 @@
+.custom-models-selector-table {
+  border-collapse: collapse;
+}
+.custom-models-selector-table tr:nth-child(2n) {
+  background-color: #f6f8fa;
+}
+.custom-models-selector-table td,
+.custom-models-selector-table th {
+  border: 1px solid #dfe2e5;
+  padding: 6px 13px;
+  text-align: center;
+  line-height: 20px;
+}
+
+.custom-models-selector-table td {
+  text-align: left;
+}
+
+.custom-models-selector-table tr td:nth-child(4) {
+  text-align: center;
+}
+
+.el-radio-group.multi-line label.el-radio {
+  display: block;
+  margin-left: 0px;
+  margin-bottom: 5px;
+}

supervisely/app/widgets/experiment_selector/template.html

@@ -0,0 +1,82 @@
+<link rel="stylesheet" href="./sly/css/app/widgets/custom_models_selector/style.css"/>
+
+<div
+    {% if widget._changes_handled == true %}
+    @change="post('/{{{widget.widget_id}}}/value_changed')"
+    {% endif %}
+    >
+
+    <div v-if="Object.keys(data.{{{widget.widget_id}}}.rowsHtml).length === 0"> You don't have any custom models</div>
+    <div v-else>
+
+        <div v-if="data.{{{widget.widget_id}}}.taskTypes.length > 1">
+            <sly-field
+            title="Task Type"
+            >
+                <el-radio-group
+                    class="multi-line mt10"
+                    :value="state.{{{widget.widget_id}}}.selectedTaskType"
+    
+                    {% if widget._task_type_changes_handled == true %}
+                        @input="(evt) => {state.{{{widget.widget_id}}}.selectedTaskType = evt; state.{{{widget.widget_id}}}.selectedRow = 0; post('/{{{widget.widget_id}}}/task_type_changed')}"
+                    {% else %}
+                        @input="(evt) => {state.{{{widget.widget_id}}}.selectedTaskType = evt; state.{{{widget.widget_id}}}.selectedRow = 0;}"
+                    {% endif %}
+                    >
+    
+                    <el-radio 
+                        v-for="(item, idx) in {{{widget._task_types}}}"
+                        :key="item"
+                        :label="item"
+                        >
+                            {{ item }}
+                    </el-radio>
+                </el-radio-group>
+        </sly-field>
+    </div>
+
+        <div>
+        
+    <table class="custom-models-selector-table">
+        <thead>
+        <tr>
+            <th v-for="col in data.{{{widget.widget_id}}}.columns">
+                <div> {{col}} </div>
+            </th>
+        </tr> 
+        </thead>
+        <tbody>
+        <tr v-for="row, ridx in data.{{{widget.widget_id}}}.rowsHtml[state.{{{widget.widget_id}}}.selectedTaskType]">
+            <td v-for="col, vidx in row">
+                <div v-if="vidx === 0" style="display: flex;">
+                    <el-radio
+                    style="display: flex;"
+                    v-model="state.{{{widget.widget_id}}}.selectedRow"
+                    :label="ridx"
+                    >&#8205;</el-radio>
+                    
+                    <sly-html-compiler :params="{ridx: ridx, vidx: vidx}" :template="col" :data="data" :state="state"></sly-html-compiler>
+                    
+                </div>
+
+                <div v-else>
+
+                    <sly-html-compiler :params="{ridx: ridx, vidx: vidx}" :template="col" :data="data" :state="state">
+                    </sly-html-compiler>
+
+                </div>
+
+            </td>
+        </tr>
+        </tbody>
+    </table>
+</div>
+    <div class="mt10" v-if="{{{widget.show_custom_checkpoint_path}}}"
+    >
+    {{{widget.show_custom_checkpoint_path_checkbox}}}
+    <div class="mt10">
+    {{{widget.custom_tab_widgets}}}
+</div>
+</div>
+</div>
+</div>

supervisely/app/widgets/pretrained_models_selector/pretrained_models_selector.py

@@ -136,6 +136,9 @@ class PretrainedModelsSelector(Widget):
 
     def get_selected_model_params(self, model_name_column: str = "Model") -> Union[Dict, None]:
         selected_model = self.get_selected_row()
+        if selected_model is None:
+            return {}
+
         model_name = selected_model.get(model_name_column)
         if model_name is None:
             raise ValueError(
@@ -172,24 +175,43 @@ class PretrainedModelsSelector(Widget):
         if widget_actual_state is not None and widget_actual_data is not None:
             return widget_actual_state["selectedRow"]
 
-    def set_active_arch_type(self, arch_type: str):
+    def set_active_arch_type(self, arch_type: str) -> None:
         if arch_type not in self._arch_types:
             raise ValueError(f'Architecture type "{arch_type}" does not exist')
         StateJson()[self.widget_id]["selectedArchType"] = arch_type
         StateJson().send_changes()
 
-    def set_active_task_type(self, task_type: str):
+    def set_active_task_type(self, task_type: str) -> None:
         if task_type not in self._task_types:
             raise ValueError(f'Task type "{task_type}" does not exist')
         StateJson()[self.widget_id]["selectedTaskType"] = task_type
         StateJson().send_changes()
 
-    def set_active_row(self, row_index: int):
+    def set_active_row(self, row_index: int) -> None:
         if row_index < 0:
             raise ValueError(f'Row with index "{row_index}" does not exist')
         StateJson()[self.widget_id]["selectedRow"] = row_index
         StateJson().send_changes()
 
+    def set_by_model_name(self, model_name: str) -> None:
+        for task_type in self._table_data:
+            for arch_type in self._table_data[task_type]:
+                for idx, model in enumerate(self._table_data[task_type][arch_type]):
+                    model_meta = model.get("meta", {})
+                    if model_meta.get("model_name") == model_name:
+                        self.set_active_task_type(task_type)
+                        self.set_active_arch_type(arch_type)
+                        self.set_active_row(idx)
+                        return
+
+    def get_by_model_name(self, model_name: str) -> Union[Dict, None]:
+        for task_type in self._table_data:
+            for arch_type in self._table_data[task_type]:
+                for idx, model in enumerate(self._table_data[task_type][arch_type]):
+                    model_meta = model.get("meta", {})
+                    if model_meta.get("model_name") == model_name:
+                        return model
+
     def _filter_and_sort_models(self, models: List[Dict], sort_models: bool = True) -> Dict:
         filtered_models = {}
 

supervisely/app/widgets/random_splits_table/random_splits_table.py

@@ -1,15 +1,16 @@
 from typing import Dict, Optional
 
-from supervisely.app import StateJson
+from supervisely.app import DataJson, StateJson
 from supervisely.app.widgets import Widget
 
+
 class RandomSplitsTable(Widget):
     def __init__(
-        self, 
-        items_count: int, 
+        self,
+        items_count: int,
         start_train_percent: Optional[int] = 80,
         disabled: Optional[bool] = False,
-        widget_id: Optional[int] = None
+        widget_id: Optional[int] = None,
     ):
         self._disabled = disabled
         if 1 <= start_train_percent <= 99:
@@ -23,33 +24,56 @@ class RandomSplitsTable(Widget):
         ]
         self._items_count = items_count
         train_count = int(items_count / 100 * start_train_percent)
-        self._count = {
-            "total": items_count,
-            "train": train_count,
-            "val": items_count - train_count
-        }
+        self._count = {"total": items_count, "train": train_count, "val": items_count - train_count}
 
         self._percent = {
             "total": 100,
             "train": start_train_percent,
-            "val": 100 - start_train_percent
+            "val": 100 - start_train_percent,
         }
 
         super().__init__(widget_id=widget_id, file_path=__file__)
 
-
     def get_json_data(self):
         return {
             "table_data": self._table_data,
             "items_count": self._items_count,
-            "disabled": self._disabled
+            "disabled": self._disabled,
         }
 
     def get_json_state(self):
-        return {
-            "count": self._count,
-            "percent": self._percent
-        }
+        return {"count": self._count, "percent": self._percent}
 
     def get_splits_counts(self) -> Dict[str, int]:
-        return StateJson()[self.widget_id]["count"]
+        return StateJson()[self.widget_id]["count"]
+
+    def set_train_split_percent(self, percent: int):
+        if 1 <= percent <= 99:
+            self._percent["train"] = percent
+            self._percent["val"] = 100 - percent
+            self._count["train"] = int(self._items_count / 100 * percent)
+            self._count["val"] = self._items_count - self._count["train"]
+
+            StateJson()[self.widget_id]["count"] = self._count
+            StateJson()[self.widget_id]["percent"] = self._percent
+            StateJson().send_changes()
+        else:
+            raise ValueError("percent must be in range [1; 99].")
+
+    def get_train_split_percent(self) -> Dict[str, int]:
+        return StateJson()[self.widget_id]["percent"]["train"]
+
+    def set_val_split_percent(self, percent: int):
+        if 1 <= percent <= 99:
+            self._percent["val"] = percent
+            self._percent["train"] = 100 - percent
+            self._count["val"] = int(self._items_count / 100 * percent)
+            self._count["train"] = self._items_count - self._count["val"]
+            StateJson()[self.widget_id]["count"] = self._count
+            StateJson()[self.widget_id]["percent"] = self._percent
+            StateJson().send_changes()
+        else:
+            raise ValueError("percent must be in range [1; 99].")
+
+    def get_val_split_percent(self) -> Dict[str, int]:
+        return StateJson()[self.widget_id]["percent"]["val"]

supervisely/app/widgets/select_dataset_tree/select_dataset_tree.py

@@ -119,11 +119,18 @@ class SelectDatasetTree(Widget):
         self._append_to_body = append_to_body
 
         # Extract values from Enum to match the .type property of the ProjectInfo object.
-        self._project_types = (
-            [project_type.value for project_type in allowed_project_types]
-            if allowed_project_types is not None
-            else None
-        )
+
+        if allowed_project_types is not None:
+            if all(allowed_project_types) is isinstance(allowed_project_types, ProjectType):
+                self._project_types = (
+                    [project_type.value for project_type in allowed_project_types]
+                    if allowed_project_types is not None
+                    else None
+                )
+            elif all(allowed_project_types) is isinstance(allowed_project_types, str):
+                self._project_types = allowed_project_types
+            else:
+                self._project_types = None
 
         # Widget components.
         self._select_team = None

supervisely/app/widgets/train_val_splits/train_val_splits.py

@@ -10,7 +10,6 @@ from supervisely.app.widgets import (
     Field,
     NotificationBox,
     RadioTabs,
-    SelectDataset,
     SelectString,
     SelectTagMeta,
     Widget,
@@ -18,11 +17,15 @@ from supervisely.app.widgets import (
 from supervisely.app.widgets.random_splits_table.random_splits_table import (
     RandomSplitsTable,
 )
+from supervisely.app.widgets.select_dataset_tree.select_dataset_tree import (
+    SelectDatasetTree,
+)
 from supervisely.io.fs import remove_dir
 from supervisely.project import get_project_class
 from supervisely.project.pointcloud_episode_project import PointcloudEpisodeProject
 from supervisely.project.pointcloud_project import PointcloudProject
 from supervisely.project.project import ItemInfo, Project
+from supervisely.project.project_type import ProjectType
 from supervisely.project.video_project import VideoProject
 from supervisely.project.volume_project import VolumeProject
 
@@ -65,8 +68,8 @@ class TrainValSplits(Widget):
         self._train_tag_select: SelectTagMeta = None
         self._val_tag_select: SelectTagMeta = None
         self._untagged_select: SelectString = None
-        self._train_ds_select: Union[SelectDataset, SelectString] = None
-        self._val_ds_select: Union[SelectDataset, SelectString] = None
+        self._train_ds_select: Union[SelectDatasetTree, SelectString] = None
+        self._val_ds_select: Union[SelectDatasetTree, SelectString] = None
         self._split_methods = []
 
         contents = []
@@ -97,6 +100,7 @@ class TrainValSplits(Widget):
         super().__init__(widget_id=widget_id, file_path=__file__)
 
     def _get_random_content(self):
+        items_count = 0
         if self._project_id is not None:
             items_count = self._project_info.items_count
         elif self._project_fs is not None:
@@ -163,12 +167,37 @@ class TrainValSplits(Widget):
             box_type="info",
         )
         if self._project_id is not None:
-            self._train_ds_select = SelectDataset(
-                project_id=self._project_id, multiselect=True, compact=True, show_label=False
+            self._train_ds_select = SelectDatasetTree(
+                multiselect=True,
+                flat=True,
+                select_all_datasets=False,
+                allowed_project_types=[self._project_type],
+                always_open=False,
+                compact=True,
+                team_is_selectable=False,
+                workspace_is_selectable=False,
+                append_to_body=True,
             )
-            self._val_ds_select = SelectDataset(
-                project_id=self._project_id, multiselect=True, compact=True, show_label=False
+
+            self._val_ds_select = SelectDatasetTree(
+                multiselect=True,
+                flat=True,
+                select_all_datasets=False,
+                allowed_project_types=[self._project_type],
+                always_open=False,
+                compact=True,
+                team_is_selectable=False,
+                workspace_is_selectable=False,
+                append_to_body=True,
             )
+
+            # old implementation
+            # self._train_ds_select = SelectDataset(
+            #     project_id=self._project_id, multiselect=True, compact=True, show_label=False
+            # )
+            # self._val_ds_select = SelectDataset(
+            #     project_id=self._project_id, multiselect=True, compact=True, show_label=False
+            # )
         elif self._project_fs is not None:
             ds_names = [ds.name for ds in self._project_fs.datasets]
             self._train_ds_select = SelectString(ds_names, multiple=True)
@@ -196,6 +225,7 @@ class TrainValSplits(Widget):
     def get_splits(self) -> Tuple[List[ItemInfo], List[ItemInfo]]:
         split_method = self._content.get_active_tab()
         tmp_project_dir = None
+        train_set, val_set = [], []
         if self._project_fs is None:
             tmp_project_dir = os.path.join(get_data_dir(), rand_str(15))
             self._project_class.download(self._api, self._project_id, tmp_project_dir)
@@ -226,11 +256,11 @@ class TrainValSplits(Widget):
 
         elif split_method == "Based on datasets":
             if self._project_id is not None:
-                self._train_ds_select: SelectDataset
-                self._val_ds_select: SelectDataset
+                self._train_ds_select: SelectDatasetTree
+                self._val_ds_select: SelectDatasetTree
                 train_ds_ids = self._train_ds_select.get_selected_ids()
                 val_ds_ids = self._val_ds_select.get_selected_ids()
-                ds_infos = self._api.dataset.get_list(self._project_id)
+                ds_infos = [dataset for _, dataset in self._api.dataset.tree(self._project_id)]
                 train_ds_names, val_ds_names = [], []
                 for ds_info in ds_infos:
                     if ds_info.id in train_ds_ids:
@@ -251,6 +281,65 @@ class TrainValSplits(Widget):
             remove_dir(tmp_project_dir)
         return train_set, val_set
 
+    def set_split_method(self, split_method: Literal["random", "tags", "datasets"]):
+        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"
+        self._content.set_active_tab(split_method)
+        StateJson().send_changes()
+        DataJson().send_changes()
+
+    def get_split_method(self) -> str:
+        return self._content.get_active_tab()
+
+    def set_random_splits(
+        self, split: Literal["train", "training", "val", "validation"], percent: int
+    ):
+        self._content.set_active_tab("Random")
+        if split == "train" or split == "training":
+            self._random_splits_table.set_train_split_percent(percent)
+        elif split == "val" or split == "validation":
+            self._random_splits_table.set_val_split_percent(percent)
+        else:
+            raise ValueError("Split value must be 'train', 'training', 'val' or 'validation'")
+
+    def get_train_split_percent(self) -> List[int]:
+        return self._random_splits_table.get_train_split_percent()
+
+    def get_val_split_percent(self) -> List[int]:
+        return 100 - self._random_splits_table.get_train_split_percent()
+
+    def set_tags_splits(
+        self, train_tag: str, val_tag: str, untagged_action: Literal["train", "val", "ignore"]
+    ):
+        self._content.set_active_tab("Based on item tags")
+        self._train_tag_select.set_name(train_tag)
+        self._val_tag_select.set_name(val_tag)
+        self._untagged_select.set_value(untagged_action)
+
+    def get_train_tag(self) -> str:
+        return self._train_tag_select.get_selected_name()
+
+    def get_val_tag(self) -> str:
+        return self._val_tag_select.get_selected_name()
+
+    def set_datasets_splits(self, train_datasets: List[int], val_datasets: List[int]):
+        self._content.set_active_tab("Based on datasets")
+        self._train_ds_select.set_dataset_ids(train_datasets)
+        self._val_ds_select.set_dataset_ids(val_datasets)
+
+    def get_train_dataset_ids(self) -> List[int]:
+        return self._train_ds_select.get_selected_ids()
+
+    def get_val_dataset_ids(self) -> List[int]:
+        return self._val_ds_select.get_selected_ids()
+
+    def get_untagged_action(self) -> str:
+        return self._untagged_select.get_value()
+
     def disable(self):
         self._content.disable()
         self._random_splits_table.disable()

supervisely/app/widgets/tree_select/tree_select.py

@@ -200,6 +200,8 @@ class TreeSelect(Widget):
         :rtype: Union[List[TreeSelect.Item], TreeSelect.Item]
         """
         res = StateJson()[self.widget_id]["value"]
+        if res is None:
+            return None
         if isinstance(res, list):
             return [TreeSelect.Item.from_json(item) for item in res]
         return TreeSelect.Item.from_json(res)

supervisely/nn/__init__.py

@@ -10,4 +10,6 @@ from supervisely.nn.prediction_dto import (
     PredictionMask,
     PredictionSegmentation,
 )
-from supervisely.nn.task_type import TaskType
+from supervisely.nn.task_type import TaskType
+from supervisely.nn.utils import ModelSource, RuntimeType
+from supervisely.nn.experiments import ExperimentInfo, get_experiment_infos

supervisely/nn/artifacts/artifacts.py

@@ -59,6 +59,7 @@ class BaseTrainArtifacts:
         self._weights_ext: str = None
         self._config_file: str = None
         self._pattern: str = None
+        self._available_task_types: List[str] = []
 
     @property
     def team_id(self) -> int:
@@ -516,3 +517,12 @@ class BaseTrainArtifacts:
         train_infos = self.sort_train_infos(train_infos, sort)
         logger.debug(f"Listing time: '{format(end_time - start_time, '.6f')}' sec")
         return train_infos
+
+    def get_available_task_types(self) -> List[str]:
+        """
+        Get available task types.
+
+        :return: The list of available task types.
+        :rtype: List[str]
+        """
+        return self._available_task_types

supervisely/nn/artifacts/detectron2.py

@@ -1,5 +1,6 @@
 from os.path import join
 from re import compile as re_compile
+from typing import List
 
 from supervisely.nn.artifacts.artifacts import BaseTrainArtifacts
 
@@ -16,6 +17,7 @@ class Detectron2(BaseTrainArtifacts):
         self._weights_ext = ".pth"
         self._config_file = "model_config.yaml"
         self._pattern = re_compile(r"^/detectron2/\d+_[^/]+/?$")
+        self._available_task_types: List[str] = ["instance segmentation"]
 
     def get_task_id(self, artifacts_folder: str) -> str:
         parts = artifacts_folder.split("/")

supervisely/nn/artifacts/hrda.py

@@ -1,3 +1,5 @@
+from typing import List
+
 from supervisely.nn.artifacts.artifacts import BaseTrainArtifacts
 
 
@@ -13,6 +15,7 @@ class HRDA(BaseTrainArtifacts):
         # self._task_type = "semantic segmentation"
         # self._weights_ext = ".pth"
         # self._config_file = "config.py"
+        # self._available_task_types: List[str] = ["semantic segmentation"]
 
     def get_task_id(self, artifacts_folder: str) -> str:
         raise NotImplementedError

supervisely/nn/artifacts/mmclassification.py

@@ -1,5 +1,6 @@
 from os.path import join
 from re import compile as re_compile
+from typing import List
 
 from supervisely.nn.artifacts.artifacts import BaseTrainArtifacts
 
@@ -14,6 +15,7 @@ class MMClassification(BaseTrainArtifacts):
         self._task_type = "classification"
         self._weights_ext = ".pth"
         self._pattern = re_compile(r"^/mmclassification/\d+_[^/]+/?$")
+        self._available_task_types: List[str] = ["classification"]
 
     def get_task_id(self, artifacts_folder: str) -> str:
         parts = artifacts_folder.split("/")

supervisely/nn/artifacts/mmdetection.py

@@ -1,10 +1,11 @@
+import random
+import string
 from os.path import join
 from re import compile as re_compile
+from typing import List
 
 from supervisely.io.fs import silent_remove
 from supervisely.nn.artifacts.artifacts import BaseTrainArtifacts
-import string
-import random
 
 
 class MMDetection(BaseTrainArtifacts):
@@ -19,6 +20,7 @@ class MMDetection(BaseTrainArtifacts):
         self._info_file = "info/ui_state.json"
         self._config_file = "config.py"
         self._pattern = re_compile(r"^/mmdetection/\d+_[^/]+/?$")
+        self._available_task_types: List[str] = ["object detection", "instance segmentation"]
 
     def get_task_id(self, artifacts_folder: str) -> str:
         parts = artifacts_folder.split("/")
@@ -39,7 +41,7 @@ class MMDetection(BaseTrainArtifacts):
         task_type = "undefined"
         for file_info in self._get_file_infos():
             if file_info.path == info_path:
-                json_data = self._fetch_json_from_url(file_info.full_storage_url)
+                json_data = self._fetch_json_from_path(file_info.path)
                 task_type = json_data.get("task", "undefined")
                 break
         return task_type
@@ -62,6 +64,7 @@ class MMDetection3(BaseTrainArtifacts):
         self._weights_ext = ".pth"
         self._config_file = "config.py"
         self._pattern = re_compile(r"^/mmdetection-3/\d+_[^/]+/?$")
+        self._available_task_types: List[str] = ["object detection", "instance segmentation"]
 
     def get_task_id(self, artifacts_folder: str) -> str:
         parts = artifacts_folder.split("/")

supervisely/nn/artifacts/mmsegmentation.py

@@ -1,5 +1,6 @@
 from os.path import join
 from re import compile as re_compile
+from typing import List
 
 from supervisely.nn.artifacts.artifacts import BaseTrainArtifacts
 
@@ -15,6 +16,7 @@ class MMSegmentation(BaseTrainArtifacts):
         self._weights_ext = ".pth"
         self._config_file = "config.py"
         self._pattern = re_compile(r"^/mmsegmentation/\d+_[^/]+/?$")
+        self._available_task_types: List[str] = ["instance segmentation"]
 
     def get_task_id(self, artifacts_folder: str) -> str:
         return artifacts_folder.split("/")[2].split("_")[0]

supervisely/nn/artifacts/ritm.py

@@ -1,5 +1,6 @@
 from os.path import join
 from re import compile as re_compile
+from typing import List
 
 from supervisely.nn.artifacts.artifacts import BaseTrainArtifacts
 
@@ -15,6 +16,7 @@ class RITM(BaseTrainArtifacts):
         self._info_file = "info/ui_state.json"
         self._weights_ext = ".pth"
         self._pattern = re_compile(r"^/RITM_training/\d+_[^/]+/?$")
+        self._available_task_types: List[str] = ["interactive segmentation"]
 
     def get_task_id(self, artifacts_folder: str) -> str:
         parts = artifacts_folder.split("/")
@@ -35,7 +37,7 @@ class RITM(BaseTrainArtifacts):
         task_type = "undefined"
         for file_info in self._get_file_infos():
             if file_info.path == info_path:
-                json_data = self._fetch_json_from_url(file_info.full_storage_url)
+                json_data = self._fetch_json_from_path(file_info.path)
                 task_type = json_data.get("segmentationType", "undefined")
                 if task_type is not None:
                     task_type = task_type.lower()

supervisely/nn/artifacts/rtdetr.py

@@ -1,5 +1,6 @@
 from os.path import join
 from re import compile as re_compile
+from typing import List
 
 from supervisely.nn.artifacts.artifacts import BaseTrainArtifacts
 
@@ -15,6 +16,7 @@ class RTDETR(BaseTrainArtifacts):
         self._weights_ext = ".pth"
         self._config_file = "config.yml"
         self._pattern = re_compile(r"^/RT-DETR/[^/]+/\d+/?$")
+        self._available_task_types: List[str] = ["object detection"]
 
     def get_task_id(self, artifacts_folder: str) -> str:
         return artifacts_folder.split("/")[-1]

supervisely/nn/artifacts/unet.py

@@ -1,5 +1,6 @@
 from os.path import join
 from re import compile as re_compile
+from typing import List
 
 from supervisely.nn.artifacts.artifacts import BaseTrainArtifacts
 
@@ -15,6 +16,7 @@ class UNet(BaseTrainArtifacts):
         self._weights_ext = ".pth"
         self._config_file = "train_args.json"
         self._pattern = re_compile(r"^/unet/\d+_[^/]+/?$")
+        self._available_task_types: List[str] = ["semantic segmentation"]
 
     def get_task_id(self, artifacts_folder: str) -> str:
         parts = artifacts_folder.split("/")

supervisely/nn/artifacts/yolov5.py

@@ -1,5 +1,6 @@
 from os.path import join
 from re import compile as re_compile
+from typing import List
 
 from supervisely.nn.artifacts.artifacts import BaseTrainArtifacts
 
@@ -15,6 +16,7 @@ class YOLOv5(BaseTrainArtifacts):
         self._weights_ext = ".pt"
         self._config_file = None
         self._pattern = re_compile(r"^/yolov5_train/[^/]+/\d+/?$")
+        self._available_task_types: List[str] = ["object detection"]
 
     def get_task_id(self, artifacts_folder: str) -> str:
         return artifacts_folder.split("/")[-1]
@@ -43,3 +45,4 @@ class YOLOv5v2(YOLOv5):
         self._weights_ext = ".pt"
         self._config_file = None
         self._pattern = re_compile(r"^/yolov5_2.0_train/[^/]+/\d+/?$")
+        self._available_task_types: List[str] = ["object detection"]