Prechádzať zdrojové kódy

Merge branch 'merge/list-save'

Michael Pitidis 13 rokov pred
rodič
commit
60410eff65
6 zmenil súbory, kde vykonal 425 pridanie a 104 odobranie
  1. 156 68
      canvas.py
  2. 23 5
      labelDialog.py
  3. 46 0
      labelFile.py
  4. 171 19
      labelme.py
  5. 3 0
      lib.py
  6. 26 12
      shape.py

+ 156 - 68
canvas.py

@@ -11,53 +11,64 @@ class Canvas(QWidget):
     scrollRequest = pyqtSignal(int, int)
     newShape = pyqtSignal(QPoint)
 
+    SELECT, EDIT = range(2)
+
     epsilon = 9.0 # TODO: Tune value
 
     def __init__(self, *args, **kwargs):
         super(Canvas, self).__init__(*args, **kwargs)
-        self.startLabeling=False # has to click new label buttoon to starting drawing new polygons
+        # Initialise local state.
+        self.mode = self.SELECT
         self.shapes = []
         self.current = None
         self.selectedShape=None # save the selected shape here
         self.selectedShapeCopy=None
-        self.line_color = QColor(0, 0, 255)
-        self.line = Shape(line_color=self.line_color)
+        self.lineColor = QColor(0, 0, 255)
+        self.line = Shape(line_color=self.lineColor)
         self.mouseButtonIsPressed=False #when it is true and shape is selected , move the shape with the mouse move event
-        self.prevPoint=QPoint()
+        self.prevPoint = QPointF()
+        self.offsets = QPointF(), QPointF()
         self.scale = 1.0
         self.pixmap = None
-
+        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
+
+    def setEditing(self, value=True):
+        self.mode = self.EDIT if value else self.SELECT
+
     def mouseMoveEvent(self, ev):
         """Update line with last point and current coordinates."""
+        pos = self.transformPos(ev.posF())
 
-        if (ev.buttons() & 2):  # wont work , as ev.buttons doesn't work well , or I haven't known how to use it :) to use right click
-            print ev.button()
+        # Polygon copy moving.
+        if Qt.RightButton & ev.buttons():
             if self.selectedShapeCopy:
                 if self.prevPoint:
-                    point=QPoint(self.prevPoint)
-                    dx= ev.x()-point.x()
-                    dy=ev.y()-point.y()
-                    self.selectedShapeCopy.moveBy(dx,dy)
+                    self.selectedShapeCopy.moveBy(pos - self.prevPoint)
                     self.repaint()
-                self.prevPoint=ev.pos()
+                self.prevPoint = pos
             elif self.selectedShape:
-                newShape=Shape()
-                for point in self.selectedShape.points:
-                    newShape.addPoint(point)
-                self.selectedShapeCopy=newShape
+                self.selectedShapeCopy = self.selectedShape.copy()
                 self.repaint()
             return
 
         # Polygon drawing.
-        if self.current and self.startLabeling:
-            pos = self.transformPos(ev.posF())
-            color = self.line_color
+        if self.current and self.editing():
+            color = self.lineColor
             if self.outOfPixmap(pos):
                 # Don't allow the user to draw outside the pixmap.
                 # Project the point to the pixmap's edges.
-                pos = self.intersectionPoint(pos)
+                pos = self.intersectionPoint(self.current[-1], pos)
             elif len(self.current) > 1 and self.closeEnough(pos, self.current[0]):
                 # Attract line to starting point and colorise to alert the user:
                 # TODO: I would also like to highlight the pixel somehow.
@@ -66,67 +77,124 @@ class Canvas(QWidget):
             self.line[1] = pos
             self.line.line_color = color
             self.repaint()
+
+        # Polygon moving.
+        elif Qt.LeftButton & ev.buttons() and self.selectedShape and self.prevPoint:
+            self.boundedMoveShape(pos)
+            self.repaint()
             return
 
