diff --git a/sloth/items/inserters.py b/sloth/items/inserters.py index 0f7ce57..5d9a2be 100644 --- a/sloth/items/inserters.py +++ b/sloth/items/inserters.py @@ -109,6 +109,70 @@ class FixedRatioRectItemInserter(RectItemInserter): event.accept() +class NPointFaceInserter(ItemInserter): + landmarks = [ + ("leoc", "left eye outer corner"), + ("leic", "left eye inner corner"), + ("reic", "right eye inner corner"), + ("reoc", "right eye outer corner"), + ("nt", "nose tip"), + ("mlc", "left mouth corner"), + ("mrc", "right mouth corner"), + ] + + def __init__(self, labeltool, scene, default_properties=None): + ItemInserter.__init__(self, labeltool, scene, default_properties) + self._reset() + + def _reset(self): + self._state = 0 + self._items = [] + self._values = {} + self._image_item = None + self._scene.setMessage("Labeling %s" % self.landmarks[self._state][1]) + pass + + def _cleanup(self): + for item in self._items: + self._scene.removeItem(item) + self._scene.clearMessage() + + def _insertItem(self): + item = {} + item.update(self._default_properties) + item.update(self._values) + self._image_item.addAnnotation(item) + + def mousePressEvent(self, event, image_item): + if self._image_item is None: + self._image_item = image_item + else: + assert(self._image_item == image_item) + + pos = event.scenePos() + if event.buttons() & Qt.LeftButton: + self._values[self.landmarks[self._state][0] + "x"] = pos.x() + self._values[self.landmarks[self._state][0] + "y"] = pos.y() + item = QGraphicsEllipseItem(QRectF(pos.x()-2, pos.y()-2, 5, 5)) + item.setPen(Qt.red) + self._scene.addItem(item) + self._items.append(item) + else: + self._values[self.landmarks[self._state][0] + "x"] = -1 + self._values[self.landmarks[self._state][0] + "y"] = -1 + + if self._state == len(self.landmarks)-1: + self._cleanup() + self._insertItem() + self.inserterFinished.emit() + else: + self._state += 1 + self._scene.setMessage("Labeling %s" % self.landmarks[self._state][1]) + + def abort(self): + self._cleanup() + self.inserterFinished.emit() + # TODO class PolygonItemInserter(ItemInserter): def __init__(self, scene, mode=None): diff --git a/sloth/items/items.py b/sloth/items/items.py index a124756..1380f3e 100644 --- a/sloth/items/items.py +++ b/sloth/items/items.py @@ -337,3 +337,101 @@ class ControlItem(QGraphicsItem): color = QColor('black') color.setAlpha(200) painter.fillRect(self.boundingRect(), color) + +class NPointFacePointItem(QGraphicsEllipseItem): + def __init__(self, landmark, *args, **kwargs): + self._landmark = landmark + QGraphicsEllipseItem.__init__(self, *args, **kwargs) + self.setFlags(QGraphicsItem.ItemIsMovable | + QGraphicsItem.ItemSendsGeometryChanges | + QGraphicsItem.ItemSendsScenePositionChanges) + + def landmark(self): + return self._landmark + + def itemChange(self, change, value): + if change == QGraphicsItem.ItemPositionHasChanged: + parent = self.parentItem() + if parent is not None: + parent.landmarkChanged(self, value) + return QAbstractGraphicsShapeItem.itemChange(self, change, value) + + def setColor(self, color): + self.setPen(color) + self.update() + +class NPointFaceItem(BaseItem): + landmarks = [ + ("leoc", "left eye outer corner"), + ("leic", "left eye inner corner"), + ("reic", "right eye inner corner"), + ("reoc", "right eye outer corner"), + ("nt", "nose tip"), + ("mlc", "left mouth corner"), + ("mrc", "right mouth corner"), + ] + + def __init__(self, model_item=None, parent=None): + self._children = {} + BaseItem.__init__(self, model_item, parent) + self.dataChange() + self.changeColor() + + def updateModel(self): + changes = {} + for lm, lmstr in self.landmarks: + lmx, lmy = (lm+"x", lm+"y") + if self._model_item[lmx] != self._children[lm].scenePos().x(): + changes[lmx] = self._children[lm].scenePos().x() + if self._model_item[lmy] != self._children[lm].scenePos().y(): + changes[lmy] = self._children[lm].scenePos().y() + if changes: + self._model_item.update(changes) + + def dataChange(self): + for lm, lmstr in self.landmarks: + lmx, lmy = (lm+"x", lm+"y") + landmark_present = True + if lmx not in self._model_item or lmy not in self._model_item: + landmark_present = False + elif self._model_item[lmx] < 0 or self._model_item[lmy] < 0: + landmark_present = False + + if landmark_present: + px, py = (self._model_item[lmx], self._model_item[lmy]) + if lm in self._children: + # Update item position if it is different + if px != self._children[lm].scenePos().x() or py != self._children[lm].scenePos().y(): + self._children[lm].setPos(px, py) + else: + self._children[lm] = NPointFacePointItem(lm, QRectF(-2, -2, 5, 5), self) + self._children[lm].setPos(px, py) + else: + # Landmark is not present + if lm in self._children: + # Remove landmark from scene + pass + + def setColor(self, *args, **kwargs): + for c in self._children.values(): + c.setColor(*args, **kwargs) + BaseItem.setColor(self, *args, **kwargs) + + def landmarkChanged(self, item, value): + self.prepareGeometryChange() + self.updateModel() + + def boundingRect(self): + br = self.childrenBoundingRect() + offset = 0.2 * br.height() + return br.adjusted(-offset, -offset, +offset, +offset) + + 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) + painter.drawRect(self.boundingRect()) +