Add the missing PolygonItem

Add the functionality to add labels of polygon shape. That means:

- Fix the non-working PolygonItemInserter class
- Add the missing PolygonItem class
- Add the "polygon" class to the default configuration

sloth/conf/default_config.py

@@ -61,6 +61,15 @@ LABELS = (
         'hotkey':   'p',
         'text':     'Point',
     },
+    {
+        'attributes': {
+            'class':    'polygon',
+        },
+        'inserter': 'sloth.items.PolygonItemInserter',
+        'item':     'sloth.items.PolygonItem',
+        'hotkey':   'o',
+        'text':     'Polygon',
+    },
 )
 
 # HOTKEYS

sloth/items/inserters.py

@@ -321,31 +321,61 @@ class NPointFaceInserter(SequenceItemInserter):
                                    "Now at: " + self.inserters[self._state][2])
 
 
-# TODO
 class PolygonItemInserter(ItemInserter):
-    def __init__(self, scene, mode=None):
-        ItemInserter.__init__(self, scene, mode)
-        self._current_item = None
+    def __init__(self, labeltool, scene, default_properties=None,
+                 prefix="", commit=True):
+        ItemInserter.__init__(self, labeltool, scene, default_properties,
+                              prefix, commit)
+        self._item = None
 
     def mousePressEvent(self, event, image_item):
         pos = event.scenePos()
-        if self._current_item is None:
+        if self._item is None:
             item = QGraphicsPolygonItem(QPolygonF([pos]))
-            self._current_item = item
+            self._item = item
+            self._item.setPen(self.pen())
             self._scene.addItem(item)
+            self._current_image_item = image_item
         else:
-            polygon = self._current_item.polygon()
+            polygon = self._item.polygon()
             polygon.append(pos)
-            self._current_item.setPolygon(polygon)
+            self._item.setPolygon(polygon)
 
         event.accept()
 
     def mouseMoveEvent(self, event, image_item):
-        if self._current_item is not None:
+        if self._item is not None:
             pos = event.scenePos()
-            polygon = self._current_item.polygon()
+            polygon = self._item.polygon()
             assert polygon.size() > 0
             polygon[-1] = pos
-            self._current_item.setPolygon(polygon)
+            self._item.setPolygon(polygon)
 
         event.accept()
+
+    def abort(self):
+        # XXX Is it abuse to handle the end of inserting here in abort()?
+        if self._item is not None:
+            polygon = self._item.polygon()
+            assert polygon.size() > 0
+
+            # The last point of the polygon is the point the user would add
+            # to the polygon when pressing the mouse button. At this point,
+            # we want to throw it away.
+            polygon.remove(polygon.size()-1)
+            assert polygon.size() > 0
+
+            self._ann.update({self._prefix + 'xn':
+                                  ";".join([str(p.x()) for p in polygon]),
+                              self._prefix + 'yn':
+                                  ";".join([str(p.y()) for p in polygon])})
+            self._ann.update(self._default_properties)
+            if self._commit:
+                self._current_image_item.addAnnotation(self._ann)
+            self._scene.removeItem(self._item)
+            self.annotationFinished.emit()
+            self._init_pos = None
+            self._item = None
+            self._current_image_item = None
+
+        self.inserterFinished.emit()

sloth/items/items.py

@@ -757,3 +757,77 @@ class NPointFaceItem(GroupItem):
             pen.setStyle(Qt.DashLine)
         painter.setPen(pen)
         painter.drawRect(self.boundingRect())
+
+class PolygonItem(BaseItem):
+    # XXX prefix="pointlist"
+    def __init__(self, model_item=None, prefix="pointlist", parent=None):
+        BaseItem.__init__(self, model_item, prefix, parent)
+
+        # Make it non-movable for now
+        self.setFlags(QGraphicsItem.ItemIsSelectable |
+                      QGraphicsItem.ItemSendsGeometryChanges |
+                      QGraphicsItem.ItemSendsScenePositionChanges)
+        self._polygon = None
+
+        self._updatePolygon(self._dataToPolygon(self._model_item))
+        LOG.debug("Constructed polygon %s for model item %s" %
+                  (self._polygon, model_item))
+
+    def __call__(self, model_item=None, parent=None):
+        item = PolygonItem(model_item, parent)
+        item.setPen(self.pen())
+        item.setBrush(self.brush())
+        return item
+
+    def _dataToPolygon(self, model_item):
+        if model_item is None:
+            return QPolygonF()
+
+        try:
+            polygon = QPolygonF()
+            xn = [float(x) for x in model_item["xn"].split(";")]
+            yn = [float(y) for y in model_item["yn"].split(";")]
+            for x, y in zip(xn, yn):
+              polygon.append(QPointF(x, y))
+
+            return polygon
+
+        except KeyError as e:
+            LOG.debug("PolygonItem: Could not find expected key in item: "
+                      + str(e) + ". Check your config!")
+            self.setValid(False)
+            return QPolygonF()
+
+    def _updatePolygon(self, polygon):
+        if polygon == self._polygon:
+            return
+
+        self.prepareGeometryChange()
+        self._polygon = polygon
+        self.setPos(QPointF(0, 0))
+
+    def boundingRect(self):
+        xn = [p.x() for p in self._polygon]
+        yn = [p.y() for p in self._polygon]
+        xmin = min(xn)
+        xmax = max(xn)
+        ymin = min(yn)
+        ymax = max(yn)
+        return QRectF(xmin, ymin, xmax - xmin, ymax - ymin)
+
+    def paint(self, painter, option, widget=None):
+        BaseItem.paint(self, painter, option, widget)
+
+        pen = self.pen()
+        if self.isSelected():
+            pen.setStyle(Qt.DashLine)
+        painter.setPen(pen)
+
+        for k in range(-1, len(self._polygon)-1):
+            p1 = self._polygon[k]
+            p2 = self._polygon[k+1]
+            painter.drawLine(p1, p2)
+
+    def dataChange(self):
+        polygon = self._dataToPolygon(self._model_item)
+        self._updatePolygon(polygon)