diff --git a/sloth/annotations/container.py b/sloth/annotations/container.py index 8f5e09a..577f355 100644 --- a/sloth/annotations/container.py +++ b/sloth/annotations/container.py @@ -14,7 +14,7 @@ try: import yaml except: pass - +import okapy class AnnotationContainerFactory: def __init__(self, containers): diff --git a/sloth/annotations/model.py b/sloth/annotations/model.py index 63c2926..023c539 100644 --- a/sloth/annotations/model.py +++ b/sloth/annotations/model.py @@ -1,267 +1,347 @@ """ The annotationmodel module contains the classes for the AnnotationModel. """ -from PyQt4.QtGui import * -from PyQt4.QtCore import * -from functools import partial +from PyQt4.QtGui import QTreeView, QSortFilterProxyModel, QAbstractItemView +from PyQt4.QtCore import QModelIndex, QPersistentModelIndex, QAbstractItemModel, QVariant, Qt, pyqtSignal import os.path -import okapy -import okapy.videoio as okv -TypeRole, DataRole, ImageRole = [Qt.UserRole + i + 1 for i in range(3)] +ItemRole, TypeRole, DataRole, ImageRole = [Qt.UserRole + i + 1 for i in range(4)] class ModelItem: - def __init__(self, model, parent=None): - self.model_ = model - self.parent_ = parent - self.children_ = [] + def __init__(self): + self._children = [] + self._pindex = [] + self._model = None + self._parent = None + self._columns = 1 - def children(self, index=None): - if index is None: - return self.children_ - else: - # return tuple child, index of the child - return [(child, index.child(row, 0)) for row, child in enumerate(self.children_)] + def children(self): + return self._children def model(self): - return self.model_ + return self._model def parent(self): - return self.parent_ + assert self._parent != self + return self._parent - def rowOfChild(self, item): - try: - return self.children_.index(item) - except: - return -1 + def data(self, role=Qt.DisplayRole, column=0): + if role == ItemRole: + return QVariant(self) + else: + return QVariant() - def data(self, index, role): - return QVariant() + def setData(self, value, role=Qt.DisplayRole, column=0): + return False + + def getPosOfChild(self, item): + return self._children.index(item) + + def getChildAt(self, pos): + return self._children[pos] + + def getPreviousSibling(self): + p = self.parent() + if p is not None: + row = p.getPosOfChild(self) + if row > 0: + return p.getChildAt(row-1) + return None + + def getNextSibling(self): + p = self.parent() + if p is not None: + row = p.getPosOfChild(self) + if row < len(p.children()) - 2: + return p.getChildAt(row+1) + return None + + def _attachToModel(self, model, indices): + assert self.model() is None + assert not self._pindex + assert self.parent() is not None + assert self.parent().model() is not None + + self._model = model + + for i in range(self.model().columnCount()): + if i < self._columns: + ind = indices[i] + else: + ind = QModelIndex() + self._pindex.append(QPersistentModelIndex(ind)) + + # Recurse + for i in range(len(self.children())): + item = self.children()[i] + cindices = [self.model().createIndex(i, j, item) for j in range(item._columns)] + item._attachToModel(model, cindices) + + def pindex(self, column=0): + assert self._pindex + return self._pindex[column] + + def index(self, column=0): + assert self._pindex + return QModelIndex(self._pindex[column]) + + def appendChild(self, item): + assert isinstance(item, ModelItem) + assert item.model() is None + assert item.parent() is None + + if self.model() is not None: + next_row = len(self._children) + self.model().beginInsertRows(self.index(), next_row, next_row) + + item._parent = self + self.children().append(item) + + if self.model() is not None: + indices = [self.model().createIndex(next_row, i, item) for i in range(item._columns)] + item._attachToModel(self.model(), indices) + self.model().endInsertRows() + + def appendChildren(self, items): + for item in items: + assert isinstance(item, ModelItem) + assert item.model() is None + assert item.parent() is None + + if self.model() is not None: + next_row = len(self._children) + self.model().beginInsertRows(self.index(), next_row, next_row + len(items) - 1) + + for item in items: + item._parent = self + self.children().append(item) + + if self.model() is not None: + for i in range(len(items)): + item = items[i] + indices = [self.model().createIndex(next_row+i, j, item) for j in range(item._columns)] + item._attachToModel(self.model(), indices) + + self.model().endInsertRows() + + def deleteAllChildren(self): + for child in self._children: + child.deleteAllChildren() + + self._model.beginRemoveRows(self.index(), 0, len(self._children) - 1) + self._children = [] + self._model.endRemoveRows() + + def delete(self): + if self.parent() is None: + raise RuntimeError("Trying to delete orphan") + else: + self.parent().deleteChild(self) + + def deleteChild(self, arg): + if isinstance(arg, ModelItem): + self.deleteChild(self._children.index(arg)) + else: + if arg < 0 or arg >= len(self._children): + raise IndexError("child index out of range") + self._children[arg].deleteAllChildren() + self._model.beginRemoveRows(self.index(), arg, arg) + del self._children[arg] + self._model.endRemoveRows() class RootModelItem(ModelItem): - def __init__(self, model, files): - ModelItem.__init__(self, model, None) - self.files_ = files + def __init__(self, model): + ModelItem.__init__(self) + self._model = model + self._pindex = [QPersistentModelIndex() for i in range(model.columnCount())] - for file in files: - fmi = FileModelItem.create(self.model(), file, self) - self.children_.append(fmi) + def appendChild(self, item): + if isinstance(item, FileModelItem): + ModelItem.appendChild(self, item) + else: + raise TypeError("Only FileModelItems can be attached to RootModelItem") - def addFile(self, file): - fmi = FileModelItem.create(self.model(), file, self) - next = len(self.children_) - index = QModelIndex() - self.model().beginInsertRows(index, next, next) - self.children_.append(fmi) - self.files_.append(file) - self.model().endInsertRows() - self.model().emit(SIGNAL("dataChanged(QModelIndex,QModelIndex)"), index, index) + def appendFileItem(self, fileinfo): + item = FileModelItem.create(fileinfo) + self.appendChild(item) + + def appendFileItems(self, fileinfos): + items = [FileModelItem.create(fi) for fi in fileinfos] + self.appendChildren(items) class FileModelItem(ModelItem): - def __init__(self, model, file, parent): - ModelItem.__init__(self, model, parent) - self.file_ = file + def __init__(self, fileinfo): + ModelItem.__init__(self) + self._fileinfo = fileinfo def filename(self): - return self.file_['filename'] + return self._fileinfo['filename'] - def fullpath(self): - return os.path.join(self.model().basedir(), self.filename()) - - def data(self, index, role): - if role == Qt.DisplayRole and index.column() == 0: + def data(self, role=Qt.DisplayRole, column=0): + if role == Qt.DisplayRole and column == 0: return os.path.basename(self.filename()) - return ModelItem.data(self, index, role) + return ModelItem.data(self, role, column) @staticmethod - def create(model, file, parent): - if file['type'] == 'image': - return ImageFileModelItem(model, file, parent) - elif file['type'] == 'video': - return VideoFileModelItem(model, file, parent) + def create(fileinfo): + if fileinfo['type'] == 'image': + return ImageFileModelItem(fileinfo) + elif fileinfo['type'] == 'video': + return VideoFileModelItem(fileinfo) -class ImageFileModelItem(FileModelItem): - def __init__(self, model, file, parent): - FileModelItem.__init__(self, model, file, parent) +class ImageModelItem(ModelItem): + def __init__(self, annotations): + ModelItem.__init__(self) + for ann in annotations: + self.addAnnotation(ann) - for ann in file['annotations']: - ami = AnnotationModelItem(self.model(), ann, self) - self.children_.append(ami) + def appendChild(self, item): + if isinstance(item, AnnotationModelItem): + ModelItem.appendChild(self, item) + else: + raise TypeError("Only AnnotationModelItems can be attached to ImageModelItem") def addAnnotation(self, ann): - self.file_['annotations'].append(ann) - ami = AnnotationModelItem(self.model(), ann, self) - self.children_.append(ami) + self.appendChild(AnnotationModelItem(ann)) - def updateAnnotation(self, index, ann): - child_found = False - for child in self.children_: + def removeAnnotation(self, pos): + self.deleteChild(pos) + + def updateAnnotation(self, ann): + for child in self._children: if child.type() == ann['type']: if (child.has_key('id') and ann.has_key('id') and child.value('id') == ann['id']) or (not child.has_key('id') and not ann.has_key('id')): ann[None] = None - child.setData(index, QVariant(ann), DataRole) - child_found = True - break - if not child_found: - raise Exception("No ImageFileModelItem found that could be updated!") + child.setData(QVariant(ann), DataRole, 1) + return + raise Exception("No AnnotationModelItem found that could be updated!") - def removeAnnotation(self, pos): - del self.file_['annotations'][pos] - del self.children_[pos] +class ImageFileModelItem(FileModelItem, ImageModelItem): + def __init__(self, fileinfo): + annotations = fileinfo.get("annotations", []) + if fileinfo.has_key("annotations"): + del fileinfo["annotations"] + FileModelItem.__init__(self, fileinfo) + ImageModelItem.__init__(self, annotations) - def data(self, index, role): - if role == ImageRole: - return okapy.loadImage(self.fullpath()) - elif role == DataRole: - return self.file_ - return FileModelItem.data(self, index, role) + def data(self, role=Qt.DisplayRole, column=0): + if role == DataRole: + return self._fileinfo + return FileModelItem.data(self, role) class VideoFileModelItem(FileModelItem): - _cached_vs_filename = None - _cached_vs = None + def __init__(self, fileinfo): + frameinfos = fileinfo.get("frames", []) + if fileinfo.has_key("frames"): + del fileinfo["frames"] + FileModelItem.__init__(self, fileinfo) - def __init__(self, model, file, parent): - FileModelItem.__init__(self, model, file, parent) + for frameinfo in frameinfos: + self.appendChild(FrameModelItem(frameinfo)) - for frame in file['frames']: - fmi = FrameModelItem(self.model(), 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.FFMPEGIndexedVideoSource(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, model, frame, parent): - ModelItem.__init__(self, model, parent) - self.frame_ = frame - - for ann in frame['annotations']: - ami = AnnotationModelItem(ann, self) - self.children_.append(ami) +class FrameModelItem(ImageModelItem): + def __init__(self, frameinfo): + if frameinfo.has_key("annotations"): + ImageModelItem.__init__(self, frameinfo["annotations"]) + del frameinfo["annotations"] + self._frameinfo = frameinfo def framenum(self): - return int(self.frame_.get('num', -1)) + return int(self._frameinfo.get('num', -1)) def timestamp(self): - return float(self.frame_.get('timestamp', -1)) + return float(self._frameinfo.get('timestamp', -1)) - def addAnnotation(self, ann): - self.frame_['annotations'].append(ann) - ami = AnnotationModelItem(ann, self) - self.children_.append(ami) - - def updateAnnotation(self, index, ann): - child_found = False - for child in self.children_: - if child.type() == ann['type']: - if (child.has_key('id') and ann.has_key('id') and child.value('id') == ann['id']) or (not child.has_key('id') and not ann.has_key('id')): - ann[None] = None - child.setData(index, QVariant(ann), DataRole) - child_found = True - break - if not child_found: - raise Exception("No FrameModelItem found that could be updated!") - - 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: + def data(self, role=Qt.DisplayRole, column=0): + if role == Qt.DisplayRole and column == 0: return "%d / %.3f" % (self.framenum(), self.timestamp()) - elif role == ImageRole: - return self.parent().getFrame(self.frame_['num']) - return QVariant() + return ImageModelItem.data(self, role, column) class AnnotationModelItem(ModelItem): - def __init__(self, model, annotation, parent): - ModelItem.__init__(self, model, parent) - self.annotation_ = annotation + def __init__(self, annotation): + ModelItem.__init__(self) + 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 + self._annotation[None] = None for key, value in annotation.iteritems(): if key == None: continue - self.children_.append(KeyValueModelItem(model, key, self)) + self.appendChild(KeyValueModelItem(key)) def type(self): - return self.annotation_['type'] + return self._annotation['type'] - def setData(self, index, data, role): + def setData(self, value, role, column=0): 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 self._annotation + value = value.toPyObject() + print value, type(value) + print self._annotation + for key, val in value.iteritems(): + print key, val + if not key in self._annotation: print "not in annotation: ", 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] + self._annotation[key] = val + self.appendChild(KeyValueModelItem(key)) - for key in self.annotation_.keys(): - if not key in data: - #TODO beginRemoveRows, delete child, etc. - del self.annotation_[key] + for key in self._annotation.keys(): + if not key in value: + for child in [e for e in self.children() if e.key() == key]: + self.deleteChild(child) + 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)) + self._annotation[key] = value[key] + if self.model() is not None: + for child in [e for e in self.children() if e.key() == key]: + self.model().dataChanged.emit(child.index(1), child.index(1)) + + print "new annotation:", self._annotation return True return False - def data(self, index, role): - if role == Qt.DisplayRole and index.column() == 0: + def data(self, role=Qt.DisplayRole, column=0): + print "Annotation:", self._annotation + if role == Qt.DisplayRole and column == 0: return self.type() elif role == TypeRole: return self.type() elif role == DataRole: - #print "data():", self.annotation_ - return self.annotation_ + return self._annotation + return ModelItem.data(self, role, column) - return QVariant() - - def setValue(self, key, value, index): - self.annotation_[key] = value - index.model().dataChanged.emit(index, index.sibling(index.row(), 0)) + def setValue(self, key, value): + self._annotation[key] = value + if self.model() is not None: + self.model().dataChanged.emit(self.index(), self.index()) def value(self, key): - return self.annotation_[key] + return self._annotation[key] def has_key(self, key): - return self.annotation_.has_key(key) + return self._annotation.has_key(key) class KeyValueModelItem(ModelItem): - def __init__(self, model, key, parent): - ModelItem.__init__(self, model, parent) - self.key_ = key + def __init__(self, key): + ModelItem.__init__(self) + self._key = key + self._columns = 2 - def data(self, index, role): + def key(self): + return self._key + + def data(self, role=Qt.DisplayRole, column=0): + print "KeyValue:", self._key if role == Qt.DisplayRole: - if index.column() == 0: - return self.key_ - elif index.column() == 1: - return self.parent().value(self.key_) + if column == 0: + return self._key + elif column == 1: + return self.parent().value(self._key) else: return QVariant() + else: + return ModelItem.data(self, role, column) class AnnotationModel(QAbstractItemModel): # signals @@ -269,73 +349,12 @@ class AnnotationModel(QAbstractItemModel): def __init__(self, annotations, parent=None): QAbstractItemModel.__init__(self, parent) - self.annotations_ = annotations - self.root_ = RootModelItem(self, 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=QModelIndex()): - 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=Qt.DisplayRole): - 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) + self._annotations = annotations + self._dirty = False + self._root = RootModelItem(self) + self._root.appendFileItems(annotations) + # QAbstractItemModel overloads def columnCount(self, index=QModelIndex()): return 2 @@ -344,108 +363,34 @@ class AnnotationModel(QAbstractItemModel): return len(item.children()) def parent(self, index): + if index is None: + return QModelIndex() item = self.itemFromIndex(index) parent = item.parent() if parent is None: return QModelIndex() - grandparent = parent.parent() - if grandparent is None: + return parent.index() + + def index(self, row, column, parent_idx=QModelIndex()): + parent = self.itemFromIndex(parent_idx) + if row >= len(parent.children()): return QModelIndex() - row = grandparent.rowOfChild(parent) - assert row != -1 - return self.createIndex(row, 0, parent) + return parent.children()[row].index(column) - def mapToSource(self, index): - return index - - def flags(self, index): - return Qt.ItemIsEnabled | Qt.ItemIsSelectable + def data(self, index, role=Qt.DisplayRole): if not index.isValid(): - return Qt.ItemIsEnabled - index = QModelIndex(index) # explicitly convert from QPersistentModelIndex + return QVariant() item = self.itemFromIndex(index) - return item.flags(index) + return item.data(role, index.column()) def setData(self, index, value, role=Qt.EditRole): if not index.isValid(): return False - index = QModelIndex(index) # explicitly convert from QPersistentModelIndex + item = self.itemFromIndex(index) + return item.setData(value, role, index.column()) - #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 updateAnnotation(self, imageidx, ann={}, **kwargs): - ann.update(kwargs) - print "updateAnnotation", ann - imageidx = QModelIndex(imageidx) # explicitly convert from QPersistentModelIndex - item = self.itemFromIndex(imageidx) - assert isinstance(item, FrameModelItem) or isinstance(item, ImageFileModelItem) - - item.updateAnnotation(imageidx, ann) - 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 flags(self, index): + return Qt.ItemIsEnabled | Qt.ItemIsSelectable def headerData(self, section, orientation, role): if orientation == Qt.Horizontal and role == Qt.DisplayRole: @@ -453,38 +398,21 @@ class AnnotationModel(QAbstractItemModel): 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() + # Own methods + def dirty(self): + return self._dirty - 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): - # TODO bool parameter to disable wrap around - """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 + # TODO: This might need to be updated from within the ModelItems when they change + def setDirty(self, dirty=True): + if dirty != self._dirty: + self._dirty = dirty + self.dirtyChanged.emit(self._dirty) + def itemFromIndex(self, index): + index = QModelIndex(index) # explicitly convert from QPersistentModelIndex + if index.isValid(): + return index.internalPointer() + return self._root ####################################################################################### @@ -520,6 +448,7 @@ class AnnotationSortFilterProxyModel(QSortFilterProxyModel): def insertFile(self, filename): return self.sourceModel().insertFile(filename) + ####################################################################################### # view ####################################################################################### @@ -535,18 +464,13 @@ class AnnotationTreeView(QTreeView): 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) + self.expanded.connect(self.onExpanded) def resizeColumns(self): for column in range(self.model().columnCount(QModelIndex())): self.resizeColumnToContents(column) - def expanded(self): + def onExpanded(self): self.resizeColumns() def setModel(self, model): @@ -556,11 +480,7 @@ class AnnotationTreeView(QTreeView): 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) + self.model().itemFromIndex(self.currentindex()).delete() ## it is important to use the keyPressEvent of QAbstractItemView, not QTreeView QAbstractItemView.keyPressEvent(self, event) @@ -569,76 +489,3 @@ class AnnotationTreeView(QTreeView): 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_()) - diff --git a/sloth/core/labeltool.py b/sloth/core/labeltool.py index c470278..c3f8577 100755 --- a/sloth/core/labeltool.py +++ b/sloth/core/labeltool.py @@ -17,7 +17,10 @@ class LabelTool(QObject): statusMessage = pyqtSignal(QString) annotationsLoaded = pyqtSignal() pluginLoaded = pyqtSignal(QAction) - currentIndexChanged = pyqtSignal(QModelIndex) + # This still emits a QModelIndex, because Qt cannot handle emiting + # a derived class instead of a base class, i.e. ImageFileModelItem + # instead of ModelItem + currentImageChanged = pyqtSignal(QModelIndex) def __init__(self, argv, parent=None): QObject.__init__(self, parent) @@ -32,7 +35,8 @@ class LabelTool(QObject): # Instatiate container factory self.container_factory_ = AnnotationContainerFactory(config.CONTAINERS) self.container_ = AnnotationContainer() - self.current_index_ = None + self._current_image = None + self._model = None # Load annotation file if len(args) > 0: @@ -68,18 +72,18 @@ class LabelTool(QObject): ###___________________________________________________________________________________________ def loadAnnotations(self, fname): fname = str(fname) # convert from QString - try: - self.container_ = self.container_factory_.create(fname) - self.container_.load(fname) - msg = "Successfully loaded %s (%d files, %d annotations)" % \ - (fname, self.container_.numFiles(), self.container_.numAnnotations()) - self._model = AnnotationModel(self.container_.annotations()) - if self.container_.filename() is not None: - self._model.setBasedir(os.path.dirname(self.container_.filename())) - else: - self._model.setBasedir("") - except Exception, e: - msg = "Error: Loading failed (%s)" % str(e) + #try: + self.container_ = self.container_factory_.create(fname) + self.container_.load(fname) + msg = "Successfully loaded %s (%d files, %d annotations)" % \ + (fname, self.container_.numFiles(), self.container_.numAnnotations()) + self._model = AnnotationModel(self.container_.annotations()) + #if self.container_.filename() is not None: + #self._model.setBasedir(os.path.dirname(self.container_.filename())) + #else: + #self._model.setBasedir("") + #except Exception, e: + #msg = "Error: Loading failed (%s)" % str(e) self.statusMessage.emit(msg) self.annotationsLoaded.emit() @@ -120,7 +124,7 @@ class LabelTool(QObject): def clearAnnotations(self): self.container_.clear() self._model = AnnotationModel(self.container_.annotations()) - self._model.setBasedir("") + #self._model.setBasedir("") self.statusMessage.emit('') self.annotationsLoaded.emit() @@ -136,15 +140,15 @@ class LabelTool(QObject): def gotoNext(self): # TODO move this to the scene - if self._model is not None and self.current_index_ is not None: - next_index = self._model.getNextIndex(self.current_index_) - self.setCurrentIndex(next_index) + if self._model is not None and self._current_image is not None: + next_image = self._current_image.getNextSibling() + self.setCurrentImage(next_image) def gotoPrevious(self): # TODO move this to the scene - if self._model is not None and self.current_index_ is not None: - prev_index = self._model.getPreviousIndex(self.current_index_) - self.setCurrentIndex(prev_index) + if self._model is not None and self._current_image is not None: + prev_image = self._current_image.getPreviousSibling() + self.setCurrentImage(prev_image) def updateModified(self): """update all GUI elements which depend on the state of the model, @@ -155,15 +159,23 @@ class LabelTool(QObject): #self.setWindowModified(self.annotations.dirty()) pass - def currentIndex(self): - return self.current_index_ + def currentImage(self): + return self._current_image - def setCurrentIndex(self, index): - assert index.isValid() - newindex = index.model().imageIndex(index) - if newindex.isValid() and newindex != self.current_index_: - self.current_index_ = newindex - self.currentIndexChanged.emit(self.current_index_) + def setCurrentImage(self, image): + if isinstance(image, QModelIndex): + image = self._model.itemFromIndex(image) + while (image is not None) and (not isinstance(image, ImageModelItem)): + image = image.parent() + if image is None: + raise RuntimeError("Tried to set current image to item that has no Image or Frame as parent!") + if image != self._current_image: + self._current_image = image + self.currentImageChanged.emit(self._current_image.index()) + + def getImage(self, item): + # TODO: Also handle video frames + return self.container_.loadImage(item.filename()) def getAnnotationFilePatterns(self): return self.container_factory_.patterns() diff --git a/sloth/gui/annotationscene.py b/sloth/gui/annotationscene.py index 5941198..7068719 100644 --- a/sloth/gui/annotationscene.py +++ b/sloth/gui/annotationscene.py @@ -13,7 +13,7 @@ class AnnotationScene(QGraphicsScene): # TODO signal itemadded - def __init__(self, items=None, inserters=None, parent=None): + def __init__(self, labeltool, items=None, inserters=None, parent=None): super(AnnotationScene, self).__init__(parent) self.model_ = None @@ -22,6 +22,7 @@ class AnnotationScene(QGraphicsScene): self.debug_ = True self.message_ = "" self.last_key_ = None + self.labeltool_ = labeltool self.itemfactory_ = Factory(items) self.inserterfactory_ = Factory(inserters) @@ -79,7 +80,8 @@ class AnnotationScene(QGraphicsScene): return assert self.root_.model() == self.model_ - self.image_ = self.root_.data(ImageRole).toPyObject() + item = self.model_.itemFromIndex(root) + self.image_ = self.labeltool_.getImage(item) self.pixmap_ = QPixmap(okapy.guiqt.toQImage(self.image_)) item = QGraphicsPixmapItem(self.pixmap_) item.setZValue(-1) @@ -272,7 +274,6 @@ class AnnotationScene(QGraphicsScene): 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 diff --git a/sloth/gui/labeltool.py b/sloth/gui/labeltool.py index 5b53584..0a8d25b 100755 --- a/sloth/gui/labeltool.py +++ b/sloth/gui/labeltool.py @@ -46,19 +46,19 @@ class MainWindow(QMainWindow): self.setWindowTitle("%s - Unnamed[*]" % APP_NAME) self.treeview.setModel(self.labeltool.model()) self.scene.setModel(self.labeltool.model()) - self.treeview.selectionModel().currentChanged.connect(self.labeltool.setCurrentIndex) + self.treeview.selectionModel().currentChanged.connect(self.labeltool.setCurrentImage) - def onCurrentIndexChanged(self, new_index): - self.scene.setRoot(new_index) + def onCurrentImageChanged(self, index): + self.scene.setRoot(index) + new_image = self.labeltool.currentImage() # TODO: This info should be obtained from AnnotationModel or LabelTool - item = self.labeltool.model().itemFromIndex(new_index) - if isinstance(item, FrameModelItem): + if isinstance(new_image, FrameModelItem): self.controls.setFrameNumAndTimestamp(item.framenum(), item.timestamp()) - elif isinstance(item, ImageFileModelItem): - self.controls.setFilename(os.path.basename(item.filename())) - if new_index != self.treeview.currentIndex(): - self.treeview.setCurrentIndex(new_index) + elif isinstance(new_image, ImageFileModelItem): + self.controls.setFilename(os.path.basename(new_image.filename())) + if new_image.index() != self.treeview.currentIndex(): + self.treeview.setCurrentIndex(new_image.index()) def initShortcuts(self): # TODO clean up, make configurable @@ -91,7 +91,7 @@ class MainWindow(QMainWindow): def setupGui(self): self.ui = uic.loadUi(os.path.join(GUIDIR, "labeltool.ui"), self) - self.scene = AnnotationScene(items=config.ITEMS, inserters=config.INSERTERS) + self.scene = AnnotationScene(self.labeltool, items=config.ITEMS, inserters=config.INSERTERS) self.view = GraphicsView(self) self.view.setSizePolicy(QSizePolicy.MinimumExpanding, QSizePolicy.MinimumExpanding) self.view.setScene(self.scene) @@ -141,7 +141,7 @@ class MainWindow(QMainWindow): self.labeltool.pluginLoaded. connect(self.onPluginLoaded) self.labeltool.statusMessage. connect(self.onStatusMessage) self.labeltool.annotationsLoaded. connect(self.onAnnotationsLoaded) - self.labeltool.currentIndexChanged.connect(self.onCurrentIndexChanged) + self.labeltool.currentImageChanged.connect(self.onCurrentImageChanged) def loadApplicationSettings(self): settings = QSettings()