sloth/annotationscene.py

from PyQt4.QtGui import *
from PyQt4.QtCore import *
from sceneitems import *
from annotationmodel import ImageRole
import math
import okapy
from okapy.guiqt.utilities import toQImage

class ItemInserter:
    def __init__(self, scene, mode=None):
        self.scene_ = scene
        self.mode_  = mode

    def setScene(self, scene):
        self.scene_ = scene
    def scene(self):
        return self.scene_

    def setMode(self, mode):
        self.mode_  = mode
    def mode(self):
        return self.mode

    def mousePressEvent(self, event, index):
        event.accept()

    def mouseReleaseEvent(self, event, index):
        event.accept()

    def mouseMoveEvent(self, event, index):
        event.accept()

    def keyPressEvent(self, event, index):
        event.ignore()

class PointItemInserter(ItemInserter):
    def mousePressEvent(self, event, index):
        pos = event.scenePos()
        ann = {'type': 'point',
               'x': event.pos().x(), 'y': event.pos().y()}
        index.model().addAnnotation(index, ann)
        event.accept()

class RectItemInserter(ItemInserter):
    def __init__(self, scene, mode=None):
        ItemInserter.__init__(self, scene, mode)
        self.current_item_ = None
        self.init_pos_ = None

    def mousePressEvent(self, event, index):
        pos = event.scenePos()
        item = QGraphicsRectItem(QRectF(pos.x(), pos.y(), 0, 0))
        item.setPen(Qt.red)
        self.current_item_ = item
        self.init_pos_     = pos
        self.scene().addItem(item)
        event.accept()

    def mouseMoveEvent(self, event, index):
        if self.current_item_ is not None:
            assert self.init_pos_ is not None
            rect = QRectF(self.init_pos_, event.scenePos()).normalized()
            self.current_item_.setRect(rect)

        event.accept()

    def mouseReleaseEvent(self, event, index):
        if self.current_item_ is not None:
            if self.current_item_.rect().width() > 1 and \
               self.current_item_.rect().height() > 1:
                rect = self.current_item_.rect()
                ann = {'type': 'rect',
                       'x': rect.x(), 'y': rect.y(),
                       'width': rect.width(), 'height': rect.height()}
                index.model().addAnnotation(index, ann)
            self.scene().removeItem(self.current_item_)
            self.current_item_ = None
            self.init_pos_ = None

        event.accept()

class FixedRatioRectItemInserter(RectItemInserter):
    def __init__(self, scene, mode=None):
        RectItemInserter.__init__(self, scene, mode)
        self.ratio_ = 1
        if mode is not None:
            self.ratio_ = float(mode.get('_ratio', 1))

    def setMode(self, mode):
        if mode is not None:
            self.ratio_ = float(mode.get('_ratio', 1))
        RectItemInserter.setMode(self, mode)

    def mouseMoveEvent(self, event, index):
        if self.current_item_ is not None:
            new_geometry = QRectF(self.current_item_.rect().topLeft(), event.scenePos())
            dx = new_geometry.width()
            dy = new_geometry.height()
            d = math.sqrt(dx*dx + dy*dy)
            r = self.ratio_
            k = math.sqrt(r*r+1)
            h = d/k
            w = d*r/k
            new_geometry.setWidth(w)
            new_geometry.setHeight(h)
            self.current_item_.setRect(new_geometry.normalized())

        event.accept()

class PolygonItemInserter(ItemInserter):
    def __init__(self, scene, mode=None):
        ItemInserter.__init__(self, scene, mode)
        self.current_item_ = None

    def mousePressEvent(self, event, index):
        pos = event.scenePos()
        if self.current_item_ is None:
            item = QGraphicsPolygonItem(QPolygonF([pos]))
            self.current_item_ = item
            self.scene().addItem(item)
        else:
            polygon = self.current_item_.polygon()
            polygon.append(pos)
            self.current_item_.setPolygon(polygon)

        event.accept()

    def mouseMoveEvent(self, event, index):
        if self.current_item_ is not None:
            pos = event.scenePos()
            polygon = self.current_item_.polygon()
            assert polygon.size() > 0
            polygon[-1] = pos
            self.current_item_.setPolygon(polygon)

        event.accept()

