sloth/annotations/model.py

"""
The annotationmodel module contains the classes for the AnnotationModel.
"""
from PyQt4.QtGui import *
from PyQt4.QtCore import *
from functools import partial
import os.path
import okapy

TypeRole, DataRole, ImageRole = [Qt.UserRole + i + 1 for i in range(3)]

class ModelItem:
    def __init__(self, parent=None):
        self.parent_   = parent
        self.children_ = []

    def children(self):
        return self.children_

    def parent(self):
        return self.parent_

    def rowOfChild(self, item):
        try:
            return self.children_.index(item)
        except:
            return -1

    def data(self, index, role):
        return QVariant()

class RootModelItem(ModelItem):
    def __init__(self, files):
        ModelItem.__init__(self, None)
        self.files_ = files

        for file in files:
            fmi = FileModelItem.create(file, self)
            self.children_.append(fmi)

class FileModelItem(ModelItem):
    def __init__(self, file, parent):
        ModelItem.__init__(self, parent)
        self.file_ = file

    def filename(self):
        return self.file_['filename']

    def fullpath(self, index):
        return os.path.join(index.model().basedir(), self.filename())

    def data(self, index, role):
        if role == Qt.DisplayRole and index.column() == 0:
            return os.path.basename(self.filename())
        return QVariant()

    @staticmethod
    def create(file, parent):
        if file['type'] == 'image':
            return ImageFileModelItem(file, parent)
        elif file['type'] == 'video':
            return VideoFileModelItem(file, parent)

class ImageFileModelItem(FileModelItem):
    def __init__(self, file, parent):
        FileModelItem.__init__(self, file, parent)

        for ann in file['annotations']:
            ami = AnnotationModelItem(ann, self)
            self.children_.append(ami)

    def addAnnotation(self, ann):
        self.file_['annotations'].append(ann)
        ami = AnnotationModelItem(ann, self)
        self.children_.append(ami)

    def removeAnnotation(self, pos):
        del self.file_['annotations'][pos]
        del self.children_[pos]

    def data(self, index, role):
        if role == ImageRole:
            return okapy.loadImage(self.fullpath(index))
        return FileModelItem.data(self, index, role)

class VideoFileModelItem(FileModelItem):
    _cached_vs_filename = None
    _cached_vs          = None

    def __init__(self, file, parent):
        FileModelItem.__init__(self, file, parent)

        for frame in file['frames']:
            fmi = FrameModelItem(frame, self)
            self.children_.append(fmi)

    def updateCachedVideoSource(self):
        # have only one cached video source at a time for now
        # TODO: for labeling multiple synchronized videos this should
        # be modified, otherwise it might be awfully slow
        VideoFileModelItem._cached_vs = okv.ImageSequence()
        VideoFileModelItem._cached_vs.open(self.fullpath())
        VideoFileModelItem._cached_vs_filename = self.fullpath()

    def getFrame(self, frame):
        if VideoFileModelItem._cached_vs_filename != self.fullpath():
            self.updateCachedVideoSource()

        VideoFileModelItem._cached_vs.getFrame(frame)
        return VideoFileModelItem._cached_vs.getImage()

class FrameModelItem(ModelItem):
    def __init__(self, frame, parent):
        ModelItem.__init__(self, parent)
        self.frame_ = frame

        for ann in frame['annotations']:
            ami = AnnotationModelItem(ann, self)
            self.children_.append(ami)

    def framenum(self):
        return int(self.frame_.get('num', -1))

    def timestamp(self):
        return float(self.frame_.get('timestamp', -1))

    def addAnnotation(self, ann):
        self.frame_['annotations'].append(ann)
        ami = AnnotationModelItem(ann, self)
        self.children_.append(ami)

    def removeAnnotation(self, pos):
        del self.frame_['annotations'][pos]
        del self.children_[pos]

    def data(self, index, role):
        if role == Qt.DisplayRole and index.column() == 0:
            return "%d / %.3f" % (self.framenum(), self.timestamp())
        elif role == ImageRole:
            return self.parent().getFrame(self.frame_['num'])
        return QVariant()

