mirror of
https://github.com/wassname/sloth.git
synced 2026-07-06 05:16:56 +08:00
integrate annotation model and scene
This commit is contained in:
Martin Baeuml
+39
-16
@@ -3,6 +3,8 @@ from PyQt4.QtCore import *
|
||||
from functools import partial
|
||||
import os.path
|
||||
|
||||
DataRole = Qt.UserRole + 1
|
||||
|
||||
class ModelItem:
|
||||
def __init__(self, parent=None):
|
||||
self.parent_ = parent
|
||||
@@ -79,40 +81,55 @@ class FrameModelItem(ModelItem):
|
||||
return QVariant()
|
||||
|
||||
class AnnotationModelItem(ModelItem):
|
||||
def __init__(self, annotations, parent):
|
||||
def __init__(self, annotation, parent):
|
||||
ModelItem.__init__(self, parent)
|
||||
self.annotations_ = annotations
|
||||
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
|
||||
|
||||
for key, value in annotations.iteritems():
|
||||
self.children_.append(KeyValueModelItem(key, value, self))
|
||||
for key, value in annotation.iteritems():
|
||||
if key == None:
|
||||
continue
|
||||
self.children_.append(KeyValueModelItem(key, self))
|
||||
|
||||
def type(self):
|
||||
return self.annotations_['type']
|
||||
return self.annotation_['type']
|
||||
|
||||
def setData(self, index, value, role):
|
||||
if role == DataRole:
|
||||
self.annotation_ = value.toPyObject()
|
||||
#print "setData", self.annotation_
|
||||
index.model().dataChanged.emit(index, index.sibling(index.row(), 1))
|
||||
return True
|
||||
return False
|
||||
|
||||
def data(self, index, role):
|
||||
if role == Qt.DisplayRole:
|
||||
if index.column() == 0:
|
||||
return self.type()
|
||||
else:
|
||||
return QVariant()
|
||||
if role == Qt.DisplayRole and index.column() == 0:
|
||||
return self.type()
|
||||
elif role == DataRole:
|
||||
#print "data():", self.annotation_
|
||||
return self.annotation_
|
||||
|
||||
return QVariant()
|
||||
|
||||
def value(self, key):
|
||||
return self.annotation_[key]
|
||||
|
||||
class KeyValueModelItem(ModelItem):
|
||||
def __init__(self, key, value, parent):
|
||||
def __init__(self, key, parent):
|
||||
ModelItem.__init__(self, parent)
|
||||
self.key_ = key
|
||||
self.value_ = value
|
||||
self.key_ = key
|
||||
|
||||
def data(self, index, role):
|
||||
if role == Qt.DisplayRole:
|
||||
if index.column() == 0:
|
||||
return self.key_
|
||||
elif index.column() == 1:
|
||||
return self.value_
|
||||
return self.parent().value(self.key_)
|
||||
else:
|
||||
return QVariant()
|
||||
|
||||
|
||||
|
||||
class AnnotationModel(QAbstractItemModel):
|
||||
def __init__(self, annotations, parent=None):
|
||||
QAbstractItemModel.__init__(self, parent)
|
||||
@@ -192,6 +209,9 @@ class AnnotationModel(QAbstractItemModel):
|
||||
assert row != -1
|
||||
return self.createIndex(row, 0, parent)
|
||||
|
||||
def mapToSource(self, index):
|
||||
return index
|
||||
|
||||
def flags(self, index):
|
||||
return Qt.ItemIsEnabled
|
||||
if not index.isValid():
|
||||
@@ -401,6 +421,9 @@ def someAnnotations():
|
||||
annotations.append({'type': 'point',
|
||||
'x': '30',
|
||||
'y': '30'})
|
||||
annotations.append({'type': 'point',
|
||||
'x': '100',
|
||||
'y': '100'})
|
||||
return annotations
|
||||
|
||||
def defaultAnnotations():
|
||||
|
||||
+179
-52
@@ -113,17 +113,16 @@ class AnnotationScene(QGraphicsScene):
|
||||
|
||||
def __init__(self, parent=None):
|
||||
super(AnnotationScene, self).__init__(parent)
|
||||
|
||||
self.model_ = None
|
||||
self.mode_ = None
|
||||
self.inserters_ = {}
|
||||
self.inserter_ = None
|
||||
self.debug_ = True
|
||||
self.message_ = ""
|
||||
|
||||
self.setBackgroundBrush(Qt.darkGray)
|
||||
|
||||
self.reset()
|
||||
self.setSceneRect(0,0, 640, 480)
|
||||
self.addRect(QRectF(0, 0, 640, 480), brush=Qt.white)
|
||||
|
||||
self.mode_ = None
|
||||
self.inserters_ = {}
|
||||
self.inserter_ = None
|
||||
self.debug_ = True
|
||||
|
||||
self.addItemInserter('point', PointItemInserter(self))
|
||||
self.addItemInserter('rect', RectItemInserter(self))
|
||||
self.addItemInserter('poly', PolygonItemInserter(self))
|
||||
@@ -133,8 +132,78 @@ class AnnotationScene(QGraphicsScene):
|
||||
#self.setMode({'type': 'rect'})
|
||||
self.setMode({'type': 'poly'})
|
||||
|
||||
def reset(self):
|
||||
self.reset()
|
||||
self.setSceneRect(0,0, 640, 480)
|
||||
self.addRect(QRectF(0, 0, 640, 480), brush=Qt.white)
|
||||
|
||||
#
|
||||
# getters/setters
|
||||
#______________________________________________________________________________________________________
|
||||
def model(self):
|
||||
return self.model_
|
||||
|
||||
def setModel(self, model):
|
||||
if model == self.model_:
|
||||
# same model as the current one
|
||||
# reset caches anyway, invalidate root
|
||||
self.reset()
|
||||
return
|
||||
|
||||
# disconnect old signals
|
||||
if self.model_ is not None:
|
||||
self.disconnect(self.model_, SIGNAL('dataChanged(QModelIndex,QModelIndex)'), self.dataChanged)
|
||||
self.disconnect(self.model_, SIGNAL('rowsInserted(QModelIndex,int,int)'), self.rowsInserted)
|
||||
self.disconnect(self.model_, SIGNAL('rowsAboutToBeRemoved(QModelIndex,int,int)'), self.rowsAboutToBeRemoved)
|
||||
self.disconnect(self.model_, SIGNAL('rowsRemoved(QModelIndex,int,int)'), self.rowsRemoved)
|
||||
self.disconnect(self.model_, SIGNAL('modelReset()'), self.reset)
|
||||
self.model_ = model
|
||||
|
||||
# connect new signals
|
||||
if self.model_ is not None:
|
||||
self.connect(self.model_, SIGNAL('dataChanged(QModelIndex,QModelIndex)'), self.dataChanged)
|
||||
self.connect(self.model_, SIGNAL('rowsInserted(QModelIndex,int,int)'), self.rowsInserted)
|
||||
self.connect(self.model_, SIGNAL('rowsAboutToBeRemoved(QModelIndex,int,int)'), self.rowsAboutToBeRemoved)
|
||||
self.connect(self.model_, SIGNAL('rowsRemoved(QModelIndex,int,int)'), self.rowsRemoved)
|
||||
self.connect(self.model_, SIGNAL('modelReset()'), self.reset)
|
||||
|
||||
# reset caches, invalidate root
|
||||
self.reset()
|
||||
|
||||
def root(self):
|
||||
return self.root_
|
||||
|
||||
def setRoot(self, root):
|
||||
"""
|
||||
Set the index of the model which denotes the current image to be
|
||||
displayed by the scene. This can be either the index to a frame in a
|
||||
video, or to an image.
|
||||
"""
|
||||
self.root_ = root
|
||||
self.clear()
|
||||
if not root.isValid():
|
||||
return
|
||||
|
||||
assert self.root_.model() == self.model_
|
||||
self.setSceneRect(0, 0, 640, 480)
|
||||
self.addRect(QRectF(0, 0, 640, 480))
|
||||
|
||||
num_items = self.model_.rowCount(self.root_)
|
||||
self.insertItems(0, num_items)
|
||||
|
||||
def insertItems(self, first, last):
|
||||
assert self.model_ is not None
|
||||
assert self.root_.isValid()
|
||||
|
||||
for row in range(first, last+1):
|
||||
child = self.root_.child(row, 0)
|
||||
item = AnnotationGraphicsItem.createItem(child)
|
||||
if item is not None:
|
||||
#checked = (child.data(Qt.CheckStateRole).toInt()[0] == Qt.Checked)
|
||||
#item.setVisible(checked)
|
||||
self.addItem(item)
|
||||
|
||||
def mode(self):
|
||||
return self.mode_
|
||||
|
||||
def setMode(self, mode):
|
||||
print "setMode :", mode
|
||||
@@ -150,50 +219,108 @@ class AnnotationScene(QGraphicsScene):
|
||||
self.inserter_ = self.inserters_[self.mode_['type']]
|
||||
self.inserter_.setMode(self.mode_)
|
||||
|
||||
def addItemInserter(self, type, inserter):
|
||||
if type in self.inserters_:
|
||||
raise Exception("Type %s already has an inserter" % type)
|
||||
|
||||
self.inserters_[type] = inserter
|
||||
|
||||
def removeItemInserter(self, type):
|
||||
if type in self.inserters_:
|
||||
del self.inserters_[type]
|
||||
|
||||
def mousePressEvent(self, event):
|
||||
if self.debug_:
|
||||
print "mousePressEvent", self.sceneRect().contains(event.scenePos()), event.scenePos()
|
||||
if not self.sceneRect().contains(event.scenePos()):
|
||||
# ignore events outside the scene rect
|
||||
return
|
||||
elif self.inserter_ is not None:
|
||||
# insert mode
|
||||
self.inserter_.mousePressEvent(event, None)
|
||||
else:
|
||||
# selection mode
|
||||
QGraphicsScene.mousePressEvent(self, event)
|
||||
|
||||
def mouseReleaseEvent(self, event):
|
||||
if self.debug_:
|
||||
print "mouseReleaseEvent", self.sceneRect().contains(event.scenePos()), event.scenePos()
|
||||
if self.inserter_ is not None:
|
||||
# insert mode
|
||||
self.inserter_.mouseReleaseEvent(event, None)
|
||||
else:
|
||||
# selection mode
|
||||
QGraphicsScene.mouseReleaseEvent(self, event)
|
||||
|
||||
def mouseMoveEvent(self, event):
|
||||
if self.debug_:
|
||||
print "mouseMoveEvent", self.sceneRect().contains(event.scenePos()), event.scenePos()
|
||||
if self.inserter_ is not None:
|
||||
# insert mode
|
||||
self.inserter_.mouseMoveEvent(event, None)
|
||||
else:
|
||||
# selection mode
|
||||
QGraphicsScene.mouseMoveEvent(self, event)
|
||||
#
|
||||
# common methods
|
||||
#______________________________________________________________________________________________________
|
||||
def reset(self):
|
||||
self.clear()
|
||||
self.setRoot(QModelIndex())
|
||||
self.clearMessage()
|
||||
|
||||
def addItem(self, item):
|
||||
QGraphicsScene.addItem(self, item)
|
||||
# TODO emit signal itemAdded
|
||||
|
||||
def addItemInserter(self, type, inserter, replace=False):
|
||||
type = type.lower()
|
||||
if type in self.inserters_ and not replace:
|
||||
raise Exception("Type %s already has an inserter" % type)
|
||||
|
||||
self.inserters_[type] = inserter
|
||||
|
||||
def removeItemInserter(self, type):
|
||||
type = type.lower()
|
||||
if type in self.inserters_:
|
||||
del self.inserters_[type]
|
||||
|
||||
#
|
||||
# slots for signals from the model
|
||||
# this is the implemenation of the scene as a view of the model
|
||||
#______________________________________________________________________________________________________
|
||||
def dataChanged(self, indexFrom, indexTo):
|
||||
if self.root_ != indexFrom.parent() or self.root_ != indexTo.parent():
|
||||
return
|
||||
|
||||
for row in range(indexFrom.row(), indexTo.row()+1):
|
||||
item = self.itemFromIndex(indexFrom.sibling(row, 0))
|
||||
if item is not None:
|
||||
item.dataChanged()
|
||||
|
||||
def rowsInserted(self, index, first, last):
|
||||
if self.root_ != index:
|
||||
return
|
||||
|
||||
self.insertItems(first, last)
|
||||
|
||||
|
||||
def rowsAboutToBeRemoved(self, index, first, last):
|
||||
if self.root_ != index:
|
||||
return
|
||||
|
||||
for row in range(first, last+1):
|
||||
item = self.itemFromIndex(index.child(row, 0))
|
||||
if item is not None:
|
||||
self.removeItem(item)
|
||||
|
||||
def rowsRemoved(self, index, first, last):
|
||||
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
|
||||
if hasattr(item, 'index') and item.index() == index:
|
||||
return item
|
||||
return None
|
||||
|
||||
#
|
||||
# message handling and displaying
|
||||
#______________________________________________________________________________________________________
|
||||
def setMessage(self, message):
|
||||
if self.message_ is not None:
|
||||
self.clearMessage()
|
||||
|
||||
if message is None or message == "":
|
||||
return
|
||||
|
||||
# TODO don't use text item at all, just draw the text in drawForeground
|
||||
self.message_ = message
|
||||
self.message_text_item_ = QGraphicsSimpleTextItem(message)
|
||||
self.message_text_item_.setPos(20, 20)
|
||||
self.invalidate(QRectF(), QGraphicsScene.ForegroundLayer)
|
||||
|
||||
def clearMessage(self):
|
||||
if self.message_ is not None:
|
||||
self.message_text_item_ = None
|
||||
self.message_ = None
|
||||
self.invalidate(QRectF(), QGraphicsScene.ForegroundLayer)
|
||||
|
||||
def drawForeground(self, painter, rect):
|
||||
QGraphicsScene.drawForeground(self, painter, rect)
|
||||
|
||||
if self.message_ is not None:
|
||||
assert self.message_text_item_ is not None
|
||||
|
||||
painter.setTransform(QTransform())
|
||||
painter.setBrush(QColor('lightGray'))
|
||||
painter.setPen(QPen(QBrush(QColor('black')), 2))
|
||||
|
||||
br = self.message_text_item_.boundingRect()
|
||||
|
||||
painter.drawRoundedRect(QRectF(10, 10, br.width()+20, br.height()+20), 10.0, 10.0)
|
||||
painter.setTransform(QTransform.fromTranslate(20, 20))
|
||||
painter.setPen(QPen(QColor('black'), 1))
|
||||
|
||||
self.message_text_item_.paint(painter, QStyleOptionGraphicsItem(), None)
|
||||
|
||||
|
||||
+19
-9
@@ -3,6 +3,7 @@ import sys, os
|
||||
from PyQt4.QtGui import *
|
||||
from PyQt4.QtCore import *
|
||||
from annotationscene import *
|
||||
from annotationmodel import *
|
||||
|
||||
class MainWindow(QMainWindow):
|
||||
def __init__(self, argv, parent=None):
|
||||
@@ -15,9 +16,6 @@ class MainWindow(QMainWindow):
|
||||
buttonPoint = QPushButton("Point")
|
||||
buttonPoint.clicked.connect(self.clickedPoint)
|
||||
vlayout.addWidget(buttonPoint)
|
||||
buttonLine = QPushButton("Line")
|
||||
buttonLine.clicked.connect(self.clickedLine)
|
||||
vlayout.addWidget(buttonLine)
|
||||
buttonRectangle = QPushButton("Rectangle")
|
||||
buttonRectangle.clicked.connect(self.clickedRectangle)
|
||||
vlayout.addWidget(buttonRectangle)
|
||||
@@ -28,9 +26,11 @@ class MainWindow(QMainWindow):
|
||||
|
||||
hlayout = QHBoxLayout()
|
||||
self.view_ = QGraphicsView()
|
||||
self.annotree_ = AnnotationTreeView()
|
||||
hlayout.addLayout(vlayout)
|
||||
hlayout.addWidget(self.view_, 1)
|
||||
|
||||
hlayout.addWidget(self.annotree_)
|
||||
|
||||
self.scene_ = AnnotationScene(self)
|
||||
self.view_.setScene(self.scene_)
|
||||
|
||||
@@ -38,27 +38,37 @@ class MainWindow(QMainWindow):
|
||||
central.setLayout(hlayout)
|
||||
self.setCentralWidget(central)
|
||||
|
||||
def setModel(self, model):
|
||||
self.annotree_.setModel(model)
|
||||
self.scene_.setModel(model)
|
||||
|
||||
file_index = model.index(0, 0, QModelIndex())
|
||||
frame_index = model.index(0, 0, file_index)
|
||||
self.scene_.setRoot(frame_index)
|
||||
|
||||
def clickedSelect(self):
|
||||
self.scene_.setMode(None)
|
||||
|
||||
def clickedPoint(self):
|
||||
self.scene_.setMode({'type': 'point'})
|
||||
self.scene_.setMode({'type': 'point'})
|
||||
|
||||
def clickedRectangle(self):
|
||||
self.scene_.setMode({'type': 'rect'})
|
||||
|
||||
def clickedLine(self):
|
||||
self.scene_.setMode(LINE)
|
||||
|
||||
def clickedPolygon(self):
|
||||
self.scene_.setMode({'type': 'polygon'})
|
||||
|
||||
|
||||
def main():
|
||||
import sys
|
||||
app = QApplication(sys.argv)
|
||||
|
||||
annotations = defaultAnnotations()
|
||||
model = AnnotationModel(annotations)
|
||||
|
||||
wnd = MainWindow(sys.argv[1:])
|
||||
wnd.resize(800,600)
|
||||
wnd.resize(800, 600)
|
||||
wnd.setModel(model)
|
||||
wnd.show()
|
||||
|
||||
return app.exec_()
|
||||
|
||||
Reference in New Issue
Block a user