Browse Source

Add label specific content

Michał Woś 3 years ago
parent
commit
25bbb82784
4 changed files with 32 additions and 5 deletions
  1. 8 2
      labelme/app.py
  2. 2 0
      labelme/label_file.py
  3. 2 0
      labelme/shape.py
  4. 20 3
      labelme/widgets/label_dialog.py

+ 8 - 2
labelme/app.py

@@ -1055,10 +1055,11 @@ class MainWindow(QtWidgets.QMainWindow):
         shape = item.shape()
         if shape is None:
             return
-        text, flags, group_id = self.labelDialog.popUp(
+        text, flags, group_id, content = self.labelDialog.popUp(
             text=shape.label,
             flags=shape.flags,
             group_id=shape.group_id,
+            content=shape.content
         )
         if text is None:
             return
@@ -1073,6 +1074,7 @@ class MainWindow(QtWidgets.QMainWindow):
         shape.label = text
         shape.flags = flags
         shape.group_id = group_id
+        shape.content = content
 
         self._update_shape_color(shape)
         if shape.group_id is None:
@@ -1204,6 +1206,7 @@ class MainWindow(QtWidgets.QMainWindow):
             points = shape["points"]
             shape_type = shape["shape_type"]
             flags = shape["flags"]
+            content = shape["content"]
             group_id = shape["group_id"]
             other_data = shape["other_data"]
 
@@ -1215,6 +1218,7 @@ class MainWindow(QtWidgets.QMainWindow):
                 label=label,
                 shape_type=shape_type,
                 group_id=group_id,
+                content=content,
             )
             for x, y in points:
                 shape.addPoint(QtCore.QPointF(x, y))
@@ -1251,6 +1255,7 @@ class MainWindow(QtWidgets.QMainWindow):
                     label=s.label.encode("utf-8") if PY2 else s.label,
                     points=[(p.x(), p.y()) for p in s.points],
                     group_id=s.group_id,
+                    content=s.content,
                     shape_type=s.shape_type,
                     flags=s.flags,
                 )
@@ -1346,7 +1351,7 @@ class MainWindow(QtWidgets.QMainWindow):
         group_id = None
         if self._config["display_label_popup"] or not text:
             previous_text = self.labelDialog.edit.text()
-            text, flags, group_id = self.labelDialog.popUp(text)
+            text, flags, group_id, content = self.labelDialog.popUp(text)
             if not text:
                 self.labelDialog.edit.setText(previous_text)
 
@@ -1362,6 +1367,7 @@ class MainWindow(QtWidgets.QMainWindow):
             self.labelList.clearSelection()
             shape = self.canvas.setLastLabel(text, flags)
             shape.group_id = group_id
+            shape.content = content
             self.addLabel(shape)
             self.actions.editMode.setEnabled(True)
             self.actions.undoLastPoint.setEnabled(False)

+ 2 - 0
labelme/label_file.py

@@ -83,6 +83,7 @@ class LabelFile(object):
             "group_id",
             "shape_type",
             "flags",
+            "content",
         ]
         try:
             with open(filename, "r") as f:
@@ -124,6 +125,7 @@ class LabelFile(object):
                     points=s["points"],
                     shape_type=s.get("shape_type", "polygon"),
                     flags=s.get("flags", {}),
+                    content=s.get("content"),
                     group_id=s.get("group_id"),
                     other_data={
                         k: v for k, v in s.items() if k not in shape_keys

+ 2 - 0
labelme/shape.py

@@ -52,6 +52,7 @@ class Shape(object):
         shape_type=None,
         flags=None,
         group_id=None,
+        content=None,
     ):
         self.label = label
         self.group_id = group_id
@@ -60,6 +61,7 @@ class Shape(object):
         self.selected = False
         self.shape_type = shape_type
         self.flags = flags
+        self.content = content
         self.other_data = {}
 
         self._highlightIndex = None

+ 20 - 3
labelme/widgets/label_dialog.py

@@ -103,6 +103,11 @@ class LabelDialog(QtWidgets.QDialog):
         self.resetFlags()
         layout.addItem(self.flagsLayout)
         self.edit.textChanged.connect(self.updateFlags)
+        #text edit
+        self.textEdit = QtWidgets.QTextEdit()
+        self.textEdit.setPlaceholderText('Label content')
+        layout.addWidget(self.textEdit)
+        self.resize(300,200) 
         self.setLayout(layout)
         # completion
         completer = QtWidgets.QCompleter()
@@ -200,7 +205,16 @@ class LabelDialog(QtWidgets.QDialog):
             return int(group_id)
         return None
 
-    def popUp(self, text=None, move=True, flags=None, group_id=None):
+    def getContent(self):
+        content = self.textEdit.toPlainText()
+        if content:
+            return content
+        return None
+
+    def setContent(self, content):
+        self.textEdit.setPlainText(content)
+
+    def popUp(self, text=None, move=True, flags=None, group_id=None, content=None):
         if self._fit_to_content["row"]:
             self.labelList.setMinimumHeight(
                 self.labelList.sizeHintForRow(0) * self.labelList.count() + 2
@@ -212,6 +226,9 @@ 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 content is None:
+            content=""
+        self.setContent(content)
         if flags:
             self.setFlags(flags)
         else:
@@ -233,6 +250,6 @@ class LabelDialog(QtWidgets.QDialog):
         if move:
             self.move(QtGui.QCursor.pos())
         if self.exec_():
-            return self.edit.text(), self.getFlags(), self.getGroupId()
+            return self.edit.text(), self.getFlags(), self.getGroupId(), self.getContent()
         else:
-            return None, None, None
+            return None, None, None, None