wassname/sloth
1
0
Fork 0
mirror of https://github.com/wassname/sloth.git synced 2026-07-01 04:51:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

merge with remote master

Browse Source
This commit is contained in:
Martin Baeuml
2011-06-14 19:10:51 +02:00
parent 08909f0183 5f19bf6cc7
commit 7be783f9db
6 changed files with 367 additions and 510 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sloth/annotations/container.py
+1 -1
View File
@@ -14,7 +14,7 @@ try:
import yaml
except:
pass
import okapy
class AnnotationContainerFactory:
def __init__(self, containers):
sloth/annotations/model.py
+300 -455
View File
@@ -1,267 +1,345 @@
"""
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):
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):
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 +347,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 +361,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 +396,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 +446,7 @@ class AnnotationSortFilterProxyModel(QSortFilterProxyModel):
def insertFile(self, filename):
return self.sourceModel().insertFile(filename)
#######################################################################################
# view
#######################################################################################
@@ -535,18 +462,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 +478,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 +487,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_())
sloth/core/labeltool.py
+28 -21
View File
@@ -43,7 +43,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)
# TODO clean up --> prefix all members with _
def __init__(self, parent=None):
@@ -61,7 +64,7 @@ class LabelTool(QObject):
self.container_factory_ = None
self.container_ = AnnotationContainer()
self.current_index_ = None
self._current_image = None
self._model = None
def main_help_text(self):
@@ -188,10 +191,6 @@ class LabelTool(QObject):
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)
@@ -223,7 +222,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()
@@ -239,15 +238,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,
@@ -258,15 +257,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()
sloth/gui/annotationscene.py
+14 -9
View File
@@ -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)
@@ -73,15 +74,20 @@ class AnnotationScene(QGraphicsScene):
displayed by the scene. This can be either the index to a frame in a
video, or to an image.
"""
self.image_item_ = None
self.image_ = None
self.pixmap_ = None
self.root_ = root
self.clear()
if not root.isValid():
return
assert self.root_.model() == self.model_
self.image_ = self.root_.data(ImageRole).toPyObject()
self.pixmap_ = QPixmap(okapy.guiqt.toQImage(self.image_))
item = QGraphicsPixmapItem(self.pixmap_)
self.image_item_ = self.model_.itemFromIndex(root)
self.image_ = self.labeltool_.getImage(self.image_item_)
self.pixmap_ = QPixmap(okapy.guiqt.toQImage(self.image_))
item = QGraphicsPixmapItem(self.pixmap_)
item.setZValue(-1)
self.setSceneRect(0, 0, self.pixmap_.width(), self.pixmap_.height())
self.addItem(item)
@@ -146,7 +152,7 @@ class AnnotationScene(QGraphicsScene):
# ignore events outside the scene rect
return
# insert mode
self.inserter_.mousePressEvent(event, self.root_)
self.inserter_.mousePressEvent(event, self.image_item_)
else:
# selection mode
QGraphicsScene.mousePressEvent(self, event)
@@ -156,7 +162,7 @@ class AnnotationScene(QGraphicsScene):
print "mouseReleaseEvent", self.sceneRect().contains(event.scenePos()), event.scenePos()
if self.inserter_ is not None:
# insert mode
self.inserter_.mouseReleaseEvent(event, self.root_)
self.inserter_.mouseReleaseEvent(event, self.image_item_)
else:
# selection mode
QGraphicsScene.mouseReleaseEvent(self, event)
@@ -166,7 +172,7 @@ class AnnotationScene(QGraphicsScene):
# print "mouseMoveEvent", self.sceneRect().contains(event.scenePos()), event.scenePos()
if self.inserter_ is not None:
# insert mode
self.inserter_.mouseMoveEvent(event, self.root_)
self.inserter_.mouseMoveEvent(event, self.image_item_)
else:
# selection mode
QGraphicsScene.mouseMoveEvent(self, event)
@@ -215,7 +221,7 @@ class AnnotationScene(QGraphicsScene):
if self.inserter_ is not None:
# insert mode
self.inserter_.keyPressEvent(event, self.root_)
self.inserter_.keyPressEvent(event, self.image_item_)
else:
# selection mode
if event.key() == Qt.Key_Delete:
@@ -272,7 +278,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
sloth/gui/labeltool.py
+11 -11
View File
@@ -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()
sloth/items/inserters.py
+13 -13
View File
@@ -17,27 +17,27 @@ class ItemInserter:
def mode(self):
return self.mode_
def mousePressEvent(self, event, index):
def mousePressEvent(self, event, image_item):
event.accept()
def mouseReleaseEvent(self, event, index):
def mouseReleaseEvent(self, event, image_item):
event.accept()
def mouseMoveEvent(self, event, index):
def mouseMoveEvent(self, event, image_item):
event.accept()
def keyPressEvent(self, event, index):
def keyPressEvent(self, event, image_item):
event.ignore()
def allowOutOfSceneEvents(self):
return False
class PointItemInserter(ItemInserter):
def mousePressEvent(self, event, index):
def mousePressEvent(self, event, image_item):
pos = event.scenePos()
ann = {'type': 'point',
'x': pos.x(), 'y': pos.y()}
index.model().addAnnotation(index, ann)
image_item.addAnnotation(ann)
event.accept()
class RectItemInserter(ItemInserter):
@@ -46,7 +46,7 @@ class RectItemInserter(ItemInserter):
self.current_item_ = None
self.init_pos_ = None
def mousePressEvent(self, event, index):
def mousePressEvent(self, event, image_item):
pos = event.scenePos()
item = QGraphicsRectItem(QRectF(pos.x(), pos.y(), 0, 0))
item.setPen(Qt.red)
@@ -55,7 +55,7 @@ class RectItemInserter(ItemInserter):
self.scene().addItem(item)
event.accept()
def mouseMoveEvent(self, event, index):
def mouseMoveEvent(self, event, image_item):
if self.current_item_ is not None:
assert self.init_pos_ is not None
rect = QRectF(self.init_pos_, event.scenePos()).normalized()
@@ -63,7 +63,7 @@ class RectItemInserter(ItemInserter):
event.accept()
def mouseReleaseEvent(self, event, index):
def mouseReleaseEvent(self, event, image_item):
if self.current_item_ is not None:
if self.current_item_.rect().width() > 1 and \
self.current_item_.rect().height() > 1:
@@ -72,7 +72,7 @@ class RectItemInserter(ItemInserter):
'x': rect.x(), 'y': rect.y(),
'width': rect.width(), 'height': rect.height()}
ann.update(self.mode())
index.model().addAnnotation(index, ann)
image_item.addAnnotation(ann)
self.scene().removeItem(self.current_item_)
self.current_item_ = None
self.init_pos_ = None
@@ -94,7 +94,7 @@ class FixedRatioRectItemInserter(RectItemInserter):
self.ratio_ = float(mode.get('_ratio', 1))
RectItemInserter.setMode(self, mode)
def mouseMoveEvent(self, event, index):
def mouseMoveEvent(self, event, image_item):
if self.current_item_ is not None:
new_geometry = QRectF(self.current_item_.rect().topLeft(), event.scenePos())
dx = new_geometry.width()
@@ -115,7 +115,7 @@ class PolygonItemInserter(ItemInserter):
ItemInserter.__init__(self, scene, mode)
self.current_item_ = None
def mousePressEvent(self, event, index):
def mousePressEvent(self, event, image_item):
pos = event.scenePos()
if self.current_item_ is None:
item = QGraphicsPolygonItem(QPolygonF([pos]))
@@ -128,7 +128,7 @@ class PolygonItemInserter(ItemInserter):
event.accept()
def mouseMoveEvent(self, event, index):
def mouseMoveEvent(self, event, image_item):
if self.current_item_ is not None:
pos = event.scenePos()
polygon = self.current_item_.polygon()