Pārlūkot izejas kodu

Merge branch 'merge/list-save'

Michael Pitidis 13 gadi atpakaļ
vecāks
revīzija
60410eff65
6 mainītis faili ar 425 papildinājumiem un 104 dzēšanām
  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)
     scrollRequest = pyqtSignal(int, int)
     newShape = pyqtSignal(QPoint)
     newShape = pyqtSignal(QPoint)
 
 
+    SELECT, EDIT = range(2)
+
     epsilon = 9.0 # TODO: Tune value
     epsilon = 9.0 # TODO: Tune value
 
 
     def __init__(self, *args, **kwargs):
     def __init__(self, *args, **kwargs):
         super(Canvas, self).__init__(*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.shapes = []
         self.current = None
         self.current = None
         self.selectedShape=None # save the selected shape here
         self.selectedShape=None # save the selected shape here
         self.selectedShapeCopy=None
         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.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.scale = 1.0
         self.pixmap = None
         self.pixmap = None
-
+        self.visible = {}
+        self._hideBackround = False
+        self.hideBackround = False
+        # Set widget options.
+        self.setMouseTracking(True)
         self.setFocusPolicy(Qt.WheelFocus)
         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):
     def mouseMoveEvent(self, ev):
         """Update line with last point and current coordinates."""
         """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.selectedShapeCopy:
                 if self.prevPoint:
                 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.repaint()
-                self.prevPoint=ev.pos()
+                self.prevPoint = pos
             elif self.selectedShape:
             elif self.selectedShape:
-                newShape=Shape()
-                for point in self.selectedShape.points:
-                    newShape.addPoint(point)
-                self.selectedShapeCopy=newShape
+                self.selectedShapeCopy = self.selectedShape.copy()
                 self.repaint()
                 self.repaint()
             return
             return
 
 
         # Polygon drawing.
         # 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):
             if self.outOfPixmap(pos):
                 # Don't allow the user to draw outside the pixmap.
                 # Don't allow the user to draw outside the pixmap.
                 # Project the point to the pixmap's edges.
                 # 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]):
             elif len(self.current) > 1 and self.closeEnough(pos, self.current[0]):
                 # Attract line to starting point and colorise to alert the user:
                 # Attract line to starting point and colorise to alert the user:
                 # TODO: I would also like to highlight the pixel somehow.
                 # TODO: I would also like to highlight the pixel somehow.
@@ -66,67 +77,124 @@ class Canvas(QWidget):
             self.line[1] = pos
             self.line[1] = pos
             self.line.line_color = color
             self.line.line_color = color
             self.repaint()
             self.repaint()
+
+        # Polygon moving.
+        elif Qt.LeftButton & ev.buttons() and self.selectedShape and self.prevPoint:
+            self.boundedMoveShape(pos)
+            self.repaint()
             return
             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):
     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:
                 if self.current:
                     self.current.addPoint(self.line[1])
                     self.current.addPoint(self.line[1])
                     self.line[0] = self.current[-1]
                     self.line[0] = self.current[-1]
                     if self.current.isClosed():
                     if self.current.isClosed():
                         self.finalise(ev)
                         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.current = Shape()
-                    self.line.points = [pos, pos]
                     self.current.addPoint(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):
     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)
             self.finalise(ev)
 
 
     def selectShape(self, point):
     def selectShape(self, point):
         """Select the first shape created which contains this point."""
         """Select the first shape created which contains this point."""
         self.deSelectShape()
         self.deSelectShape()
-        for shape in self.shapes:
+        for shape in reversed(self.shapes):
             if shape.containsPoint(point):
             if shape.containsPoint(point):
                 shape.selected = True
                 shape.selected = True
                 self.selectedShape = shape
                 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):
     def deSelectShape(self):
         if self.selectedShape:
         if self.selectedShape:
             self.selectedShape.selected = False
             self.selectedShape.selected = False
-            self.repaint()
+            self.selectedShape = None
+            self.setHiding(False)
+
     def deleteSelected(self):
     def deleteSelected(self):
         if self.selectedShape:
         if self.selectedShape:
              self.shapes.remove(self.selectedShape)
              self.shapes.remove(self.selectedShape)
-             self.selectedShape=None
-             #print self.selectedShape()
+             self.selectedShape = None
              self.repaint()
              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):
     def paintEvent(self, event):
         if not self.pixmap:
         if not self.pixmap:
             return super(Canvas, self).paintEvent(event)
             return super(Canvas, self).paintEvent(event)
@@ -139,14 +207,15 @@ class Canvas(QWidget):
         p.translate(self.offsetToCenter())
         p.translate(self.offsetToCenter())
 
 
         p.drawPixmap(0, 0, self.pixmap)
         p.drawPixmap(0, 0, self.pixmap)
+        Shape.scale = self.scale
         for shape in self.shapes:
         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:
         if self.current:
             self.current.paint(p)
             self.current.paint(p)
             self.line.paint(p)
             self.line.paint(p)
         if self.selectedShapeCopy:
         if self.selectedShapeCopy:
             self.selectedShapeCopy.paint(p)
             self.selectedShapeCopy.paint(p)
-        
 
 
         p.end()
         p.end()
 
 
@@ -167,15 +236,13 @@ class Canvas(QWidget):
         w, h = self.pixmap.width(), self.pixmap.height()
         w, h = self.pixmap.width(), self.pixmap.height()
         return not (0 <= p.x() <= w and 0 <= p.y() <= h)
         return not (0 <= p.x() <= w and 0 <= p.y() <= h)
 
 
-
     def finalise(self, ev):
     def finalise(self, ev):
         assert self.current
         assert self.current
         self.current.fill = True
         self.current.fill = True
         self.shapes.append(self.current)
         self.shapes.append(self.current)
         self.current = None
         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.repaint()
         self.newShape.emit(self.mapToGlobal(ev.pos()))
         self.newShape.emit(self.mapToGlobal(ev.pos()))
 
 
@@ -185,7 +252,7 @@ class Canvas(QWidget):
         #print "d %.2f, m %d, %.2f" % (d, m, d - m)
         #print "d %.2f, m %d, %.2f" % (d, m, d - m)
         return distance(p1 - p2) < self.epsilon
         return distance(p1 - p2) < self.epsilon
 
 
-    def intersectionPoint(self, mousePos):
+    def intersectionPoint(self, p1, p2):
         # Cycle through each image edge in clockwise fashion,
         # Cycle through each image edge in clockwise fashion,
         # and find the one intersecting the current line segment.
         # and find the one intersecting the current line segment.
         # http://paulbourke.net/geometry/lineline2d/
         # http://paulbourke.net/geometry/lineline2d/
@@ -194,8 +261,8 @@ class Canvas(QWidget):
                   (size.width(), 0),
                   (size.width(), 0),
                   (size.width(), size.height()),
                   (size.width(), size.height()),
                   (0, 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))
         d, i, (x, y) = min(self.intersectingEdges((x1, y1), (x2, y2), points))
         x3, y3 = points[i]
         x3, y3 = points[i]
         x4, y4 = points[(i+1)%4]
         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(min(max(0, x2), max(x3, x4)), y3)
         return QPointF(x, y)
         return QPointF(x, y)
 
 
-
     def intersectingEdges(self, (x1, y1), (x2, y2), points):
     def intersectingEdges(self, (x1, y1), (x2, y2), points):
         """For each edge formed by `points', yield the intersection
         """For each edge formed by `points', yield the intersection
         with the line segment `(x1,y1) - (x2,y2)`, if it exists.
         with the line segment `(x1,y1) - (x2,y2)`, if it exists.
