From 4e278bb5110ab37c8d3d6a40d3feb45af97dbbbd Mon Sep 17 00:00:00 2001 From: Antoine Broyelle Date: Tue, 2 Oct 2018 22:01:36 +0200 Subject: [PATCH] feature: draw square bounding boxes --- labelImg.py | 26 ++++++++++++++++++++++++-- libs/canvas.py | 30 ++++++++++++++++++++++++++++-- libs/constants.py | 1 + 3 files changed, 53 insertions(+), 4 deletions(-) diff --git a/labelImg.py b/labelImg.py index 8e83a2fe3..018e86e7d 100755 --- a/labelImg.py +++ b/labelImg.py @@ -184,6 +184,7 @@ def __init__(self, defaultFilename=None, defaultPrefdefClassFile=None, defaultSa self.canvas = Canvas(parent=self) self.canvas.zoomRequest.connect(self.zoomRequest) + self.canvas.setDrawingShapeToSquare(settings.get(SETTING_DRAW_SQUARE, False)) scroll = QScrollArea() scroll.setWidget(self.canvas) @@ -332,6 +333,13 @@ def __init__(self, defaultFilename=None, defaultPrefdefClassFile=None, defaultSa self.labelList.customContextMenuRequested.connect( self.popLabelListMenu) + # Draw squares/rectangles + self.drawSquaresOption = QAction('Draw Squares', self) + self.drawSquaresOption.setShortcut('Ctrl+Shift+R') + self.drawSquaresOption.setCheckable(True) + self.drawSquaresOption.setChecked(settings.get(SETTING_DRAW_SQUARE, False)) + self.drawSquaresOption.triggered.connect(self.toogleDrawSquare) + # Store actions for further handling. self.actions = struct(save=save, save_format=save_format, saveAs=saveAs, open=open, close=close, resetAll = resetAll, lineColor=color1, create=create, delete=delete, edit=edit, copy=copy, @@ -344,7 +352,7 @@ def __init__(self, defaultFilename=None, defaultPrefdefClassFile=None, defaultSa open, opendir, save, saveAs, close, resetAll, quit), beginner=(), advanced=(), editMenu=(edit, copy, delete, - None, color1), + None, color1, self.drawSquaresOption), beginnerContext=(create, edit, copy, delete), advancedContext=(createMode, editMode, edit, copy, delete, shapeLineColor, shapeFillColor), @@ -480,6 +488,15 @@ def xbool(x): if self.filePath and os.path.isdir(self.filePath): self.openDirDialog(dirpath=self.filePath) + def keyReleaseEvent(self, event): + if event.key() == Qt.Key_Control: + self.canvas.setDrawingShapeToSquare(False) + + def keyPressEvent(self, event): + if event.key() == Qt.Key_Control: + # Draw rectangle if Ctrl is pressed + self.canvas.setDrawingShapeToSquare(True) + ## Support Functions ## def set_format(self, save_format): if save_format == FORMAT_PASCALVOC: @@ -1099,6 +1116,7 @@ def closeEvent(self, event): settings[SETTING_AUTO_SAVE] = self.autoSaving.isChecked() settings[SETTING_SINGLE_CLASS] = self.singleClassMode.isChecked() settings[SETTING_PAINT_LABEL] = self.paintLabelsOption.isChecked() + settings[SETTING_DRAW_SQUARE] = self.drawSquaresOption.isChecked() settings.save() ## User Dialogs ## @@ -1171,6 +1189,7 @@ def importDirImages(self, dirpath): if not self.mayContinue() or not dirpath: return + self.lastOpenDir = dirpath self.dirname = dirpath self.filePath = None @@ -1183,7 +1202,7 @@ def importDirImages(self, dirpath): def verifyImg(self, _value=False): # Proceding next image without dialog if having any label - if self.filePath is not None: + if self.filePath is not None: try: self.labelFile.toggleVerify() except AttributeError: @@ -1416,6 +1435,9 @@ def togglePaintLabelsOption(self): for shape in self.canvas.shapes: shape.paintLabel = paintLabelsOptionChecked + def toogleDrawSquare(self): + self.canvas.setDrawingShapeToSquare(self.drawSquaresOption.isChecked()) + def inverted(color): return QColor(*[255 - v for v in color.getRgb()]) diff --git a/libs/canvas.py b/libs/canvas.py index 6e79af5c3..d026fb9f9 100644 --- a/libs/canvas.py +++ b/libs/canvas.py @@ -61,6 +61,7 @@ def __init__(self, *args, **kwargs): self.setMouseTracking(True) self.setFocusPolicy(Qt.WheelFocus) self.verified = False + self.drawSquare = False def setDrawingColor(self, qColor): self.drawingLineColor = qColor @@ -126,7 +127,18 @@ def mouseMoveEvent(self, ev): color = self.current.line_color self.overrideCursor(CURSOR_POINT) self.current.highlightVertex(0, Shape.NEAR_VERTEX) - self.line[1] = pos + + if self.drawSquare: + initPos = self.current[0] + minX = initPos.x() + minY = initPos.y() + min_size = min(abs(pos.x() - minX), abs(pos.y() - minY)) + directionX = -1 if pos.x() - minX < 0 else 1 + directionY = -1 if pos.y() - minY < 0 else 1 + self.line[1] = QPointF(minX + directionX * min_size, minY + directionY * min_size) + else: + self.line[1] = pos + self.line.line_color = color self.prevPoint = QPointF() self.current.highlightClear() @@ -320,7 +332,18 @@ def boundedMoveVertex(self, pos): if self.outOfPixmap(pos): pos = self.intersectionPoint(point, pos) - shiftPos = pos - point + if self.drawSquare: + opposite_point_index = (index + 2) % 4 + opposite_point = shape[opposite_point_index] + + min_size = min(abs(pos.x() - opposite_point.x()), abs(pos.y() - opposite_point.y())) + directionX = -1 if pos.x() - opposite_point.x() < 0 else 1 + directionY = -1 if pos.y() - opposite_point.y() < 0 else 1 + shiftPos = QPointF(opposite_point.x() + directionX * min_size - point.x(), + opposite_point.y() + directionY * min_size - point.y()) + else: + shiftPos = pos - point + shape.moveVertexBy(index, shiftPos) lindex = (index + 1) % 4 @@ -678,3 +701,6 @@ def resetState(self): self.restoreCursor() self.pixmap = None self.update() + + def setDrawingShapeToSquare(self, status): + self.drawSquare = status diff --git a/libs/constants.py b/libs/constants.py index 7e12d682e..1c1430a0a 100755 --- a/libs/constants.py +++ b/libs/constants.py @@ -14,3 +14,4 @@ SETTING_SINGLE_CLASS = 'singleclass' FORMAT_PASCALVOC='PascalVOC' FORMAT_YOLO='YOLO' +SETTING_DRAW_SQUARE = 'draw/square'