class AnnotationModelItem(ModelItem):
    def __init__(self, annotation, parent):
        ModelItem.__init__(self, parent)
        self.annotation_ = annotation
        # dummy key/value so that pyqt does not convert the dict
        # into a QVariantMap while communicating with the Views
        self.annotation_[None] = None

        for key, value in annotation.iteritems():
            if key == None:
                continue
            self.children_.append(KeyValueModelItem(key, self))

    def type(self):
        return self.annotation_['type']

    def setData(self, index, data, role):
        if role == DataRole:
            print self.annotation_
            data = data.toPyObject()
            print data, type(data)
            print self.annotation_
            for key, value in data.iteritems():
                print key, value
                if not key in self.annotation_:
                    print "not in annoation: ", key
                    next = len(self.children_)
                    index.model().beginInsertRows(index, next, next)
                    self.children_.append(KeyValueModelItem(key, self))
                    index.model().endInsertRows()
                    index.model().emit(SIGNAL("dataChanged(QModelIndex,QModelIndex)"), index, index)
                    self.annotation_[key] = data[key]

            for key in self.annotation_.keys():
                if not key in data:
                    #TODO beginRemoveRows, delete child, etc.
                    del self.annotation[key]
                else:
                    self.annotation_[key] = data[key]
            print "new annotation:", self.annotation_
            index.model().dataChanged.emit(index, index.sibling(index.row(), 0))
            return True
        return False

    def data(self, index, role):
        if role == Qt.DisplayRole and index.column() == 0:
            return self.type()
        elif role == TypeRole:
            return self.type()
        elif role == DataRole:
            #print "data():", self.annotation_
            return self.annotation_

        return QVariant()

    def value(self, key):
        return self.annotation_[key]

class KeyValueModelItem(ModelItem):
    def __init__(self, key, parent):
        ModelItem.__init__(self, parent)
        self.key_  = key

    def data(self, index, role):
        if role == Qt.DisplayRole:
            if index.column() == 0:
                return self.key_
            elif index.column() == 1:
                return self.parent().value(self.key_)
            else:
                return QVariant()