-        if self.selectedShape:
-            if self.prevPoint:
-                    point=QPoint(self.prevPoint)
-                   # print point.x()
-                    dx= ev.x()-point.x()
-                    dy=ev.y()-point.y()
-                    self.selectedShape.moveBy(dx,dy)
-                    self.repaint()
-            self.prevPoint=ev.pos()
-            
+        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):
-        if ev.button() == 1:
-            if self.startLabeling:
+        pos = self.transformPos(ev.posF())
+        if ev.button() == Qt.LeftButton:
+            if self.editing():
                 if self.current:
                     self.current.addPoint(self.line[1])
                     self.line[0] = self.current[-1]
                     if self.current.isClosed():
                         self.finalise(ev)
-                    self.repaint()
-                else:
-                    pos = self.transformPos(ev.posF())
-                    if self.outOfPixmap(pos):
-                        return
+                elif not self.outOfPixmap(pos):
                     self.current = Shape()
-                    self.line.points = [pos, pos]
                     self.current.addPoint(pos)
-                    self.setMouseTracking(True)
-            else: # not in adding new label mode
-                self.selectShape(ev.pos())
-                self.prevPoint=ev.pos()
-       
-                
+                    self.line.points = [pos, pos]
+                    self.setHiding()
+            else:
+                self.selectShape(pos)
+                self.prevPoint = pos
+            self.repaint()
+        #elif ev.button() == Qt.RightButton and not self.editing():
+        #    self.selectShape(pos)
+        #    self.prevPoint = pos
+
+    def hideBackroundShapes(self, value):
+        self.hideBackround = value
+        if self.selectedShape:
+            # Only hide other shapes if there is a current selection.
+            # Otherwise the user will not be able to select a shape.
+            self.setHiding(True)
+            self.repaint()
+
+    def setHiding(self, enable=True):
+        self._hideBackround = self.hideBackround if enable else False
 
     def mouseDoubleClickEvent(self, ev):
-        if self.current and self.startLabeling:
-            # Add first point in the list so that it is consistent
-            # with shapes created the normal way.
-            self.current.addPoint(self.current[0])
+        if self.current and self.editing():
+            # Shapes need to have at least 3 vertices.
+            if len(self.current) < 4:
+                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):
         """Select the first shape created which contains this point."""
         self.deSelectShape()
-        for shape in self.shapes:
+        for shape in reversed(self.shapes):
             if shape.containsPoint(point):
                 shape.selected = True
                 self.selectedShape = shape
-                return self.repaint()
+                self.calculateOffsets(shape, point)
+                self.setHiding()
+                self.repaint()
+                return
+
+    def calculateOffsets(self, shape, point):
+        rect = shape.boundingRect()
+        x1 = rect.x() - point.x()
+        y1 = rect.y() - point.y()
+        x2 = (rect.x() + rect.width()) - point.x()
+        y2 = (rect.y() + rect.height()) - point.y()
+        self.offsets = QPointF(x1, y1), QPointF(x2, y2)
+
+    def boundedMoveShape(self, pos):
+        if self.outOfPixmap(pos):
+            return # No need to move
+        o1 = pos + self.offsets[0]
+        if self.outOfPixmap(o1):
+            pos -= QPointF(min(0, o1.x()), min(0, o1.y()))
+        o2 = pos + self.offsets[1]
+        if self.outOfPixmap(o2):
+            pos += QPointF(min(0, self.pixmap.width() - o2.x()),
+                           min(0, self.pixmap.height()- o2.y()))
+        # The next line tracks the new position of the cursor
+        # relative to the shape, but also results in making it
+        # a bit "shaky" when nearing the border and allows it to
+        # go outside of the shape's area for some reason. XXX
+        #self.calculateOffsets(self.selectedShape, pos)
+        self.selectedShape.moveBy(pos - self.prevPoint)
+        self.prevPoint = pos
 
     def deSelectShape(self):
         if self.selectedShape:
             self.selectedShape.selected = False
-            self.repaint()
+            self.selectedShape = None
+            self.setHiding(False)
+
     def deleteSelected(self):
         if self.selectedShape:
              self.shapes.remove(self.selectedShape)
