From 5e1eb5215239743fed63521ffb34808cf44694aa Mon Sep 17 00:00:00 2001 From: Martin Baeuml Date: Tue, 23 Nov 2010 22:59:52 +0100 Subject: [PATCH] add first skeleton of AnnotationModel --- annotationmodel.py | 773 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 773 insertions(+) create mode 100644 annotationmodel.py diff --git a/annotationmodel.py b/annotationmodel.py new file mode 100644 index 0000000..d87edee --- /dev/null +++ b/annotationmodel.py @@ -0,0 +1,773 @@ +from PyQt4.QtGui import * +from PyQt4.QtCore import * +from functools import partial +import os.path + +class KeyValueModelItem(): + def __init__(self, key, value=""): + self.key_ = key + self.value_ = value + +class ModelItem: + def __init__(self, parent=None): + self.parent_ = parent + self.children_ = [] + + def children(self): + return self.children_ + +class AnnotationModelItem(ModelItem): + def __init__(self, ann, parent): + ModelItem.__init__(self, parent) + +class RootModelItem(ModelItem): + def __init__(self, files): + ModelItem.__init__(self, None) + + for file in files: + fmi = FileModelItem(file, self) + self.children_.append(fmi) + +class FileModelItem(ModelItem): + def __init__(self, file, parent): + ModelItem.__init__(self, parent) + + for frame in file['frames']: + fmi = FrameModelItem(frame, self) + self.children_.append(fmi) + +class FrameModelItem(ModelItem): + def __init__(self, frame, parent): + ModelItem.__init__(self, parent) + + for ann in frame['annotations']: + ami = AnnotationModelItem(ann, self) + self.children_.append(ami) + + + +class AnnotationModelItemOld(object): + def __init__(self, data, model, parent=None): + self.children_ = [] + self.type_ = type + self.data_ = data + self.parent_ = parent + self.model_ = model + self.visible_ = True + + def parent(self): + return self.parent_ + + def model(self): + return self.model_ + + def children(self): + return self.children_ + + def rowOfChild(self, item): + for row, child in enumerate(self.children_): + if child is item: + return row + return -1 + + def data(self, index, role): + return QVariant() + + def flags(self, index): + return Qt.ItemIsEnabled | Qt.ItemIsSelectable + + def removeRows(self, position, rows): + return False + + def visible(self): return self.visible_ + def set_visible(self, visible): self.visible_ = visible + + def isCheckable(self, column): return False + +class PropertyDelegate(QItemDelegate): + def __init__(self, model, parent=None): + super(PropertyDelegate, self).__init__(parent) + self.model_ = model + + def createEditor(self, parent, option, index): + if index.column() == 0: + return None + item = self.model_.itemFromIndex(index) + return item.createEditor(parent, option, index) + + def setEditorData(self, editor, index): + item = self.model.getItem(index) + item.setEditorData(editor, index) + + def setModelData(self, editor, model, index): + item = self.model.getItem(index) + item.setModelData(editor, model, index) + +# def updateEditorGeometry(self, editor, option, index): +# editor.setGeometry(option.rect) + + def paint(self, painter, option, index): + item = self.model.getItem(index) + if isinstance(item, ColorProperty) and index.column() == 1: + item.paint(painter, option, index) + return + QtGui.QItemDelegate.paint(self, painter, option, index) + +class Property(AnnotationModelItem): + def __init__(self, key, data, model, parent=None): + super(Property, self).__init__(data, model, parent) + self.key_ = key + + def data(self, index, role): + if index.column() == 0: + return QVariant(self.key_) + elif index.column() == 1: + return QVariant(str(self.data_[self.key_])) + return QVariant("") + + def flags(self, index): + if index.column() == 1: + return super(Property, self).flags(index) | Qt.ItemIsEditable + return super(Property, self).flags(index) + + def setData(self, index, value, role): + if role == Qt.EditRole: + return self.parent().setData(index, value, role) + return False + + def createEditor(self, parent, option, index): + return None + + def setEditorData(self, editor, index): + return + + def setModelData(self, editor, model, index): + return + +class StringProperty(Property): + def __init__(self, key, data, model, parent): + super(StringProperty, self).__init__(key, data, model, parent) + + def createEditor(self, parent, option, index): + editor = QLineEdit(parent) + return editor + + def setEditorData(self, editor, index): + editor.setText(str(self.data_[self.key_])) + + def setModelData(self, editor, model, index): + self.data_[self.key_] = str(editor.text()) + +class RootAnnotationModelItem(AnnotationModelItem): + def __init__(self, data, model): + super(RootAnnotationModelItem, self).__init__(data, model) + for row, file in enumerate(self.data_.files): + self.children_.append(FileAnnotationModelItem(file, model, self)) + + def removeRows(self, position, rows): + self.data_.files = self.data_.files[:position] + self.data_.files[position+rows:] + self.children_ = self.children_[:position] + self.children_[position+rows:] + return True + +class FileAnnotationModelItem(AnnotationModelItem): + def __init__(self, data, model, parent): + super(FileAnnotationModelItem, self).__init__(data, model, parent) + for row, annotation in enumerate(self.data_.annotations): + if annotation.type == 'rect': + self.children_.append(RectAnnotationModelItem(annotation, model, self)) + elif annotation.type == 'mask': + self.children_.append(MaskAnnotationModelItem(annotation, model, self)) + elif annotation.type == 'point': + self.children_.append(PointAnnotationModelItem(annotation, model, self)) + + def data(self, index, role): + if role == DataRole: + return QVariant(self.data_.filename) + + if index.column() == 0: + return QVariant(os.path.split(self.data_.filename)[1]) + else: + return QVariant("") + + def removeRows(self, position, rows): + self.data_.annotations = self.data_.annotations[:position] + self.data_.annotations[position+rows:] + self.children_ = self.children_[:position] + self.children_[position+rows:] + return True + +class RectAnnotationModelItem(AnnotationModelItem): + def __init__(self, data, model, parent): + super(RectAnnotationModelItem, self).__init__(data, model, parent) + self.children_.append(StringProperty('id', data, model, self)) + self.children_.append(StringProperty('class', data, model, self)) + self.children_.append(StringProperty('x', data, model, self)) + self.children_.append(StringProperty('y', data, model, self)) + self.children_.append(StringProperty('width', data, model, self)) + self.children_.append(StringProperty('height', data, model, self)) + + def data(self, index, role): + if role == GraphicsItemRole: + return QVariant(AnnotationGraphicsRectItem(self, index)) + elif role == DataRole: + return QVariant(QRectF(self.data_.x, self.data_.y, self.data_.width, self.data_.height)) + elif role == Qt.DisplayRole: + if index.column() == 0: + return QVariant("Rect") + elif index.column() == 1: + s = "id: '%s', class: '%s', (%d, %d, %d, %d)" % \ + (self.data_['id'], self.data_['class'], int(self.data_['x']), int(self.data_['y']), + int(self.data_['width']), int(self.data_['height'])) + return QVariant(s) + assert False + + def setData(self, index, value, role): + if role == Qt.EditRole: + assert index.model().itemFromIndex(index.parent()) == self + check_pairs = [('id', value.toString), ('class', value.toString), + ('x', value.toString), ('y', value.toString), + ('height', value.toString), ('width', value.toString)] + key, f = check_pairs[index.row()] + if self.data_[key] != f(): + self.data_[key] = str(f()) + index.model().emit(SIGNAL('dataChanged(QModelIndex,QModelIndex)'), index.parent(), index.parent()) + index.model().emit(SIGNAL('dataChanged(QModelIndex,QModelIndex)'), index, index) + return True + return False + + if role == DataRole: + rect = value.toRectF().toRect() + modified = False + check_pairs = [('x', rect.topLeft().x), ('y', rect.topLeft().y), + ('width', rect.size().width), ('height', rect.size().height)] + for key, f in check_pairs: + if self.data_[key] != f(): + self.data_[key] = f() + modified = True + + return modified + return False + + def flags(self, index): + if index.column() == 0: + return Qt.ItemIsUserCheckable | super(RectAnnotationModelItem, self).flags(index) + return super(RectAnnotationModelItem, self).flags(index) + + def isCheckable(self, column): + if column == 0: + return True + return False + +class MaskAnnotationModelItem(AnnotationModelItem): + def __init__(self, data, model, parent): + super(MaskAnnotationModelItem, self).__init__(data, model, parent) + self.children_.append(StringProperty('id', data, model, self)) + self.children_.append(StringProperty('class', data, model, self)) + self.children_.append(StringProperty('filename', data, model, self)) + self.image_ = None + self.shape_ = None + self.dirty_ = False + + def set_dirty(self, dirty=True): self.dirty_ = dirty + def dirty(self): return self.dirty_ + + def data(self, index, role): + if role == GraphicsItemRole: + return QVariant(AnnotationGraphicsMaskItem(self, index)) + elif role == DataRole: + return QVariant(self.image()) + elif role == Qt.DisplayRole: + if index.column() == 0: + return QVariant("Mask") + elif index.column() == 1: + s = "id: '%s', class: '%s'" % (self.data_['id'], self.data_['class']) + return QVariant(s) + + def setData(self, index, value, role): + if role == Qt.EditRole: + assert index.model().itemFromIndex(index.parent()) == self + check_pairs = [('id', value.toString), ('class', value.toString)] + key, f = check_pairs[index.row()] + if self.data_[key] != f(): + self.data_[key] = str(f()) + index.model().emit(SIGNAL('dataChanged(QModelIndex,QModelIndex)'), index.parent(), index.parent()) + index.model().emit(SIGNAL('dataChanged(QModelIndex,QModelIndex)'), index, index) + return True + return False + + if role == DataRole: + assert value.type() == QVariant.Image + self.image_ = QImage(value) + self.set_dirty() + index.model().emit(SIGNAL('dataChanged(QModelIndex,QModelIndex)'), index, index) + # emit dataChanged (?) shouldn't be necessary + return True + return False + + def flags(self, index): + if index.column() == 0: + return Qt.ItemIsUserCheckable | super(MaskAnnotationModelItem, self).flags(index) + return super(MaskAnnotationModelItem, self).flags(index) + + def isCheckable(self, column): + if column == 0: + return True + return False + + def filename(self): + if os.path.isabs(self.data_.filename): + return self.data_.filename + else: + return os.path.join(self.model().baseDir(), self.data_.filename) + + def image(self): + if self.image_ is None: + if os.path.exists(self.filename()): + self.image_ = QImage(self.filename()).convertToFormat(QImage.Format_MonoLSB) + else: + # find parent image + parent = self.parent() + while (parent is not None): + if isinstance(parent, FileAnnotationModelItem): + break + if parent is None: + print >>sys.stderr, "Error: mask annotation item does not have a File parent" + else: + print 'Creating new mask for', self.filename() + im = QImage(parent.data_.filename) + self.image_ = QImage(im.width(), im.height(), QImage.Format_MonoLSB) + self.image_.fill(0) + self.set_dirty() + + # strange qt convention: + # white == Qt.color0 == 0 is background + # black == Qt.color1 == 1 is foreground + # thus we switch the pixel values and reassign the colors + self.image_.invertPixels() + return self.image_ + + def writeback(self): + assert self.filename() is not None + if not self.dirty(): return + tmpimage = QImage(self.image()) # important to wrap this in a QImage constructor, otherwise no real copy made + tmpimage.invertPixels() + tmpimage.save(self.filename(), os.path.splitext(self.filename())[1][1:]) + +class PointAnnotationModelItem(AnnotationModelItem): + def __init__(self, data, model, parent): + super(PointAnnotationModelItem, self).__init__(data, model, parent) + self.children_.append(StringProperty('id', data, model, self)) + self.children_.append(StringProperty('class', data, model, self)) + self.children_.append(StringProperty('x', data, model, self)) + self.children_.append(StringProperty('y', data, model, self)) + + def data(self, index, role): + if role == GraphicsItemRole: + return QVariant(AnnotationGraphicsPointItem(self, index)) + elif role == DataRole: + return QVariant(QPointF(self.data_.x, self.data_.y)) + elif role == Qt.DisplayRole: + if index.column() == 0: + return QVariant("Point") + elif index.column() == 1: + s = "id: '%s', class: '%s', (%d, %d)" % \ + (self.data_['id'], self.data_['class'], int(self.data_['x']), int(self.data_['y'])) + return QVariant(s) + return QVariant() + + def setData(self, index, value, role): + if role == Qt.EditRole: + assert index.model().itemFromIndex(index.parent()) == self + check_pairs = [('id', value.toString), ('class', value.toString), + ('x', value.toString), ('y', value.toString)] + key, f = check_pairs[index.row()] + if self.data_[key] != f(): + self.data_[key] = str(f()) + index.model().emit(SIGNAL('dataChanged(QModelIndex,QModelIndex)'), index.parent(), index.parent()) + index.model().emit(SIGNAL('dataChanged(QModelIndex,QModelIndex)'), index, index) + return True + return False + + if role == DataRole: + point = value.toPointF().toPoint() + modified = False + check_pairs = [('x', point.x), ('y', point.y)] + for key, f in check_pairs: + if self.data_[key] != f(): + self.data_[key] = f() + modified = True + return modified + + return False + + def flags(self, index): + if index.column() == 0: + return Qt.ItemIsUserCheckable | super(PointAnnotationModelItem, self).flags(index) + return super(PointAnnotationModelItem, self).flags(index) + + def isCheckable(self, column): + if column == 0: + return True + return False + +class AnnotationModel(QAbstractItemModel): + def __init__(self, annotations, parent=None): + QAbstractItemModel.__init__(self, parent) + self.annotations_ = annotations + self.root_ = RootModelItem(self.annotations_) + self.dirty_ = False + + def dirty(self): + return self.dirty_ + + def setDirty(self, dirty=True): + previous = self.dirty_ + self.dirty_ = dirty + if previous != dirty: + self.emit(SIGNAL("dirtyChanged()")) + + dirty = property(dirty, setDirty) + + 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 fileIndex(self, index): + """return index that points to the (maybe parental) file object""" + if not index.isValid(): + return QModelIndex() + index = QModelIndex(index) # explicitly convert from QPersistentModelIndex + item = self.itemFromIndex(index) + if isinstance(item, FileAnnotationModelItem): + return index + return self.fileIndex(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 flags(self, index): + 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) + if item.setData(index, value, role): + self.emit(SIGNAL("dataChanged(QModelIndex,QModelIndex)"), index, index.sibling(index.row(), 1)) + return True + + return False + + def insertRows(self, position, rows=1, index=QModelIndex()): + index = QModelIndex(index) # explicitly convert from QPersistentModelIndex + item = self.itemFromIndex(index) + if isinstance(item, RootAnnotationModelItem): + self.beginInsertRows(QModelIndex(), position, position + rows - 1) + for row in range(rows): + file = File('') + item.data_.files.insert(position + row, file) + item.children_.insert(position + row, FileAnnotationModelItem(file, item)) + self.endInsertRows() + self.emit(SIGNAL("dataChanged(QModelIndex,QModelIndex)"), index, index) + self.set_dirty(True) + return True + # TODO handle inserts of rects, points etc + + def removeRows(self, position, rows=1, index=QModelIndex()): + index = QModelIndex(index) # explicitly convert from QPersistentModelIndex + item = self.itemFromIndex(index) + self.beginRemoveRows(index, position, position + rows - 1) + data_changed = item.removeRows(position, rows) + self.endRemoveRows() + if data_changed: + self.set_dirty(True) + return True + return False + + def insertPoint(self, pos, parent, **kwargs): + parent = QModelIndex(parent) # explicitly convert from QPersistentModelIndex + item = self.itemFromIndex(parent) + #TODO move this to the item class + assert isinstance(item, FileAnnotationModelItem) + position = len(item.children_) + assert position == len(item.data_.annotations) + self.beginInsertRows(parent, position, position) + data = Annotation(type='point', x=pos.x(), y=pos.y()) + data.update(kwargs) + item.data_.annotations.append(data) + item.children_.append(PointAnnotationModelItem(data, self, item)) + self.endInsertRows() + self.set_dirty(True) + return True + + def insertRect(self, rect, parent, **kwargs): + parent = QModelIndex(parent) # explicitly convert from QPersistentModelIndex + item = self.itemFromIndex(parent) + #TODO move this to the item class + assert isinstance(item, FileAnnotationModelItem) + position = len(item.children_) + assert position == len(item.data_.annotations) + self.beginInsertRows(parent, position, position) + rect = rect.normalized().toRect() + data = Annotation(type='rect', x=rect.x(), y=rect.y(), width=rect.width(), height=rect.height()) + data.update(kwargs) + item.data_.annotations.append(data) + item.children_.append(RectAnnotationModelItem(data, self, item)) + self.endInsertRows() + self.set_dirty(True) + return True + + def insertMask(self, fname, parent, **kwargs): + parent = QModelIndex(parent) # explicitly convert from QPersistentModelIndex + item = self.itemFromIndex(parent) + #TODO move this to the item class + assert isinstance(item, FileAnnotationModelItem) + position = len(item.children_) + assert position == len(item.data_.annotations) + self.beginInsertRows(parent, position, position) + + data = Annotation(type='mask', filename=fname) + data.update(kwargs) + item.data_.annotations.append(data) + item.children_.append(MaskAnnotationModelItem(data, self, item)) + self.endInsertRows() + self.set_dirty(True) + return True + + def insertFile(self, filename): + parent = QModelIndex() + position = len(self.root_.children()) + assert position == len(self.annotations.files) + self.beginInsertRows(parent, position, position) + file = File(filename, type='image') # TODO determine type image vs. movie here + self.annotations.append(file) + self.root_.children_.append(FileAnnotationModelItem(file, self, self.root_)) + self.endInsertRows() + self.set_dirty(True) + return True + + def headerData(self, section, orientation, role): + if orientation == Qt.Horizontal and role == Qt.DisplayRole: + if section == 0: return QVariant("File/Type") + elif section == 1: return QVariant("Value") + return QVariant() + + def baseDir(self): + return self.annotations.base_dir() + + def writeback(self, root=None): + """Write back all items to disk which have been modified since + the last writeback (such as image masks). This should usually be done + everytime the underlying annotations are saved.""" + # Iterate over all items and check whether they have a writeback method. + # Then save these items to disk. + if root is None: + root = self.root_ + for child in root.children(): + if hasattr(child, 'writeback'): + child.writeback() + self.writeback(child) + +####################################################################################### +# 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.setSelectionBehavior(QTreeView.SelectItems) + self.setEditTriggers(QAbstractItemView.SelectedClicked) + self.setSortingEnabled(True) + + 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) + + ## super(...) does not work here... + ## 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 selectionModel(self): + return QAbstractItemView.selectionModel(self) + +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'}) + return annotations + +def defaultAnnotations(): + annotations = [] + for i in range(5): + file = { + 'filename': 'file%d.png' % i, + 'type': 'image', + 'frames': [] + } + file['frames'].append({'annotations': someAnnotations()}) + annotations.append(file) + for i in range(5): + file = { + 'filename': 'file%d.png' % 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.set_annotations(ann) + #wnd.show() + + sys.exit(app.exec_()) +