diff --git a/sloth/conf/default_config.py b/sloth/conf/default_config.py index 52d939f..45b999c 100644 --- a/sloth/conf/default_config.py +++ b/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 diff --git a/sloth/gui/annotationscene.py b/sloth/gui/annotationscene.py index f7ab14d..69fc47d 100644 --- a/sloth/gui/annotationscene.py +++ b/sloth/gui/annotationscene.py @@ -185,6 +185,19 @@ class AnnotationScene(QGraphicsScene): # selection mode QGraphicsScene.mousePressEvent(self, event) + def mouseDoubleClickEvent(self, event): + LOG.debug("mouseDoubleClickEvent %s %s" % (self.sceneRect().contains(event.scenePos()), event.scenePos())) + if self._inserter is not None: + if not self.sceneRect().contains(event.scenePos()) and \ + not self._inserter.allowOutOfSceneEvents(): + # ignore events outside the scene rect + return + # insert mode + self._inserter.mouseDoubleClickEvent(event, self._image_item) + else: + # selection mode + QGraphicsScene.mouseDoubleClickEvent(self, event) + def mouseReleaseEvent(self, event): LOG.debug("mouseReleaseEvent %s %s" % (self.sceneRect().contains(event.scenePos()), event.scenePos())) if self._inserter is not None: diff --git a/sloth/items/inserters.py b/sloth/items/inserters.py index 93fbd9a..3b573b3 100644 --- a/sloth/items/inserters.py +++ b/sloth/items/inserters.py @@ -38,6 +38,9 @@ class ItemInserter(QObject): def mousePressEvent(self, event, image_item): event.accept() + def mouseDoubleClickEvent(self, event, image_item): + event.accept() + def mouseReleaseEvent(self, event, image_item): event.accept() @@ -321,31 +324,102 @@ 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) - else: - polygon = self._current_item.polygon() - polygon.append(pos) - self._current_item.setPolygon(polygon) + + self._scene.setMessage("Press Enter to finish the polygon.") + + polygon = self._item.polygon() + polygon.append(pos) + self._item.setPolygon(polygon) event.accept() + def mouseDoubleClickEvent(self, event, image_item): + """Finish the polygon when the user double clicks.""" + + # No need to add the position of the click, as a single mouse + # press event added the point already. + # Even then, the last point of the polygon is duplicate as it would be + # shortly after a single mouse press. At this point, we want to throw it + # away. + polygon = self._item.polygon() + polygon.remove(polygon.size()-1) + assert polygon.size() > 0 + self._item.setPolygon(polygon) + + self._updateAnnotation() + if self._commit: + image_item.addAnnotation(self._ann) + self._scene.removeItem(self._item) + self.annotationFinished.emit() + self._item = None + self._scene.clearMessage() + + self.inserterFinished.emit() + + 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 keyPressEvent(self, event, image_item): + """ + When the user presses Enter, the polygon is finished. + """ + if event.key() == Qt.Key_Return and 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._item.setPolygon(polygon) + + self._updateAnnotation() + if self._commit: + image_item.addAnnotation(self._ann) + self._scene.removeItem(self._item) + self.annotationFinished.emit() + self._item = None + self._scene.clearMessage() + + self.inserterFinished.emit() + + def abort(self): + if self._item is not None: + self._scene.removeItem(self._item) + self._item = None + self._scene.clearMessage() + ItemInserter.abort(self) + + def _updateAnnotation(self): + polygon = self._item.polygon() + 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) diff --git a/sloth/items/items.py b/sloth/items/items.py index 64fcbf0..2b41bc3 100644 --- a/sloth/items/items.py +++ b/sloth/items/items.py @@ -757,3 +757,76 @@ class NPointFaceItem(GroupItem): pen.setStyle(Qt.DashLine) painter.setPen(pen) painter.drawRect(self.boundingRect()) + +class PolygonItem(BaseItem): + def __init__(self, model_item=None, prefix="", 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)