Skip to content
This repository has been archived by the owner on Feb 29, 2024. It is now read-only.

Commit

Permalink
feature: draw square bounding boxes
Browse files Browse the repository at this point in the history
  • Loading branch information
antoinebrl authored and tzutalin committed Oct 20, 2018
1 parent 5abf660 commit 4e278bb
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 4 deletions.
26 changes: 24 additions & 2 deletions labelImg.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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),
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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 ##

Expand Down Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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()])

Expand Down
30 changes: 28 additions & 2 deletions libs/canvas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -678,3 +701,6 @@ def resetState(self):
self.restoreCursor()
self.pixmap = None
self.update()

def setDrawingShapeToSquare(self, status):
self.drawSquare = status
1 change: 1 addition & 0 deletions libs/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,3 +14,4 @@
SETTING_SINGLE_CLASS = 'singleclass'
FORMAT_PASCALVOC='PascalVOC'
FORMAT_YOLO='YOLO'
SETTING_DRAW_SQUARE = 'draw/square'

0 comments on commit 4e278bb

Please sign in to comment.