-             self.selectedShape=None
-             #print self.selectedShape()
+             self.selectedShape = None
              self.repaint()
+
+    def copySelectedShape(self):
+        if self.selectedShape:
+            shape = self.selectedShape.copy()
+            self.shapes.append(shape)
+            self.selectedShape = shape
+            self.deSelectShape()
+            self.repaint()
+            return shape
+
+
     def paintEvent(self, event):
         if not self.pixmap:
             return super(Canvas, self).paintEvent(event)
@@ -139,14 +207,15 @@ class Canvas(QWidget):
         p.translate(self.offsetToCenter())
 
         p.drawPixmap(0, 0, self.pixmap)
+        Shape.scale = self.scale
         for shape in self.shapes:
-            shape.paint(p)
+            if (shape.selected or not self._hideBackround) and self.isVisible(shape):
+                shape.paint(p)
         if self.current:
             self.current.paint(p)
             self.line.paint(p)
         if self.selectedShapeCopy:
             self.selectedShapeCopy.paint(p)
-        
 
         p.end()
 
@@ -167,15 +236,13 @@ class Canvas(QWidget):
         w, h = self.pixmap.width(), self.pixmap.height()
         return not (0 <= p.x() <= w and 0 <= p.y() <= h)
 
-
     def finalise(self, ev):
         assert self.current
         self.current.fill = True
         self.shapes.append(self.current)
         self.current = None
-        self.startLabeling = False
-        # TODO: Mouse tracking is still useful for selecting shapes!
-        self.setMouseTracking(False)
+        self.setEditing(False)
+        self.setHiding(False)
         self.repaint()
         self.newShape.emit(self.mapToGlobal(ev.pos()))
 
@@ -185,7 +252,7 @@ class Canvas(QWidget):
         #print "d %.2f, m %d, %.2f" % (d, m, d - m)
         return distance(p1 - p2) < self.epsilon
 
-    def intersectionPoint(self, mousePos):
+    def intersectionPoint(self, p1, p2):
         # Cycle through each image edge in clockwise fashion,
         # and find the one intersecting the current line segment.
         # http://paulbourke.net/geometry/lineline2d/
@@ -194,8 +261,8 @@ class Canvas(QWidget):
                   (size.width(), 0),
                   (size.width(), size.height()),
                   (0, size.height())]
-        x1, y1 = self.current[-1].x(), self.current[-1].y()
-        x2, y2 = mousePos.x(), mousePos.y()
+        x1, y1 = p1.x(), p1.y()
+        x2, y2 = p2.x(), p2.y()
         d, i, (x, y) = min(self.intersectingEdges((x1, y1), (x2, y2), points))
         x3, y3 = points[i]
         x4, y4 = points[(i+1)%4]
@@ -207,7 +274,6 @@ class Canvas(QWidget):
                 return QPointF(min(max(0, x2), max(x3, x4)), y3)
         return QPointF(x, y)
 
