shape.py 8.5 KB
Newer Older
1
import copy
T
Tatiana Malygina 已提交
2
import math
3

4
from qtpy import QtCore
5
from qtpy import QtGui
K
Kentaro Wada 已提交
6

K
Kentaro Wada 已提交
7
import labelme.utils
M
Michael Pitidis 已提交
8

K
Kentaro Wada 已提交
9 10

# TODO(unknown):
11 12
# - [opt] Store paths instead of creating new ones at each paint.

K
Kentaro Wada 已提交
13

14 15 16 17 18 19 20
R, G, B = SHAPE_COLOR = 0, 255, 0  # green
DEFAULT_LINE_COLOR = QtGui.QColor(R, G, B, 128)                # bf hovering
DEFAULT_FILL_COLOR = QtGui.QColor(R, G, B, 128)                # hovering
DEFAULT_SELECT_LINE_COLOR = QtGui.QColor(255, 255, 255)        # selected
DEFAULT_SELECT_FILL_COLOR = QtGui.QColor(R, G, B, 155)         # selected
DEFAULT_VERTEX_FILL_COLOR = QtGui.QColor(R, G, B, 255)         # hovering
DEFAULT_HVERTEX_FILL_COLOR = QtGui.QColor(255, 255, 255, 255)  # hovering
K
Kentaro Wada 已提交
21

22

23
class Shape(object):
K
Kentaro Wada 已提交
24

K
Kentaro Wada 已提交
25
    P_SQUARE, P_ROUND = 0, 1
M
Michael Pitidis 已提交
26

K
Kentaro Wada 已提交
27
    MOVE_VERTEX, NEAR_VERTEX = 0, 1
28

K
Kentaro Wada 已提交
29
    # The following class variables influence the drawing of all shape objects.
30 31
    line_color = DEFAULT_LINE_COLOR
    fill_color = DEFAULT_FILL_COLOR
32 33 34
    select_line_color = DEFAULT_SELECT_LINE_COLOR
    select_fill_color = DEFAULT_SELECT_FILL_COLOR
    vertex_fill_color = DEFAULT_VERTEX_FILL_COLOR
35
    hvertex_fill_color = DEFAULT_HVERTEX_FILL_COLOR
36
    point_type = P_ROUND
M
Michael Pitidis 已提交
37
    point_size = 8
38
    scale = 1.0
M
Michael Pitidis 已提交
39

40
    def __init__(self, label=None, line_color=None, shape_type=None,
K
Kentaro Wada 已提交
41
                 flags=None, group_id=None):
42
        self.label = label
K
Kentaro Wada 已提交
43
        self.group_id = group_id
M
Michael Pitidis 已提交
44
        self.points = []
45
        self.fill = False
M
Michael Pitidis 已提交
46
        self.selected = False
T
Tatiana Malygina 已提交
47
        self.shape_type = shape_type
C
cmerchant 已提交
48
        self.flags = flags
49 50 51 52 53

        self._highlightIndex = None
        self._highlightMode = self.NEAR_VERTEX
        self._highlightSettings = {
            self.NEAR_VERTEX: (4, self.P_ROUND),
54
            self.MOVE_VERTEX: (1.5, self.P_SQUARE),
K
Kentaro Wada 已提交
55
        }
56 57 58

        self._closed = False

59 60 61 62 63
        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
64

65 66 67 68 69 70 71 72 73 74
        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'
S
Shohei Fujii 已提交
75 76
        if value not in ['polygon', 'rectangle', 'point',
           'line', 'circle', 'linestrip']:
77 78 79
            raise ValueError('Unexpected shape_type: {}'.format(value))
        self._shape_type = value

80 81 82
    def close(self):
        self._closed = True

M
Michael Pitidis 已提交
83
    def addPoint(self, point):
84
        if self.points and point == self.points[0]:
85 86 87
            self.close()
        else:
            self.points.append(point)
M
Michael Pitidis 已提交
88

89 90 91
    def canAddPoint(self):
        return self.shape_type in ['polygon', 'linestrip']

M
Michael Pitidis 已提交
92 93 94
    def popPoint(self):
        if self.points:
            return self.points.pop()
95
        return None
M
Michael Pitidis 已提交
96

A
AlexMarshall011 已提交
97 98 99
    def insertPoint(self, i, point):
        self.points.insert(i, point)

100 101 102
    def removePoint(self, i):
        self.points.pop(i)

M
Michael Pitidis 已提交
103
    def isClosed(self):
104
        return self._closed
M
Michael Pitidis 已提交
105

106 107 108
    def setOpen(self):
        self._closed = False

109 110 111 112 113
    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)

