Support rectangle annotation

labelme/app.py

@@ -262,6 +262,10 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         createMode = action('Create\nPolygo&ns', self.setCreateMode,
                             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,
                           shortcuts['edit_polygon'], 'edit',
                           'Move and edit polygons', enabled=True)
@@ -362,6 +366,7 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
             undoLastPoint=undoLastPoint, undo=undo,
             addPoint=addPoint,
             createMode=createMode, editMode=editMode,
+            createRectangleMode=createRectangleMode,
             shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
             zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
             fitWindow=fitWindow, fitWidth=fitWidth,
@@ -371,11 +376,12 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
             editMenu=(edit, copy, delete, None, undo, undoLastPoint,
                       None, color1, color2),
             menu=(
-                createMode, editMode, edit, copy,
+                createMode, createRectangleMode,
+                editMode, edit, copy,
                 delete, shapeLineColor, shapeFillColor,
                 undo, undoLastPoint, addPoint,
             ),
-            onLoadActive=(close, createMode, editMode),
+            onLoadActive=(close, createMode, createRectangleMode, editMode),
             onShapesPresent=(saveAs, hideAll, showAll),
         )
 
@@ -410,7 +416,8 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         self.tools = self.toolbar('Tools')
         self.actions.tool = (
             open_, opendir, openNextImg, openPrevImg, save,
-            None, createMode, copy, delete, editMode, undo, None,
+            None, createMode, createRectangleMode,
+            copy, delete, editMode, undo, None,
             zoomIn, zoom, zoomOut, fitWindow, fitWidth)
 
         self.statusBar().showMessage('%s started.' % __appname__)
@@ -486,7 +493,11 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         self.canvas.menus[0].clear()
         addActions(self.canvas.menus[0], menu)
         self.menus.edit.clear()
-        actions = (self.actions.createMode, self.actions.editMode)
+        actions = (
+            self.actions.createMode,
+            self.actions.createRectangleMode,
+            self.actions.editMode,
+        )
         addActions(self.menus.edit, actions + self.actions.editMenu)
 
     def setDirty(self):
@@ -506,6 +517,7 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         self.dirty = False
         self.actions.save.setEnabled(False)
         self.actions.createMode.setEnabled(True)
+        self.actions.createRectangleMode.setEnabled(True)
         title = __appname__
         if self.filename is not None:
             title = '{} - {}'.format(title, self.filename)
@@ -571,13 +583,24 @@ class MainWindow(QtWidgets.QMainWindow, WindowMixin):
         self.actions.undoLastPoint.setEnabled(drawing)
         self.actions.undo.setEnabled(not drawing)
 
-    def toggleDrawMode(self, edit=True):
+    def toggleDrawMode(self, edit=True, createMode='polygon'):
         self.canvas.setEditing(edit)
-        self.actions.createMode.setEnabled(edit)
+        self.canvas.createMode = createMode
+        if createMode == 'polygon':
+            self.actions.createMode.setEnabled(edit)
+            self.actions.createRectangleMode.setEnabled(not edit)
+        elif createMode == 'rectangle':
+            self.actions.createMode.setEnabled(not edit)
+            self.actions.createRectangleMode.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)
+        self.toggleDrawMode(False, createMode='polygon')
 
     def setEditMode(self):
         self.toggleDrawMode(True)

labelme/canvas.py

@@ -30,6 +30,9 @@ class Canvas(QtWidgets.QWidget):
 
     CREATE, EDIT = 0, 1
 
+    # polygon, rectangle
+    createMode = 'polygon'
+
     epsilon = 11.0
 
     def __init__(self, *args, **kwargs):
@@ -42,6 +45,9 @@ 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 = Shape(line_color=self.lineColor)
         self.prevPoint = QtCore.QPoint()
         self.prevMovePoint = QtCore.QPoint()
@@ -131,26 +137,35 @@ class Canvas(QtWidgets.QWidget):
         # Polygon drawing.
         if self.drawing():
             self.overrideCursor(CURSOR_DRAW)
-            if self.current:
-                color = self.lineColor
-                if self.outOfPixmap(pos):
-                    # Don't allow the user to draw outside the pixmap.
-                    # Project the point to the pixmap's edges.
-                    pos = self.intersectionPoint(self.current[-1], pos)
-                elif len(self.current) > 1 and \
-                        self.closeEnough(pos, self.current[0]):
-                    # Attract line to starting point and
-                    # colorise to alert the user.
-                    pos = self.current[0]
-                    color = self.current.line_color
-                    self.overrideCursor(CURSOR_POINT)
-                    self.current.highlightVertex(0, Shape.NEAR_VERTEX)
+            if not self.current:
+                return
+
+            color = self.lineColor
+            if self.outOfPixmap(pos):
+                # Don't allow the user to draw outside the pixmap.
+                # Project the point to the pixmap's edges.
+                pos = self.intersectionPoint(self.current[-1], pos)
+            elif len(self.current) > 1 and \
+                    self.closeEnough(pos, self.current[0]):
+                # Attract line to starting point and
+                # colorise to alert the user.
+                pos = self.current[0]
+                color = self.current.line_color
+                self.overrideCursor(CURSOR_POINT)
+                self.current.highlightVertex(0, Shape.NEAR_VERTEX)
+            if self.createMode == 'polygon':
                 self.line[0] = self.current[-1]
                 self.line[1] = pos
-                self.line.line_color = color
-                self.repaint()
-                self.current.highlightClear()
-            return
+            elif self.createMode == 'rectangle':
+                self.line.points = list(self.getRectangleFromLine(
+                    (self.current[0], pos)
+                ))
+                self.line.close()
+            else:
+                raise ValueError
+            self.line.line_color = color
+            self.repaint()
+            self.current.highlightClear()
 
         # Polygon copy moving.
         if QtCore.Qt.RightButton & ev.buttons():
@@ -232,6 +247,13 @@ class Canvas(QtWidgets.QWidget):
         self.hVertex = index
         self.hEdge = None
 
+    def getRectangleFromLine(self, line):
+        pt1 = line[0]
+        pt3 = line[1]
+        pt2 = QtCore.QPoint(pt3.x(), pt1.y())
+        pt4 = QtCore.QPoint(pt1.x(), pt3.y())
+        return pt1, pt2, pt3, pt4
+
     def mousePressEvent(self, ev):
         if QT5:
             pos = self.transformPos(ev.pos())
@@ -240,11 +262,20 @@ class Canvas(QtWidgets.QWidget):
         if ev.button() == QtCore.Qt.LeftButton:
             if self.drawing():
                 if self.current:
-                    self.current.addPoint(self.line[1])
-                    self.line[0] = self.current[-1]
-                    if self.current.isClosed():
+                    # Add point to existing shape.
+                    if self.createMode == 'polygon':
+                        self.current.addPoint(self.line[1])
+                        self.line[0] = self.current[-1]
+                        if self.current.isClosed():
+                            self.finalise()
+                    elif self.createMode == 'rectangle':
+                        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()
                     self.current.addPoint(pos)
                     self.line.points = [pos, pos]
@@ -582,7 +613,12 @@ class Canvas(QtWidgets.QWidget):
         assert self.shapes
         self.current = self.shapes.pop()
         self.current.setOpen()
-        self.line.points = [self.current[-1], self.current[0]]
+        if self.createMode == 'polygon':
+            self.line.points = [self.current[-1], self.current[0]]
+        elif self.createMode == 'rectangle':
+            self.current.points = self.current.points[0:1]
+        else:
+            raise ValueError
         self.drawingPolygon.emit(True)
 
     def undoLastPoint(self):

labelme/config/default_config.yaml

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