@@ -232,7 +298,6 @@ class Canvas(QWidget):
                 d = distance(m - QPointF(x2, y2))
                 d = distance(m - QPointF(x2, y2))
                 yield d, i, (x, y)
                 yield d, i, (x, y)
 
 
-
     # These two, along with a call to adjustSize are required for the
     # These two, along with a call to adjustSize are required for the
     # scroll area.
     # scroll area.
     def sizeHint(self):
     def sizeHint(self):
@@ -259,20 +324,43 @@ class Canvas(QWidget):
     def keyPressEvent(self, ev):
     def keyPressEvent(self, ev):
         if ev.key() == Qt.Key_Escape and self.current:
         if ev.key() == Qt.Key_Escape and self.current:
             self.current = None
             self.current = None
-            self.setMouseTracking(False)
             self.repaint()
             self.repaint()
 
 
     def setLastLabel(self, text):
     def setLastLabel(self, text):
         assert text
         assert text
-        print "Setting shape label to %r" % text
+        print "shape <- '%s'" % text
         self.shapes[-1].label = text
         self.shapes[-1].label = text
+        return self.shapes[-1]
 
 
     def undoLastLine(self):
     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):
     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):
 def distance(p):
     return sqrt(p.x() * p.x() + p.y() * p.y())
     return sqrt(p.x() * p.x() + p.y() * p.y())

