Bladeren bron

Added dialog for label flags when a new shape is created

cmerchant 6 jaren geleden
bovenliggende
commit
92d2b226ba
2 gewijzigde bestanden met toevoegingen van 41 en 6 verwijderingen
  1. 9 3
      labelme/app.py
  2. 32 3
      labelme/widgets/label_dialog.py

+ 9 - 3
labelme/app.py

@@ -81,6 +81,7 @@ class MainWindow(QtWidgets.QMainWindow):
             show_text_field=self._config['show_label_text_field'],
             completion=self._config['label_completion'],
             fit_to_content=self._config['fit_to_content'],
+            flags=self._config['label_flags']
         )
 
         self.labelList = LabelQListWidget()
@@ -861,7 +862,8 @@ class MainWindow(QtWidgets.QMainWindow):
         if not self.canvas.editing():
             return
         item = item if item else self.currentItem()
-        text = self.labelDialog.popUp(item.text() if item else None)
+        shape = self.labelList.get_shape_from_item(item)
+        text, flags = self.labelDialog.popUp(item.text() if item else None, flags=(shape.flags if item else None))
         if text is None:
             return
         if not self.validateLabel(text):
@@ -869,6 +871,8 @@ class MainWindow(QtWidgets.QMainWindow):
                               "Invalid label '{}' with validation type '{}'"
                               .format(text, self._config['validate_label']))
             return
+        shape.flags = flags
+        self.loadLabelFlags(flags)
         item.setText(text)
         self.setDirty()
         if not self.uniqLabelList.findItems(text, Qt.MatchExactly):
@@ -1083,7 +1087,7 @@ class MainWindow(QtWidgets.QMainWindow):
                 if len(split) > 1 and split[-1].isdigit():
                     split[-1] = str(int(split[-1]) + 1)
                     text = '-'.join(split)
-            text = self.labelDialog.popUp(text)
+            text, flags = self.labelDialog.popUp(text)
 
         if text is not None and not self.validateLabel(text):
             self.errorMessage('Invalid label',
@@ -1094,7 +1098,9 @@ class MainWindow(QtWidgets.QMainWindow):
             self.canvas.undoLastLine()
             self.canvas.shapesBackups.pop()
         else:
-            self.addLabel(self.canvas.setLastLabel(text))
+            shape = self.canvas.setLastLabel(text)
+            shape.flags = flags
+            self.addLabel(shape)
             self.actions.editMode.setEnabled(True)
             self.actions.undoLastPoint.setEnabled(False)
             self.actions.undo.setEnabled(True)

+ 32 - 3
labelme/widgets/label_dialog.py

@@ -29,7 +29,7 @@ class LabelDialog(QtWidgets.QDialog):
 
     def __init__(self, text="Enter object label", parent=None, labels=None,
                  sort_labels=True, show_text_field=True,
-                 completion='startswith', fit_to_content=None):
+                 completion='startswith', fit_to_content=None, flags=[]):
         if fit_to_content is None:
             fit_to_content = {'row': False, 'column': True}
         self._fit_to_content = fit_to_content
@@ -74,6 +74,14 @@ class LabelDialog(QtWidgets.QDialog):
         self.labelList.currentItemChanged.connect(self.labelSelected)
         self.edit.setListWidget(self.labelList)
         layout.addWidget(self.labelList)
+        # label_flags
+        self.flags = flags
+        self.label_flags = None
+        if flags:
+            self.label_flags = QtWidgets.QVBoxLayout()
+            for flag in flags:
+                self.label_flags.addWidget(QtWidgets.QCheckBox(flag, self))
+            layout.addItem(self.label_flags)
         self.setLayout(layout)
         # completion
         completer = QtWidgets.QCompleter()
@@ -122,7 +130,24 @@ class LabelDialog(QtWidgets.QDialog):
             text = text.trimmed()
         self.edit.setText(text)
 
-    def popUp(self, text=None, move=True):
+    def resetFlags(self):
+        for i in range(self.label_flags.count()):
+            item = self.label_flags.itemAt(i).widget()
+            item.setChecked(False)
+
+    def setFlags(self, flags):
+        for i in range(self.label_flags.count()):
+            item = self.label_flags.itemAt(i).widget()
+            item.setChecked(flags[item.text()])
+
+    def getFlags(self):
+        flags = {}
+        for i in range(self.label_flags.count()):
+            item = self.label_flags.itemAt(i).widget()
+            flags[item.text()] = True if item.isChecked() else False
+        return flags
+
+    def popUp(self, text=None, move=True, flags=None):
         if self._fit_to_content['row']:
             self.labelList.setMinimumHeight(
                 self.labelList.sizeHintForRow(0) * self.labelList.count() + 2
@@ -134,6 +159,10 @@ class LabelDialog(QtWidgets.QDialog):
         # if text is None, the previous label in self.edit is kept
         if text is None:
             text = self.edit.text()
+        if flags:
+            self.setFlags(flags)
+        else:
+            self.resetFlags()
         self.edit.setText(text)
         self.edit.setSelection(0, len(text))
         items = self.labelList.findItems(text, QtCore.Qt.MatchFixedString)
@@ -146,4 +175,4 @@ class LabelDialog(QtWidgets.QDialog):
         self.edit.setFocus(QtCore.Qt.PopupFocusReason)
         if move:
             self.move(QtGui.QCursor.pos())
-        return self.edit.text() if self.exec_() else None
+        return (self.edit.text(), self.getFlags()) if self.exec_() else None