114
    def paint(self, painter):
M
Michael Pitidis 已提交
115
        if self.points:
K
Kentaro Wada 已提交
116 117 118
            color = self.select_line_color \
                if self.selected else self.line_color
            pen = QtGui.QPen(color)
119 120
            # Try using integer sizes for smoother drawing(?)
            pen.setWidth(max(1, int(round(2.0 / self.scale))))
M
Michael Pitidis 已提交
121
            painter.setPen(pen)
H
Hussein 已提交
122

K
Kentaro Wada 已提交
123 124
            line_path = QtGui.QPainterPath()
            vrtx_path = QtGui.QPainterPath()
H
Hussein 已提交
125

126 127 128 129 130 131 132
            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)
T
Tatiana Malygina 已提交
133 134 135 136 137 138 139
            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)
S
Shohei Fujii 已提交
140 141 142 143 144
            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)
145 146 147 148 149 150 151 152 153 154 155 156
            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])
157

M
Michael Pitidis 已提交
158
            painter.drawPath(line_path)
159
            painter.drawPath(vrtx_path)
160
            painter.fillPath(vrtx_path, self._vertex_fill_color)
161
            if self.fill:
K
Kentaro Wada 已提交
162 163
                color = self.select_fill_color \
                    if self.selected else self.fill_color
164
                painter.fillPath(line_path, color)
M
Michael Pitidis 已提交
165

166
    def drawVertex(self, path, i):
167
        d = self.point_size / self.scale
168 169 170 171 172
        shape = self.point_type
        point = self.points[i]
        if i == self._highlightIndex:
            size, shape = self._highlightSettings[self._highlightMode]
            d *= size
173
        if self._highlightIndex is not None:
174
            self._vertex_fill_color = self.hvertex_fill_color
175
        else:
176
            self._vertex_fill_color = self.vertex_fill_color
177
        if shape == self.P_SQUARE:
K
Kentaro Wada 已提交
178
            path.addRect(point.x() - d / 2, point.y() - d / 2, d, d)
179
        elif shape == self.P_ROUND:
K
Kentaro Wada 已提交
180
            path.addEllipse(point, d / 2.0, d / 2.0)
181 182 183 184
        else:
            assert False, "unsupported vertex shape"

    def nearestVertex(self, point, epsilon):
S
serycjon 已提交
185 186
        min_distance = float('inf')
        min_i = None
187
        for i, p in enumerate(self.points):
K
Kentaro Wada 已提交
188
            dist = labelme.utils.distance(p - point)
S
serycjon 已提交
189 190 191 192
            if dist <= epsilon and dist < min_distance:
                min_distance = dist
                min_i = i
        return min_i
M
Michael Pitidis 已提交
193

A
AlexMarshall011 已提交
194 195 196 197 198
    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]]
K
Kentaro Wada 已提交
199
            dist = labelme.utils.distancetoline(point, line)
A
AlexMarshall011 已提交
200 201 202 203 204
            if dist <= epsilon and dist < min_distance:
                min_distance = dist
                post_i = i
        return post_i

M
Michael Pitidis 已提交
205
    def containsPoint(self, point):
206 207
        return self.makePath().contains(point)

T
Tatiana Malygina 已提交
208 209 210 211 212 213 214 215 216 217
    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

218
    def makePath(self):
219 220 221 222 223
        if self.shape_type == 'rectangle':
            path = QtGui.QPainterPath()
            if len(self.points) == 2:
                rectangle = self.getRectFromLine(*self.points)
                path.addRect(rectangle)
T
Tatiana Malygina 已提交
224 225 226 227 228
        elif self.shape_type == "circle":
            path = QtGui.QPainterPath()
            if len(self.points) == 2:
                rectangle = self.getCircleRectFromLine(self.points)
                path.addEllipse(rectangle)
229 230 231 232
        else:
            path = QtGui.QPainterPath(self.points[0])
            for p in self.points[1:]:
                path.lineTo(p)
233 234 235 236 237 238 239 240
        return path

    def boundingRect(self):
        return self.makePath().boundingRect()

    def moveBy(self, offset):
        self.points = [p + offset for p in self.points]

241 242 243 244 245 246 247 248 249 250
    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

H
Hussein 已提交
251
    def copy(self):
K
Kentaro Wada 已提交
252
        return copy.deepcopy(self)
253

254
    def __len__(self):
M
Michael Pitidis 已提交
255
        return len(self.points)
256 257

    def __getitem__(self, key):
M
Michael Pitidis 已提交
258 259 260 261
        return self.points[key]

    def __setitem__(self, key, value):
        self.points[key] = value