Skip to content

Commit

Permalink
Optimize label\shape.py
Browse files Browse the repository at this point in the history
  • Loading branch information
@PairZhu
PairZhu committed Jun 18, 2024
1 parent aca3e64 commit d5a7002
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 45 deletions.
91 changes: 54 additions & 37 deletions labelme/shape.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import copy
import math

import numpy as np
import skimage.measure
Expand All @@ -26,6 +25,8 @@ class Shape(object):
# Flag for all other handles on the current shape
NEAR_VERTEX = 1

PEN_WIDTH = 2

# The following class variables influence the drawing of all shape objects.
line_color = None
fill_color = None
Expand All @@ -35,7 +36,7 @@ class Shape(object):
hvertex_fill_color = None
point_type = P_ROUND
point_size = 8
scale = 1.0
scale = 1

def __init__(
self,
Expand Down Expand Up @@ -77,6 +78,9 @@ def __init__(
# is used for drawing the pending line a different color.
self.line_color = line_color

def scaleTranslate(self, point):
return QtCore.QPointF(point.x() * self.scale, point.y() * self.scale)

def setShapeRefined(self, shape_type, points, point_labels, mask=None):
self._shape_raw = (self.shape_type, self.points, self.point_labels)
self.shape_type = shape_type
Expand Down Expand Up @@ -168,19 +172,17 @@ def isClosed(self):
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.mask is None and not self.points:
return

painter.save()
painter.scale(1 / self.scale, 1 / self.scale)

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))))
pen.setWidth(self.PEN_WIDTH)
painter.setPen(pen)

if self.mask is not None:
Expand All @@ -192,19 +194,29 @@ def paint(self, painter):
)
image_to_draw[self.mask] = fill_color
qimage = QtGui.QImage.fromData(labelme.utils.img_arr_to_data(image_to_draw))
qimage = qimage.scaled(
qimage.size() * self.scale,
QtCore.Qt.IgnoreAspectRatio,
QtCore.Qt.SmoothTransformation,
)

painter.drawImage(
int(round(self.points[0].x())),
int(round(self.points[0].y())),
# Position error will increase with the scaling factor, so it needs to be corrected
self.scaleTranslate(self.points[0] + QtCore.QPointF(0.5, 0.5)),
qimage,
)

line_path = QtGui.QPainterPath()
contours = skimage.measure.find_contours(np.pad(self.mask, pad_width=1))
for contour in contours:
contour += [self.points[0].y(), self.points[0].x()]
line_path.moveTo(contour[0, 1], contour[0, 0])
line_path.moveTo(
self.scaleTranslate(QtCore.QPointF(contour[0, 1], contour[0, 0]))
)
for point in contour[1:]:
line_path.lineTo(point[1], point[0])
line_path.lineTo(
self.scaleTranslate(QtCore.QPointF(point[1], point[0]))
)
painter.drawPath(line_path)

