Prechádzať zdrojové kódy

Add create line mode

Kentaro Wada 6 rokov pred
rodič
commit
abadf01870
5 zmenil súbory, kde vykonal 67 pridanie a 35 odobranie
  1. 44 18
      labelme/app.py
  2. 21 13
      labelme/canvas.py
  3. 1 0
      labelme/config/default_config.yaml
  4. 1 0
      labelme/lib.py
  5. 0 4
      labelme/shape.py

+ 44 - 18
labelme/app.py

@@ -261,21 +261,38 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
                         shortcuts['edit_fill_color'], 'color',
                         'Choose polygon fill color')
 
-        createMode = action('Create\nPolygo&ns', self.setCreateMode,
-                            shortcuts['create_polygon'], 'objects',
-                            'Start drawing polygons', enabled=True)
+        createMode = action(
+            'Create Polygons',
+            lambda: self.toggleDrawMode(False, createMode='polygon'),
+            shortcuts['create_polygon'],
+            'objects',
+            'Start drawing polygons',
+            enabled=True,
+        )
         createRectangleMode = action(
-            'Create\nRectangle', self.setCreateRectangleMode,
-            shortcuts['create_rectangle'], 'objects',
-            'Start drawing rectangles', enabled=True)
-        editMode = action('&Edit\nPolygons', self.setEditMode,
+            'Create Rectangle',
+            lambda: self.toggleDrawMode(False, createMode='rectangle'),
+            shortcuts['create_rectangle'],
+            'objects',
+            'Start drawing rectangles',
+            enabled=True,
+        )
+        createLineMode = action(
+            'Create Line',
+            lambda: self.toggleDrawMode(False, createMode='line'),
+            shortcuts['create_line'],
+            'objects',
+            'Start drawing lines',
+            enabled=True,
+        )
+        editMode = action('Edit Polygons', self.setEditMode,
                           shortcuts['edit_polygon'], 'edit',
                           'Move and edit polygons', enabled=True)
 
-        delete = action('Delete\nPolygon', self.deleteSelectedShape,
+        delete = action('Delete Polygon', self.deleteSelectedShape,
                         shortcuts['delete_polygon'], 'cancel',
                         'Delete', enabled=False)
-        copy = action('&Duplicate\nPolygon', self.copySelectedShape,
+        copy = action('Duplicate Polygon', self.copySelectedShape,
                       shortcuts['duplicate_polygon'], 'copy',
                       'Create a duplicate of the selected polygon',
                       enabled=False)
@@ -369,6 +386,7 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
             addPoint=addPoint,
             createMode=createMode, editMode=editMode,
             createRectangleMode=createRectangleMode,
+            createLineMode=createLineMode,
             shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
             zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
             fitWindow=fitWindow, fitWidth=fitWidth,
@@ -377,13 +395,20 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
             tool=(),
             editMenu=(edit, copy, delete, None, undo, undoLastPoint,
                       None, color1, color2),
+            # menu shown at right click
             menu=(
-                createMode, createRectangleMode,
+                createMode, createRectangleMode, createLineMode,
                 editMode, edit, copy,
                 delete, shapeLineColor, shapeFillColor,
                 undo, undoLastPoint, addPoint,
             ),
-            onLoadActive=(close, createMode, createRectangleMode, editMode),
+            onLoadActive=(
+                close,
+                createMode,
+                createRectangleMode,
+                createLineMode,
+                editMode,
+            ),
             onShapesPresent=(saveAs, hideAll, showAll),
         )
 
@@ -426,7 +451,6 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
             save,
             None,
             createMode,
-            createRectangleMode,
             editMode,
             copy,
             delete,
@@ -515,6 +539,7 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         actions = (
             self.actions.createMode,
             self.actions.createRectangleMode,
+            self.actions.createLineMode,
             self.actions.editMode,
         )
         addActions(self.menus.edit, actions + self.actions.editMenu)
@@ -537,6 +562,7 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         self.actions.save.setEnabled(False)
         self.actions.createMode.setEnabled(True)
         self.actions.createRectangleMode.setEnabled(True)
+        self.actions.createLineMode.setEnabled(True)
         title = __appname__
         if self.filename is not None:
             title = '{} - {}'.format(title, self.filename)
@@ -608,19 +634,19 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         if createMode == 'polygon':
             self.actions.createMode.setEnabled(edit)
             self.actions.createRectangleMode.setEnabled(not edit)
+            self.actions.createLineMode.setEnabled(not edit)
         elif createMode == 'rectangle':
             self.actions.createMode.setEnabled(not edit)
             self.actions.createRectangleMode.setEnabled(edit)
+            self.actions.createLineMode.setEnabled(not edit)
+        elif createMode == 'line':
+            self.actions.createMode.setEnabled(not edit)
+            self.actions.createRectangleMode.setEnabled(not edit)
+            self.actions.createLineMode.setEnabled(edit)
         else:
             raise ValueError
         self.actions.editMode.setEnabled(not edit)
 
-    def setCreateRectangleMode(self):
-        self.toggleDrawMode(False, createMode='rectangle')
-
-    def setCreateMode(self):
-        self.toggleDrawMode(False, createMode='polygon')
-
     def setEditMode(self):
         self.toggleDrawMode(True)
 

+ 21 - 13
labelme/canvas.py

@@ -30,8 +30,8 @@ class Canvas(QtWidgets.QWidget):
 
     CREATE, EDIT = 0, 1
 
-    # polygon, rectangle
-    createMode = 'polygon'
+    # polygon, rectangle, or line
+    _createMode = 'polygon'
 
     def __init__(self, *args, **kwargs):
         self.epsilon = kwargs.pop('epsilon', 11.0)
@@ -44,9 +44,10 @@ class Canvas(QtWidgets.QWidget):
         self.selectedShape = None  # save the selected shape here
         self.selectedShapeCopy = None
         self.lineColor = QtGui.QColor(0, 0, 255)
-        # self.line represents
-        # if createMode == 'polygon': edge from last point to current
-        # if createMode == 'rectangle': the rectangle
+        # self.line represents:
+        #   - createMode == 'polygon': edge from last point to current
+        #   - createMode == 'rectangle': diagonal line of the rectangle
+        #   - createMode == 'line': the line
         self.line = Shape(line_color=self.lineColor)
         self.prevPoint = QtCore.QPoint()
         self.prevMovePoint = QtCore.QPoint()
@@ -68,6 +69,16 @@ class Canvas(QtWidgets.QWidget):
         self.setMouseTracking(True)
         self.setFocusPolicy(QtCore.Qt.WheelFocus)
 
+    @property
+    def createMode(self):
+        return self._createMode
+
+    @createMode.setter
+    def createMode(self, value):
+        if value not in ['polygon', 'rectangle', 'line']:
+            raise ValueError('Unsupported createMode: %s' % value)
+        self._createMode = value
+
     def storeShapes(self):
         shapesBackup = []
         for shape in self.shapes:
@@ -160,8 +171,9 @@ class Canvas(QtWidgets.QWidget):
                     (self.current[0], pos)
                 ))
                 self.line.close()
