13a700ad53
When labelImg was closed on the second monitor, and later opened when it is disabled, it is impossible to see the app window.
1485 lines
58 KiB
Python
Executable File
1485 lines
58 KiB
Python
Executable File
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
import codecs
|
|
import distutils.spawn
|
|
import os.path
|
|
import platform
|
|
import re
|
|
import sys
|
|
import subprocess
|
|
|
|
from functools import partial
|
|
from collections import defaultdict
|
|
|
|
try:
|
|
from PyQt5.QtGui import *
|
|
from PyQt5.QtCore import *
|
|
from PyQt5.QtWidgets import *
|
|
except ImportError:
|
|
# needed for py3+qt4
|
|
# Ref:
|
|
# http://pyqt.sourceforge.net/Docs/PyQt4/incompatible_apis.html
|
|
# http://stackoverflow.com/questions/21217399/pyqt4-qtcore-qvariant-object-instead-of-a-string
|
|
if sys.version_info.major >= 3:
|
|
import sip
|
|
sip.setapi('QVariant', 2)
|
|
from PyQt4.QtGui import *
|
|
from PyQt4.QtCore import *
|
|
|
|
import resources
|
|
# Add internal libs
|
|
from libs.constants import *
|
|
from libs.lib import struct, newAction, newIcon, addActions, fmtShortcut, generateColorByText
|
|
from libs.settings import Settings
|
|
from libs.shape import Shape, DEFAULT_LINE_COLOR, DEFAULT_FILL_COLOR
|
|
from libs.canvas import Canvas
|
|
from libs.zoomWidget import ZoomWidget
|
|
from libs.labelDialog import LabelDialog
|
|
from libs.colorDialog import ColorDialog
|
|
from libs.labelFile import LabelFile, LabelFileError
|
|
from libs.toolBar import ToolBar
|
|
from libs.pascal_voc_io import PascalVocReader
|
|
from libs.pascal_voc_io import XML_EXT
|
|
from libs.yolo_io import YoloReader
|
|
from libs.yolo_io import TXT_EXT
|
|
from libs.ustr import ustr
|
|
from libs.version import __version__
|
|
|
|
__appname__ = 'labelImg'
|
|
|
|
# Utility functions and classes.
|
|
|
|
def have_qstring():
|
|
'''p3/qt5 get rid of QString wrapper as py3 has native unicode str type'''
|
|
return not (sys.version_info.major >= 3 or QT_VERSION_STR.startswith('5.'))
|
|
|
|
def util_qt_strlistclass():
|
|
return QStringList if have_qstring() else list
|
|
|
|
|
|
class WindowMixin(object):
|
|
|
|
def menu(self, title, actions=None):
|
|
menu = self.menuBar().addMenu(title)
|
|
if actions:
|
|
addActions(menu, actions)
|
|
return menu
|
|
|
|
def toolbar(self, title, actions=None):
|
|
toolbar = ToolBar(title)
|
|
toolbar.setObjectName(u'%sToolBar' % title)
|
|
# toolbar.setOrientation(Qt.Vertical)
|
|
toolbar.setToolButtonStyle(Qt.ToolButtonTextUnderIcon)
|
|
if actions:
|
|
addActions(toolbar, actions)
|
|
self.addToolBar(Qt.LeftToolBarArea, toolbar)
|
|
return toolbar
|
|
|
|
|
|
# PyQt5: TypeError: unhashable type: 'QListWidgetItem'
|
|
class HashableQListWidgetItem(QListWidgetItem):
|
|
|
|
def __init__(self, *args):
|
|
super(HashableQListWidgetItem, self).__init__(*args)
|
|
|
|
def __hash__(self):
|
|
return hash(id(self))
|
|
|
|
|
|
class MainWindow(QMainWindow, WindowMixin):
|
|
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
|
|
|
|
def __init__(self, defaultFilename=None, defaultPrefdefClassFile=None, defaultSaveDir=None):
|
|
super(MainWindow, self).__init__()
|
|
self.setWindowTitle(__appname__)
|
|
|
|
# Load setting in the main thread
|
|
self.settings = Settings()
|
|
self.settings.load()
|
|
settings = self.settings
|
|
|
|
# Save as Pascal voc xml
|
|
self.defaultSaveDir = defaultSaveDir
|
|
self.usingPascalVocFormat = True
|
|
self.usingYoloFormat = False
|
|
|
|
# For loading all image under a directory
|
|
self.mImgList = []
|
|
self.dirname = None
|
|
self.labelHist = []
|
|
self.lastOpenDir = None
|
|
|
|
# Whether we need to save or not.
|
|
self.dirty = False
|
|
|
|
self._noSelectionSlot = False
|
|
self._beginner = True
|
|
self.screencastViewer = self.getAvailableScreencastViewer()
|
|
self.screencast = "https://youtu.be/p0nR2YsCY_U"
|
|
|
|
# Load predefined classes to the list
|
|
self.loadPredefinedClasses(defaultPrefdefClassFile)
|
|
|
|
# Main widgets and related state.
|
|
self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
|
|
|
|
self.itemsToShapes = {}
|
|
self.shapesToItems = {}
|
|
self.prevLabelText = ''
|
|
|
|
listLayout = QVBoxLayout()
|
|
listLayout.setContentsMargins(0, 0, 0, 0)
|
|
|
|
# Create a widget for using default label
|
|
self.useDefaultLabelCheckbox = QCheckBox(u'Use default label')
|
|
self.useDefaultLabelCheckbox.setChecked(False)
|
|
self.defaultLabelTextLine = QLineEdit()
|
|
useDefaultLabelQHBoxLayout = QHBoxLayout()
|
|
useDefaultLabelQHBoxLayout.addWidget(self.useDefaultLabelCheckbox)
|
|
useDefaultLabelQHBoxLayout.addWidget(self.defaultLabelTextLine)
|
|
useDefaultLabelContainer = QWidget()
|
|
useDefaultLabelContainer.setLayout(useDefaultLabelQHBoxLayout)
|
|
|
|
# Create a widget for edit and diffc button
|
|
self.diffcButton = QCheckBox(u'difficult')
|
|
self.diffcButton.setChecked(False)
|
|
self.diffcButton.stateChanged.connect(self.btnstate)
|
|
self.editButton = QToolButton()
|
|
self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
|
|
|
|
# Add some of widgets to listLayout
|
|
listLayout.addWidget(self.editButton)
|
|
listLayout.addWidget(self.diffcButton)
|
|
listLayout.addWidget(useDefaultLabelContainer)
|
|
|
|
# Create and add a widget for showing current label items
|
|
self.labelList = QListWidget()
|
|
labelListContainer = QWidget()
|
|
labelListContainer.setLayout(listLayout)
|
|
self.labelList.itemActivated.connect(self.labelSelectionChanged)
|
|
self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
|
|
self.labelList.itemDoubleClicked.connect(self.editLabel)
|
|
# Connect to itemChanged to detect checkbox changes.
|
|
self.labelList.itemChanged.connect(self.labelItemChanged)
|
|
listLayout.addWidget(self.labelList)
|
|
|
|
self.dock = QDockWidget(u'Box Labels', self)
|
|
self.dock.setObjectName(u'Labels')
|
|
self.dock.setWidget(labelListContainer)
|
|
|
|
# Tzutalin 20160906 : Add file list and dock to move faster
|
|
self.fileListWidget = QListWidget()
|
|
self.fileListWidget.itemDoubleClicked.connect(self.fileitemDoubleClicked)
|
|
filelistLayout = QVBoxLayout()
|
|
filelistLayout.setContentsMargins(0, 0, 0, 0)
|
|
filelistLayout.addWidget(self.fileListWidget)
|
|
fileListContainer = QWidget()
|
|
fileListContainer.setLayout(filelistLayout)
|
|
self.filedock = QDockWidget(u'File List', self)
|
|
self.filedock.setObjectName(u'Files')
|
|
self.filedock.setWidget(fileListContainer)
|
|
|
|
self.zoomWidget = ZoomWidget()
|
|
self.colorDialog = ColorDialog(parent=self)
|
|
|
|
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)
|
|
scroll.setWidgetResizable(True)
|
|
self.scrollBars = {
|
|
Qt.Vertical: scroll.verticalScrollBar(),
|
|
Qt.Horizontal: scroll.horizontalScrollBar()
|
|
}
|
|
self.scrollArea = scroll
|
|
self.canvas.scrollRequest.connect(self.scrollRequest)
|
|
|
|
self.canvas.newShape.connect(self.newShape)
|
|
self.canvas.shapeMoved.connect(self.setDirty)
|
|
self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
|
|
self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
|
|
|
|
self.setCentralWidget(scroll)
|
|
self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
|
|
# Tzutalin 20160906 : Add file list and dock to move faster
|
|
self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
|
|
self.filedock.setFeatures(QDockWidget.DockWidgetFloatable)
|
|
|
|
self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
|
|
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
|
|
|
|
# Actions
|
|
action = partial(newAction, self)
|
|
quit = action('&Quit', self.close,
|
|
'Ctrl+Q', 'quit', u'Quit application')
|
|
|
|
open = action('&Open', self.openFile,
|
|
'Ctrl+O', 'open', u'Open image or label file')
|
|
|
|
opendir = action('&Open Dir', self.openDirDialog,
|
|
'Ctrl+u', 'open', u'Open Dir')
|
|
|
|
changeSavedir = action('&Change Save Dir', self.changeSavedirDialog,
|
|
'Ctrl+r', 'open', u'Change default saved Annotation dir')
|
|
|
|
openAnnotation = action('&Open Annotation', self.openAnnotationDialog,
|
|
'Ctrl+Shift+O', 'open', u'Open Annotation')
|
|
|
|
openNextImg = action('&Next Image', self.openNextImg,
|
|
'd', 'next', u'Open Next')
|
|
|
|
openPrevImg = action('&Prev Image', self.openPrevImg,
|
|
'a', 'prev', u'Open Prev')
|
|
|
|
verify = action('&Verify Image', self.verifyImg,
|
|
'space', 'verify', u'Verify Image')
|
|
|
|
save = action('&Save', self.saveFile,
|
|
'Ctrl+S', 'save', u'Save labels to file', enabled=False)
|
|
|
|
save_format = action('&PascalVOC', self.change_format,
|
|
'Ctrl+', 'format_voc', u'Change save format', enabled=True)
|
|
|
|
saveAs = action('&Save As', self.saveFileAs,
|
|
'Ctrl+Shift+S', 'save-as', u'Save labels to a different file', enabled=False)
|
|
|
|
close = action('&Close', self.closeFile, 'Ctrl+W', 'close', u'Close current file')
|
|
|
|
resetAll = action('&ResetAll', self.resetAll, None, 'resetall', u'Reset all')
|
|
|
|
color1 = action('Box Line Color', self.chooseColor1,
|
|
'Ctrl+L', 'color_line', u'Choose Box line color')
|
|
|
|
createMode = action('Create\nRectBox', self.setCreateMode,
|
|
'w', 'new', u'Start drawing Boxs', enabled=False)
|
|
editMode = action('&Edit\nRectBox', self.setEditMode,
|
|
'Ctrl+J', 'edit', u'Move and edit Boxs', enabled=False)
|
|
|
|
create = action('Create\nRectBox', self.createShape,
|
|
'w', 'new', u'Draw a new Box', enabled=False)
|
|
delete = action('Delete\nRectBox', self.deleteSelectedShape,
|
|
'Delete', 'delete', u'Delete', enabled=False)
|
|
copy = action('&Duplicate\nRectBox', self.copySelectedShape,
|
|
'Ctrl+D', 'copy', u'Create a duplicate of the selected Box',
|
|
enabled=False)
|
|
|
|
advancedMode = action('&Advanced Mode', self.toggleAdvancedMode,
|
|
'Ctrl+Shift+A', 'expert', u'Switch to advanced mode',
|
|
checkable=True)
|
|
|
|
hideAll = action('&Hide\nRectBox', partial(self.togglePolygons, False),
|
|
'Ctrl+H', 'hide', u'Hide all Boxs',
|
|
enabled=False)
|
|
showAll = action('&Show\nRectBox', partial(self.togglePolygons, True),
|
|
'Ctrl+A', 'hide', u'Show all Boxs',
|
|
enabled=False)
|
|
|
|
help = action('&Tutorial', self.showTutorialDialog, None, 'help', u'Show demos')
|
|
showInfo = action('&Information', self.showInfoDialog, None, 'help', u'Information')
|
|
|
|
zoom = QWidgetAction(self)
|
|
zoom.setDefaultWidget(self.zoomWidget)
|
|
self.zoomWidget.setWhatsThis(
|
|
u"Zoom in or out of the image. Also accessible with"
|
|
" %s and %s from the canvas." % (fmtShortcut("Ctrl+[-+]"),
|
|
fmtShortcut("Ctrl+Wheel")))
|
|
self.zoomWidget.setEnabled(False)
|
|
|
|
zoomIn = action('Zoom &In', partial(self.addZoom, 10),
|
|
'Ctrl++', 'zoom-in', u'Increase zoom level', enabled=False)
|
|
zoomOut = action('&Zoom Out', partial(self.addZoom, -10),
|
|
'Ctrl+-', 'zoom-out', u'Decrease zoom level', enabled=False)
|
|
zoomOrg = action('&Original size', partial(self.setZoom, 100),
|
|
'Ctrl+=', 'zoom', u'Zoom to original size', enabled=False)
|
|
fitWindow = action('&Fit Window', self.setFitWindow,
|
|
'Ctrl+F', 'fit-window', u'Zoom follows window size',
|
|
checkable=True, enabled=False)
|
|
fitWidth = action('Fit &Width', self.setFitWidth,
|
|
'Ctrl+Shift+F', 'fit-width', u'Zoom follows window width',
|
|
checkable=True, enabled=False)
|
|
# Group zoom controls into a list for easier toggling.
|
|
zoomActions = (self.zoomWidget, zoomIn, zoomOut,
|
|
zoomOrg, fitWindow, fitWidth)
|
|
self.zoomMode = self.MANUAL_ZOOM
|
|
self.scalers = {
|
|
self.FIT_WINDOW: self.scaleFitWindow,
|
|
self.FIT_WIDTH: self.scaleFitWidth,
|
|
# Set to one to scale to 100% when loading files.
|
|
self.MANUAL_ZOOM: lambda: 1,
|
|
}
|
|
|
|
edit = action('&Edit Label', self.editLabel,
|
|
'Ctrl+E', 'edit', u'Modify the label of the selected Box',
|
|
enabled=False)
|
|
self.editButton.setDefaultAction(edit)
|
|
|
|
shapeLineColor = action('Shape &Line Color', self.chshapeLineColor,
|
|
icon='color_line', tip=u'Change the line color for this specific shape',
|
|
enabled=False)
|
|
shapeFillColor = action('Shape &Fill Color', self.chshapeFillColor,
|
|
icon='color', tip=u'Change the fill color for this specific shape',
|
|
enabled=False)
|
|
|
|
labels = self.dock.toggleViewAction()
|
|
labels.setText('Show/Hide Label Panel')
|
|
labels.setShortcut('Ctrl+Shift+L')
|
|
|
|
# Lavel list context menu.
|
|
labelMenu = QMenu()
|
|
addActions(labelMenu, (edit, delete))
|
|
self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
|
|
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,
|
|
createMode=createMode, editMode=editMode, advancedMode=advancedMode,
|
|
shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
|
|
zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
|
|
fitWindow=fitWindow, fitWidth=fitWidth,
|
|
zoomActions=zoomActions,
|
|
fileMenuActions=(
|
|
open, opendir, save, saveAs, close, resetAll, quit),
|
|
beginner=(), advanced=(),
|
|
editMenu=(edit, copy, delete,
|
|
None, color1, self.drawSquaresOption),
|
|
beginnerContext=(create, edit, copy, delete),
|
|
advancedContext=(createMode, editMode, edit, copy,
|
|
delete, shapeLineColor, shapeFillColor),
|
|
onLoadActive=(
|
|
close, create, createMode, editMode),
|
|
onShapesPresent=(saveAs, hideAll, showAll))
|
|
|
|
self.menus = struct(
|
|
file=self.menu('&File'),
|
|
edit=self.menu('&Edit'),
|
|
view=self.menu('&View'),
|
|
help=self.menu('&Help'),
|
|
recentFiles=QMenu('Open &Recent'),
|
|
labelList=labelMenu)
|
|
|
|
# Auto saving : Enable auto saving if pressing next
|
|
self.autoSaving = QAction("Auto Saving", self)
|
|
self.autoSaving.setCheckable(True)
|
|
self.autoSaving.setChecked(settings.get(SETTING_AUTO_SAVE, False))
|
|
# Sync single class mode from PR#106
|
|
self.singleClassMode = QAction("Single Class Mode", self)
|
|
self.singleClassMode.setShortcut("Ctrl+Shift+S")
|
|
self.singleClassMode.setCheckable(True)
|
|
self.singleClassMode.setChecked(settings.get(SETTING_SINGLE_CLASS, False))
|
|
self.lastLabel = None
|
|
# Add option to enable/disable labels being painted at the top of bounding boxes
|
|
self.paintLabelsOption = QAction("Paint Labels", self)
|
|
self.paintLabelsOption.setShortcut("Ctrl+Shift+P")
|
|
self.paintLabelsOption.setCheckable(True)
|
|
self.paintLabelsOption.setChecked(settings.get(SETTING_PAINT_LABEL, False))
|
|
self.paintLabelsOption.triggered.connect(self.togglePaintLabelsOption)
|
|
|
|
addActions(self.menus.file,
|
|
(open, opendir, changeSavedir, openAnnotation, self.menus.recentFiles, save, save_format, saveAs, close, resetAll, quit))
|
|
addActions(self.menus.help, (help, showInfo))
|
|
addActions(self.menus.view, (
|
|
self.autoSaving,
|
|
self.singleClassMode,
|
|
self.paintLabelsOption,
|
|
labels, advancedMode, None,
|
|
hideAll, showAll, None,
|
|
zoomIn, zoomOut, zoomOrg, None,
|
|
fitWindow, fitWidth))
|
|
|
|
self.menus.file.aboutToShow.connect(self.updateFileMenu)
|
|
|
|
# Custom context menu for the canvas widget:
|
|
addActions(self.canvas.menus[0], self.actions.beginnerContext)
|
|
addActions(self.canvas.menus[1], (
|
|
action('&Copy here', self.copyShape),
|
|
action('&Move here', self.moveShape)))
|
|
|
|
self.tools = self.toolbar('Tools')
|
|
self.actions.beginner = (
|
|
open, opendir, changeSavedir, openNextImg, openPrevImg, verify, save, save_format, None, create, copy, delete, None,
|
|
zoomIn, zoom, zoomOut, fitWindow, fitWidth)
|
|
|
|
self.actions.advanced = (
|
|
open, opendir, changeSavedir, openNextImg, openPrevImg, save, save_format, None,
|
|
createMode, editMode, None,
|
|
hideAll, showAll)
|
|
|
|
self.statusBar().showMessage('%s started.' % __appname__)
|
|
self.statusBar().show()
|
|
|
|
# Application state.
|
|
self.image = QImage()
|
|
self.filePath = ustr(defaultFilename)
|
|
self.recentFiles = []
|
|
self.maxRecent = 7
|
|
self.lineColor = None
|
|
self.fillColor = None
|
|
self.zoom_level = 100
|
|
self.fit_window = False
|
|
# Add Chris
|
|
self.difficult = False
|
|
|
|
## Fix the compatible issue for qt4 and qt5. Convert the QStringList to python list
|
|
if settings.get(SETTING_RECENT_FILES):
|
|
if have_qstring():
|
|
recentFileQStringList = settings.get(SETTING_RECENT_FILES)
|
|
self.recentFiles = [ustr(i) for i in recentFileQStringList]
|
|
else:
|
|
self.recentFiles = recentFileQStringList = settings.get(SETTING_RECENT_FILES)
|
|
|
|
size = settings.get(SETTING_WIN_SIZE, QSize(600, 500))
|
|
position = QPoint(0, 0)
|
|
saved_position = settings.get(SETTING_WIN_POSE, position)
|
|
# Fix the multiple monitors issue
|
|
for i in range(QApplication.desktop().screenCount()):
|
|
if QApplication.desktop().availableGeometry(i).contains(saved_position):
|
|
position = saved_position
|
|
break
|
|
self.resize(size)
|
|
self.move(position)
|
|
saveDir = ustr(settings.get(SETTING_SAVE_DIR, None))
|
|
self.lastOpenDir = ustr(settings.get(SETTING_LAST_OPEN_DIR, None))
|
|
if self.defaultSaveDir is None and saveDir is not None and os.path.exists(saveDir):
|
|
self.defaultSaveDir = saveDir
|
|
self.statusBar().showMessage('%s started. Annotation will be saved to %s' %
|
|
(__appname__, self.defaultSaveDir))
|
|
self.statusBar().show()
|
|
|
|
self.restoreState(settings.get(SETTING_WIN_STATE, QByteArray()))
|
|
Shape.line_color = self.lineColor = QColor(settings.get(SETTING_LINE_COLOR, DEFAULT_LINE_COLOR))
|
|
Shape.fill_color = self.fillColor = QColor(settings.get(SETTING_FILL_COLOR, DEFAULT_FILL_COLOR))
|
|
self.canvas.setDrawingColor(self.lineColor)
|
|
# Add chris
|
|
Shape.difficult = self.difficult
|
|
|
|
def xbool(x):
|
|
if isinstance(x, QVariant):
|
|
return x.toBool()
|
|
return bool(x)
|
|
|
|
if xbool(settings.get(SETTING_ADVANCE_MODE, False)):
|
|
self.actions.advancedMode.setChecked(True)
|
|
self.toggleAdvancedMode()
|
|
|
|
# Populate the File menu dynamically.
|
|
self.updateFileMenu()
|
|
|
|
# Since loading the file may take some time, make sure it runs in the background.
|
|
if self.filePath and os.path.isdir(self.filePath):
|
|
self.queueEvent(partial(self.importDirImages, self.filePath or ""))
|
|
elif self.filePath:
|
|
self.queueEvent(partial(self.loadFile, self.filePath or ""))
|
|
|
|
# Callbacks:
|
|
self.zoomWidget.valueChanged.connect(self.paintCanvas)
|
|
|
|
self.populateModeActions()
|
|
|
|
# Display cursor coordinates at the right of status bar
|
|
self.labelCoordinates = QLabel('')
|
|
self.statusBar().addPermanentWidget(self.labelCoordinates)
|
|
|
|
# Open Dir if deafult file
|
|
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:
|
|
self.actions.save_format.setText(FORMAT_PASCALVOC)
|
|
self.actions.save_format.setIcon(newIcon("format_voc"))
|
|
self.usingPascalVocFormat = True
|
|
self.usingYoloFormat = False
|
|
LabelFile.suffix = XML_EXT
|
|
|
|
elif save_format == FORMAT_YOLO:
|
|
self.actions.save_format.setText(FORMAT_YOLO)
|
|
self.actions.save_format.setIcon(newIcon("format_yolo"))
|
|
self.usingPascalVocFormat = False
|
|
self.usingYoloFormat = True
|
|
LabelFile.suffix = TXT_EXT
|
|
|
|
def change_format(self):
|
|
if self.usingPascalVocFormat: self.set_format(FORMAT_YOLO)
|
|
elif self.usingYoloFormat: self.set_format(FORMAT_PASCALVOC)
|
|
|
|
def noShapes(self):
|
|
return not self.itemsToShapes
|
|
|
|
def toggleAdvancedMode(self, value=True):
|
|
self._beginner = not value
|
|
self.canvas.setEditing(True)
|
|
self.populateModeActions()
|
|
self.editButton.setVisible(not value)
|
|
if value:
|
|
self.actions.createMode.setEnabled(True)
|
|
self.actions.editMode.setEnabled(False)
|
|
self.dock.setFeatures(self.dock.features() | self.dockFeatures)
|
|
else:
|
|
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
|
|
|
|
def populateModeActions(self):
|
|
if self.beginner():
|
|
tool, menu = self.actions.beginner, self.actions.beginnerContext
|
|
else:
|
|
tool, menu = self.actions.advanced, self.actions.advancedContext
|
|
self.tools.clear()
|
|
addActions(self.tools, tool)
|
|
self.canvas.menus[0].clear()
|
|
addActions(self.canvas.menus[0], menu)
|
|
self.menus.edit.clear()
|
|
actions = (self.actions.create,) if self.beginner()\
|
|
else (self.actions.createMode, self.actions.editMode)
|
|
addActions(self.menus.edit, actions + self.actions.editMenu)
|
|
|
|
def setBeginner(self):
|
|
self.tools.clear()
|
|
addActions(self.tools, self.actions.beginner)
|
|
|
|
def setAdvanced(self):
|
|
self.tools.clear()
|
|
addActions(self.tools, self.actions.advanced)
|
|
|
|
def setDirty(self):
|
|
self.dirty = True
|
|
self.actions.save.setEnabled(True)
|
|
|
|
def setClean(self):
|
|
self.dirty = False
|
|
self.actions.save.setEnabled(False)
|
|
self.actions.create.setEnabled(True)
|
|
|
|
def toggleActions(self, value=True):
|
|
"""Enable/Disable widgets which depend on an opened image."""
|
|
for z in self.actions.zoomActions:
|
|
z.setEnabled(value)
|
|
for action in self.actions.onLoadActive:
|
|
action.setEnabled(value)
|
|
|
|
def queueEvent(self, function):
|
|
QTimer.singleShot(0, function)
|
|
|
|
def status(self, message, delay=5000):
|
|
self.statusBar().showMessage(message, delay)
|
|
|
|
def resetState(self):
|
|
self.itemsToShapes.clear()
|
|
self.shapesToItems.clear()
|
|
self.labelList.clear()
|
|
self.filePath = None
|
|
self.imageData = None
|
|
self.labelFile = None
|
|
self.canvas.resetState()
|
|
self.labelCoordinates.clear()
|
|
|
|
def currentItem(self):
|
|
items = self.labelList.selectedItems()
|
|
if items:
|
|
return items[0]
|
|
return None
|
|
|
|
def addRecentFile(self, filePath):
|
|
if filePath in self.recentFiles:
|
|
self.recentFiles.remove(filePath)
|
|
elif len(self.recentFiles) >= self.maxRecent:
|
|
self.recentFiles.pop()
|
|
self.recentFiles.insert(0, filePath)
|
|
|
|
def beginner(self):
|
|
return self._beginner
|
|
|
|
def advanced(self):
|
|
return not self.beginner()
|
|
|
|
def getAvailableScreencastViewer(self):
|
|
osName = platform.system()
|
|
|
|
if osName == 'Windows':
|
|
return ['C:\\Program Files\\Internet Explorer\\iexplore.exe']
|
|
elif osName == 'Linux':
|
|
return ['xdg-open']
|
|
elif osName == 'Darwin':
|
|
return ['open', '-a', 'Safari']
|
|
|
|
## Callbacks ##
|
|
def showTutorialDialog(self):
|
|
subprocess.Popen(self.screencastViewer + [self.screencast])
|
|
|
|
def showInfoDialog(self):
|
|
msg = u'Name:{0} \nApp Version:{1} \n{2} '.format(__appname__, __version__, sys.version_info)
|
|
QMessageBox.information(self, u'Information', msg)
|
|
|
|
def createShape(self):
|
|
assert self.beginner()
|
|
self.canvas.setEditing(False)
|
|
self.actions.create.setEnabled(False)
|
|
|
|
def toggleDrawingSensitive(self, drawing=True):
|
|
"""In the middle of drawing, toggling between modes should be disabled."""
|
|
self.actions.editMode.setEnabled(not drawing)
|
|
if not drawing and self.beginner():
|
|
# Cancel creation.
|
|
print('Cancel creation.')
|
|
self.canvas.setEditing(True)
|
|
self.canvas.restoreCursor()
|
|
self.actions.create.setEnabled(True)
|
|
|
|
def toggleDrawMode(self, edit=True):
|
|
self.canvas.setEditing(edit)
|
|
self.actions.createMode.setEnabled(edit)
|
|
self.actions.editMode.setEnabled(not edit)
|
|
|
|
def setCreateMode(self):
|
|
assert self.advanced()
|
|
self.toggleDrawMode(False)
|
|
|
|
def setEditMode(self):
|
|
assert self.advanced()
|
|
self.toggleDrawMode(True)
|
|
self.labelSelectionChanged()
|
|
|
|
def updateFileMenu(self):
|
|
currFilePath = self.filePath
|
|
|
|
def exists(filename):
|
|
return os.path.exists(filename)
|
|
menu = self.menus.recentFiles
|
|
menu.clear()
|
|
files = [f for f in self.recentFiles if f !=
|
|
currFilePath and exists(f)]
|
|
for i, f in enumerate(files):
|
|
icon = newIcon('labels')
|
|
action = QAction(
|
|
icon, '&%d %s' % (i + 1, QFileInfo(f).fileName()), self)
|
|
action.triggered.connect(partial(self.loadRecent, f))
|
|
menu.addAction(action)
|
|
|
|
def popLabelListMenu(self, point):
|
|
self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
|
|
|
|
def editLabel(self):
|
|
if not self.canvas.editing():
|
|
return
|
|
item = self.currentItem()
|
|
text = self.labelDialog.popUp(item.text())
|
|
if text is not None:
|
|
item.setText(text)
|
|
item.setBackground(generateColorByText(text))
|
|
self.setDirty()
|
|
|
|
# Tzutalin 20160906 : Add file list and dock to move faster
|
|
def fileitemDoubleClicked(self, item=None):
|
|
currIndex = self.mImgList.index(ustr(item.text()))
|
|
if currIndex < len(self.mImgList):
|
|
filename = self.mImgList[currIndex]
|
|
if filename:
|
|
self.loadFile(filename)
|
|
|
|
# Add chris
|
|
def btnstate(self, item= None):
|
|
""" Function to handle difficult examples
|
|
Update on each object """
|
|
if not self.canvas.editing():
|
|
return
|
|
|
|
item = self.currentItem()
|
|
if not item: # If not selected Item, take the first one
|
|
item = self.labelList.item(self.labelList.count()-1)
|
|
|
|
difficult = self.diffcButton.isChecked()
|
|
|
|
try:
|
|
shape = self.itemsToShapes[item]
|
|
except:
|
|
pass
|
|
# Checked and Update
|
|
try:
|
|
if difficult != shape.difficult:
|
|
shape.difficult = difficult
|
|
self.setDirty()
|
|
else: # User probably changed item visibility
|
|
self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
|
|
except:
|
|
pass
|
|
|
|
# React to canvas signals.
|
|
def shapeSelectionChanged(self, selected=False):
|
|
if self._noSelectionSlot:
|
|
self._noSelectionSlot = False
|
|
else:
|
|
shape = self.canvas.selectedShape
|
|
if shape:
|
|
self.shapesToItems[shape].setSelected(True)
|
|
else:
|
|
self.labelList.clearSelection()
|
|
self.actions.delete.setEnabled(selected)
|
|
self.actions.copy.setEnabled(selected)
|
|
self.actions.edit.setEnabled(selected)
|
|
self.actions.shapeLineColor.setEnabled(selected)
|
|
self.actions.shapeFillColor.setEnabled(selected)
|
|
|
|
def addLabel(self, shape):
|
|
shape.paintLabel = self.paintLabelsOption.isChecked()
|
|
item = HashableQListWidgetItem(shape.label)
|
|
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
|
|
item.setCheckState(Qt.Checked)
|
|
item.setBackground(generateColorByText(shape.label))
|
|
self.itemsToShapes[item] = shape
|
|
self.shapesToItems[shape] = item
|
|
self.labelList.addItem(item)
|
|
for action in self.actions.onShapesPresent:
|
|
action.setEnabled(True)
|
|
|
|
def remLabel(self, shape):
|
|
if shape is None:
|
|
# print('rm empty label')
|
|
return
|
|
item = self.shapesToItems[shape]
|
|
self.labelList.takeItem(self.labelList.row(item))
|
|
del self.shapesToItems[shape]
|
|
del self.itemsToShapes[item]
|
|
|
|
def loadLabels(self, shapes):
|
|
s = []
|
|
for label, points, line_color, fill_color, difficult in shapes:
|
|
shape = Shape(label=label)
|
|
for x, y in points:
|
|
shape.addPoint(QPointF(x, y))
|
|
shape.difficult = difficult
|
|
shape.close()
|
|
s.append(shape)
|
|
|
|
if line_color:
|
|
shape.line_color = QColor(*line_color)
|
|
else:
|
|
shape.line_color = generateColorByText(label)
|
|
|
|
if fill_color:
|
|
shape.fill_color = QColor(*fill_color)
|
|
else:
|
|
shape.fill_color = generateColorByText(label)
|
|
|
|
self.addLabel(shape)
|
|
|
|
self.canvas.loadShapes(s)
|
|
|
|
def saveLabels(self, annotationFilePath):
|
|
annotationFilePath = ustr(annotationFilePath)
|
|
if self.labelFile is None:
|
|
self.labelFile = LabelFile()
|
|
self.labelFile.verified = self.canvas.verified
|
|
|
|
def format_shape(s):
|
|
return dict(label=s.label,
|
|
line_color=s.line_color.getRgb(),
|
|
fill_color=s.fill_color.getRgb(),
|
|
points=[(p.x(), p.y()) for p in s.points],
|
|
# add chris
|
|
difficult = s.difficult)
|
|
|
|
shapes = [format_shape(shape) for shape in self.canvas.shapes]
|
|
# Can add differrent annotation formats here
|
|
try:
|
|
if self.usingPascalVocFormat is True:
|
|
if ustr(annotationFilePath[-4:]) != ".xml":
|
|
annotationFilePath += XML_EXT
|
|
print ('Img: ' + self.filePath + ' -> Its xml: ' + annotationFilePath)
|
|
self.labelFile.savePascalVocFormat(annotationFilePath, shapes, self.filePath, self.imageData,
|
|
self.lineColor.getRgb(), self.fillColor.getRgb())
|
|
elif self.usingYoloFormat is True:
|
|
if annotationFilePath[-4:] != ".txt":
|
|
annotationFilePath += TXT_EXT
|
|
print ('Img: ' + self.filePath + ' -> Its txt: ' + annotationFilePath)
|
|
self.labelFile.saveYoloFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.labelHist,
|
|
self.lineColor.getRgb(), self.fillColor.getRgb())
|
|
else:
|
|
self.labelFile.save(annotationFilePath, shapes, self.filePath, self.imageData,
|
|
self.lineColor.getRgb(), self.fillColor.getRgb())
|
|
return True
|
|
except LabelFileError as e:
|
|
self.errorMessage(u'Error saving label data', u'<b>%s</b>' % e)
|
|
return False
|
|
|
|
def copySelectedShape(self):
|
|
self.addLabel(self.canvas.copySelectedShape())
|
|
# fix copy and delete
|
|
self.shapeSelectionChanged(True)
|
|
|
|
def labelSelectionChanged(self):
|
|
item = self.currentItem()
|
|
if item and self.canvas.editing():
|
|
self._noSelectionSlot = True
|
|
self.canvas.selectShape(self.itemsToShapes[item])
|
|
shape = self.itemsToShapes[item]
|
|
# Add Chris
|
|
self.diffcButton.setChecked(shape.difficult)
|
|
|
|
def labelItemChanged(self, item):
|
|
shape = self.itemsToShapes[item]
|
|
label = item.text()
|
|
if label != shape.label:
|
|
shape.label = item.text()
|
|
shape.line_color = generateColorByText(shape.label)
|
|
self.setDirty()
|
|
else: # User probably changed item visibility
|
|
self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
|
|
|
|
# Callback functions:
|
|
def newShape(self):
|
|
"""Pop-up and give focus to the label editor.
|
|
|
|
position MUST be in global coordinates.
|
|
"""
|
|
if not self.useDefaultLabelCheckbox.isChecked() or not self.defaultLabelTextLine.text():
|
|
if len(self.labelHist) > 0:
|
|
self.labelDialog = LabelDialog(
|
|
parent=self, listItem=self.labelHist)
|
|
|
|
# Sync single class mode from PR#106
|
|
if self.singleClassMode.isChecked() and self.lastLabel:
|
|
text = self.lastLabel
|
|
else:
|
|
text = self.labelDialog.popUp(text=self.prevLabelText)
|
|
self.lastLabel = text
|
|
else:
|
|
text = self.defaultLabelTextLine.text()
|
|
|
|
# Add Chris
|
|
self.diffcButton.setChecked(False)
|
|
if text is not None:
|
|
self.prevLabelText = text
|
|
generate_color = generateColorByText(text)
|
|
shape = self.canvas.setLastLabel(text, generate_color, generate_color)
|
|
self.addLabel(shape)
|
|
if self.beginner(): # Switch to edit mode.
|
|
self.canvas.setEditing(True)
|
|
self.actions.create.setEnabled(True)
|
|
else:
|
|
self.actions.editMode.setEnabled(True)
|
|
self.setDirty()
|
|
|
|
if text not in self.labelHist:
|
|
self.labelHist.append(text)
|
|
else:
|
|
# self.canvas.undoLastLine()
|
|
self.canvas.resetAllLines()
|
|
|
|
def scrollRequest(self, delta, orientation):
|
|
units = - delta / (8 * 15)
|
|
bar = self.scrollBars[orientation]
|
|
bar.setValue(bar.value() + bar.singleStep() * units)
|
|
|
|
def setZoom(self, value):
|
|
self.actions.fitWidth.setChecked(False)
|
|
self.actions.fitWindow.setChecked(False)
|
|
self.zoomMode = self.MANUAL_ZOOM
|
|
self.zoomWidget.setValue(value)
|
|
|
|
def addZoom(self, increment=10):
|
|
self.setZoom(self.zoomWidget.value() + increment)
|
|
|
|
def zoomRequest(self, delta):
|
|
# get the current scrollbar positions
|
|
# calculate the percentages ~ coordinates
|
|
h_bar = self.scrollBars[Qt.Horizontal]
|
|
v_bar = self.scrollBars[Qt.Vertical]
|
|
|
|
# get the current maximum, to know the difference after zooming
|
|
h_bar_max = h_bar.maximum()
|
|
v_bar_max = v_bar.maximum()
|
|
|
|
# get the cursor position and canvas size
|
|
# calculate the desired movement from 0 to 1
|
|
# where 0 = move left
|
|
# 1 = move right
|
|
# up and down analogous
|
|
cursor = QCursor()
|
|
pos = cursor.pos()
|
|
relative_pos = QWidget.mapFromGlobal(self, pos)
|
|
|
|
cursor_x = relative_pos.x()
|
|
cursor_y = relative_pos.y()
|
|
|
|
w = self.scrollArea.width()
|
|
h = self.scrollArea.height()
|
|
|
|
# the scaling from 0 to 1 has some padding
|
|
# you don't have to hit the very leftmost pixel for a maximum-left movement
|
|
margin = 0.1
|
|
move_x = (cursor_x - margin * w) / (w - 2 * margin * w)
|
|
move_y = (cursor_y - margin * h) / (h - 2 * margin * h)
|
|
|
|
# clamp the values from 0 to 1
|
|
move_x = min(max(move_x, 0), 1)
|
|
move_y = min(max(move_y, 0), 1)
|
|
|
|
# zoom in
|
|
units = delta / (8 * 15)
|
|
scale = 10
|
|
self.addZoom(scale * units)
|
|
|
|
# get the difference in scrollbar values
|
|
# this is how far we can move
|
|
d_h_bar_max = h_bar.maximum() - h_bar_max
|
|
d_v_bar_max = v_bar.maximum() - v_bar_max
|
|
|
|
# get the new scrollbar values
|
|
new_h_bar_value = h_bar.value() + move_x * d_h_bar_max
|
|
new_v_bar_value = v_bar.value() + move_y * d_v_bar_max
|
|
|
|
h_bar.setValue(new_h_bar_value)
|
|
v_bar.setValue(new_v_bar_value)
|
|
|
|
def setFitWindow(self, value=True):
|
|
if value:
|
|
self.actions.fitWidth.setChecked(False)
|
|
self.zoomMode = self.FIT_WINDOW if value else self.MANUAL_ZOOM
|
|
self.adjustScale()
|
|
|
|
def setFitWidth(self, value=True):
|
|
if value:
|
|
self.actions.fitWindow.setChecked(False)
|
|
self.zoomMode = self.FIT_WIDTH if value else self.MANUAL_ZOOM
|
|
self.adjustScale()
|
|
|
|
def togglePolygons(self, value):
|
|
for item, shape in self.itemsToShapes.items():
|
|
item.setCheckState(Qt.Checked if value else Qt.Unchecked)
|
|
|
|
def loadFile(self, filePath=None):
|
|
"""Load the specified file, or the last opened file if None."""
|
|
self.resetState()
|
|
self.canvas.setEnabled(False)
|
|
if filePath is None:
|
|
filePath = self.settings.get(SETTING_FILENAME)
|
|
|
|
# Make sure that filePath is a regular python string, rather than QString
|
|
filePath = ustr(filePath)
|
|
|
|
unicodeFilePath = ustr(filePath)
|
|
# Tzutalin 20160906 : Add file list and dock to move faster
|
|
# Highlight the file item
|
|
if unicodeFilePath and self.fileListWidget.count() > 0:
|
|
index = self.mImgList.index(unicodeFilePath)
|
|
fileWidgetItem = self.fileListWidget.item(index)
|
|
fileWidgetItem.setSelected(True)
|
|
|
|
if unicodeFilePath and os.path.exists(unicodeFilePath):
|
|
if LabelFile.isLabelFile(unicodeFilePath):
|
|
try:
|
|
self.labelFile = LabelFile(unicodeFilePath)
|
|
except LabelFileError as e:
|
|
self.errorMessage(u'Error opening file',
|
|
(u"<p><b>%s</b></p>"
|
|
u"<p>Make sure <i>%s</i> is a valid label file.")
|
|
% (e, unicodeFilePath))
|
|
self.status("Error reading %s" % unicodeFilePath)
|
|
return False
|
|
self.imageData = self.labelFile.imageData
|
|
self.lineColor = QColor(*self.labelFile.lineColor)
|
|
self.fillColor = QColor(*self.labelFile.fillColor)
|
|
self.canvas.verified = self.labelFile.verified
|
|
else:
|
|
# Load image:
|
|
# read data first and store for saving into label file.
|
|
self.imageData = read(unicodeFilePath, None)
|
|
self.labelFile = None
|
|
self.canvas.verified = False
|
|
|
|
image = QImage.fromData(self.imageData)
|
|
if image.isNull():
|
|
self.errorMessage(u'Error opening file',
|
|
u"<p>Make sure <i>%s</i> is a valid image file." % unicodeFilePath)
|
|
self.status("Error reading %s" % unicodeFilePath)
|
|
return False
|
|
self.status("Loaded %s" % os.path.basename(unicodeFilePath))
|
|
self.image = image
|
|
self.filePath = unicodeFilePath
|
|
self.canvas.loadPixmap(QPixmap.fromImage(image))
|
|
if self.labelFile:
|
|
self.loadLabels(self.labelFile.shapes)
|
|
self.setClean()
|
|
self.canvas.setEnabled(True)
|
|
self.adjustScale(initial=True)
|
|
self.paintCanvas()
|
|
self.addRecentFile(self.filePath)
|
|
self.toggleActions(True)
|
|
|
|
# Label xml file and show bound box according to its filename
|
|
# if self.usingPascalVocFormat is True:
|
|
if self.defaultSaveDir is not None:
|
|
basename = os.path.basename(
|
|
os.path.splitext(self.filePath)[0])
|
|
xmlPath = os.path.join(self.defaultSaveDir, basename + XML_EXT)
|
|
txtPath = os.path.join(self.defaultSaveDir, basename + TXT_EXT)
|
|
|
|
"""Annotation file priority:
|
|
PascalXML > YOLO
|
|
"""
|
|
if os.path.isfile(xmlPath):
|
|
self.loadPascalXMLByFilename(xmlPath)
|
|
elif os.path.isfile(txtPath):
|
|
self.loadYOLOTXTByFilename(txtPath)
|
|
else:
|
|
xmlPath = os.path.splitext(filePath)[0] + XML_EXT
|
|
txtPath = os.path.splitext(filePath)[0] + TXT_EXT
|
|
if os.path.isfile(xmlPath):
|
|
self.loadPascalXMLByFilename(xmlPath)
|
|
elif os.path.isfile(txtPath):
|
|
self.loadYOLOTXTByFilename(txtPath)
|
|
|
|
self.setWindowTitle(__appname__ + ' ' + filePath)
|
|
|
|
# Default : select last item if there is at least one item
|
|
if self.labelList.count():
|
|
self.labelList.setCurrentItem(self.labelList.item(self.labelList.count()-1))
|
|
self.labelList.item(self.labelList.count()-1).setSelected(True)
|
|
|
|
self.canvas.setFocus(True)
|
|
return True
|
|
return False
|
|
|
|
def resizeEvent(self, event):
|
|
if self.canvas and not self.image.isNull()\
|
|
and self.zoomMode != self.MANUAL_ZOOM:
|
|
self.adjustScale()
|
|
super(MainWindow, self).resizeEvent(event)
|
|
|
|
def paintCanvas(self):
|
|
assert not self.image.isNull(), "cannot paint null image"
|
|
self.canvas.scale = 0.01 * self.zoomWidget.value()
|
|
self.canvas.adjustSize()
|
|
self.canvas.update()
|
|
|
|
def adjustScale(self, initial=False):
|
|
value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
|
|
self.zoomWidget.setValue(int(100 * value))
|
|
|
|
def scaleFitWindow(self):
|
|
"""Figure out the size of the pixmap in order to fit the main widget."""
|
|
e = 2.0 # So that no scrollbars are generated.
|
|
w1 = self.centralWidget().width() - e
|
|
h1 = self.centralWidget().height() - e
|
|
a1 = w1 / h1
|
|
# Calculate a new scale value based on the pixmap's aspect ratio.
|
|
w2 = self.canvas.pixmap.width() - 0.0
|
|
h2 = self.canvas.pixmap.height() - 0.0
|
|
a2 = w2 / h2
|
|
return w1 / w2 if a2 >= a1 else h1 / h2
|
|
|
|
def scaleFitWidth(self):
|
|
# The epsilon does not seem to work too well here.
|
|
w = self.centralWidget().width() - 2.0
|
|
return w / self.canvas.pixmap.width()
|
|
|
|
def closeEvent(self, event):
|
|
if not self.mayContinue():
|
|
event.ignore()
|
|
settings = self.settings
|
|
# If it loads images from dir, don't load it at the begining
|
|
if self.dirname is None:
|
|
settings[SETTING_FILENAME] = self.filePath if self.filePath else ''
|
|
else:
|
|
settings[SETTING_FILENAME] = ''
|
|
|
|
settings[SETTING_WIN_SIZE] = self.size()
|
|
settings[SETTING_WIN_POSE] = self.pos()
|
|
settings[SETTING_WIN_STATE] = self.saveState()
|
|
settings[SETTING_LINE_COLOR] = self.lineColor
|
|
settings[SETTING_FILL_COLOR] = self.fillColor
|
|
settings[SETTING_RECENT_FILES] = self.recentFiles
|
|
settings[SETTING_ADVANCE_MODE] = not self._beginner
|
|
if self.defaultSaveDir and os.path.exists(self.defaultSaveDir):
|
|
settings[SETTING_SAVE_DIR] = ustr(self.defaultSaveDir)
|
|
else:
|
|
settings[SETTING_SAVE_DIR] = ""
|
|
|
|
if self.lastOpenDir and os.path.exists(self.lastOpenDir):
|
|
settings[SETTING_LAST_OPEN_DIR] = self.lastOpenDir
|
|
else:
|
|
settings[SETTING_LAST_OPEN_DIR] = ""
|
|
|
|
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 ##
|
|
|
|
def loadRecent(self, filename):
|
|
if self.mayContinue():
|
|
self.loadFile(filename)
|
|
|
|
def scanAllImages(self, folderPath):
|
|
extensions = ['.%s' % fmt.data().decode("ascii").lower() for fmt in QImageReader.supportedImageFormats()]
|
|
images = []
|
|
|
|
for root, dirs, files in os.walk(folderPath):
|
|
for file in files:
|
|
if file.lower().endswith(tuple(extensions)):
|
|
relativePath = os.path.join(root, file)
|
|
path = ustr(os.path.abspath(relativePath))
|
|
images.append(path)
|
|
images.sort(key=lambda x: x.lower())
|
|
return images
|
|
|
|
def changeSavedirDialog(self, _value=False):
|
|
if self.defaultSaveDir is not None:
|
|
path = ustr(self.defaultSaveDir)
|
|
else:
|
|
path = '.'
|
|
|
|
dirpath = ustr(QFileDialog.getExistingDirectory(self,
|
|
'%s - Save annotations to the directory' % __appname__, path, QFileDialog.ShowDirsOnly
|
|
| QFileDialog.DontResolveSymlinks))
|
|
|
|
if dirpath is not None and len(dirpath) > 1:
|
|
self.defaultSaveDir = dirpath
|
|
|
|
self.statusBar().showMessage('%s . Annotation will be saved to %s' %
|
|
('Change saved folder', self.defaultSaveDir))
|
|
self.statusBar().show()
|
|
|
|
def openAnnotationDialog(self, _value=False):
|
|
if self.filePath is None:
|
|
self.statusBar().showMessage('Please select image first')
|
|
self.statusBar().show()
|
|
return
|
|
|
|
path = os.path.dirname(ustr(self.filePath))\
|
|
if self.filePath else '.'
|
|
if self.usingPascalVocFormat:
|
|
filters = "Open Annotation XML file (%s)" % ' '.join(['*.xml'])
|
|
filename = ustr(QFileDialog.getOpenFileName(self,'%s - Choose a xml file' % __appname__, path, filters))
|
|
if filename:
|
|
if isinstance(filename, (tuple, list)):
|
|
filename = filename[0]
|
|
self.loadPascalXMLByFilename(filename)
|
|
|
|
def openDirDialog(self, _value=False, dirpath=None):
|
|
if not self.mayContinue():
|
|
return
|
|
|
|
defaultOpenDirPath = dirpath if dirpath else '.'
|
|
if self.lastOpenDir and os.path.exists(self.lastOpenDir):
|
|
defaultOpenDirPath = self.lastOpenDir
|
|
else:
|
|
defaultOpenDirPath = os.path.dirname(self.filePath) if self.filePath else '.'
|
|
|
|
targetDirPath = ustr(QFileDialog.getExistingDirectory(self,
|
|
'%s - Open Directory' % __appname__, defaultOpenDirPath,
|
|
QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks))
|
|
self.importDirImages(targetDirPath)
|
|
|
|
def importDirImages(self, dirpath):
|
|
if not self.mayContinue() or not dirpath:
|
|
return
|
|
|
|
|
|
self.lastOpenDir = dirpath
|
|
self.dirname = dirpath
|
|
self.filePath = None
|
|
self.fileListWidget.clear()
|
|
self.mImgList = self.scanAllImages(dirpath)
|
|
self.openNextImg()
|
|
for imgPath in self.mImgList:
|
|
item = QListWidgetItem(imgPath)
|
|
self.fileListWidget.addItem(item)
|
|
|
|
def verifyImg(self, _value=False):
|
|
# Proceding next image without dialog if having any label
|
|
if self.filePath is not None:
|
|
try:
|
|
self.labelFile.toggleVerify()
|
|
except AttributeError:
|
|
# If the labelling file does not exist yet, create if and
|
|
# re-save it with the verified attribute.
|
|
self.saveFile()
|
|
if self.labelFile != None:
|
|
self.labelFile.toggleVerify()
|
|
else:
|
|
return
|
|
|
|
self.canvas.verified = self.labelFile.verified
|
|
self.paintCanvas()
|
|
self.saveFile()
|
|
|
|
def openPrevImg(self, _value=False):
|
|
# Proceding prev image without dialog if having any label
|
|
if self.autoSaving.isChecked():
|
|
if self.defaultSaveDir is not None:
|
|
if self.dirty is True:
|
|
self.saveFile()
|
|
else:
|
|
self.changeSavedirDialog()
|
|
return
|
|
|
|
if not self.mayContinue():
|
|
return
|
|
|
|
if len(self.mImgList) <= 0:
|
|
return
|
|
|
|
if self.filePath is None:
|
|
return
|
|
|
|
currIndex = self.mImgList.index(self.filePath)
|
|
if currIndex - 1 >= 0:
|
|
filename = self.mImgList[currIndex - 1]
|
|
if filename:
|
|
self.loadFile(filename)
|
|
|
|
def openNextImg(self, _value=False):
|
|
# Proceding prev image without dialog if having any label
|
|
if self.autoSaving.isChecked():
|
|
if self.defaultSaveDir is not None:
|
|
if self.dirty is True:
|
|
self.saveFile()
|
|
else:
|
|
self.changeSavedirDialog()
|
|
return
|
|
|
|
if not self.mayContinue():
|
|
return
|
|
|
|
if len(self.mImgList) <= 0:
|
|
return
|
|
|
|
filename = None
|
|
if self.filePath is None:
|
|
filename = self.mImgList[0]
|
|
else:
|
|
currIndex = self.mImgList.index(self.filePath)
|
|
if currIndex + 1 < len(self.mImgList):
|
|
filename = self.mImgList[currIndex + 1]
|
|
|
|
if filename:
|
|
self.loadFile(filename)
|
|
|
|
def openFile(self, _value=False):
|
|
if not self.mayContinue():
|
|
return
|
|
path = os.path.dirname(ustr(self.filePath)) if self.filePath else '.'
|
|
formats = ['*.%s' % fmt.data().decode("ascii").lower() for fmt in QImageReader.supportedImageFormats()]
|
|
filters = "Image & Label files (%s)" % ' '.join(formats + ['*%s' % LabelFile.suffix])
|
|
filename = QFileDialog.getOpenFileName(self, '%s - Choose Image or Label file' % __appname__, path, filters)
|
|
if filename:
|
|
if isinstance(filename, (tuple, list)):
|
|
filename = filename[0]
|
|
self.loadFile(filename)
|
|
|
|
def saveFile(self, _value=False):
|
|
if self.defaultSaveDir is not None and len(ustr(self.defaultSaveDir)):
|
|
if self.filePath:
|
|
imgFileName = os.path.basename(self.filePath)
|
|
savedFileName = os.path.splitext(imgFileName)[0]
|
|
savedPath = os.path.join(ustr(self.defaultSaveDir), savedFileName)
|
|
self._saveFile(savedPath)
|
|
else:
|
|
imgFileDir = os.path.dirname(self.filePath)
|
|
imgFileName = os.path.basename(self.filePath)
|
|
savedFileName = os.path.splitext(imgFileName)[0]
|
|
savedPath = os.path.join(imgFileDir, savedFileName)
|
|
self._saveFile(savedPath if self.labelFile
|
|
else self.saveFileDialog())
|
|
|
|
def saveFileAs(self, _value=False):
|
|
assert not self.image.isNull(), "cannot save empty image"
|
|
self._saveFile(self.saveFileDialog())
|
|
|
|
def saveFileDialog(self):
|
|
caption = '%s - Choose File' % __appname__
|
|
filters = 'File (*%s)' % LabelFile.suffix
|
|
openDialogPath = self.currentPath()
|
|
dlg = QFileDialog(self, caption, openDialogPath, filters)
|
|
dlg.setDefaultSuffix(LabelFile.suffix[1:])
|
|
dlg.setAcceptMode(QFileDialog.AcceptSave)
|
|
filenameWithoutExtension = os.path.splitext(self.filePath)[0]
|
|
dlg.selectFile(filenameWithoutExtension)
|
|
dlg.setOption(QFileDialog.DontUseNativeDialog, False)
|
|
if dlg.exec_():
|
|
fullFilePath = ustr(dlg.selectedFiles()[0])
|
|
return os.path.splitext(fullFilePath)[0] # Return file path without the extension.
|
|
return ''
|
|
|
|
def _saveFile(self, annotationFilePath):
|
|
if annotationFilePath and self.saveLabels(annotationFilePath):
|
|
self.setClean()
|
|
self.statusBar().showMessage('Saved to %s' % annotationFilePath)
|
|
self.statusBar().show()
|
|
|
|
def closeFile(self, _value=False):
|
|
if not self.mayContinue():
|
|
return
|
|
self.resetState()
|
|
self.setClean()
|
|
self.toggleActions(False)
|
|
self.canvas.setEnabled(False)
|
|
self.actions.saveAs.setEnabled(False)
|
|
|
|
def resetAll(self):
|
|
self.settings.reset()
|
|
self.close()
|
|
proc = QProcess()
|
|
proc.startDetached(os.path.abspath(__file__))
|
|
|
|
def mayContinue(self):
|
|
return not (self.dirty and not self.discardChangesDialog())
|
|
|
|
def discardChangesDialog(self):
|
|
yes, no = QMessageBox.Yes, QMessageBox.No
|
|
msg = u'You have unsaved changes, proceed anyway?'
|
|
return yes == QMessageBox.warning(self, u'Attention', msg, yes | no)
|
|
|
|
def errorMessage(self, title, message):
|
|
return QMessageBox.critical(self, title,
|
|
'<p><b>%s</b></p>%s' % (title, message))
|
|
|
|
def currentPath(self):
|
|
return os.path.dirname(self.filePath) if self.filePath else '.'
|
|
|
|
def chooseColor1(self):
|
|
color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
|
|
default=DEFAULT_LINE_COLOR)
|
|
if color:
|
|
self.lineColor = color
|
|
Shape.line_color = color
|
|
self.canvas.setDrawingColor(color)
|
|
self.canvas.update()
|
|
self.setDirty()
|
|
|
|
def deleteSelectedShape(self):
|
|
self.remLabel(self.canvas.deleteSelected())
|
|
self.setDirty()
|
|
if self.noShapes():
|
|
for action in self.actions.onShapesPresent:
|
|
action.setEnabled(False)
|
|
|
|
def chshapeLineColor(self):
|
|
color = self.colorDialog.getColor(self.lineColor, u'Choose line color',
|
|
default=DEFAULT_LINE_COLOR)
|
|
if color:
|
|
self.canvas.selectedShape.line_color = color
|
|
self.canvas.update()
|
|
self.setDirty()
|
|
|
|
def chshapeFillColor(self):
|
|
color = self.colorDialog.getColor(self.fillColor, u'Choose fill color',
|
|
default=DEFAULT_FILL_COLOR)
|
|
if color:
|
|
self.canvas.selectedShape.fill_color = color
|
|
self.canvas.update()
|
|
self.setDirty()
|
|
|
|
def copyShape(self):
|
|
self.canvas.endMove(copy=True)
|
|
self.addLabel(self.canvas.selectedShape)
|
|
self.setDirty()
|
|
|
|
def moveShape(self):
|
|
self.canvas.endMove(copy=False)
|
|
self.setDirty()
|
|
|
|
def loadPredefinedClasses(self, predefClassesFile):
|
|
if os.path.exists(predefClassesFile) is True:
|
|
with codecs.open(predefClassesFile, 'r', 'utf8') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if self.labelHist is None:
|
|
self.labelHist = [line]
|
|
else:
|
|
self.labelHist.append(line)
|
|
|
|
def loadPascalXMLByFilename(self, xmlPath):
|
|
if self.filePath is None:
|
|
return
|
|
if os.path.isfile(xmlPath) is False:
|
|
return
|
|
|
|
self.set_format(FORMAT_PASCALVOC)
|
|
|
|
tVocParseReader = PascalVocReader(xmlPath)
|
|
shapes = tVocParseReader.getShapes()
|
|
self.loadLabels(shapes)
|
|
self.canvas.verified = tVocParseReader.verified
|
|
|
|
def loadYOLOTXTByFilename(self, txtPath):
|
|
if self.filePath is None:
|
|
return
|
|
if os.path.isfile(txtPath) is False:
|
|
return
|
|
|
|
self.set_format(FORMAT_YOLO)
|
|
tYoloParseReader = YoloReader(txtPath, self.image)
|
|
shapes = tYoloParseReader.getShapes()
|
|
print (shapes)
|
|
self.loadLabels(shapes)
|
|
self.canvas.verified = tYoloParseReader.verified
|
|
|
|
def togglePaintLabelsOption(self):
|
|
paintLabelsOptionChecked = self.paintLabelsOption.isChecked()
|
|
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()])
|
|
|
|
|
|
def read(filename, default=None):
|
|
try:
|
|
with open(filename, 'rb') as f:
|
|
return f.read()
|
|
except:
|
|
return default
|
|
|
|
|
|
def get_main_app(argv=[]):
|
|
"""
|
|
Standard boilerplate Qt application code.
|
|
Do everything but app.exec_() -- so that we can test the application in one thread
|
|
"""
|
|
app = QApplication(argv)
|
|
app.setApplicationName(__appname__)
|
|
app.setWindowIcon(newIcon("app"))
|
|
# Tzutalin 201705+: Accept extra agruments to change predefined class file
|
|
# Usage : labelImg.py image predefClassFile saveDir
|
|
win = MainWindow(argv[1] if len(argv) >= 2 else None,
|
|
argv[2] if len(argv) >= 3 else os.path.join(
|
|
os.path.dirname(sys.argv[0]),
|
|
'data', 'predefined_classes.txt'),
|
|
argv[3] if len(argv) >= 4 else None)
|
|
win.show()
|
|
return app, win
|
|
|
|
|
|
def main():
|
|
'''construct main app and run it'''
|
|
app, _win = get_main_app(sys.argv)
|
|
return app.exec_()
|
|
|
|
if __name__ == '__main__':
|
|
sys.exit(main())
|