Browse Source

Add circle shapes support

Tatiana Malygina 6 năm trước cách đây
mục cha
commit
bd747d869f
5 tập tin đã thay đổi với 72 bổ sung8 xóa
  1. 24 0
      labelme/app.py
  2. 1 0
      labelme/config/default_config.yaml
  3. 25 1
      labelme/shape.py
  4. 14 4
      labelme/utils/shape.py
  5. 8 3
      labelme/widgets/canvas.py

+ 24 - 0
labelme/app.py

@@ -229,6 +229,14 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
             'Start drawing rectangles',
             enabled=True,
         )
+        createCircleMode = action(
+            'Create Circle',
+            lambda: self.toggleDrawMode(False, createMode='circle'),
+            shortcuts['create_circle'],
+            'objects',
+            'Start drawing circles',
+            enabled=True,
+        )
         createLineMode = action(
             'Create Line',
             lambda: self.toggleDrawMode(False, createMode='line'),
@@ -352,6 +360,7 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
             addPoint=addPoint,
             createMode=createMode, editMode=editMode,
             createRectangleMode=createRectangleMode,
+            createCircleMode=createCircleMode,
             createLineMode=createLineMode,
             createPointMode=createPointMode,
             shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
@@ -366,6 +375,7 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
             menu=(
                 createMode,
                 createRectangleMode,
+                createCircleMode,
                 createLineMode,
                 createPointMode,
                 editMode,
@@ -382,6 +392,7 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
                 close,
                 createMode,
                 createRectangleMode,
+                createCircleMode,
                 createLineMode,
                 createPointMode,
                 editMode,
@@ -534,6 +545,7 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         actions = (
             self.actions.createMode,
             self.actions.createRectangleMode,
+            self.actions.createCircleMode,
             self.actions.createLineMode,
             self.actions.createPointMode,
             self.actions.editMode,
@@ -558,6 +570,7 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         self.actions.save.setEnabled(False)
         self.actions.createMode.setEnabled(True)
         self.actions.createRectangleMode.setEnabled(True)
+        self.actions.createCircleMode.setEnabled(True)
         self.actions.createLineMode.setEnabled(True)
         self.actions.createPointMode.setEnabled(True)
         title = __appname__
@@ -630,29 +643,40 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         if edit:
             self.actions.createMode.setEnabled(True)
             self.actions.createRectangleMode.setEnabled(True)
+            self.actions.createCircleMode.setEnabled(True)
             self.actions.createLineMode.setEnabled(True)
             self.actions.createPointMode.setEnabled(True)
         else:
             if createMode == 'polygon':
                 self.actions.createMode.setEnabled(False)
                 self.actions.createRectangleMode.setEnabled(True)
+                self.actions.createCircleMode.setEnabled(True)
                 self.actions.createLineMode.setEnabled(True)
                 self.actions.createPointMode.setEnabled(True)
             elif createMode == 'rectangle':
                 self.actions.createMode.setEnabled(True)
                 self.actions.createRectangleMode.setEnabled(False)
+                self.actions.createCircleMode.setEnabled(True)
                 self.actions.createLineMode.setEnabled(True)
                 self.actions.createPointMode.setEnabled(True)
             elif createMode == 'line':
                 self.actions.createMode.setEnabled(True)
                 self.actions.createRectangleMode.setEnabled(True)
+                self.actions.createCircleMode.setEnabled(True)
                 self.actions.createLineMode.setEnabled(False)
                 self.actions.createPointMode.setEnabled(True)
             elif createMode == 'point':
                 self.actions.createMode.setEnabled(True)
                 self.actions.createRectangleMode.setEnabled(True)
+                self.actions.createCircleMode.setEnabled(True)
                 self.actions.createLineMode.setEnabled(True)
                 self.actions.createPointMode.setEnabled(False)
+            elif createMode == "circle":
+                self.actions.createMode.setEnabled(True)
+                self.actions.createRectangleMode.setEnabled(True)
+                self.actions.createCircleMode.setEnabled(False)
+                self.actions.createLineMode.setEnabled(True)
+                self.actions.createPointMode.setEnabled(True)
             else:
                 raise ValueError('Unsupported createMode: %s' % createMode)
         self.actions.editMode.setEnabled(not edit)

+ 1 - 0
labelme/config/default_config.yaml

@@ -59,6 +59,7 @@ shortcuts:
   add_point: Ctrl+Shift+P
   create_polygon: Ctrl+N
   create_rectangle: Ctrl+R
+  create_circle: null
   create_line: null
   create_point: null
   edit_polygon: Ctrl+J

+ 25 - 1
labelme/shape.py

@@ -1,4 +1,5 @@
 import copy
+import math
 
 from qtpy import QtCore
 from qtpy import QtGui
@@ -40,6 +41,7 @@ class Shape(object):
         self.points = []
         self.fill = False
         self.selected = False
+        self.shape_type = shape_type
 
         self._highlightIndex = None
         self._highlightMode = self.NEAR_VERTEX
@@ -66,7 +68,7 @@ class Shape(object):
     def shape_type(self, value):
         if value is None:
             value = 'polygon'
-        if value not in ['polygon', 'rectangle', 'point', 'line']:
+        if value not in ['polygon', 'rectangle', 'point', 'line', 'circle']:
             raise ValueError('Unexpected shape_type: {}'.format(value))
         self._shape_type = value
 
@@ -117,6 +119,13 @@ class Shape(object):
                     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)
             else:
                 line_path.moveTo(self.points[0])
                 # Uncommenting the following line will draw 2 paths
@@ -180,12 +189,27 @@ class Shape(object):
     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:]:

+ 14 - 4
labelme/utils/shape.py

@@ -1,15 +1,24 @@
+from math import pow
+from math import sqrt
 import numpy as np
+
 import PIL.Image
 import PIL.ImageDraw
 
 from labelme import logger
 
 
-def polygons_to_mask(img_shape, polygons):
+def polygons_to_mask(img_shape, polygons, shape_type=None):
     mask = np.zeros(img_shape[:2], dtype=np.uint8)
     mask = PIL.Image.fromarray(mask)
-    xy = list(map(tuple, polygons))
-    PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
+    draw = PIL.ImageDraw.Draw(mask)
+    if shape_type == "circle" and len(polygons) == 2:
+        ((cx, cy), (px, py)) = polygons
+        d = sqrt(pow(cx - px, 2) + pow(cy - py, 2))
+        draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1)
+    else:
+        xy = list(map(tuple, polygons))
+        draw.polygon(xy=xy, outline=1, fill=1)
     mask = np.array(mask, dtype=bool)
     return mask
 