class AnnotationScene(QGraphicsScene):

    # TODO signal itemadded

    def __init__(self, parent=None):
        super(AnnotationScene, self).__init__(parent)

        self.model_     = None
        self.mode_      = None
        self.inserters_ = {}
        self.inserter_  = None
        self.debug_     = True
        self.message_   = ""
        self.last_key_  = None

        self.setBackgroundBrush(Qt.darkGray)

        self.addItemInserter('point',      PointItemInserter(self))
        self.addItemInserter('rect',       RectItemInserter(self))
        self.addItemInserter('ratiorect',  FixedRatioRectItemInserter(self))
        self.addItemInserter('poly',       PolygonItemInserter(self))
        self.addItemInserter('polygon',    PolygonItemInserter(self))

        self.setMode(None)
        self.reset()

    #
    # getters/setters
    #______________________________________________________________________________________________________
    def model(self):
        return self.model_

    def setModel(self, model):
        if model == self.model_:
            # same model as the current one
            # reset caches anyway, invalidate root
            self.reset()
            return

        # disconnect old signals
        if self.model_ is not None:
            self.disconnect(self.model_, SIGNAL('dataChanged(QModelIndex,QModelIndex)'), self.dataChanged)
            self.disconnect(self.model_, SIGNAL('rowsInserted(QModelIndex,int,int)'), self.rowsInserted)
            self.disconnect(self.model_, SIGNAL('rowsAboutToBeRemoved(QModelIndex,int,int)'), self.rowsAboutToBeRemoved)
            self.disconnect(self.model_, SIGNAL('rowsRemoved(QModelIndex,int,int)'), self.rowsRemoved)
            self.disconnect(self.model_, SIGNAL('modelReset()'), self.reset)
        self.model_ = model

        # connect new signals
        if self.model_ is not None:
            self.connect(self.model_, SIGNAL('dataChanged(QModelIndex,QModelIndex)'), self.dataChanged)
            self.connect(self.model_, SIGNAL('rowsInserted(QModelIndex,int,int)'), self.rowsInserted)
            self.connect(self.model_, SIGNAL('rowsAboutToBeRemoved(QModelIndex,int,int)'), self.rowsAboutToBeRemoved)
            self.connect(self.model_, SIGNAL('rowsRemoved(QModelIndex,int,int)'), self.rowsRemoved)
            self.connect(self.model_, SIGNAL('modelReset()'), self.reset)

        # reset caches, invalidate root
        self.reset()

    def root(self):
        return self.root_

    def setRoot(self, root):
        """
        Set the index of the model which denotes the current image to be
        displayed by the scene.  This can be either the index to a frame in a
        video, or to an image.
        """
        self.root_ = root
        self.clear()
        if not root.isValid():
            return

        assert self.root_.model() == self.model_
        self.image_ = self.root_.data(ImageRole).toPyObject()
        self.pixmap_ = QPixmap(okapy.guiqt.toQImage(self.image_))
        item = QGraphicsPixmapItem(self.pixmap_)
        item.setZValue(-1)
        self.setSceneRect(0, 0, self.pixmap_.width(), self.pixmap_.height())
        self.addItem(item)

        num_items = self.model_.rowCount(self.root_)
        self.insertItems(0, num_items)
        self.update()

    def insertItems(self, first, last):
        assert self.model_ is not None
        assert self.root_.isValid()
        print "insertItems"

        for row in range(first, last+1):
            child = self.root_.child(row, 0)
            item = AnnotationGraphicsItem.createItem(child)
            if item is not None:
                #checked = (child.data(Qt.CheckStateRole).toInt()[0] == Qt.Checked)
                #item.setVisible(checked)
                self.addItem(item)

    def mode(self):
        return self.mode_

    def setMode(self, mode):
        print "setMode :", mode
        self.mode_ = mode

        if self.mode_ == None:
            self.inserter_ = None
            return

        if not self.mode_['type'] in self.inserters_:
            raise InvalidArgumentException("Invalid mode")

        self.inserter_ = self.inserters_[self.mode_['type']]
        self.inserter_.setMode(self.mode_)

    #
    # common methods
    #______________________________________________________________________________________________________
    def reset(self):
        self.clear()
        self.setRoot(QModelIndex())
        self.clearMessage()

    def addItem(self, item):
        QGraphicsScene.addItem(self, item)
        # TODO emit signal itemAdded

    def addItemInserter(self, type, inserter, replace=False):
        type = type.lower()
        if type in self.inserters_ and not replace:
            raise Exception("Type %s already has an inserter" % type)

        self.inserters_[type] = inserter

    def removeItemInserter(self, type):
        type = type.lower()
        if type in self.inserters_:
            del self.inserters_[type]

    #
    # mouse event handlers
    #______________________________________________________________________________________________________
    def mousePressEvent(self, event):
        if self.debug_:
            print "mousePressEvent", self.sceneRect().contains(event.scenePos()), event.scenePos()
        if not self.sceneRect().contains(event.scenePos()):
            # ignore events outside the scene rect
            return
        elif self.inserter_ is not None:
            # insert mode
            self.inserter_.mousePressEvent(event, self.root_)
        else:
            # selection mode
            QGraphicsScene.mousePressEvent(self, event)

    def mouseReleaseEvent(self, event):
        if self.debug_:
            print "mouseReleaseEvent", self.sceneRect().contains(event.scenePos()), event.scenePos()
        if self.inserter_ is not None:
            # insert mode
            self.inserter_.mouseReleaseEvent(event, self.root_)
        else:
            # selection mode
            QGraphicsScene.mouseReleaseEvent(self, event)

    def mouseMoveEvent(self, event):
        #if self.debug_:
        #   print "mouseMoveEvent", self.sceneRect().contains(event.scenePos()), event.scenePos()
        if self.inserter_ is not None:
            # insert mode
            self.inserter_.mouseMoveEvent(event, self.root_)
        else:
            # selection mode
            QGraphicsScene.mouseMoveEvent(self, event)

    #
    # key event handlers
    #______________________________________________________________________________________________________
    def keyPressEvent(self, event):
        if self.debug_:
           print "keyPressEvent", event

        if self.model_ is None or not self.root_.isValid():
            event.ignore()
            return

        if self.inserter_ is not None:
            # insert mode
            self.inserter_.keyPressEvent(event, self.root_)
        else:
            # selection mode
            if event.key() == Qt.Key_Delete:
                for item in self.selectedItems():
                    index = item.index()
                    index.model().removeAnnotation(index)
                event.accept()
                return

            if ord('0') <= event.key() <= ord('9'):
                if self.last_key_ is None:
                    self.last_key_ = event.key()
                else:
                    id = chr(self.last_key_) + chr(event.key())
                    print "id=", id
                    for item in self.selectedItems():
                        index = item.index()
                        data = dict(index.data(DataRole).toPyObject().iteritems())
                        if data['type'] == 'rect':
                            data['id'] = id
                        index.model().setData(index, QVariant(data), DataRole)
                    self.last_key_ = None

        event.ignore()

    #
    # slots for signals from the model
    # this is the implemenation of the scene as a view of the model
    #______________________________________________________________________________________________________
    def dataChanged(self, indexFrom, indexTo):
        if self.root_ != indexFrom.parent() or self.root_ != indexTo.parent():
            return

        for row in range(indexFrom.row(), indexTo.row()+1):
            item = self.itemFromIndex(indexFrom.sibling(row, 0))
            if item is not None:
                item.dataChanged()

    def rowsInserted(self, index, first, last):
        if self.root_ != index:
            return

        self.insertItems(first, last)


    def rowsAboutToBeRemoved(self, index, first, last):
        if self.root_ != index:
            return

        for row in range(first, last+1):
            item = self.itemFromIndex(index.child(row, 0))
            if item is not None:
                self.removeItem(item)

    def rowsRemoved(self, index, first, last):
        pass

    def itemFromIndex(self, index):
        index = index.model().mapToSource(index)  # TODO: solve this somehow else
        for item in self.items():
            # some graphics items will not have an index method,
            # we just skip these
            if hasattr(item, 'index') and item.index() == index:
                return item
        return None

    #
    # message handling and displaying
    #______________________________________________________________________________________________________
    def setMessage(self, message):
        if self.message_ is not None:
            self.clearMessage()

        if message is None or message == "":
            return

        # TODO don't use text item at all, just draw the text in drawForeground
        self.message_ = message
        self.message_text_item_ = QGraphicsSimpleTextItem(message)
        self.message_text_item_.setPos(20, 20)
        self.invalidate(QRectF(), QGraphicsScene.ForegroundLayer)

    def clearMessage(self):
        if self.message_ is not None:
            self.message_text_item_ = None
            self.message_ = None
            self.invalidate(QRectF(), QGraphicsScene.ForegroundLayer)

    def drawForeground(self, painter, rect):
        QGraphicsScene.drawForeground(self, painter, rect)

        if self.message_ is not None:
            assert self.message_text_item_ is not None

            painter.setTransform(QTransform())
            painter.setBrush(QColor('lightGray'))
            painter.setPen(QPen(QBrush(QColor('black')), 2))

            br = self.message_text_item_.boundingRect()

            painter.drawRoundedRect(QRectF(10, 10, br.width()+20, br.height()+20), 10.0, 10.0)
            painter.setTransform(QTransform.fromTranslate(20, 20))
            painter.setPen(QPen(QColor('black'), 1))

            self.message_text_item_.paint(painter, QStyleOptionGraphicsItem(), None)