Draw polygonal shapes freely over the picture

Refactor shape code and fix errors. The shapes are not scaled along with
the image currently.

canvas.py

@@ -1,4 +1,6 @@
 
+from math import sqrt
+
 from PyQt4.QtGui import *
 from PyQt4.QtCore import *
 
@@ -6,54 +8,71 @@ from shape import Shape
 
 class Canvas(QLabel):
     done = pyqtSignal()
-    epsilon = 7**2 # TODO: Tune value
+    epsilon = 7.0 # TODO: Tune value
 
     def __init__(self, *args, **kwargs):
         super(Canvas, self).__init__(*args, **kwargs)
-        self.points = []
-        self.shapes = [Shape('one', QColor(0, 255, 0))]
-        self.current_line = Shape('line', QColor(255, 0, 0))
+        self.shapes = []
+        self.current = None
+        self.line = Shape(line_color=QColor(0, 0, 255))
 
     def mousePressEvent(self, ev):
-        if ev.button() != 1:
-            return
-
-        if not self.points:
-            self.setMouseTracking(True)
-
-        self.points.append(ev.pos())
-        self.shapes[0].addPoint(ev.pos())
-        if self.isClosed():
-            self.done.emit()
-            print "Points:", self.points
-            self.points = []
-            self.shapes[0].setFill(True)
-            self.setMouseTracking(False)
+        if ev.button() == 1:
+            if self.current:
+                self.current.vertices.append(ev.pos())
+                if self.isClosed():
+                    self.finalise()
+                self.repaint()
+            else:
+                self.current = Shape()
+                self.current.vertices.append(ev.pos())
+                self.setMouseTracking(True)
+
+    def mouseDoubleClickEvent(self, ev):
+        if self.current:
+            self.current.vertices.append(self.current[0])
+            self.finalise()
+
+    def finalise(self):
+        assert self.current
+        self.current.fill = True
+        self.shapes.append(self.current)
+        self.current = None
+        # TODO: Mouse tracking is still useful for selecting shapes!
+        self.setMouseTracking(False)
         self.repaint()
+        self.done.emit()
 
     def isClosed(self):
+        assert self.current
+        return len(self.current) > 1\
+           and self.closeEnough()
         return len(self.points) > 1 and self.closeEnough(self.points[0], self.points[-1])
 
     def mouseMoveEvent(self, ev):
-        self.current_line.points = [self.points[-1], self.pos()]
-        self.repaint()
-       #print "moving", ev.pos()
-
-
+        """Update line with last point and current coordinates."""
+        if self.current:
+            self.line.vertices = (self.current[-1], ev.pos())
+            self.repaint()
 
-    def closeEnough(self, p1, p2):
-        def dist(p):
-            return p.x() * p.x() + p.y() * p.y()
-        print p1, p2
-        print abs(dist(p1) - dist(p2)), self.epsilon
-        return abs(dist(p1) - dist(p2)) < self.epsilon
+    def closeEnough(self):
+        assert self.current
+        def distance(p):
+            return sqrt(p.x() * p.x() + p.y() * p.y())
+        p1, p2 = self.current.vertices[0], self.current.vertices[-1]
+        d = distance(p1 - p2)
+        m = (p1-p2).manhattanLength()
+        print "d %.2f, m %d, %.2f" % (d, m, d - m)
+        return distance(p1 - p2) < self.epsilon
 
     def paintEvent(self, event):
         super(Canvas, self).paintEvent(event)
+        qp = QPainter()
+        qp.begin(self)
         for shape in self.shapes:
-            qp = QPainter()
-            qp.begin(self)
-            shape.drawShape(qp)
-            self.current_line.drawShape(qp)
-            qp.end()
+            shape.paint(qp)
+        if self.current:
+            self.current.paint(qp)
+            self.line.paint(qp)
+        qp.end()
 

shape.py

@@ -6,34 +6,42 @@ from PyQt4.QtCore import *
 
 class Shape(object):
 
-    def __init__(self ,label,color):
-        self.label=label
-        self.points=[]
-        self.color=color
-        self.fill=False
+    def __init__(self, label=None,
+            line_color=QColor(0, 255, 0, 128),
+            fill_color=QColor(255, 0, 0, 128)):
 
-    def getLabel(self):
-        return label
+        self.label = label
+        self.line_color = line_color
+        self.fill_color = fill_color
 
-    def setLabel(self,label):
-        self.label=label
+        self.vertices = []
+        self.fill = False
 
-    def setFill(self,fillB):
-        self.fill=fillB
+    def addVertex(self, vertex):
+        self.vertices.append(vertex)
 
-    def addPoint(self,point):
-        self.points.append(point)
+    def popVertex(self):
+        if self.vertices:
+            return self.vertices.pop()
+        return None
 
-    def drawShape(self,painter):
-
-        if len(self.points) >0 :
-            pen=QPen(self.color)
+    def paint(self, painter):
+        if self.vertices:
+            pen = QPen(self.line_color)
             painter.setPen(pen)
-            prePoint=self.points[0]
-            path=QPainterPath()
-            path.moveTo(prePoint.x(),prePoint.y())
-            for point in self.points:
-                path.lineTo(point.x(),point.y())
+            path = QPainterPath()
+            p0 = self.vertices[0]
+            path.moveTo(p0.x(), p0.y())
+            for v in self.vertices[1:]:
+                path.lineTo(v.x(), v.y())
             painter.drawPath(path)
-        if self.fill:
-            painter.fillPath(path,Qt.red)
+
+            if self.fill:
+                painter.fillPath(path, self.fill_color)
+
+    def __len__(self):
+        return len(self.vertices)
+
+    def __getitem__(self, key):
+        return self.vertices[key]
+