+ 23 - 5
labelDialog.py

@@ -2,10 +2,13 @@
 from PyQt4.QtGui import *
 from PyQt4.QtGui import *
 from PyQt4.QtCore import *
 from PyQt4.QtCore import *
 
 
-from lib import newButton
+from lib import newButton, labelValidator
 
 
 BB = QDialogButtonBox
 BB = QDialogButtonBox
 
 
+# FIXME:
+# - Use the validator when accepting the dialog.
+
 class LabelDialog(QDialog):
 class LabelDialog(QDialog):
     OK, UNDO, DELETE = range(3)
     OK, UNDO, DELETE = range(3)
 
 
@@ -14,6 +17,8 @@ class LabelDialog(QDialog):
         self.action = self.OK
         self.action = self.OK
         self.edit = QLineEdit()
         self.edit = QLineEdit()
         self.edit.setText(text)
         self.edit.setText(text)
+        self.edit.setValidator(labelValidator())
+        self.edit.editingFinished.connect(self.postProcess)
         layout = QHBoxLayout()
         layout = QHBoxLayout()
         layout.addWidget(self.edit)
         layout.addWidget(self.edit)
         delete = newButton('Delete', icon='delete', slot=self.delete)
         delete = newButton('Delete', icon='delete', slot=self.delete)
@@ -22,24 +27,37 @@ class LabelDialog(QDialog):
         bb.addButton(BB.Ok)
         bb.addButton(BB.Ok)
         bb.addButton(undo, BB.RejectRole)
         bb.addButton(undo, BB.RejectRole)
         bb.addButton(delete, 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)
         layout.addWidget(bb)
         self.setLayout(layout)
         self.setLayout(layout)
 
 
     def undo(self):
     def undo(self):
         self.action = self.UNDO
         self.action = self.UNDO
+        self.reject()
 
 
     def delete(self):
     def delete(self):
         self.action = self.DELETE
         self.action = self.DELETE
+        self.reject()
 
 
     def text(self):
     def text(self):
         return self.edit.text()
         return self.edit.text()
 
 
     def popUp(self, position):
     def popUp(self, position):
-        self.move(position)
+        # It actually works without moving...
+        #self.move(position)
         self.edit.setText(u"Enter label")
         self.edit.setText(u"Enter label")
         self.edit.setSelection(0, len(self.text()))
         self.edit.setSelection(0, len(self.text()))
         self.edit.setFocus(Qt.PopupFocusReason)
         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
 import resources
 
 
-from lib import newAction, addActions
+from lib import newAction, addActions, labelValidator
 from shape import Shape
 from shape import Shape
 from canvas import Canvas
 from canvas import Canvas
 from zoomWidget import ZoomWidget
 from zoomWidget import ZoomWidget
 from labelDialog import LabelDialog
 from labelDialog import LabelDialog
+from labelFile import LabelFile
 
 
 
 
 __appname__ = 'labelme'
 __appname__ = 'labelme'
 
 
+# FIXME
+# - [low] Label validation/postprocessing breaks with TAB.
+
 # TODO:
 # 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".
 # - Zoom is too "steppy".
 
 
+
 ### Utility functions and classes.
 ### Utility functions and classes.
 
 
 class WindowMixin(object):
 class WindowMixin(object):
@@ -56,10 +65,22 @@ class MainWindow(QMainWindow, WindowMixin):
 
 
         # Main widgets.
         # Main widgets.
         self.label = LabelDialog(parent=self)
         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.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 = Canvas()
         #self.canvas.setAlignment(Qt.AlignCenter)
         #self.canvas.setAlignment(Qt.AlignCenter)
+
         self.canvas.setContextMenuPolicy(Qt.ActionsContextMenu)
         self.canvas.setContextMenuPolicy(Qt.ActionsContextMenu)
         self.canvas.zoomRequest.connect(self.zoomRequest)
         self.canvas.zoomRequest.connect(self.zoomRequest)
 
 
@@ -75,6 +96,7 @@ class MainWindow(QMainWindow, WindowMixin):
         self.canvas.newShape.connect(self.newShape)
         self.canvas.newShape.connect(self.newShape)
 
 
         self.setCentralWidget(scroll)
         self.setCentralWidget(scroll)
+        self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
 
 
         # Actions
         # Actions
         action = partial(newAction, self)
         action = partial(newAction, self)