if self.points:
Expand All @@ -215,22 +227,29 @@ def paint(self, painter):
if self.shape_type in ["rectangle", "mask"]:
assert len(self.points) in [1, 2]
if len(self.points) == 2:
rectangle = self.getRectFromLine(*self.points)
rectangle = QtCore.QRectF(
self.scaleTranslate(self.points[0]),
self.scaleTranslate(self.points[1]),
)
line_path.addRect(rectangle)
if self.shape_type == "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)
raidus = labelme.utils.distance(
self.scaleTranslate(self.points[0] - self.points[1])
)
line_path.addEllipse(
self.scaleTranslate(self.points[0]), raidus, raidus
)
for i in range(len(self.points)):
self.drawVertex(vrtx_path, i)
elif self.shape_type == "linestrip":
line_path.moveTo(self.points[0])
line_path.moveTo(self.scaleTranslate(self.points[0]))
for i, p in enumerate(self.points):
line_path.lineTo(p)
line_path.lineTo(self.scaleTranslate(p))
self.drawVertex(vrtx_path, i)
elif self.shape_type == "points":
assert len(self.points) == len(self.point_labels)
Expand All @@ -240,17 +259,17 @@ def paint(self, painter):
else:
self.drawVertex(negative_vrtx_path, i)
else:
line_path.moveTo(self.points[0])
line_path.moveTo(self.scaleTranslate(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)
line_path.lineTo(self.scaleTranslate(p))
self.drawVertex(vrtx_path, i)
if self.isClosed():
line_path.lineTo(self.points[0])
line_path.lineTo(self.scaleTranslate(self.points[0]))

painter.drawPath(line_path)
if vrtx_path.length() > 0:
Expand All @@ -265,10 +284,12 @@ def paint(self, painter):
painter.drawPath(negative_vrtx_path)
painter.fillPath(negative_vrtx_path, QtGui.QColor(255, 0, 0, 255))

painter.restore()

def drawVertex(self, path, i):
d = self.point_size / self.scale
d = self.point_size
shape = self.point_type
point = self.points[i]
point = self.scaleTranslate(self.points[i])
if i == self._highlightIndex:
size, shape = self._highlightSettings[self._highlightMode]
d *= size
Expand All @@ -286,7 +307,9 @@ def drawVertex(self, path, i):
def nearestVertex(self, point, epsilon):
min_distance = float("inf")
min_i = None
point = QtCore.QPointF(point.x() * self.scale, point.y() * self.scale)
for i, p in enumerate(self.points):
p = QtCore.QPointF(p.x() * self.scale, p.y() * self.scale)
dist = labelme.utils.distance(p - point)
if dist <= epsilon and dist < min_distance:
min_distance = dist
Expand All @@ -296,8 +319,13 @@ def nearestVertex(self, point, epsilon):
def nearestEdge(self, point, epsilon):
min_distance = float("inf")
post_i = None
point = QtCore.QPointF(point.x() * self.scale, point.y() * self.scale)
for i in range(len(self.points)):
line = [self.points[i - 1], self.points[i]]
start = self.points[i - 1]
end = self.points[i]
start = QtCore.QPointF(start.x() * self.scale, start.y() * self.scale)
end = QtCore.QPointF(end.x() * self.scale, end.y() * self.scale)
line = [start, end]
dist = labelme.utils.distancetoline(point, line)
if dist <= epsilon and dist < min_distance:
min_distance = dist
Expand All @@ -319,27 +347,16 @@ def containsPoint(self, point):
return self.mask[y, x]
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 in ["rectangle", "mask"]:
path = QtGui.QPainterPath()
if len(self.points) == 2:
rectangle = self.getRectFromLine(*self.points)
path.addRect(rectangle)
path.addRect(QtCore.QRectF(self.points[0], self.points[1]))
elif self.shape_type == "circle":
path = QtGui.QPainterPath()
if len(self.points) == 2:
rectangle = self.getCircleRectFromLine(self.points)
path.addEllipse(rectangle)
raidus = labelme.utils.distance(self.points[0] - self.points[1])
path.addEllipse(self.points[0], raidus, raidus)
else:
path = QtGui.QPainterPath(self.points[0])
for p in self.points[1:]:
Expand Down
20 changes: 12 additions & 8 deletions labelme/widgets/canvas.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -331,8 +331,8 @@ def mouseMoveEvent(self, ev):
for shape in reversed([s for s in self.shapes if self.isVisible(s)]):
# Look for a nearby vertex to highlight. If that fails,
# check if we happen to be inside a shape.
index = shape.nearestVertex(pos, self.epsilon / self.scale)
index_edge = shape.nearestEdge(pos, self.epsilon / self.scale)
index = shape.nearestVertex(pos, self.epsilon)
index_edge = shape.nearestEdge(pos, self.epsilon)
if index is not None:
if self.selectedVertex():
self.hShape.highlightClear()
Expand Down Expand Up @@ -439,9 +439,11 @@ def mousePressEvent(self, ev):
elif not self.outOfPixmap(pos):
# Create new shape.
self.current = Shape(
shape_type="points"
if self.createMode in ["ai_polygon", "ai_mask"]
else self.createMode
shape_type=(
"points"
if self.createMode in ["ai_polygon", "ai_mask"]
else self.createMode
)
)
self.current.addPoint(pos, label=0 if is_shift_pressed else 1)
if self.createMode == "point":
Expand Down Expand Up @@ -933,9 +935,11 @@ def wheelEvent(self, ev):
else:
self.scrollRequest.emit(
ev.delta(),
QtCore.Qt.Horizontal
if (QtCore.Qt.ShiftModifier == int(mods))
else QtCore.Qt.Vertical,
(
QtCore.Qt.Horizontal
if (QtCore.Qt.ShiftModifier == int(mods))
else QtCore.Qt.Vertical
),
)
else:
self.scrollRequest.emit(ev.delta(), QtCore.Qt.Horizontal)
Expand Down

0 comments on commit d5a7002

Please sign in to comment.