mirror of
https://github.com/wassname/sloth.git
synced 2026-06-30 17:08:02 +08:00
Merge pull request #1 from mika-fischer/model-refactoring
Model refactoring
This commit is contained in:
Mika Fischer
@@ -14,7 +14,7 @@ try:
|
||||
import yaml
|
||||
except:
|
||||
pass
|
||||
|
||||
import okapy
|
||||
|
||||
class AnnotationContainerFactory:
|
||||
def __init__(self, containers):
|
||||
|
||||
+302
-455
@@ -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_())
|
||||
|
||||
|
||||
+41
-29
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
+11
-11
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user