-
     def intersectingEdges(self, (x1, y1), (x2, y2), points):
         """For each edge formed by `points', yield the intersection
         with the line segment `(x1,y1) - (x2,y2)`, if it exists.
@@ -232,7 +298,6 @@ class Canvas(QWidget):
                 d = distance(m - QPointF(x2, y2))
                 yield d, i, (x, y)
 
-
     # These two, along with a call to adjustSize are required for the
     # scroll area.
     def sizeHint(self):
@@ -259,20 +324,43 @@ class Canvas(QWidget):
     def keyPressEvent(self, ev):
         if ev.key() == Qt.Key_Escape and self.current:
             self.current = None
-            self.setMouseTracking(False)
             self.repaint()
 
     def setLastLabel(self, text):
         assert text
-        print "Setting shape label to %r" % text
+        print "shape <- '%s'" % text
         self.shapes[-1].label = text
+        return self.shapes[-1]
 
     def undoLastLine(self):
-        print "Undoing last line"
+        assert self.shapes
+        self.current = self.shapes.pop()
+        self.current.fill = False
+        pos = self.current.popPoint()
+        self.line.points = [self.current[-1], pos]
+        self.setEditing()
 
     def deleteLastShape(self):
-        print "Deleting last shape"
+        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())
 
 def distance(p):
     return sqrt(p.x() * p.x() + p.y() * p.y())

+ 23 - 5
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,24 +27,37 @@ class LabelDialog(QDialog):
         bb.addButton(BB.Ok)
         bb.addButton(undo, BB.RejectRole)
         bb.addButton(delete, BB.RejectRole)
-        bb.accepted.connect(self.accept)
-        bb.rejected.connect(self.reject)
+        bb.accepted.connect(self.validate)
         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()
 
     def popUp(self, position):
-        self.move(position)
+        # It actually works without moving...
+        #self.move(position)
         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()
 

+ 46 - 0
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
+

+ 171 - 19
labelme.py

@@ -13,18 +13,27 @@ 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:
+# - 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.
 
 class WindowMixin(object):
@@ -56,10 +65,22 @@ class MainWindow(QMainWindow, WindowMixin):
 
         # Main widgets.
         self.label = LabelDialog(parent=self)
+        self.labels = {}
+        self.highlighted = None
+        self.labelList = QListWidget()
+        self.dock = QDockWidget(u'Labels', self)
+        self.dock.setObjectName(u'Labels')
+        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)
+
         self.canvas.setContextMenuPolicy(Qt.ActionsContextMenu)
         self.canvas.zoomRequest.connect(self.zoomRequest)
 
@@ -75,6 +96,7 @@ class MainWindow(QMainWindow, WindowMixin):
         self.canvas.newShape.connect(self.newShape)
 
         self.setCentralWidget(scroll)
+        self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
 
         # Actions
         action = partial(newAction, self)
@@ -82,17 +104,40 @@ 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,
                 'Ctrl+N', 'new', u'Add new label')
+        copy = action('&Copy', self.copySelectedShape,
+                'Ctrl+C', 'copy', u'Copy')
         delete = action('&Delete', self.deleteSelectedShape,
                 'Ctrl+D', 'delete', u'Delete')
+        hide = action('&Hide labels', self.hideLabelsToggle,
+                'Ctrl+H', 'hide', u'Hide background labels when drawing',
+                checkable=True)
+
+        self.canvas.setContextMenuPolicy( Qt.CustomContextMenu )
+        self.canvas.customContextMenuRequested.connect(self.popContextMenu)
+
+        # Popup Menu
+        self.popMenu = QMenu(self )
+        self.popMenu.addAction( label )
+        self.popMenu.addAction(copy)
+        self.popMenu.addAction( delete )
 
+        labels = self.dock.toggleViewAction()
+        labels.setShortcut('Ctrl+L')
 
         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)
 
@@ -100,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, (labl,))
+        addActions(self.menus.view, (labels,))
 
         self.tools = self.toolbar('Tools')
-        addActions(self.tools, (open, color, None, label, delete, None,
+        addActions(self.tools, (open, save, color, None, label, delete, hide, None,
             zoom, fit_window, None, quit))
 
 
@@ -151,6 +196,56 @@ class MainWindow(QMainWindow, WindowMixin):
         # Callbacks:
         self.zoom_widget.editingFinished.connect(self.paintCanvas)
 
+    def popContextMenu(self, point):
+        self.popMenu.exec_(self.canvas.mapToGlobal(point))
+
+    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())
+
+    def highlightLabel(self, item):
+        if self.highlighted:
+            self.highlighted.fill_color = Shape.fill_color
+        shape = self.labels[item]
+        shape.fill_color = inverted(Shape.fill_color)
+        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):
@@ -160,11 +255,10 @@ class MainWindow(QMainWindow, WindowMixin):
         """
         action = self.label.popUp(position)
         if action == self.label.OK:
-            print "Setting label to %s" % self.label.text()
-            self.canvas.setLastLabel(self.label.text())
-            # TODO: Add to list of labels.
+            self.addLabel(self.canvas.setLastLabel(self.label.text()))
+            # Enable the save action.
+            self.actions.save.setEnabled(True)
         elif action == self.label.UNDO:
