Merge branch 'feature/save' into feature/label-list

canvas.py

@@ -30,13 +30,16 @@ class Canvas(QWidget):
         self.offsets = QPointF(), QPointF()
         self.scale = 1.0
         self.pixmap = None
-        # Background label hiding.
+        self.visible = {}
         self._hideBackround = False
         self.hideBackround = False
         # Set widget options.
         self.setMouseTracking(True)
         self.setFocusPolicy(Qt.WheelFocus)
 
+    def isVisible(self, shape):
+        return self.visible.get(shape, True)
+
     def editing(self):
         return self.mode == self.EDIT
 
@@ -79,6 +82,14 @@ class Canvas(QWidget):
         elif Qt.LeftButton & ev.buttons() and self.selectedShape and self.prevPoint:
             self.boundedMoveShape(pos)
             self.repaint()
+            return
+
+        self.setToolTip("Image")
+        pos = self.transformPos(ev.posF())
+        for shape in reversed(self.shapes):
+            if shape.containsPoint(pos) and self.isVisible(shape):
+                return self.setToolTip("Object '%s'" % shape.label)
+
 
     def mousePressEvent(self, ev):
         pos = self.transformPos(ev.posF())
@@ -114,11 +125,14 @@ class Canvas(QWidget):
         self._hideBackround = self.hideBackround if enable else False
 
     def mouseDoubleClickEvent(self, ev):
-        # FIXME: Don't create shape with 2 vertices only.
         if self.current and self.editing():
-            # Add first point in the list so that it is consistent
+            # Shapes need to have at least 3 vertices.
-            # with shapes created the normal way.
+            if len(self.current) < 4:
-            self.current.addPoint(self.current[0])
+                return
+            # Replace the last point with the starting point.
+            # We have to do this because the mousePressEvent handler
+            # adds that point before this handler is called!
+            self.current[-1] = self.current[0]
             self.finalise(ev)
 
     def selectShape(self, point):
@@ -195,7 +209,7 @@ class Canvas(QWidget):
         p.drawPixmap(0, 0, self.pixmap)
         Shape.scale = self.scale
         for shape in self.shapes:
-            if shape.selected or not self._hideBackround:
+            if (shape.selected or not self._hideBackround) and self.isVisible(shape):
                 shape.paint(p)
         if self.current:
             self.current.paint(p)
@@ -330,6 +344,20 @@ class Canvas(QWidget):
         assert self.shapes
         self.shapes.pop()
 
+    def loadPixmap(self, pixmap):
+        self.pixmap = pixmap
+        self.shapes = []
+        self.repaint()
+
+    def loadShapes(self, shapes):
+        self.shapes = list(shapes)
+        self.current = None
+        self.repaint()
+
+    def setShapeVisible(self, shape, value):
+        self.visible[shape] = value
+        self.repaint()
+
 
 def pp(p):
     return '%.2f, %.2f' % (p.x(), p.y())

labelDialog.py

@@ -2,10 +2,13 @@
 from PyQt4.QtGui import *
 from PyQt4.QtCore import *
 
-from lib import newButton
+from lib import newButton, labelValidator
 
 BB = QDialogButtonBox
 
+# FIXME:
+# - Use the validator when accepting the dialog.
+
 class LabelDialog(QDialog):
     OK, UNDO, DELETE = range(3)
 
@@ -14,6 +17,8 @@ class LabelDialog(QDialog):
         self.action = self.OK
         self.edit = QLineEdit()
         self.edit.setText(text)
+        self.edit.setValidator(labelValidator())
+        self.edit.editingFinished.connect(self.postProcess)
         layout = QHBoxLayout()
         layout.addWidget(self.edit)
         delete = newButton('Delete', icon='delete', slot=self.delete)
@@ -22,16 +27,17 @@ class LabelDialog(QDialog):
         bb.addButton(BB.Ok)
         bb.addButton(undo, BB.RejectRole)
         bb.addButton(delete, BB.RejectRole)
