Selaa lähdekoodia

Fix editLabel crash due to the novel label name in self.uniqLabelList

Kentaro Wada 2 vuotta sitten
vanhempi
commit
a37aa20515
2 muutettua tiedostoa jossa 18 lisäystä ja 9 poistoa
  1. 11 5
      labelme/app.py
  2. 7 4
      labelme/widgets/unique_label_qlist_widget.py

+ 11 - 5
labelme/app.py

@@ -1083,10 +1083,11 @@ class MainWindow(QtWidgets.QMainWindow):
         else:
             item.setText("{} ({})".format(shape.label, shape.group_id))
         self.setDirty()
-        if not self.uniqLabelList.findItemsByLabel(shape.label):
-            item = QtWidgets.QListWidgetItem()
-            item.setData(Qt.UserRole, shape.label)
+        if self.uniqLabelList.findItemByLabel(shape.label) is None:
+            item = self.uniqLabelList.createItemFromLabel(shape.label)
             self.uniqLabelList.addItem(item)
+            rgb = self._get_rgb_by_label(shape.label)
+            self.uniqLabelList.setItemLabel(item, shape.label, rgb)
 
     def fileSearchChanged(self):
         self.importDirImages(
@@ -1136,7 +1137,7 @@ class MainWindow(QtWidgets.QMainWindow):
             text = "{} ({})".format(shape.label, shape.group_id)
         label_list_item = LabelListWidgetItem(text, shape)
         self.labelList.addItem(label_list_item)
-        if not self.uniqLabelList.findItemsByLabel(shape.label):
+        if self.uniqLabelList.findItemByLabel(shape.label) is None:
             item = self.uniqLabelList.createItemFromLabel(shape.label)
             self.uniqLabelList.addItem(item)
             rgb = self._get_rgb_by_label(shape.label)
@@ -1163,7 +1164,12 @@ class MainWindow(QtWidgets.QMainWindow):
 
     def _get_rgb_by_label(self, label):
         if self._config["shape_color"] == "auto":
-            item = self.uniqLabelList.findItemsByLabel(label)[0]
+            item = self.uniqLabelList.findItemByLabel(label)
+            if item is None:
+                item = self.uniqLabelList.createItemFromLabel(label)
+                self.uniqLabelList.addItem(item)
+                rgb = self._get_rgb_by_label(label)
+                self.uniqLabelList.setItemLabel(item, label, rgb)
             label_id = self.uniqLabelList.indexFromItem(item).row() + 1
             label_id += self._config["shift_auto_shape_color"]
             return LABEL_COLORMAP[label_id % len(LABEL_COLORMAP)]

+ 7 - 4
labelme/widgets/unique_label_qlist_widget.py

@@ -14,15 +14,18 @@ class UniqueLabelQListWidget(EscapableQListWidget):
         if not self.indexAt(event.pos()).isValid():
             self.clearSelection()
 
-    def findItemsByLabel(self, label):
-        items = []
+    def findItemByLabel(self, label):
         for row in range(self.count()):
             item = self.item(row)
             if item.data(Qt.UserRole) == label:
-                items.append(item)
-        return items
+                return item
 
     def createItemFromLabel(self, label):
+        if self.findItemByLabel(label):
+            raise ValueError(
+                "Item for label '{}' already exists".format(label)
+            )
+
         item = QtWidgets.QListWidgetItem()
         item.setData(Qt.UserRole, label)
         return item