@@ -24,6 +33,7 @@ def shapes_to_label(img_shape, shapes, label_name_to_value, type='class'):
     for shape in shapes:
         polygons = shape['points']
         label = shape['label']
+        shape_type = shape.get('shape_type', None)
         if type == 'class':
             cls_name = label
         elif type == 'instance':
@@ -32,7 +42,7 @@ def shapes_to_label(img_shape, shapes, label_name_to_value, type='class'):
                 instance_names.append(label)
             ins_id = len(instance_names) - 1
         cls_id = label_name_to_value[cls_name]
-        mask = polygons_to_mask(img_shape[:2], polygons)
+        mask = polygons_to_mask(img_shape[:2], polygons, shape_type)
         cls[mask] = cls_id
         if type == 'instance':
             ins[mask] = ins_id

+ 8 - 3
labelme/widgets/canvas.py

@@ -84,7 +84,7 @@ class Canvas(QtWidgets.QWidget):
 
     @createMode.setter
     def createMode(self, value):
-        if value not in ['polygon', 'rectangle', 'line', 'point']:
+        if value not in ['polygon', 'rectangle', 'circle', 'line', 'point']:
             raise ValueError('Unsupported createMode: %s' % value)
         self._createMode = value
 
@@ -180,6 +180,9 @@ class Canvas(QtWidgets.QWidget):
             elif self.createMode == 'rectangle':
                 self.line.points = [self.current[0], pos]
                 self.line.close()
+            elif self.createMode == 'circle':
+                self.line.points = [self.current[0], pos]
+                self.line.shape_type = "circle"
             elif self.createMode == 'line':
                 self.line.points = [self.current[0], pos]
                 self.line.close()
@@ -285,7 +288,7 @@ class Canvas(QtWidgets.QWidget):
                         self.line[0] = self.current[-1]
                         if self.current.isClosed():
                             self.finalise()
-                    elif self.createMode in ['rectangle', 'line']:
+                    elif self.createMode in ['rectangle', 'circle', 'line']:
                         assert len(self.current.points) == 1
                         self.current.points = self.line.points
                         self.finalise()
@@ -296,6 +299,8 @@ class Canvas(QtWidgets.QWidget):
                     if self.createMode == 'point':
                         self.finalise()
                     else:
+                        if self.createMode == 'circle':
+                            self.current.shape_type = 'circle'
                         self.line.points = [pos, pos]
                         self.setHiding()
                         self.drawingPolygon.emit(True)
@@ -641,7 +646,7 @@ class Canvas(QtWidgets.QWidget):
         self.current.setOpen()
         if self.createMode == 'polygon':
             self.line.points = [self.current[-1], self.current[0]]
-        elif self.createMode in ['rectangle', 'line']:
+        elif self.createMode in ['rectangle', 'line', 'circle']:
             self.current.points = self.current.points[0:1]
         elif self.createMode == 'point':
             self.current = None