-        bb.accepted.connect(self.accept)
+        bb.accepted.connect(self.validate)
-        bb.rejected.connect(self.reject)
         layout.addWidget(bb)
         self.setLayout(layout)
 
     def undo(self):
         self.action = self.UNDO
+        self.reject()
 
     def delete(self):
         self.action = self.DELETE
+        self.reject()
 
     def text(self):
         return self.edit.text()
@@ -42,5 +48,16 @@ class LabelDialog(QDialog):
         self.edit.setText(u"Enter label")
         self.edit.setSelection(0, len(self.text()))
         self.edit.setFocus(Qt.PopupFocusReason)
-        return self.OK if self.exec_() else self.action
+        return self.OK if self.exec_() == QDialog.Accepted else self.action
+
+    def validate(self):
+        if self.edit.text().trimmed():
+            self.accept()
+
+    def postProcess(self):
+        self.edit.setText(self.edit.text().trimmed())
+
+    def keyPressEvent(self, ev):
+        if ev.key() == Qt.Key_Escape:
+            self.undo()
 

labelFile.py

@@ -0,0 +1,46 @@
+
+import json
+import os.path
+
+from base64 import b64encode, b64decode
+
+class LabelFileError(Exception):
+    pass
+
+class LabelFile(object):
+    suffix = '.lif'
+
+    def __init__(self):
+        self.shapes = ()
+        self.imagePath = None
+        self.imageData = None
+
+    def load(self, filename):
+        try:
+            with open(filename, 'rb') as f:
+                data = json.load(f)
+                imagePath = data['imagePath']
+                imageData = b64decode(data['imageData'])
+                shapes = ((s['label'], s['points']) for s in data['shapes'])
+                # Only replace data after everything is loaded.
+                self.shapes = shapes
+                self.imagePath = imagePath
+                self.imageData = imageData
+        except Exception, e:
+            raise LabelFileError(e)
+
+    def save(self, filename, shapes, imagePath, imageData):
+        try:
+            with open(filename, 'wb') as f:
+                json.dump(dict(
+                    shapes=[dict(label=l, points=p) for (l, p) in shapes],
+                    imagePath=imagePath,
+                    imageData=b64encode(imageData)),
+                    f, ensure_ascii=True, indent=2)
+        except Exception, e:
+            raise LabelFileError(e)
+
+    @staticmethod
+    def isLabelFile(filename):
+        return os.path.splitext(filename)[1].lower() == LabelFile.suffix
+

labelme.py

@@ -13,18 +13,26 @@ from PyQt4.QtCore import *
 
 import resources
 
-from lib import newAction, addActions
+from lib import newAction, addActions, labelValidator
 from shape import Shape
 from canvas import Canvas
 from zoomWidget import ZoomWidget
 from labelDialog import LabelDialog
+from labelFile import LabelFile
 
 
 __appname__ = 'labelme'
 
+# FIXME
+# - [low] Label validation/postprocessing breaks with TAB.
+
 # TODO:
-# - Zoom is too "steppy".
 # - Add a new column in list widget with checkbox to show/hide shape.
+# - Make sure the `save' action is disabled when no labels are
+#   present in the image, e.g. when all of them are deleted.
+# - [easy] Add button to Hide/Show all labels.
+# - Zoom is too "steppy".
+
 
 ### Utility functions and classes.
 
@@ -65,7 +73,10 @@ class MainWindow(QMainWindow, WindowMixin):
         self.dock.setWidget(self.labelList)
         self.zoom_widget = ZoomWidget()
 
+        self.labelList.setItemDelegate(LabelDelegate())
         self.labelList.itemActivated.connect(self.highlightLabel)
+        # Connect to itemChanged to detect checkbox changes.
+        self.labelList.itemChanged.connect(self.labelItemChanged)
 
         self.canvas = Canvas()
         #self.canvas.setAlignment(Qt.AlignCenter)
@@ -93,6 +104,8 @@ class MainWindow(QMainWindow, WindowMixin):
                 'Ctrl+Q', 'quit', u'Exit application')
         open = action('&Open', self.openFile,
                 'Ctrl+O', 'open', u'Open file')
