feature: draw square bounding boxes
This commit is contained in:
parent
5abf660222
commit
4e278bb511
26
labelImg.py
26
labelImg.py
@ -184,6 +184,7 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
|
||||
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 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
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 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
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 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
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 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
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 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
if not self.mayContinue() or not dirpath:
|
||||
return
|
||||
|
||||
|
||||
self.lastOpenDir = dirpath
|
||||
self.dirname = dirpath
|
||||
self.filePath = None
|
||||
@ -1183,7 +1202,7 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
|
||||
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 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
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()])
|
||||
|
||||
|
||||
@ -61,6 +61,7 @@ class Canvas(QWidget):
|
||||
self.setMouseTracking(True)
|
||||
self.setFocusPolicy(Qt.WheelFocus)
|
||||
self.verified = False
|
||||
self.drawSquare = False
|
||||
|
||||
def setDrawingColor(self, qColor):
|
||||
self.drawingLineColor = qColor
|
||||
@ -126,7 +127,18 @@ class Canvas(QWidget):
|
||||
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 @@ class Canvas(QWidget):
|
||||
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 @@ class Canvas(QWidget):
|
||||
self.restoreCursor()
|
||||
self.pixmap = None
|
||||
self.update()
|
||||
|
||||
def setDrawingShapeToSquare(self, status):
|
||||
self.drawSquare = status
|
||||
|
||||
@ -14,3 +14,4 @@ SETTING_AUTO_SAVE = 'autosave'
|
||||
SETTING_SINGLE_CLASS = 'singleclass'
|
||||
FORMAT_PASCALVOC='PascalVOC'
|
||||
FORMAT_YOLO='YOLO'
|
||||
SETTING_DRAW_SQUARE = 'draw/square'
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user