code improvements

supervisely-ecosystem · Sep 24, 2024 · f18220e · f18220e
1 parent 3e0eb91
commit f18220e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 10 deletions.
diff --git a/src/ui/dtl/actions/neural_networks/deploy/deploy.py b/src/ui/dtl/actions/neural_networks/deploy/deploy.py
@@ -127,7 +127,6 @@ def agent_selector_sidebar_save_btn_cb():
             # public options
             model_selector_sidebar_public_model_table,
             model_selector_runtime_selector_sidebar,
-            model_selector_sidebar_public_container,
             # sidebar
             model_selector_sidebar_model_source_tabs,
             model_selector_sidebar_save_btn,

diff --git a/src/ui/dtl/actions/neural_networks/deploy/layout/model_selector.py b/src/ui/dtl/actions/neural_networks/deploy/layout/model_selector.py
@@ -9,8 +9,7 @@
     Select,
     Field,
 )
-
-
+from supervisely.nn.inference import RuntimeType
 import src.globals as g
 from src.ui.dtl.utils import (
     get_text_font_size,
@@ -21,6 +20,10 @@
 )
 
 
+def get_available_runtimes():
+    return [RuntimeType.ONNXRUNTIME, RuntimeType.PYTORCH, RuntimeType.TENSORRT]
+
+
 def create_model_selector_widgets(
     framework_name: str, pretrained_models: list, custom_models: list, custom_task_types: list = []
 ):
@@ -50,19 +53,15 @@ def create_model_selector_widgets(
         model_selector_sidebar_public_model_table.set_active_task_type("object detection")
 
     # runtime selector
-    available_runtimes = ["PyTorch", "ONNXRuntime", "TensorRT"]
     model_selector_runtime_selector_sidebar = Select(
-        [Select.Item(runtime, runtime) for runtime in available_runtimes]
+        [Select.Item(runtime, runtime) for runtime in get_available_runtimes()]
     )
     model_selector_runtime_field = Field(
         model_selector_runtime_selector_sidebar,
         "Runtime",
         "The model will be exported to the selected runtime for efficient inference (exporting to TensorRT may take about a minute).",
     )
 
-    model_selector_sidebar_public_container = Container(
-        [model_selector_sidebar_public_model_table, model_selector_runtime_field]
-    )
     # ------------------------------
 
     # CUSTOM /PUBLIC TABS
@@ -71,7 +70,7 @@ def create_model_selector_widgets(
         descriptions=["Models trained by you", f"Models trained by {framework_name} team"],
         contents=[
             model_selector_sidebar_custom_model_table,
-            model_selector_sidebar_public_container,
+            model_selector_sidebar_public_model_table,
         ],
     )
     if len(custom_models) == 0:
@@ -82,6 +81,7 @@ def create_model_selector_widgets(
     model_selector_sidebar_container = Container(
         [
             model_selector_sidebar_model_source_tabs,
+            model_selector_runtime_field,
             model_selector_sidebar_save_btn,
         ]
     )
@@ -125,7 +125,6 @@ def create_model_selector_widgets(
         # public options
         model_selector_sidebar_public_model_table,
         model_selector_runtime_selector_sidebar,
-        model_selector_sidebar_public_container,
         # sidebar
         model_selector_sidebar_model_source_tabs,
         model_selector_sidebar_save_btn,