mirror of
https://github.com/wassname/sloth.git
synced 2026-07-03 17:49:49 +08:00
60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
from items import ItemInserter, AnnotationGraphicsItem
|
|
from PyQt4.QtGui import *
|
|
from PyQt4.Qt import *
|
|
|
|
class BodyBoundingboxItemInserter(ItemInserter):
|
|
def __init__(self, scene, mode=None):
|
|
ItemInserter.__init__(self, scene, mode)
|
|
self.point_items_ = []
|
|
self.points_ = []
|
|
|
|
def mousePressEvent(self, event, index):
|
|
pos = event.scenePos()
|
|
self.points_.append((pos.x(), pos.y()))
|
|
item = QGraphicsEllipseItem(pos.x()-2, pos.y()-2, 4, 4)
|
|
item.setPen(Qt.red)
|
|
self.point_items_.append(item)
|
|
self.scene().addItem(item)
|
|
|
|
event.accept()
|
|
if self.complete(index):
|
|
self.clear()
|
|
|
|
def complete(self, index):
|
|
if len(self.points_) == 4:
|
|
ann = {'type': 'bodybbox',
|
|
'x1': self.points_[0][0], 'y1': self.points_[0][1],
|
|
'x2': self.points_[1][0], 'y2': self.points_[1][1],
|
|
'x3': self.points_[2][0], 'y3': self.points_[2][1],
|
|
'x4': self.points_[3][0], 'y4': self.points_[3][1]}
|
|
index.model().addAnnotation(index, ann)
|
|
return True
|
|
return False
|
|
|
|
def clear(self):
|
|
for item in self.point_items_:
|
|
self.scene().removeItem(item)
|
|
self.point_items_ = []
|
|
self.points_ = []
|
|
|
|
class BodyBoundingboxItem(AnnotationGraphicsItem):
|
|
def __init__(self, index, parent=None):
|
|
AnnotationGraphicsItem.__init__(self, index, parent)
|
|
|
|
self.data_ = self.index().data(DataRole).toPyObject()
|
|
self.points_ = [(self.data_['x1'], self.data_['y1']),
|
|
(self.data_['x2'], self.data_['y2']),
|
|
(self.data_['x3'], self.data_['y3']),
|
|
(self.data_['x4'], self.data_['y4'])]
|
|
|
|
def boundingRect(self):
|
|
return QRectF(QPointF(0, 0), QSizeF(20, 20))
|
|
|
|
def paint(self, painter, option, widget = None):
|
|
pen = self.pen()
|
|
if self.isSelected():
|
|
pen.setStyle(Qt.DashLine)
|
|
painter.setPen(pen)
|
|
painter.drawRect(self.boundingRect())
|
|
|