class AnnotationModel(QAbstractItemModel):
    # signals
    dirtyChanged = pyqtSignal(bool, name='dirtyChanged')

    def __init__(self, annotations, parent=None):
        QAbstractItemModel.__init__(self, parent)
        self.annotations_ = annotations
        self.root_        = RootModelItem(self.annotations_)
        self.dirty_       = False
        self.basedir_     = ""

    def dirty(self):
        return self.dirty_

    def setDirty(self, dirty=True):
        previous = self.dirty_
        self.dirty_ = dirty
        if previous != dirty:
            self.dirtyChanged.emit(dirty)

    def basedir(self):
        return self.basedir_

    def setBasedir(self, dir):
        print "setBasedir: \"" + dir + "\"" 
        self.basedir_ = dir

    def itemFromIndex(self, index):
        index = QModelIndex(index)  # explicitly convert from QPersistentModelIndex
        if index.isValid():
            return index.internalPointer()
        return self.root_

    def index(self, row, column, parent_idx):
        parent_item = self.itemFromIndex(parent_idx)
        if row >= len(parent_item.children()):
            return QModelIndex()
        child_item = parent_item.children()[row]
        return self.createIndex(row, column, child_item)

    def imageIndex(self, index):
        """return index that points to the (maybe parental) image/frame object"""
        if not index.isValid():
            return QModelIndex()

        index = QModelIndex(index)  # explicitly convert from QPersistentModelIndex
        item = self.itemFromIndex(index)
        if isinstance(item, ImageFileModelItem) or \
           isinstance(item, FrameModelItem):
            return index

        # try with next hierarchy up
        return self.imageIndex(index.parent())

    def data(self, index, role):
        if not index.isValid():
            return QVariant()
        index = QModelIndex(index)  # explicitly convert from QPersistentModelIndex

        #if role == Qt.CheckStateRole:
            #item = self.itemFromIndex(index)
            #if item.isCheckable(index.column()):
                #return QVariant(Qt.Checked if item.visible() else Qt.Unchecked)
            #return QVariant()

        #if role != Qt.DisplayRole and role != GraphicsItemRole and role != DataRole:
            #return QVariant()

        ## non decorational behaviour

        item = self.itemFromIndex(index)
        return item.data(index, role)

    def columnCount(self, index):
        return 2

    def rowCount(self, index):
        item = self.itemFromIndex(index)
        return len(item.children())

    def parent(self, index):
        item = self.itemFromIndex(index)
        parent = item.parent()
        if parent is None:
            return QModelIndex()
        grandparent = parent.parent()
        if grandparent is None:
            return QModelIndex()
        row = grandparent.rowOfChild(parent)
        assert row != -1
        return self.createIndex(row, 0, parent)

    def mapToSource(self, index):
        return index

    def flags(self, index):
        return Qt.ItemIsEnabled
        if not index.isValid():
            return Qt.ItemIsEnabled
        index = QModelIndex(index)  # explicitly convert from QPersistentModelIndex
        item = self.itemFromIndex(index)
        return item.flags(index)

    def setData(self, index, value, role):
        if not index.isValid():
            return False
        index = QModelIndex(index)  # explicitly convert from QPersistentModelIndex

        #if role == Qt.EditRole:
            #item = self.itemFromIndex(index)
            #item.data_ = value
            #self.emit(SIGNAL("dataChanged(QModelIndex,QModelIndex)"), index, index)
            #return True

        if role == Qt.CheckStateRole:
            item = self.itemFromIndex(index)
            checked = (value.toInt()[0] == Qt.Checked)
            item.set_visible(checked)
            self.emit(SIGNAL("dataChanged(QModelIndex,QModelIndex)"), index, index)
            return True

        if role == Qt.EditRole:
            item = self.itemFromIndex(index)
            return item.setData(index, value, role)

        if role == DataRole:
            item = self.itemFromIndex(index)
            print "setData", value.toPyObject()
            if item.setData(index, value, role):
                self.setDirty(True)
                # TODO check why this is needed (should be done by item.setData() anyway)
                self.emit(SIGNAL("dataChanged(QModelIndex,QModelIndex)"), index, index.sibling(index.row(), 1))
            return True

        return False

    def addAnnotation(self, imageidx, ann={}, **kwargs):
        ann.update(kwargs)
        print "addAnnotation", ann
        imageidx = QModelIndex(imageidx)  # explicitly convert from QPersistentModelIndex
        item = self.itemFromIndex(imageidx)
        assert isinstance(item, FrameModelItem) or isinstance(item, ImageFileModelItem)

        next = len(item.children())
        self.beginInsertRows(imageidx, next, next)
        item.addAnnotation(ann)
        self.endInsertRows()
        self.setDirty(True)

        self.emit(SIGNAL("dataChanged(QModelIndex,QModelIndex)"), imageidx, imageidx)

        return True

    def removeAnnotation(self, annidx):
        annidx = QModelIndex(annidx)  # explicitly convert from QPersistentModelIndex
        item = self.itemFromIndex(annidx)
        assert isinstance(item, AnnotationModelItem)

        parent = item.parent_
        parentidx = annidx.parent()
        assert isinstance(parent, FrameModelItem) or isinstance(parent, ImageFileModelItem)

        pos = parent.rowOfChild(item)
        self.beginRemoveRows(parentidx, pos, pos)
        parent.removeAnnotation(pos)
        self.endRemoveRows()
        self.setDirty(True)

        return True

    def headerData(self, section, orientation, role):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            if section == 0:   return QVariant("File/Type/Key")
            elif section == 1: return QVariant("Value")
        return QVariant()

    def getNextIndex(self, index):
        """returns index of next *image* or *frame*"""
        if not index.isValid():
            return QModelIndex()

        assert index == self.imageIndex(index)
        num_images = self.rowCount(index.parent())
        if index.row() < num_images - 1:
            return index.sibling(index.row()+1, 0)

        return index

    def getPreviousIndex(self, index):
        """returns index of previous *image* or *frame*"""
        if not index.isValid():
            return QModelIndex()

        assert index == self.imageIndex(index)
        if index.row() > 0:
            return index.sibling(index.row()-1, 0)

        return index

    def asDictList(self):
        """return annotations as python list of dictionary"""
        # TODO
        annotations = []
        if self.root_ is not None:
            for child in self.root_.children_:
                pass



#######################################################################################
# proxy model
#######################################################################################