+        save = action('&Save', self.saveFile,
+                'Ctrl+S', 'save', u'Save file')
         color = action('&Color', self.chooseColor,
                 'Ctrl+C', 'color', u'Choose line color')
         label = action('&New Item', self.newLabel,
@@ -120,6 +133,11 @@ class MainWindow(QMainWindow, WindowMixin):
         zoom = QWidgetAction(self)
         zoom.setDefaultWidget(self.zoom_widget)
 
+        # Store actions for further handling.
+        self.actions = struct(save=save, open=open, color=color,
+                label=label, delete=delete, zoom=zoom)
+        save.setEnabled(False)
+
         fit_window = action('&Fit Window', self.setFitWindow,
                 'Ctrl+F', 'fit',  u'Fit image to window', checkable=True)
 
@@ -127,13 +145,13 @@ class MainWindow(QMainWindow, WindowMixin):
                 file=self.menu('&File'),
                 edit=self.menu('&Image'),
                 view=self.menu('&View'))
-        addActions(self.menus.file, (open, quit))
+        addActions(self.menus.file, (open, save, quit))
         addActions(self.menus.edit, (label, color, fit_window))
 
         addActions(self.menus.view, (labels,))
 
         self.tools = self.toolbar('Tools')
-        addActions(self.tools, (open, color, None, label, delete, hide, None,
+        addActions(self.tools, (open, save, color, None, label, delete, hide, None,
             zoom, fit_window, None, quit))
 
 
@@ -183,9 +201,29 @@ class MainWindow(QMainWindow, WindowMixin):
 
     def addLabel(self, shape):
         item = QListWidgetItem(shape.label)
+        item.setFlags(item.flags() | Qt.ItemIsUserCheckable | Qt.ItemIsEditable)
+        item.setCheckState(Qt.Checked)
         self.labels[item] = shape
         self.labelList.addItem(item)
 
+
+    def loadLabels(self, shapes):
+        s = []
+        for label, points in shapes:
+            shape = Shape(label=label)
+            shape.fill = True
+            for x, y in points:
+                shape.addPoint(QPointF(x, y))
+            s.append(shape)
+            self.addLabel(shape)
+        self.canvas.loadShapes(s)
+
+    def saveLabels(self, filename):
+        lf = LabelFile()
+        shapes = [(unicode(shape.label), [(p.x(), p.y()) for p in shape.points])\
+                for shape in self.canvas.shapes]
+        lf.save(filename, shapes, unicode(self.filename), self.imageData)
+
     def copySelectedShape(self):
         self.addLabel(self.copySelectedShape())
 
@@ -197,6 +235,18 @@ class MainWindow(QMainWindow, WindowMixin):
         self.highlighted = shape
         self.canvas.repaint()
 
+    def labelItemChanged(self, item):
+        shape = self.labels[item]
+        label = unicode(item.text())
+        if label != shape.label:
+            self.stateChanged()
+            shape.label = unicode(item.text())
+        else: # User probably changed item visibility
+            self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
+
+    def stateChanged(self):
+        self.actions.save.setEnabled(True)
+
     ## Callback functions:
     def newShape(self, position):
         """Pop-up and give focus to the label editor.
@@ -206,6 +256,8 @@ class MainWindow(QMainWindow, WindowMixin):
         action = self.label.popUp(position)
         if action == self.label.OK:
             self.addLabel(self.canvas.setLastLabel(self.label.text()))
+            # Enable the save action.
+            self.actions.save.setEnabled(True)
         elif action == self.label.UNDO:
             self.canvas.undoLastLine()
         elif action == self.label.DELETE:
@@ -240,17 +292,31 @@ class MainWindow(QMainWindow, WindowMixin):
         """Load the specified file, or the last opened file if None."""
         if filename is None:
             filename = self.settings['filename']
-        # FIXME: Load the actual file here.
+        filename = unicode(filename)
         if QFile.exists(filename):
-            # Load image
+            if LabelFile.isLabelFile(filename):
-            image = QImage(filename)
+                # TODO: Error handling.
+                lf = LabelFile()
+                lf.load(filename)
+                self.labelFile = lf
+                self.imageData = lf.imageData
+            else:
+                # Load image:
+                # read data first and store for saving into label file.
+                self.imageData = read(filename, None)
+                self.labelFile = None
+            image = QImage.fromData(self.imageData)
             if image.isNull():
                 message = "Failed to read %s" % filename
             else:
                 message = "Loaded %s" % os.path.basename(unicode(filename))
                 self.image = image
                 self.filename = filename
-                self.loadPixmap()
+                self.labels = {}
+                self.labelList.clear()
+                self.canvas.loadPixmap(QPixmap.fromImage(image))
+                if self.labelFile:
+                    self.loadLabels(self.labelFile.shapes)
             self.statusBar().showMessage(message)
 
     def resizeEvent(self, event):
@@ -258,10 +324,6 @@ class MainWindow(QMainWindow, WindowMixin):
             self.paintCanvas()
         super(MainWindow, self).resizeEvent(event)
 
-    def loadPixmap(self):
-        assert not self.image.isNull(), "cannot load null image"
-        self.canvas.pixmap = QPixmap.fromImage(self.image)
-
     def paintCanvas(self):
         assert not self.image.isNull(), "cannot paint null image"
         self.canvas.scale = self.fitSize() if self.fit_window\
@@ -305,11 +367,25 @@ class MainWindow(QMainWindow, WindowMixin):
                 if self.filename else '.'
         formats = ['*.%s' % unicode(fmt).lower()\
                 for fmt in QImageReader.supportedImageFormats()]
+        filters = 'Image files (%s)\nLabel files (*%s)'\
+                % (' '.join(formats), LabelFile.suffix)
         filename = unicode(QFileDialog.getOpenFileName(self,
-            '%s - Choose Image', path, 'Image files (%s)' % ' '.join(formats)))
+            '%s - Choose Image', path, filters))
         if filename:
             self.loadFile(filename)
 
+    def saveFile(self):
+        assert not self.image.isNull(), "cannot save empty image"
+        # XXX: What if user wants to remove label file?
+        assert self.labels, "cannot save empty labels"
+        path = os.path.dirname(unicode(self.filename))\
+                if self.filename else '.'
+        formats = ['*%s' % LabelFile.suffix]
+        filename = unicode(QFileDialog.getSaveFileName(self,
+            '%s - Choose File', path, 'Label files (%s)' % ''.join(formats)))
+        if filename:
+            self.saveLabels(filename)
+
     def check(self):
         # TODO: Prompt user to save labels etc.
         return True
@@ -329,6 +405,24 @@ class MainWindow(QMainWindow, WindowMixin):
         self.canvas.deleteSelected()
 
 
+class LabelDelegate(QItemDelegate):
+    def __init__(self, parent=None):
+        super(LabelDelegate, self).__init__(parent)
+        self.validator = labelValidator()
+
+    # FIXME: Validation and trimming are completely broken if the
+    # user navigates away from the editor with something like TAB.
+    def createEditor(self, parent, option, index):
+        """Make sure the user cannot enter empty labels.
+        Also remove trailing whitespace."""
+        edit = super(LabelDelegate, self).createEditor(parent, option, index)
+        if isinstance(edit, QLineEdit):
+            edit.setValidator(self.validator)
+            def strip():
+                edit.setText(edit.text().trimmed())
+            edit.editingFinished.connect(strip)
+        return edit
+
 class Settings(object):
     """Convenience dict-like wrapper around QSettings."""
     def __init__(self, types=None):
@@ -359,6 +453,13 @@ def inverted(color):
     return QColor(*[255 - v for v in color.getRgb()])
 
 
+def read(filename, default=None):
+    try:
+        with open(filename, 'rb') as f:
+            return f.read()
+    except:
+        return default
+
 class struct(object):
     def __init__(self, **kwargs):
         self.__dict__.update(kwargs)

lib.py

@@ -39,3 +39,6 @@ def addActions(widget, actions):
         else:
             widget.addAction(action)
 
+def labelValidator():
+    return QRegExpValidator(QRegExp(r'^[^ \t].+'), None)
+