-            else:
-                raise ValueError
+            elif self.createMode == 'line':
+                self.line.points = [self.current[0], pos]
+                self.line.close()
             self.line.line_color = color
             self.repaint()
             self.current.highlightClear()
@@ -268,12 +280,10 @@ class Canvas(QtWidgets.QWidget):
                         self.line[0] = self.current[-1]
                         if self.current.isClosed():
                             self.finalise()
-                    elif self.createMode == 'rectangle':
+                    elif self.createMode in ['rectangle', 'line']:
                         assert len(self.current.points) == 1
                         self.current.points = self.line.points
                         self.finalise()
-                    else:
-                        raise ValueError
                 elif not self.outOfPixmap(pos):
                     # Create new shape.
                     self.current = Shape()
@@ -615,10 +625,8 @@ class Canvas(QtWidgets.QWidget):
         self.current.setOpen()
         if self.createMode == 'polygon':
             self.line.points = [self.current[-1], self.current[0]]
-        elif self.createMode == 'rectangle':
+        elif self.createMode in ['rectangle', 'line']:
             self.current.points = self.current.points[0:1]
-        else:
-            raise ValueError
         self.drawingPolygon.emit(True)
 
     def undoLastPoint(self):

+ 1 - 0
labelme/config/default_config.yaml

@@ -30,6 +30,7 @@ shortcuts:
   add_point: Ctrl+Shift+P
   create_polygon: Ctrl+N
   create_rectangle: Ctrl+R
+  create_line: null
   edit_polygon: Ctrl+J
   delete_polygon: Delete
   duplicate_polygon: Ctrl+D

+ 1 - 0
labelme/lib.py

@@ -30,6 +30,7 @@ def newAction(parent, text, slot=None, shortcut=None, icon=None,
     """Create a new action and assign callbacks, shortcuts, etc."""
     a = QtWidgets.QAction(text, parent)
     if icon is not None:
+        a.setIconText(text.replace(' ', '\n'))
         a.setIcon(newIcon(icon))
     if shortcut is not None:
         if isinstance(shortcut, (list, tuple)):

+ 0 - 4
labelme/shape.py

@@ -2,7 +2,6 @@ from qtpy import QtGui
 
 from labelme.lib import distance
 from labelme.lib import distancetoline
-from labelme import logger
 
 
 # TODO(unknown):
@@ -56,9 +55,6 @@ class Shape(object):
             self.line_color = line_color
 
     def close(self):
-        if len(self.points) <= 2:
-            logger.error('Polygon should be created with points >2')
-            return
         self._closed = True
 
     def addPoint(self, point):