mirror of
https://github.com/wassname/sloth.git
synced 2026-07-04 05:36:22 +08:00
Martin Baeuml
832 lines
30 KiB
Python
832 lines
30 KiB
Python
from PyQt4.QtGui import *
|
|
from PyQt4.QtCore import *
|
|
from functools import partial
|
|
import os.path
|
|
|
|
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):
|
|
for row, child in enumerate(self.children_):
|
|
if child is item:
|
|
return row
|
|
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(file, self)
|
|
self.children_.append(fmi)
|
|
|
|
class FileModelItem(ModelItem):
|
|
def __init__(self, file, parent):
|
|
ModelItem.__init__(self, parent)
|
|
self.file_ = file
|
|
|
|
for frame in file['frames']:
|
|
fmi = FrameModelItem(frame, self)
|
|
self.children_.append(fmi)
|
|
|
|
def filename(self):
|
|
return self.file_['filename']
|
|
|
|
def data(self, index, role):
|
|
if role == Qt.DisplayRole and index.column() == 0:
|
|
return self.filename()
|
|
return QVariant()
|
|
|
|
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 data(self, index, role):
|
|
if role == Qt.DisplayRole and index.column() == 0:
|
|
return "%d / %.3f" % (self.framenum(), self.timestamp())
|
|
return QVariant()
|
|
|
|
class AnnotationModelItem(ModelItem):
|
|
def __init__(self, annotations, parent):
|
|
ModelItem.__init__(self, parent)
|
|
self.annotations_ = annotations
|
|
|
|
for key, value in annotations.iteritems():
|
|
self.children_.append(KeyValueModelItem(key, value, self))
|
|
|
|
def type(self):
|
|
return self.annotations_['type']
|
|
|
|
def data(self, index, role):
|
|
if role == Qt.DisplayRole:
|
|
if index.column() == 0:
|
|
return self.type()
|
|
else:
|
|
return QVariant()
|
|
|
|
class KeyValueModelItem(ModelItem):
|
|
def __init__(self, key, value, parent):
|
|
ModelItem.__init__(self, parent)
|
|
self.key_ = key
|
|
self.value_ = value
|
|
|
|
def data(self, index, role):
|
|
if role == Qt.DisplayRole:
|
|
if index.column() == 0:
|
|
return self.key_
|
|
elif index.column() == 1:
|
|
return self.value_
|
|
else:
|
|
return QVariant()
|
|
|
|
|
|
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):
|
|
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)
|
|
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.setModel(model)
|
|
wnd.show()
|
|
|
|
sys.exit(app.exec_())
|
|
|