@@ -82,17 +104,40 @@ class MainWindow(QMainWindow, WindowMixin):
                 'Ctrl+Q', 'quit', u'Exit application')
                 'Ctrl+Q', 'quit', u'Exit application')
         open = action('&Open', self.openFile,
         open = action('&Open', self.openFile,
                 'Ctrl+O', 'open', u'Open file')
                 'Ctrl+O', 'open', u'Open file')
+        save = action('&Save', self.saveFile,
+                'Ctrl+S', 'save', u'Save file')
         color = action('&Color', self.chooseColor,
         color = action('&Color', self.chooseColor,
                 'Ctrl+C', 'color', u'Choose line color')
                 'Ctrl+C', 'color', u'Choose line color')
         label = action('&New Item', self.newLabel,
         label = action('&New Item', self.newLabel,
                 'Ctrl+N', 'new', u'Add new label')
                 'Ctrl+N', 'new', u'Add new label')
+        copy = action('&Copy', self.copySelectedShape,
+                'Ctrl+C', 'copy', u'Copy')
         delete = action('&Delete', self.deleteSelectedShape,
         delete = action('&Delete', self.deleteSelectedShape,
                 'Ctrl+D', 'delete', u'Delete')
                 '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 = QWidgetAction(self)
         zoom.setDefaultWidget(self.zoom_widget)
         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,
         fit_window = action('&Fit Window', self.setFitWindow,
                 'Ctrl+F', 'fit',  u'Fit image to window', checkable=True)
                 'Ctrl+F', 'fit',  u'Fit image to window', checkable=True)
 
 
@@ -100,13 +145,13 @@ class MainWindow(QMainWindow, WindowMixin):
                 file=self.menu('&File'),
                 file=self.menu('&File'),
                 edit=self.menu('&Image'),
                 edit=self.menu('&Image'),
                 view=self.menu('&View'))
                 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.edit, (label, color, fit_window))
 
 
-        #addActions(self.menus.view, (labl,))
+        addActions(self.menus.view, (labels,))
 
 
         self.tools = self.toolbar('Tools')
         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))
             zoom, fit_window, None, quit))
 
 
 
 
@@ -151,6 +196,56 @@ class MainWindow(QMainWindow, WindowMixin):
         # Callbacks:
         # Callbacks:
         self.zoom_widget.editingFinished.connect(self.paintCanvas)
         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:
     ## Callback functions:
     def newShape(self, position):
     def newShape(self, position):
@@ -160,11 +255,10 @@ class MainWindow(QMainWindow, WindowMixin):
         """
         """
         action = self.label.popUp(position)
         action = self.label.popUp(position)
         if action == self.label.OK:
         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:
         elif action == self.label.UNDO:
-            print "Undo last line"
             self.canvas.undoLastLine()
             self.canvas.undoLastLine()
         elif action == self.label.DELETE:
         elif action == self.label.DELETE:
             self.canvas.deleteLastShape()
             self.canvas.deleteLastShape()
@@ -191,21 +285,38 @@ class MainWindow(QMainWindow, WindowMixin):
     def queueEvent(self, function):
     def queueEvent(self, function):
         QTimer.singleShot(0, function)
         QTimer.singleShot(0, function)
 
 
+    def hideLabelsToggle(self, value):
+        self.canvas.hideBackroundShapes(value)
+
     def loadFile(self, filename=None):
     def loadFile(self, filename=None):
         """Load the specified file, or the last opened file if None."""
         """Load the specified file, or the last opened file if None."""
         if filename is None:
         if filename is None:
             filename = self.settings['filename']
             filename = self.settings['filename']
-        # FIXME: Load the actual file here.
+        filename = unicode(filename)
         if QFile.exists(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():
             if image.isNull():
                 message = "Failed to read %s" % filename
                 message = "Failed to read %s" % filename
             else:
             else:
                 message = "Loaded %s" % os.path.basename(unicode(filename))
                 message = "Loaded %s" % os.path.basename(unicode(filename))
                 self.image = image
                 self.image = image
                 self.filename = filename
                 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)
             self.statusBar().showMessage(message)
 
 
     def resizeEvent(self, event):
     def resizeEvent(self, event):
@@ -213,10 +324,6 @@ class MainWindow(QMainWindow, WindowMixin):
             self.paintCanvas()
             self.paintCanvas()
         super(MainWindow, self).resizeEvent(event)
         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):
     def paintCanvas(self):
         assert not self.image.isNull(), "cannot paint null image"
         assert not self.image.isNull(), "cannot paint null image"
         self.canvas.scale = self.fitSize() if self.fit_window\
         self.canvas.scale = self.fitSize() if self.fit_window\
