Преглед изворни кода

Set default ai model by config

Kentaro Wada пре 1 година
родитељ
комит
1c23cf0d43
2 измењених фајлова са 14 додато и 2 уклоњено
  1. 11 2
      labelme/app.py
  2. 3 0
      labelme/config/default_config.yaml

+ 11 - 2
labelme/app.py

@@ -785,8 +785,17 @@ class MainWindow(QtWidgets.QMainWindow):
         selectAiModel.defaultWidget().layout().addWidget(
             self._selectAiModelComboBox
         )
-        self._selectAiModelComboBox.addItems([model.name for model in MODELS])
-        self._selectAiModelComboBox.setCurrentIndex(1)
+        model_names = [model.name for model in MODELS]
+        self._selectAiModelComboBox.addItems(model_names)
+        if self._config["ai"]["default"] in model_names:
+            model_index = model_names.index(self._config["ai"]["default"])
+        else:
+            logger.warning(
+                "Default AI model is not found: %r",
+                self._config["ai"]["default"],
+            )
+            model_index = 0
+        self._selectAiModelComboBox.setCurrentIndex(model_index)
         self._selectAiModelComboBox.currentIndexChanged.connect(
             lambda: self.canvas.initializeAiModel(
                 name=self._selectAiModelComboBox.currentText()

+ 3 - 0
labelme/config/default_config.yaml

@@ -30,6 +30,9 @@ shape:
   hvertex_fill_color: [255, 255, 255, 255]
   point_size: 8
 
+ai:
+  default: 'SegmentAnything (balanced)'
+
 # main
 flag_dock:
   show: true