class AnnotationSortFilterProxyModel(QSortFilterProxyModel):
    """Adds sorting and filtering support to the AnnotationModel without basically
    any implementation effort.  Special functions such as ``insertPoint()`` just
    call the source models respective functions."""
    def __init__(self, parent=None):
        super(AnnotationSortFilterProxyModel, self).__init__(parent)

    def fileIndex(self, index):
        fi = self.sourceModel().fileIndex(self.mapToSource(index))
        return self.mapFromSource(fi)

    def itemFromIndex(self, index):
        return self.sourceModel().itemFromIndex(self.mapToSource(index))

    def baseDir(self):
        return self.sourceModel().baseDir()

    def insertPoint(self, pos, parent, **kwargs):
        return self.sourceModel().insertPoint(pos, self.mapToSource(parent), **kwargs)

    def insertRect(self, rect, parent, **kwargs):
        return self.sourceModel().insertRect(rect, self.mapToSource(parent), **kwargs)

    def insertMask(self, fname, parent, **kwargs):
        return self.sourceModel().insertMask(fname, self.mapToSource(parent), **kwargs)

    def insertFile(self, filename):
        return self.sourceModel().insertFile(filename)

#######################################################################################
# view
#######################################################################################

class AnnotationTreeView(QTreeView):
    def __init__(self, parent=None):
        super(AnnotationTreeView, self).__init__(parent)

        self.setUniformRowHeights(True)
        self.setSelectionMode(QTreeView.SingleSelection)
        self.setSelectionBehavior(QTreeView.SelectItems)
        self.setAllColumnsShowFocus(True)
        self.setAlternatingRowColors(True)
        self.setEditTriggers(QAbstractItemView.SelectedClicked)
        self.setSortingEnabled(True)
        self.setStyleSheet("""
            QTreeView { selection-color: blue; show-decoration-selected: 1; }
            QTreeView::item:alternate { background-color: #EEEEEE; }
        """)

        self.connect(self, SIGNAL("expanded(QModelIndex)"), self.expanded)

    def resizeColumns(self):
        for column in range(self.model().columnCount(QModelIndex())):
            self.resizeColumnToContents(column)

    def expanded(self):
        self.resizeColumns()

    def setModel(self, model):
        QTreeView.setModel(self, model)
        self.resizeColumns()

    def keyPressEvent(self, event):
        ## handle deletions of items
        if event.key() == Qt.Key_Delete:
            index = self.currentIndex()
            if not index.isValid():
                return
            parent = self.model().parent(index)
            self.model().removeRow(index.row(), parent)

        ## it is important to use the keyPressEvent of QAbstractItemView, not QTreeView
        QAbstractItemView.keyPressEvent(self, event)

    def rowsInserted(self, index, start, end):
        QTreeView.rowsInserted(self, index, start, end)
        self.resizeColumns()
#        self.setCurrentIndex(index.child(end, 0))


def someAnnotations():
    annotations = []
    annotations.append({'type': 'rect',
                        'x': '10',
                        'y': '20',
                        'w': '40',
                        'h': '60'})
    annotations.append({'type': 'rect',
                        'x': '80',
                        'y': '20',
                        'w': '40',
                        'h': '60'})
    annotations.append({'type': 'point',
                        'x': '30',
                        'y': '30'})
    annotations.append({'type': 'point',
                        'x': '100',
                        'y': '100'})
    return annotations

def defaultAnnotations():
    annotations = []
    import os, glob
    if os.path.exists('/cvhci/data/multimedia/bigbangtheory/still_images/s1e1/'):
        images = glob.glob('/cvhci/data/multimedia/bigbangtheory/still_images/s1e1/*.png')
        images.sort()
        for fname in images:
            file = {
                'filename': fname,
                'type': 'image',
                'annotations': someAnnotations()
            }
            annotations.append(file)

    for i in range(5):
        file = {
            'filename': 'file%d.png' % i,
            'type': 'image',
            'annotations': someAnnotations()
        }
        annotations.append(file)
    for i in range(5):
        file = {
            'filename': 'file%d.avi' % i,
            'type':     'video',
            'frames': [],
        }
        for j in range(5):
            frame = {
                'num':       '%d' % j,
                'timestamp': '123456.789',
                'annotations': someAnnotations()
            }
            file['frames'].append(frame)
        annotations.append(file)
    return annotations


if __name__ == '__main__':
    import sys
    app = QApplication(sys.argv)
    annotations = defaultAnnotations()

    model = AnnotationModel(annotations)

    wnd = AnnotationTreeView()
    wnd.setModel(model)
    wnd.show()

    sys.exit(app.exec_())