@@ -245,8 +352,10 @@ class MainWindow(QMainWindow, WindowMixin):
         s['window/position'] = self.pos()
         s['window/position'] = self.pos()
         s['window/state'] = self.saveState()
         s['window/state'] = self.saveState()
         s['line/color'] = self.color
         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):
     def updateFileMenu(self):
         """Populate menu with recent files."""
         """Populate menu with recent files."""
 
 
@@ -258,11 +367,25 @@ class MainWindow(QMainWindow, WindowMixin):
                 if self.filename else '.'
                 if self.filename else '.'
         formats = ['*.%s' % unicode(fmt).lower()\
         formats = ['*.%s' % unicode(fmt).lower()\
                 for fmt in QImageReader.supportedImageFormats()]
                 for fmt in QImageReader.supportedImageFormats()]
+        filters = 'Image files (%s)\nLabel files (*%s)'\
+                % (' '.join(formats), LabelFile.suffix)
         filename = unicode(QFileDialog.getOpenFileName(self,
         filename = unicode(QFileDialog.getOpenFileName(self,
-            '%s - Choose Image', path, 'Image files (%s)' % ' '.join(formats)))
+            '%s - Choose Image', path, filters))
         if filename:
         if filename:
             self.loadFile(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):
     def check(self):
         # TODO: Prompt user to save labels etc.
         # TODO: Prompt user to save labels etc.
         return True
         return True
@@ -276,12 +399,30 @@ class MainWindow(QMainWindow, WindowMixin):
 
 
     def newLabel(self):
     def newLabel(self):
         self.canvas.deSelectShape()
         self.canvas.deSelectShape()
-        self.canvas.startLabeling=True
+        self.canvas.setEditing()
 
 
     def deleteSelectedShape(self):
     def deleteSelectedShape(self):
         self.canvas.deleteSelected()
         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):
 class Settings(object):
     """Convenience dict-like wrapper around QSettings."""
     """Convenience dict-like wrapper around QSettings."""
     def __init__(self, types=None):
     def __init__(self, types=None):
@@ -308,6 +449,17 @@ class Settings(object):
         return value
         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):
 class struct(object):
     def __init__(self, **kwargs):
     def __init__(self, **kwargs):
         self.__dict__.update(kwargs)
         self.__dict__.update(kwargs)

+ 3 - 0
lib.py

@@ -39,3 +39,6 @@ def addActions(widget, actions):
         else:
         else:
             widget.addAction(action)
             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 *
 from PyQt4.QtCore import *
 
 
 # FIXME:
 # FIXME:
-# - Don't scale vertices.
+# - Add support for highlighting vertices.
 
 
 class Shape(object):
 class Shape(object):
     P_SQUARE, P_ROUND = range(2)
     P_SQUARE, P_ROUND = range(2)
@@ -17,6 +17,7 @@ class Shape(object):
     select_color = QColor(255, 255, 255)
     select_color = QColor(255, 255, 255)
     point_type = P_SQUARE
     point_type = P_SQUARE
     point_size = 8
     point_size = 8
+    scale = 1.0
 
 
     def __init__(self, label=None, line_color=None):
     def __init__(self, label=None, line_color=None):
         self.label = label
         self.label = label
@@ -45,6 +46,8 @@ class Shape(object):
     def paint(self, painter):
     def paint(self, painter):
         if self.points:
         if self.points:
             pen = QPen(self.select_color if self.selected else self.line_color)
             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)
             painter.setPen(pen)
 
 
             line_path = QPainterPath()
             line_path = QPainterPath()
@@ -64,24 +67,35 @@ class Shape(object):
                 painter.fillPath(line_path, self.fill_color)
                 painter.fillPath(line_path, self.fill_color)
 
 
     def drawVertex(self, path, point):
     def drawVertex(self, path, point):
-        d = self.point_size
+        d = self.point_size / self.scale
         if self.point_type == self.P_SQUARE:
         if self.point_type == self.P_SQUARE:
             path.addRect(point.x() - d/2, point.y() - d/2, d, d)
             path.addRect(point.x() - d/2, point.y() - d/2, d, d)
         else:
         else:
             path.addEllipse(point, d/2.0, d/2.0)
             path.addEllipse(point, d/2.0, d/2.0)
 
 
     def containsPoint(self, point):
     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:]:
         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):
     def __len__(self):
         return len(self.points)
         return len(self.points)