import copy
import math
from qtpy import QtCore
from qtpy import QtGui
from labelme.logger import logger
import labelme.utils
# TODO(unknown):
# - [opt] Store paths instead of creating new ones at each paint.
DEFAULT_LINE_COLOR = QtGui.QColor(0, 255, 0, 128) # bf hovering
DEFAULT_FILL_COLOR = QtGui.QColor(0, 255, 0, 128) # hovering
DEFAULT_SELECT_LINE_COLOR = QtGui.QColor(255, 255, 255) # selected
DEFAULT_SELECT_FILL_COLOR = QtGui.QColor(0, 255, 0, 155) # selected
DEFAULT_VERTEX_FILL_COLOR = QtGui.QColor(0, 255, 0, 255) # hovering
DEFAULT_HVERTEX_FILL_COLOR = QtGui.QColor(255, 255, 255, 255) # hovering
class Shape(object):
# Render handles as squares
P_SQUARE = 0
# Render handles as circles
P_ROUND = 1
# Flag for the handles we would move if dragging
MOVE_VERTEX = 0
# Flag for all other handles on the curent shape
NEAR_VERTEX = 1
# The following class variables influence the drawing of all shape objects.
line_color = DEFAULT_LINE_COLOR
fill_color = DEFAULT_FILL_COLOR
select_line_color = DEFAULT_SELECT_LINE_COLOR
select_fill_color = DEFAULT_SELECT_FILL_COLOR
vertex_fill_color = DEFAULT_VERTEX_FILL_COLOR
hvertex_fill_color = DEFAULT_HVERTEX_FILL_COLOR
point_type = P_ROUND
point_size = 8
scale = 1.0
def __init__(
self,
label=None,
line_color=None,
shape_type=None,
flags=None,
group_id=None,
):
self.label = label
self.group_id = group_id
self.points = []
self.fill = False
self.selected = False
self.shape_type = shape_type
self.flags = flags
self.other_data = {}
self._highlightIndex = None
self._highlightMode = self.NEAR_VERTEX
self._highlightSettings = {
self.NEAR_VERTEX: (4, self.P_ROUND),
self.MOVE_VERTEX: (1.5, self.P_SQUARE),
}
self._closed = False
if line_color is not None:
# Override the class line_color attribute
# with an object attribute. Currently this
# is used for drawing the pending line a different color.
self.line_color = line_color
self.shape_type = shape_type
@property
def shape_type(self):
return self._shape_type
@shape_type.setter
def shape_type(self, value):
if value is None:
value = "polygon"
if value not in [
"polygon",
"rectangle",
"point",
"line",
"circle",
"linestrip",
]:
raise ValueError("Unexpected shape_type: {}".format(value))
self._shape_type = value
def close(self):
self._closed = True
def addPoint(self, point):
if self.points and point == self.points[0]:
self.close()
else:
self.points.append(point)
def canAddPoint(self):
return self.shape_type in ["polygon", "linestrip"]
def popPoint(self):
if self.points:
return self.points.pop()
return None
def insertPoint(self, i, point):
self.points.insert(i, point)
def removePoint(self, i):
if not self.canAddPoint():
logger.warning(
"Cannot remove point from: shape_type=%r",
self.shape_type,
)
return
if self.shape_type == "polygon" and len(self.points) <= 3:
logger.warning(
"Cannot remove point from: shape_type=%r, len(points)=%d",
self.shape_type,
len(self.points),
)
return
if self.shape_type == "linestrip" and len(self.points) <= 2:
logger.warning(
"Cannot remove point from: shape_type=%r, len(points)=%d",
self.shape_type,
len(self.points),
)
return
self.points.pop(i)
def isClosed(self):
return self._closed
def setOpen(self):
self._closed = False
def getRectFromLine(self, pt1, pt2):
x1, y1 = pt1.x(), pt1.y()
x2, y2 = pt2.x(), pt2.y()
return QtCore.QRectF(x1, y1, x2 - x1, y2 - y1)
def paint(self, painter):
if self.points:
color = (
self.select_line_color if self.selected else self.line_color
)
pen = QtGui.QPen(color)
# Try using integer sizes for smoother drawing(?)
pen.setWidth(max(1, int(round(2.0 / self.scale))))
painter.setPen(pen)
line_path = QtGui.QPainterPath()
vrtx_path = QtGui.QPainterPath()
if self.shape_type == "rectangle":
assert len(self.points) in [1, 2]
if len(self.points) == 2:
rectangle = self.getRectFromLine(*self.points)
line_path.addRect(rectangle)
for i in range(len(self.points)):
self.drawVertex(vrtx_path, i)
elif self.shape_type == "circle":
assert len(self.points) in [1, 2]
if len(self.points) == 2:
rectangle = self.getCircleRectFromLine(self.points)
line_path.addEllipse(rectangle)
for i in range(len(self.points)):
self.drawVertex(vrtx_path, i)
elif self.shape_type == "linestrip":
line_path.moveTo(self.points[0])
for i, p in enumerate(self.points):
line_path.lineTo(p)
self.drawVertex(vrtx_path, i)
else:
line_path.moveTo(self.points[0])
# Uncommenting the following line will draw 2 paths
# for the 1st vertex, and make it non-filled, which
# may be desirable.
# self.drawVertex(vrtx_path, 0)
for i, p in enumerate(self.points):
line_path.lineTo(p)
self.drawVertex(vrtx_path, i)
if self.isClosed():
line_path.lineTo(self.points[0])
painter.drawPath(line_path)
painter.drawPath(vrtx_path)
painter.fillPath(vrtx_path, self._vertex_fill_color)
if self.fill:
color = (
self.select_fill_color
if self.selected
else self.fill_color
)
painter.fillPath(line_path, color)
def drawVertex(self, path, i):
d = self.point_size / self.scale
shape = self.point_type
point = self.points[i]
if i == self._highlightIndex:
size, shape = self._highlightSettings[self._highlightMode]
d *= size
if self._highlightIndex is not None:
self._vertex_fill_color = self.hvertex_fill_color
else:
self._vertex_fill_color = self.vertex_fill_color
if shape == self.P_SQUARE:
path.addRect(point.x() - d / 2, point.y() - d / 2, d, d)
elif shape == self.P_ROUND:
path.addEllipse(point, d / 2.0, d / 2.0)
else:
assert False, "unsupported vertex shape"
def nearestVertex(self, point, epsilon):
min_distance = float("inf")
min_i = None
for i, p in enumerate(self.points):
dist = labelme.utils.distance(p - point)
if dist <= epsilon and dist < min_distance:
min_distance = dist
min_i = i
return min_i
def nearestEdge(self, point, epsilon):
min_distance = float("inf")
post_i = None
for i in range(len(self.points)):
line = [self.points[i - 1], self.points[i]]
dist = labelme.utils.distancetoline(point, line)
if dist <= epsilon and dist < min_distance:
min_distance = dist
post_i = i
return post_i
def containsPoint(self, point):
return self.makePath().contains(point)
def getCircleRectFromLine(self, line):
"""Computes parameters to draw with `QPainterPath::addEllipse`"""
if len(line) != 2:
return None
(c, point) = line
r = line[0] - line[1]
d = math.sqrt(math.pow(r.x(), 2) + math.pow(r.y(), 2))
rectangle = QtCore.QRectF(c.x() - d, c.y() - d, 2 * d, 2 * d)
return rectangle
def makePath(self):
if self.shape_type == "rectangle":
path = QtGui.QPainterPath()
if len(self.points) == 2:
rectangle = self.getRectFromLine(*self.points)
path.addRect(rectangle)
elif self.shape_type == "circle":
path = QtGui.QPainterPath()
if len(self.points) == 2:
rectangle = self.getCircleRectFromLine(self.points)
path.addEllipse(rectangle)
else:
path = QtGui.QPainterPath(self.points[0])
for p in self.points[1:]:
path.lineTo(p)
return path
def boundingRect(self):
return self.makePath().boundingRect()
def moveBy(self, offset):
self.points = [p + offset for p in self.points]
def moveVertexBy(self, i, offset):
self.points[i] = self.points[i] + offset
def highlightVertex(self, i, action):
"""Highlight a vertex appropriately based on the current action
Args:
i (int): The vertex index
action (int): The action
(see Shape.NEAR_VERTEX and Shape.MOVE_VERTEX)
"""
self._highlightIndex = i
self._highlightMode = action
def highlightClear(self):
"""Clear the highlighted point"""
self._highlightIndex = None
def copy(self):
return copy.deepcopy(self)
def __len__(self):
return len(self.points)
def __getitem__(self, key):
return self.points[key]
def __setitem__(self, key, value):
self.points[key] = value