Create shape.py

DataAnnotation/labelme/labelme/shape.py

+import copy
+import math
+
+from qtpy import QtCore
+from qtpy import QtGui
+
+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)
+DEFAULT_FILL_COLOR = QtGui.QColor(255, 0, 0, 128)
+DEFAULT_SELECT_LINE_COLOR = QtGui.QColor(255, 255, 255)
+DEFAULT_SELECT_FILL_COLOR = QtGui.QColor(0, 128, 255, 155)
+DEFAULT_VERTEX_FILL_COLOR = QtGui.QColor(0, 255, 0, 255)
+DEFAULT_HVERTEX_FILL_COLOR = QtGui.QColor(255, 0, 0)
+
+
+class Shape(object):
+
+    P_SQUARE, P_ROUND = 0, 1
+
+    MOVE_VERTEX, NEAR_VERTEX = 0, 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):
+        self.label = label
+        self.points = []
+        self.fill = False
+        self.selected = False
+        self.shape_type = shape_type
+        self.flags = flags
+
+        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 popPoint(self):
+        if self.points:
+            return self.points.pop()
+        return None
+
+    def insertPoint(self, i, point):
+        self.points.insert(i, point)
+
+    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 = Shape.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):
+        self._highlightIndex = i
+        self._highlightMode = action
+
+    def highlightClear(self):
+        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