247 lines
9.9 KiB
Python
247 lines
9.9 KiB
Python
# -*- coding: utf-8 -*-
|
||
# @Author : LG
|
||
|
||
from PyQt5 import QtCore, QtWidgets, QtGui
|
||
from sam_ann.annotation import Object
|
||
import typing
|
||
|
||
class Vertex(QtWidgets.QGraphicsPathItem):
|
||
def __init__(self, polygon, color, nohover_size=2):
|
||
super(Vertex, self).__init__()
|
||
self.polygon = polygon
|
||
self.color = color
|
||
|
||
self.nohover_size = nohover_size
|
||
self.hover_size = self.nohover_size + 2
|
||
self.line_width = 0
|
||
|
||
self.nohover = QtGui.QPainterPath()
|
||
self.nohover.addEllipse(QtCore.QRectF(-self.nohover_size//2, -self.nohover_size//2, self.nohover_size, self.nohover_size))
|
||
self.hover = QtGui.QPainterPath()
|
||
self.hover.addEllipse(QtCore.QRectF(-self.hover_size//2, -self.hover_size//2, self.hover_size, self.hover_size))
|
||
|
||
self.setPath(self.nohover)
|
||
self.setBrush(self.color)
|
||
self.setCursor(QtGui.QCursor(QtCore.Qt.CursorShape.PointingHandCursor))
|
||
self.setPen(QtGui.QPen(self.color, self.line_width))
|
||
self.setFlag(QtWidgets.QGraphicsItem.GraphicsItemFlag.ItemIsSelectable, True)
|
||
self.setFlag(QtWidgets.QGraphicsItem.GraphicsItemFlag.ItemIsMovable, True)
|
||
self.setFlag(QtWidgets.QGraphicsItem.GraphicsItemFlag.ItemSendsGeometryChanges, True)
|
||
self.setAcceptHoverEvents(True)
|
||
self.setZValue(1e5)
|
||
|
||
def itemChange(self, change: 'QtWidgets.QGraphicsItem.GraphicsItemChange', value: typing.Any):
|
||
if change == QtWidgets.QGraphicsItem.GraphicsItemChange.ItemSelectedHasChanged:
|
||
self.scene().mainwindow.actionDelete.setEnabled(self.isSelected())
|
||
if change == QtWidgets.QGraphicsItem.GraphicsItemChange.ItemPositionChange and self.isEnabled():
|
||
# 限制顶点移动到图外
|
||
if value.x() < 0:
|
||
value.setX(0)
|
||
if value.x() > self.scene().width()-1:
|
||
value.setX(self.scene().width()-1)
|
||
if value.y() < 0:
|
||
value.setY(0)
|
||
if value.y() > self.scene().height()-1:
|
||
value.setY(self.scene().height()-1)
|
||
index = self.polygon.vertexs.index(self)
|
||
self.polygon.movePoint(index, value)
|
||
|
||
return super(Vertex, self).itemChange(change, value)
|
||
|
||
def hoverEnterEvent(self, event: 'QGraphicsSceneHoverEvent'):
|
||
self.setPath(self.hover)
|
||
super(Vertex, self).hoverEnterEvent(event)
|
||
|
||
def hoverLeaveEvent(self, event: 'QGraphicsSceneHoverEvent'):
|
||
self.setPath(self.nohover)
|
||
super(Vertex, self).hoverLeaveEvent(event)
|
||
|
||
|
||
class Polygon(QtWidgets.QGraphicsPolygonItem):
|
||
def __init__(self):
|
||
super(Polygon, self).__init__(parent=None)
|
||
self.line_width = 0
|
||
self.hover_alpha = 150
|
||
self.nohover_alpha = 80
|
||
self.points = []
|
||
self.vertexs = []
|
||
self.category = ''
|
||
self.group = 0
|
||
self.iscrowd = 0
|
||
self.note = ''
|
||
|
||
self.rxmin, self.rxmax, self.rymin, self.rymax = 0, 0, 0, 0 # 用于绘画完成后，记录多边形的各边界，此处与points对应
|
||
self.color = QtGui.QColor('#ff0000')
|
||
self.is_drawing = True
|
||
|
||
self.setPen(QtGui.QPen(self.color, self.line_width))
|
||
self.setBrush(QtGui.QBrush(self.color, QtCore.Qt.BrushStyle.FDiagPattern))
|
||
|
||
self.setAcceptHoverEvents(True)
|
||
self.setFlag(QtWidgets.QGraphicsItem.GraphicsItemFlag.ItemIsSelectable, True)
|
||
self.setFlag(QtWidgets.QGraphicsItem.GraphicsItemFlag.ItemIsMovable, True)
|
||
self.setFlag(QtWidgets.QGraphicsItem.GraphicsItemFlag.ItemSendsGeometryChanges, True)
|
||
self.setZValue(1e5)
|
||
|
||
def addPoint(self, point):
|
||
self.points.append(point)
|
||
vertex = Vertex(self, self.color, self.scene().mainwindow.cfg['vertex_size'])
|
||
# 添加路径点
|
||
self.scene().addItem(vertex)
|
||
self.vertexs.append(vertex)
|
||
vertex.setPos(point)
|
||
|
||
def movePoint(self, index, point):
|
||
if not 0 <= index < len(self.points):
|
||
return
|
||
self.points[index] = self.mapFromScene(point)
|
||
|
||
self.redraw()
|
||
if self.scene().mainwindow.load_finished and not self.is_drawing:
|
||
self.scene().mainwindow.set_saved_state(False)
|
||
|
||
def removePoint(self, index):
|
||
if not self.points:
|
||
return
|
||
self.points.pop(index)
|
||
vertex = self.vertexs.pop(index)
|
||
self.scene().removeItem(vertex)
|
||
del vertex
|
||
self.redraw()
|
||
|
||
def delete(self):
|
||
self.points.clear()
|
||
while self.vertexs:
|
||
vertex = self.vertexs.pop()
|
||
self.scene().removeItem(vertex)
|
||
del vertex
|
||
|
||
def moveVertex(self, index, point):
|
||
if not 0 <= index < len(self.vertexs):
|
||
return
|
||
vertex = self.vertexs[index]
|
||
vertex.setEnabled(False)
|
||
vertex.setPos(point)
|
||
vertex.setEnabled(True)
|
||
|
||
def itemChange(self, change: 'QGraphicsItem.GraphicsItemChange', value: typing.Any):
|
||
if change == QtWidgets.QGraphicsItem.GraphicsItemChange.ItemSelectedHasChanged and not self.is_drawing: # 选中改变
|
||
if self.isSelected():
|
||
color = QtGui.QColor('#00A0FF')
|
||
color.setAlpha(self.hover_alpha)
|
||
self.setBrush(color)
|
||
else:
|
||
self.color.setAlpha(self.nohover_alpha)
|
||
self.setBrush(self.color)
|
||
self.scene().mainwindow.annos_dock_widget.set_selected(self) # 更新label面板
|
||
|
||
if change == QtWidgets.QGraphicsItem.GraphicsItemChange.ItemPositionChange: # ItemPositionHasChanged
|
||
bias = value
|
||
l, t, b, r = self.boundingRect().left(), self.boundingRect().top(), self.boundingRect().bottom(), self.boundingRect().right()
|
||
if l + bias.x() < 0: bias.setX(-l)
|
||
if r + bias.x() > self.scene().width(): bias.setX(self.scene().width()-r)
|
||
if t + bias.y() < 0: bias.setY(-t)
|
||
if b + bias.y() > self.scene().height(): bias.setY(self.scene().height()-b)
|
||
|
||
for index, point in enumerate(self.points):
|
||
self.moveVertex(index, point+bias)
|
||
|
||
if self.scene().mainwindow.load_finished and not self.is_drawing:
|
||
self.scene().mainwindow.set_saved_state(False)
|
||
|
||
return super(Polygon, self).itemChange(change, value)
|
||
|
||
def hoverEnterEvent(self, event: 'QGraphicsSceneHoverEvent'):
|
||
if not self.is_drawing:
|
||
self.color.setAlpha(self.hover_alpha)
|
||
self.setBrush(self.color)
|
||
super(Polygon, self).hoverEnterEvent(event)
|
||
|
||
def hoverLeaveEvent(self, event: 'QGraphicsSceneHoverEvent'):
|
||
if not self.is_drawing:
|
||
self.color.setAlpha(self.nohover_alpha)
|
||
self.setBrush(self.color)
|
||
super(Polygon, self).hoverEnterEvent(event)
|
||
|
||
def mouseDoubleClickEvent(self, event: 'QGraphicsSceneMouseEvent'):
|
||
if event.button() == QtCore.Qt.MouseButton.LeftButton:
|
||
self.scene().mainwindow.category_edit_widget.polygon = self
|
||
self.scene().mainwindow.category_edit_widget.load_cfg()
|
||
self.scene().mainwindow.category_edit_widget.show()
|
||
|
||
def redraw(self):
|
||
if len(self.points) < 1:
|
||
return
|
||
xs = [p.x() for p in self.points]
|
||
ys = [p.y() for p in self.points]
|
||
self.rxmin, self.rymin, self.rxmax, self.rymax = min(xs), min(ys), max(xs), max(ys)
|
||
self.setPolygon(QtGui.QPolygonF(self.points))
|
||
|
||
def change_color(self, color):
|
||
self.color = color
|
||
self.color.setAlpha(self.nohover_alpha)
|
||
self.setPen(QtGui.QPen(self.color, self.line_width))
|
||
self.setBrush(self.color)
|
||
for vertex in self.vertexs:
|
||
vertex_color = self.color
|
||
vertex_color.setAlpha(255)
|
||
vertex.setPen(QtGui.QPen(vertex_color, self.line_width))
|
||
vertex.setBrush(vertex_color)
|
||
|
||
def set_drawed(self, category, group, iscrowd, note, color:QtGui.QColor, layer=None):
|
||
self.is_drawing = False
|
||
self.category = category
|
||
if isinstance(group, str):
|
||
group = 0 if group == '' else int(group)
|
||
self.group = group
|
||
self.iscrowd = iscrowd
|
||
self.note = note
|
||
|
||
self.color = color
|
||
self.color.setAlpha(self.nohover_alpha)
|
||
self.setPen(QtGui.QPen(self.color, self.line_width))
|
||
self.setBrush(self.color)
|
||
if layer is not None:
|
||
self.setZValue(layer)
|
||
for vertex in self.vertexs:
|
||
vertex_color = self.color
|
||
vertex_color.setAlpha(255)
|
||
vertex.setPen(QtGui.QPen(vertex_color, self.line_width))
|
||
vertex.setBrush(vertex_color)
|
||
if layer is not None:
|
||
vertex.setZValue(layer)
|
||
|
||
def calculate_area(self):
|
||
area = 0
|
||
num_points = len(self.points)
|
||
for i in range(num_points):
|
||
p1 = self.points[i]
|
||
p2 = self.points[(i + 1) % num_points]
|
||
d = p1.x() * p2.y() - p2.x() * p1.y()
|
||
area += d
|
||
return abs(area) / 2
|
||
|
||
def load_object(self, object):
|
||
segmentation = object.segmentation
|
||
for x, y in segmentation:
|
||
point = QtCore.QPointF(x, y)
|
||
self.addPoint(point)
|
||
color = self.scene().mainwindow.category_color_dict.get(object.category, '#000000')
|
||
self.set_drawed(object.category, object.group, object.iscrowd, object.note, QtGui.QColor(color), object.layer) # ...
|
||
|
||
def to_object(self):
|
||
if self.is_drawing:
|
||
return None
|
||
segmentation = []
|
||
for point in self.points:
|
||
point = point + self.pos()
|
||
segmentation.append((round(point.x(), 2), round(point.y(), 2)))
|
||
xmin = self.boundingRect().x() + self.pos().x()
|
||
ymin = self.boundingRect().y() + self.pos().y()
|
||
xmax = xmin + self.boundingRect().width()
|
||
ymax = ymin + self.boundingRect().height()
|
||
|
||
object = Object(self.category, group=self.group, segmentation=segmentation,
|
||
area=self.calculate_area(), layer=self.zValue(), bbox=(xmin, ymin, xmax, ymax), iscrowd=self.iscrowd, note=self.note)
|
||
return object
|