-            print "Undo last line"
             self.canvas.undoLastLine()
         elif action == self.label.DELETE:
             self.canvas.deleteLastShape()
@@ -191,21 +285,38 @@ class MainWindow(QMainWindow, WindowMixin):
     def queueEvent(self, function):
         QTimer.singleShot(0, function)
 
+    def hideLabelsToggle(self, value):
+        self.canvas.hideBackroundShapes(value)
+
     def loadFile(self, filename=None):
         """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
-            image = QImage(filename)
+            if LabelFile.isLabelFile(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):
@@ -213,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\
@@ -245,8 +352,10 @@ class MainWindow(QMainWindow, WindowMixin):
         s['window/position'] = self.pos()
         s['window/state'] = self.saveState()
         s['line/color'] = self.color
-        #s['window/geometry'] = self.saveGeometry()
+        # ask the use for where to save the labels
+       
 
+        #s['window/geometry'] = self.saveGeometry()
     def updateFileMenu(self):
         """Populate menu with recent files."""
 
@@ -258,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
@@ -276,12 +399,30 @@ class MainWindow(QMainWindow, WindowMixin):
 
     def newLabel(self):
         self.canvas.deSelectShape()
-        self.canvas.startLabeling=True
+        self.canvas.setEditing()
 
     def deleteSelectedShape(self):
         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):
@@ -308,6 +449,17 @@ class Settings(object):
         return value
 
 
+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)

+ 3 - 0
lib.py

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

+ 26 - 12
shape.py

@@ -5,7 +5,7 @@ from PyQt4.QtGui import *
 from PyQt4.QtCore import *
 
 # FIXME:
-# - Don't scale vertices.
+# - Add support for highlighting vertices.
 
 class Shape(object):
     P_SQUARE, P_ROUND = range(2)
@@ -17,6 +17,7 @@ class Shape(object):
     select_color = QColor(255, 255, 255)
     point_type = P_SQUARE
     point_size = 8
+    scale = 1.0
 
     def __init__(self, label=None, line_color=None):
         self.label = label
@@ -45,6 +46,8 @@ class Shape(object):
     def paint(self, painter):
         if self.points:
             pen = QPen(self.select_color if self.selected else self.line_color)
+            # Try using integer sizes for smoother drawing(?)
+            pen.setWidth(max(1, int(round(2.0 / self.scale))))
             painter.setPen(pen)
 
             line_path = QPainterPath()
@@ -64,24 +67,35 @@ class Shape(object):
                 painter.fillPath(line_path, self.fill_color)
 
     def drawVertex(self, path, point):
-        d = self.point_size
+        d = self.point_size / self.scale
         if self.point_type == self.P_SQUARE:
             path.addRect(point.x() - d/2, point.y() - d/2, d, d)
         else:
             path.addEllipse(point, d/2.0, d/2.0)
 
     def containsPoint(self, point):
-        path = QPainterPath(QPointF(self.points[0]))
+        return self.makePath().contains(point)
+
+    def makePath(self):
+        path = QPainterPath(self.points[0])
         for p in self.points[1:]:
-            path.lineTo(QPointF(p))
-        return path.contains(QPointF(point))
-    def moveBy(self,dx,dy):
-        index=0
-        for point in self.points:
-           newXPos= point.x()+dx
-           newYPos=point.y()+dy
-           self.points[index]=QPoint(newXPos,newYPos)
-           index +=1
+            path.lineTo(p)
+        return path
+
+    def boundingRect(self):
+        return self.makePath().boundingRect()
+
+    def moveBy(self, offset):
+        self.points = [p + offset for p in self.points]
+
+    def copy(self):
+        shape = Shape()
+        shape.points= [p for p in self.points]
+        shape.label = "Copy of %s" % self.label
+        shape.fill = self.fill
+        shape.selected = self.selected
+        return shape
+
     def __len__(self):
         return len(self.points)