class MainWindow(QMainWindow, WindowMixin):
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
def __init__(self, defaultFilename=None, defaultPrefdefClassFile=None):
super(MainWindow, self).__init__()
self.setWindowTitle(__appname__)
self.defaultSaveDir = None
self.usingPascalVocFormat = True
self.mImgList = []
self.dirname = None
self.labelHist = []
self.lastOpenDir = None
self.dirty = False
self._noSelectionSlot = False
self._beginner = True
self.screencastViewer = "firefox"
self.screencast = "https://www.youtube.com/?gl=UA&hl=ru"
self.loadPredefinedClasses(defaultPrefdefClassFile)
self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
self.itemsToShapes = {}
self.shapesToItems = {}
self.prevLabelText = ''
listLayout = QVBoxLayout()
listLayout.setContentsMargins(0, 0, 0, 0)
self.useDefaultLabelCheckbox = QCheckBox(u'Default label')
self.useDefaultLabelCheckbox.setChecked(True)
self.defaultLabelTextLine = QLineEdit('Object')
useDefaultLabelQHBoxLayout = QHBoxLayout()
useDefaultLabelQHBoxLayout.addWidget(self.useDefaultLabelCheckbox)
useDefaultLabelQHBoxLayout.addWidget(self.defaultLabelTextLine)
useDefaultLabelContainer = QWidget()
useDefaultLabelContainer.setLayout(useDefaultLabelQHBoxLayout)
self.editButton = QToolButton()
self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
listLayout.addWidget(self.editButton)
listLayout.addWidget(useDefaultLabelContainer)
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)
self.labelList.itemChanged.connect(self.labelItemChanged)
listLayout.addWidget(self.labelList)
self.dock = QDockWidget(u'Label List', self)
self.dock.setObjectName(u'Labels')
self.dock.setWidget(labelListContainer)
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()
self.canvas.zoomRequest.connect(self.zoomRequest)
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)
self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
action = partial(newAction, self)
open = action('&Open Image', self.openFile, 'I', 'open',
u'Open image file')
opendir = action('&Open Folder', self.openDir, 'F', 'open',
u'Open directory')
openNextImg = action('&Next Image', self.openNextImg, 'D', 'next',
u'Open next image')
openPrevImg = action('&Previous Image', self.openPrevImg, 'A', 'prev',
u'Open previous image')
save = action('&Save Label',
self.saveFile,
'S',
'save',
u'Save labels to file',
enabled=False)
createMode = action('Create Label',
self.setCreateMode,
'W',
'new',
u'Create new label',
enabled=False)
editMode = action('&Edit Label',
self.setEditMode,
'E',
'edit',
u'Move and edit label',
enabled=False)
delete = action('Delete label',
self.deleteSelectedShape,
'Delete',
'delete',
u'Delete label',
enabled=False)
advancedMode = action('&Advanced Mode',
self.toggleAdvancedMode,
'Ctrl+Shift+A',
'expert',
u'Switch to advanced mode',
checkable=True)
help = action('&Tutorial', self.tutorial, 'T', 'help',
u'Show demo video')
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)
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,
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('Label Panel')
labels.setShortcut('Ctrl+Shift+L')
files = self.filedock.toggleViewAction()
files.setText('File Panel')
files.setShortcut('Ctrl+Shift+F')
labelMenu = QMenu()
addActions(labelMenu, (edit, delete))
self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
self.labelList.customContextMenuRequested.connect(
self.popLabelListMenu)
self.actions = struct(save=save,
open=open,
create=createMode,
delete=delete,
edit=edit,
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),
beginner=(),
advanced=(),
editMenu=(edit, delete),
beginnerContext=(edit, delete),
advancedContext=(save, createMode, editMode),
onLoadActive=(createMode, editMode))
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)
self.autoSaving = QAction("Auto Saving", self)
self.autoSaving.setCheckable(True)
self.singleClassMode = QAction("Single Class Mode", self)
self.singleClassMode.setShortcut("Ctrl+Shift+S")
self.singleClassMode.setCheckable(True)
self.lastLabel = None
addActions(self.menus.file, (open, opendir, openNextImg, openPrevImg))
addActions(self.menus.help, (help, ))
addActions(self.menus.view, (labels, files, zoomIn, zoomOut))
self.menus.file.aboutToShow.connect(self.updateFileMenu)
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.advanced = (None, open, opendir, None, openNextImg,
openPrevImg, None, save, createMode, editMode,
None, zoomIn, zoomOut, None)
self.statusBar().showMessage('%s started.' % __appname__)
self.statusBar().show()
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
self.difficult = False
self.settings = Settings()
self.settings.load()
settings = self.settings
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(300, 500))
position = settings.get(SETTING_WIN_POSE, QPoint(0, 0))
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 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()))
self.lineColor = QColor(
settings.get(SETTING_LINE_COLOR, Shape.line_color))
self.fillColor = QColor(
settings.get(SETTING_FILL_COLOR, Shape.fill_color))
Shape.line_color = self.lineColor
Shape.fill_color = self.fillColor
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()
self.updateFileMenu()
self.queueEvent(partial(self.loadFile, self.filePath or ""))
self.zoomWidget.valueChanged.connect(self.paintCanvas)
self.populateModeActions()
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):
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.advancedContext)
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)
def toggleActions(self, value=True):
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()
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 tutorial(self):
subprocess.Popen([self.screencastViewer, self.screencast])
def createShape(self):
assert self.beginner()
self.canvas.setEditing(False)
self.actions.create.setEnabled(False)
def toggleDrawingSensitive(self, drawing=True):
self.actions.editMode.setEnabled(not drawing)
if not drawing and self.beginner():
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, item=None):
if not self.canvas.editing():
return
item = item if item else self.currentItem()
text = self.labelDialog.popUp(item.text())
if text is not None:
item.setText(text)
self.setDirty()
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)
def btnstate(self, item=None):
if not self.canvas.editing():
return
item = self.currentItem()
if not item:
item = self.labelList.item(self.labelList.count() - 1)
try:
shape = self.itemsToShapes[item]
except:
pass
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):
item = HashableQListWidgetItem(shape.label)
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
self.itemsToShapes[item] = shape
self.shapesToItems[shape] = item
self.labelList.addItem(item)
def remLabel(self, shape):
if shape is None:
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)
self.addLabel(shape)
if line_color:
shape.line_color = QColor(*line_color)
if fill_color:
shape.fill_color = QColor(*fill_color)
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()
if s.line_color != self.lineColor else None,
fill_color=s.fill_color.getRgb()
if s.fill_color != self.fillColor else None,
points=[(p.x(), p.y()) for p in s.points],
difficult=s.difficult)
shapes = [format_shape(shape) for shape in self.canvas.shapes]
try:
if self.usingPascalVocFormat is True:
print('Img: ' + self.filePath + ' -> Its xml: ' +
annotationFilePath)
self.labelFile.savePascalVocFormat(annotationFilePath, shapes,
self.filePath,
self.imageData,
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())
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]
def labelItemChanged(self, item):
shape = self.itemsToShapes[item]
label = item.text()
if label != shape.label:
shape.label = item.text()
self.setDirty()
else:
self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
def newShape(self):
if not self.useDefaultLabelCheckbox.isChecked(
) or not self.defaultLabelTextLine.text():
if len(self.labelHist) > 0:
self.labelDialog = LabelDialog(parent=self,
listItem=self.labelHist)
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()
if text is not None:
self.prevLabelText = text
self.addLabel(self.canvas.setLastLabel(text))
if self.beginner():
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.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):
h_bar = self.scrollBars[Qt.Horizontal]
v_bar = self.scrollBars[Qt.Vertical]
h_bar_max = h_bar.maximum()
v_bar_max = v_bar.maximum()
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()
margin = 0.1
move_x = (cursor_x - margin * w) / (w - 2 * margin * w)
move_y = (cursor_y - margin * h) / (h - 2 * margin * h)
move_x = min(max(move_x, 0), 1)
move_y = min(max(move_y, 0), 1)
units = delta / (8 * 15)
scale = 10
self.addZoom(scale * units)
d_h_bar_max = h_bar.maximum() - h_bar_max
d_v_bar_max = v_bar.maximum() - v_bar_max
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):
self.resetState()
self.canvas.setEnabled(False)
if filePath is None:
filePath = self.settings.get(SETTING_FILENAME)
unicodeFilePath = ustr(filePath)
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)
else:
self.imageData = read(unicodeFilePath, None)
self.labelFile = None
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)
if self.usingPascalVocFormat is True:
if self.defaultSaveDir is not None:
basename = os.path.basename(
os.path.splitext(self.filePath)[0]) + XML_EXT
xmlPath = os.path.join(self.defaultSaveDir, basename)
self.loadPascalXMLByFilename(xmlPath)
else:
xmlPath = os.path.splitext(filePath)[0] + XML_EXT
if os.path.isfile(xmlPath):
self.loadPascalXMLByFilename(xmlPath)
self.setWindowTitle(__appname__ + ' ' + filePath)
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):
e = 2.0
w1 = self.centralWidget().width() - e
h1 = self.centralWidget().height() - e
a1 = w1 / h1
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):
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 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 is not None and len(self.defaultSaveDir) > 1:
settings[SETTING_SAVE_DIR] = ustr(self.defaultSaveDir)
else:
settings[SETTING_SAVE_DIR] = ""
if self.lastOpenDir is not None and len(self.lastOpenDir) > 1:
settings[SETTING_LAST_OPEN_DIR] = self.lastOpenDir
else:
settings[SETTING_LAST_OPEN_DIR] = ""
settings.save()
def loadRecent(self, filename):
if self.mayContinue():
self.loadFile(filename)
def scanAllImages(self, folderPath):
extensions = ['.jpg', '.png']
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)
try:
path = os.path.split(images[0])[0]
extention = os.path.splitext(images[0])[1]
images = [
int(os.path.splitext(os.path.basename(i))[0]) for i in images
]
images = sorted(images)
if platform == 'win32' or platform == 'win64':
images = [path + '\\' + str(i) + extention for i in images]
else:
images = [path + '/' + str(i) + extention for i in images]
except:
images.sort(key=lambda x: x.lower())
return images
def openAnnotation(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 openDir(self, _value=False):
if not self.mayContinue():
return
path = os.path.dirname(self.filePath) \
if self.filePath else '.'
if self.lastOpenDir is not None and len(self.lastOpenDir) > 1:
path = self.lastOpenDir
dirpath = ustr(
QFileDialog.getExistingDirectory(
self, '%s - Open Directory' % __appname__, path,
QFileDialog.ShowDirsOnly
| QFileDialog.DontResolveSymlinks))
if dirpath is not None and len(dirpath) > 1:
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 openPrevImg(self, _value=False):
if self.autoSaving.isChecked():
if self.defaultSaveDir is not None:
if self.dirty is True:
self.saveFile()
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):
if self.autoSaving.isChecked():
if self.defaultSaveDir is not None:
if self.dirty is True:
self.saveFile()
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]
try:
self.loadFile(filename)
except:
pass
def saveFile(self, _value=False):
imgFileDir = os.path.dirname(self.filePath)
imgFileName = os.path.basename(self.filePath)
savedFileName = os.path.splitext(imgFileName)[0] + XML_EXT
savedPath = os.path.join(imgFileDir, savedFileName)
self._saveFile(savedPath)
def saveFileAs(self, _value=False):
assert not self.image.isNull(), "cannot save empty image"
self._saveFile(self.saveFileDialog())
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)
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 = self.lineColor
self.canvas.update()
self.setDirty()
def chooseColor2(self):
color = self.colorDialog.getColor(self.fillColor,
u'Choose fill color',
default=DEFAULT_FILL_COLOR)
if color:
self.fillColor = color
Shape.fill_color = self.fillColor
self.canvas.update()
self.setDirty()
def deleteSelectedShape(self):
self.remLabel(self.canvas.deleteSelected())
self.setDirty()
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.lablHist = [line]
else:
self.labelHist.append(line)
def loadPascalXMLByFilename(self, xmlPath):
if self.filePath is None:
return
if os.path.isfile(xmlPath) is False:
return
tVocParseReader = PascalVocReader(xmlPath)
shapes = tVocParseReader.getShapes()
self.loadLabels(shapes)
self.canvas.verified = tVocParseReader.verified
class MainWindow(QMainWindow):
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
def __init__(self):
QMainWindow.__init__(self)
self.ui = ui.Ui_MainWindow()
self.ui.setupUi(self)
self.canvas = Canvas(self.ui.scrollArea)
self.canvas.zoomRequest.connect(self.zoomRequest)
self.canvas.scrollRequest.connect(self.scrollRequest)
self.canvas.newShape.connect(self.newShape)
self.canvas.changeLabel.connect(self.changeLabel)
self.canvas.shapeMoved.connect(self.setDirty)
self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
self.canvas.menu.addAction(self.ui.actionDelete_Shape)
self.ui.fileList.itemDoubleClicked.connect(self.fileitemDoubleClicked)
self.ui.scrollArea.setWidget(self.canvas)
self.ui.scrollArea.setWidgetResizable(True)
self.ui.dock_params.setVisible(False)
self.ui.dock_image_editing_operations.setVisible(True)
self.ui.current_label.setStyleSheet("QLabel { color : red; }")
self.ui.actionSave.setEnabled(True)
self.ui.actionSave_As.setEnabled(False)
self.updateFontSize(15)
self.scrollBars = {
Qt.Vertical: self.ui.scrollArea.verticalScrollBar(),
Qt.Horizontal: self.ui.scrollArea.horizontalScrollBar()
}
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,
}
self.mImgList = []
self.dirname = None
self.lastOpenDir = None
self.filePath = "."
self.defaultSaveDir = None
self.zoom_value = 100
# Whether we need to save or not.
self.dirty = False
# Enble auto saving if pressing next
self.autoSaving = True
self.itemToShapes = {} #listWidgetItem: shape1, shape2, ...
self.shapeToItem = {} #shape: listWidgetItem
self.labelToItem = {} #label: listWidgetItem
self.prevLabelText = ''
self.labelInfoDict = {}
#application state
self.image_np = None
self.image = QImage()
self.recentFiles = []
self.maxRecent = 7
self.zoom_level = 100
self.zoomMode = self.MANUAL_ZOOM
self.fit_window = False
self.usingPascalVocFormat = True
self.loadPredefinedClasses(
os.path.join(os.path.dirname(__file__), "predefined_labels.json"))
self.initLabelList()
self.initUndoAction()
self.predefined_labels_num = 0
@pyqtSlot()
def on_actionOpen_triggered(self):
self.openFile()
@pyqtSlot()
def on_actionOpen_Dir_triggered(self):
opened = self.openDir()
self.ui.actionSave.setEnabled(opened)
self.ui.actionSave_As.setEnabled(opened)
@pyqtSlot()
def on_actionOpen_Recent_triggered(self):
pass
@pyqtSlot()
def on_actionSave_triggered(self):
self.saveFile()
@pyqtSlot()
def on_actionSave_As_triggered(self):
self.saveFileAs()
@pyqtSlot()
def on_actionSave_Files_triggered(self):
self.saveFiles()
@pyqtSlot()
def on_actionClose_triggered(self):
self.closeFile()
@pyqtSlot()
def on_actionQuit_triggered(self):
self.close()
@pyqtSlot()
def on_actionCreate_Shape_triggered(self):
self.createShape()
@pyqtSlot()
def on_actionDelete_Shape_triggered(self):
self.deleteSelectedShape()
@pyqtSlot()
def on_actionDelete_Label_triggered(self):
self.deleteSelectedLabel()
@pyqtSlot()
def on_actionZoom_In_triggered(self):
self.addZoom(10)
@pyqtSlot()
def on_actionZoom_Out_triggered(self):
self.addZoom(-10)
@pyqtSlot()
def on_actionZoom_Original_triggered(self):
self.setZoom(100)
@pyqtSlot()
def on_actionUpdate_Line_Width_triggered(self):
self.updateShapesWidth()
@pyqtSlot()
def on_actionUpdate_Font_Size_triggered(self):
point_size, result = QInputDialog.getInt(self, "Set Font Size",
"Font Size", 15, 0, 100, 1)
self.updateFontSize(point_size)
@pyqtSlot()
def on_actionFit_Window_triggered(self):
self.setFitWindow(self.FIT_WINDOW)
@pyqtSlot()
def on_actionFit_Width_triggered(self):
self.setFitWidth(self.FIT_WIDTH)
@pyqtSlot()
def on_actionClear_Canvas_triggered(self):
self.clearCanvas()
@pyqtSlot()
def on_actionNext_Image_triggered(self):
self.openNextImg()
@pyqtSlot()
def on_actionPrev_Image_triggered(self):
self.openPrevImg()
@pyqtSlot()
def on_actionCheck_Label_triggered(self):
if not isinstance(self.dirname, str):
QMessageBox.warning(self, 'warning', 'Please select dir')
else:
self.get_label_info(self.dirname)
@pyqtSlot(bool)
def on_actionShow_Lable_List_triggered(self, checked):
self.ui.dock_label.setVisible(checked)
@pyqtSlot(bool)
def on_actionShow_File_List_triggered(self, checked):
self.ui.dock_file_list.setVisible(checked)
@pyqtSlot(bool)
def on_actionShow_Dilate_Params_triggered(self, checked):
self.ui.dock_params.setVisible(checked)
@pyqtSlot(QListWidgetItem)
def on_labelList_itemActivated(self):
self.labelSelectionChanged()
@pyqtSlot()
def on_labelList_itemSelectionChanged(self):
self.labelSelectionChanged()
@pyqtSlot(int)
def on_zoomSlider_valueChanged(self, value):
if value == 0:
return
self.setZoom(100 + value)
self.ui.zoomValue.setValue(value)
@pyqtSlot(int)
def on_eraser_strength_valueChanged(self, value):
self.canvas.setEraserStrength(value)
@pyqtSlot()
def on_exit_image_editing_clicked(self):
if self.ui.exit_image_editing.text() == 'Exit Image Editing':
self.set_image_operations('label')
self.canvas.mode = Canvas.EDIT
self.ui.exit_image_editing.setText('Enter Image Editing')
self.ui.crop.setDisabled(True)
self.ui.reverse_crop.setDisabled(True)
self.ui.magic_wand.setDisabled(True)
self.ui.eraser_strength.setDisabled(True)
else:
self.set_image_operations('crop')
self.ui.exit_image_editing.setText('Exit Image Editing')
self.ui.crop.setDisabled(False)
self.ui.reverse_crop.setDisabled(False)
self.ui.magic_wand.setDisabled(False)
self.ui.eraser_strength.setDisabled(False)
@pyqtSlot()
def on_crop_clicked(self):
print('crop')
self.set_image_operations('crop')
self.canvas.cropping()
@pyqtSlot()
def on_reverse_crop_clicked(self):
self.set_image_operations('crop')
self.canvas.cropping(True)
@pyqtSlot(bool)
def on_magic_wand_clicked(self, checked):
if checked:
self.set_image_operations('magic_wand')
else:
self.set_image_operations('crop')
def get_label_info(self, dir_path):
labels_path = os.path.join(dir_path, 'labels')
labels_exist_path = labels_path if os.path.exists(
labels_path) else dir_path
self.label_to_num = {}
for file_name in tqdm(os.listdir(labels_exist_path), "check label"):
if os.path.splitext(file_name)[1] != XML_EXT:
continue
xml_path = os.path.join(labels_exist_path, file_name)
reader = PascalVocReader(xml_path)
for shape in reader.shapes:
if shape[0] not in self.label_to_num:
self.label_to_num.setdefault(shape[0], 1)
else:
self.label_to_num[shape[0]] += 1
list_of_dataset = os.listdir(dir_path)
self.label_to_num['origin_image_num'] = 0
self.label_to_num['labeled_image_num'] = 0
for each_item in list_of_dataset:
if os.path.splitext(each_item)[1] == '.jpg':
self.label_to_num['origin_image_num'] += 1
else:
self.check_label_exist(os.path.join(dir_path, each_item))
QMessageBox.information(None, 'info', str(self.label_to_num))
def check_label_exist(self, xml_item):
parser = ET.XMLParser(encoding="utf-8")
tree = ET.parse(xml_item, parser)
root = tree.getroot()
check_object = root.find("./object")
if check_object:
self.label_to_num['labeled_image_num'] += 1
def set_image_operations(self, value):
print('image opt', value)
if value.lower() == 'crop':
self.canvas.image_editing_status = Canvas.CROP
elif value.lower() == "magic_wand":
self.canvas.image_editing_status = Canvas.MAGIC_WAND
self.canvas.mode = Canvas.ERASING
else:
self.canvas.image_editing_status = Canvas.LABEL
def initUndoAction(self):
undoAction = self.canvas.undoStack.createUndoAction(self, 'Undo')
undoAction.setShortcuts(QKeySequence.Undo)
icon = QIcon()
icon.addPixmap(QPixmap("icons/undo.png"), QIcon.Normal, QIcon.Off)
undoAction.setIcon(icon)
self.ui.menuEdit.addAction(undoAction)
self.ui.mainToolBar.addAction(undoAction)
redoAction = self.canvas.undoStack.createRedoAction(self, 'Redo')
redoAction.setShortcuts(QKeySequence.Redo)
def noShapes(self):
return not self.itemToShapes
def toggleAdvancedMode(self, value=True):
self.canvas.setEditing(True)
self.populateModeActions()
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 setDirty(self):
self.dirty = True
def setClean(self):
self.dirty = False
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 status(self, message, delay=5000):
self.statusBar().showMessage(message, delay)
def resetState(self):
self.itemToShapes.clear()
self.shapeToItem.clear()
self.ui.labelList.clear()
self.filePath = None
self.imageData = None
self.labelFile = None
self.canvas.resetState()
def currentItem(self):
items = self.ui.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 createShape(self):
self.canvas.currentLabel = self.ui.current_label.text()
self.canvas.setEditing(False)
def toggleDrawingSensitive(self, drawing=True):
"""In the middle of drawing, toggling between modes should be disabled."""
if not drawing:
# Cancel creation.
print('Cancel creation.')
self.canvas.setEditing(True)
self.canvas.restoreCursor()
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)
def editLabel(self, item=None):
if not self.canvas.editing():
return
item = item if item else self.currentItem()
text, result = QInputDialog.getText(self, "Set Label Name",
"Label Name", QLineEdit.Normal, '')
if text is not None and result:
item.setText(text)
self.setDirty()
def editShape(self, shape):
pass
# Tzutalin 20160906 : Add file list and dock to move faster
def fileitemDoubleClicked(self, item=None):
if len(self.canvas.shapes) > 0 and self.dirty:
self.saveFile()
currIndex = self.mImgList.index(item.text())
if currIndex < len(self.mImgList):
filename = self.mImgList[currIndex]
if filename:
self.loadFile(filename)
# React to canvas signals.
def shapeSelectionChanged(self, selected=False):
self.restore_list_style()
shape = self.canvas.selectedShape
if shape is None:
return
if shape in self.shapeToItem:
blocked = self.ui.labelList.blockSignals(True)
self.shapeToItem[shape].setSelected(True)
self.shapeToItem[shape].setFont(FONT_SELECTED)
self.ui.labelList.blockSignals(blocked)
def addShapeList(self, shape):
item = self.labelToItem[shape.label]
try:
self.itemToShapes[item].append(shape)
except KeyError:
self.itemToShapes[item] = []
self.itemToShapes[item].append(shape)
self.shapeToItem[shape] = item
def addLabel(self):
text, result = QInputDialog.getText(self, "Set Label Name",
"Label Name", QLineEdit.Normal,
self.prevLabelText)
if not text:
return None
if text in self.labelInfoDict.keys():
return text #label is exist
item = HashableQListWidgetItem(text)
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
item.setSelected(True)
self.ui.labelList.addItem(item)
self.labelToItem[text] = item
self.ui.labelList.setCurrentItem(item)
self.itemToShapes[item] = []
self.labelInfoDict[text] = (
True, self.canvas.POLYGON, self.canvas.LABEL_COLORS[int(
len(self.labelInfoDict) % len(self.canvas.LABEL_COLORS))]
) #name, isClosePath, polygon, enable
return text
def updateShapesWidth(self):
lineWidth, result = QInputDialog.getInt(
self, "Set LineWidth for all Shapes", "Line Width", 1, 0, 100, 1)
self.canvas.updateShapes(lineWidth)
def updateFontSize(self, point_size):
font = QFont()
font.setPointSize(point_size)
widgets = self.findChildren(QWidget)
for w in widgets:
w.setFont(font)
self.update()
def initLabelList(self):
if self.labelInfoDict is None:
return
self.ui.labelList.clear()
for label in self.labelInfoDict.keys():
item = HashableQListWidgetItem(label)
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
item.setSelected(True)
self.ui.labelList.addItem(item)
self.labelToItem[label] = item
self.ui.labelList.setCurrentRow(0)
def removeLabel(self, shape):
if shape is None:
return
item = self.shapeToItem[shape]
self.itemToShapes[item].remove(shape)
if len(self.itemToShapes[item]) == 0:
del self.itemToShapes[item]
del self.shapeToItem[shape]
def loadLabels(self, shapes):
s = []
self.initLabelList()
idx = len(self.labelInfoDict)
for label, lineWidth, points, ellipse_points, rotate, r1, r2, center in shapes:
if label not in self.labelInfoDict.keys():
with open(
os.path.join(os.path.dirname(__file__),
"predefined_labels.json")) as f:
predefined_json = json.load(f)
new_append_predefined_dict = {
"name": "",
"isPathClosed": True,
"type": 0,
"enable": True
}
new_append_predefined_dict['name'] = label
if r1:
new_append_predefined_dict['type'] = 1
predefined_json['labels'].append(
new_append_predefined_dict)
with open(
os.path.join(os.path.dirname(__file__),
"predefined_labels.json"), "w") as f:
json.dump(predefined_json, f, indent=4)
self.canvas.LABEL_COLORS.extend(self.canvas.random_colors(N=1))
self.labelInfoDict[label] = (
True, new_append_predefined_dict['type'],
self.canvas.LABEL_COLORS[idx])
item = HashableQListWidgetItem(label)
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
item.setSelected(True)
self.ui.labelList.addItem(item)
self.labelToItem[label] = item
idx = idx + 1
shape = Shape(label=label)
shape.lineWidth = lineWidth
shape.isPathClosed = self.labelInfoDict[label][0]
shape.d_type = self.labelInfoDict[label][1]
shape.lineColor = self.labelInfoDict[label][2]
shape.rotate = rotate
shape.r1 = r1
shape.r2 = r2
shape.center = center
for x, y in points:
shape.addPoint(QPointF(x, y))
for x, y in ellipse_points:
shape.ellipse_points.append(QPointF(x, y))
if shape.d_type == self.canvas.Ellipse:
shape.sampleEllipsePoints()
shape.close()
s.append(shape)
self.addShapeList(shape)
self.canvas.loadShapes(s)
def calcShapeParent(self, s):
s.parentGuid = s.guid
for shape in self.canvas.shapes:
if shape.containsShape(s):
s.parentGuid = shape.guid
return
def saveLabels(self, annotationFilePath):
annotationFilePath = annotationFilePath
imgFileDir = os.path.dirname(annotationFilePath)
imgFileName = os.path.basename(annotationFilePath)
dataFileName = os.path.splitext(imgFileName)[0] + ".compressed"
dataPath = os.path.join(imgFileDir, dataFileName)
imgPath = os.path.join(imgFileDir,
os.path.splitext(imgFileName)[0] + ".jpg")
if self.labelFile is None:
self.labelFile = LabelFile()
self.labelFile.verified = self.canvas.verified
def format_shape(s):
self.calcShapeParent(s)
return dict(label=s.label,
guid=s.guid,
parentGuid=s.parentGuid,
lineWidth=s.lineWidth,
rotate=s.rotate,
r1=s.r1,
r2=s.r2,
center=(s.center.x(), s.center.y()),
points=[(p.x(), p.y()) for p in s.points],
ellipse_points=[(p.x(), p.y())
for p in s.ellipse_points])
shapes = [format_shape(shape) for shape in self.canvas.shapes]
# Can add differrent annotation formats here
try:
if self.usingPascalVocFormat is True:
print('Img: ' + imgPath + ' -> Its xml: ' +
annotationFilePath + ' -> label data file: ' + dataPath)
labelNames = []
for label in self.labelInfoDict.keys():
labelNames.append(label[0])
self.labelFile.setLabels(labelNames)
self.labelFile.savePascalVocFormat(annotationFilePath, shapes,
imgPath, self.imageData)
#for deeplearing, no need to save compress label data
if self.ui.actionSave_Label_Compress_Data.isChecked():
self.labelFile.saveData(
dataPath, self.canvas.shapes,
imgPath, (self.ui.cannyKSize.value(),
self.ui.cannyKSize.value()),
(self.ui.boxKSize.value(), self.ui.boxKSize.value()),
(self.ui.labelKSize.value(),
self.ui.labelKSize.value()),
self.ui.actionSave_Rf_Data.isChecked(
)) #check if save random forest data file
else:
self.labelFile.save(annotationFilePath, shapes, imgPath,
self.imageData)
return True
except LabelFileError as e:
self.errorMessage(u'Error saving label data', u'<b>%s</b>' % e)
return False
def copySelectedShape(self):
self.addShapeList(self.canvas.copySelectedShape())
# fix copy and delete
self.shapeSelectionChanged(True)
def restore_list_style(self):
for i in range(self.ui.labelList.count()):
self.ui.labelList.item(i).setFont(FONT_NORMAL)
def labelSelectionChanged(self):
self.restore_list_style()
item = self.currentItem()
if item:
if self.ui.fileList.currentItem() is not None:
self.setWindowTitle(__appname__ + ' ' +
self.ui.fileList.currentItem().text())
item.setFont(FONT_SELECTED)
self.ui.current_label.setText(item.text())
label = item.text()
labelInfo = self.labelInfoDict[label]
self.canvas.currentLabel = label
self.canvas.currentDType = labelInfo[1]
try:
shapes = []
shapes = self.itemToShapes[item]
self.canvas.selectShapes(shapes, label, labelInfo)
except KeyError:
self.canvas.deSelectShapes()
print("no shape is in current label")
def labelItemChanged(self, item):
if self.itemToShapes is None:
return
self.ui.labelList.setCurrentItem(item)
shape = self.itemToShapes[item]
if item.text() != shape.label:
shape.label = item.text()
self.setDirty()
else: # User probably changed item visibility
self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
# Callback functions:
def newShape(self):
text = self.ui.current_label.text()
if text:
self.prevLabelText = text
shape = self.canvas.setLastLabel(text)
shape.isPathClosed = self.labelInfoDict[text][0]
shape.d_type = self.labelInfoDict[text][1]
shape.label = text
shape.lineColor = self.labelInfoDict[text][2]
self.addShapeList(shape)
self.setDirty()
else:
self.canvas.resetAllLines()
def changeLabel(self, shape, change_item):
if change_item not in list(self.labelInfoDict.keys()):
QMessageBox.warning(self, 'check label',
'modified label not in Labels list')
else:
shape.label = change_item
self.setDirty()
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.zoomMode = self.MANUAL_ZOOM
self.setZoomValue(value)
def addZoom(self, increment=10):
self.setZoom(self.zoom_value + increment)
def zoomRequest(self, delta):
units = delta / (8 * 15)
scale = 10
self.addZoom(scale * units)
def setFitWindow(self, value=True):
self.zoomMode = self.FIT_WINDOW if value else self.MANUAL_ZOOM
self.adjustScale()
def setFitWidth(self, value=True):
self.zoomMode = self.FIT_WIDTH if value else self.MANUAL_ZOOM
self.adjustScale()
def clearCanvas(self):
for shape in self.canvas.shapes:
self.canvas.selectedShape = shape
self.deleteSelectedShape()
def togglePolygons(self, value):
for item, shape in self.itemToShapes.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.canvas.setEnabled(False)
if filePath is None:
filePath = self.settings.get('filename')
unicodeFilePath = filePath
# Tzutalin 20160906 : Add file list and dock to move faster
# Highlight the file item
if unicodeFilePath and self.ui.fileList.count() > 0:
index = self.mImgList.index(unicodeFilePath)
fileWidgetItem = self.ui.fileList.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
else:
# Load image:
# read data first and store for saving into label file
self.canvas.image_np = cv2.imread(
unicodeFilePath) # read as 3-channel
with open(unicodeFilePath, 'rb') as f:
self.imageData = f.read()
self.labelFile = None
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.paintCanvas()
self.addRecentFile(self.filePath)
# 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]) + XML_EXT
xmlPath = os.path.join(self.defaultSaveDir, basename)
self.loadPascalXMLByFilename(xmlPath)
else:
xmlPath = os.path.splitext(self.filePath)[0] + XML_EXT
if os.path.isfile(xmlPath):
self.loadPascalXMLByFilename(xmlPath)
self.setWindowTitle(__appname__ + ' ' + filePath)
# Default : select last item if there is at least one item
#if self.ui.labelList.count():
#self.ui.labelList.setCurrentItem(self.ui.labelList.item(self.ui.labelList.count()-1))
#self.ui.labelList.setItemSelected(self.ui.labelList.item(self.ui.labelList.count()-1), 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 adjustScale(self, initial=False):
value = self.scalers[self.FIT_WINDOW if initial else self.zoomMode]()
self.setZoomValue(value)
def setZoomValue(self, value):
self.zoom_value = value
self.paintCanvas()
def paintCanvas(self):
assert not self.image.isNull(), "cannot paint null image"
self.canvas.scale = 0.01 * self.zoom_value
self.canvas.adjustSize()
self.canvas.update()
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 loadRecent(self, filename):
if self.mayContinue():
self.loadFile(filename)
def scanAllImages(self, folderPath):
extensions = ['.jpeg', '.jpg', '.png', '.bmp']
images = []
for root, dirs, files in os.walk(folderPath):
for file in files:
if file.lower().endswith(tuple(extensions)):
relatviePath = os.path.join(root, file)
path = os.path.abspath(relatviePath)
images.append(path)
images.sort(key=lambda x: x.lower())
return images
def changeSavedir(self, _value=False):
if self.defaultSaveDir is not None:
path = self.defaultSaveDir
else:
path = '.'
dirpath = QFileDialog.getExistingDirectory(
self, '%s - Save 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 openAnnotation(self, _value=False):
if self.filePath is None:
return
path = os.path.dirname(self.filePath)\
if self.filePath else '.'
if self.usingPascalVocFormat:
filters = "Open Annotation XML file (%s)" % \
' '.join(['*.xml'])
filename = 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 openDir(self, _value=False):
if not self.mayContinue():
return False
path = os.path.dirname(self.filePath if self.filePath else '.')
if self.lastOpenDir is not None and len(self.lastOpenDir) > 1:
path = self.lastOpenDir
dirpath = QFileDialog.getExistingDirectory(
self, '%s - Open Directory' % __appname__, path,
QFileDialog.ShowDirsOnly
| QFileDialog.DontResolveSymlinks)
if dirpath is not None and len(dirpath) > 1:
self.lastOpenDir = dirpath
self.dirname = dirpath
self.filePath = None
self.ui.fileList.clear()
self.mImgList = self.scanAllImages(dirpath)
self.openNextImg()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.ui.fileList.addItem(item)
return True
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()
self.labelFile.toggleVerify()
self.canvas.verified = self.labelFile.verified
self.paintCanvas()
self.saveFile()
def openPrevImg(self, _value=False):
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 next image without dialog if having any label
print('image editing', self.canvas.image_editing_status)
if self.canvas.image_editing_status != Canvas.LABEL:
self.saveFile()
if self.autoSaving and self.dirty and self.filePath:
print("save file when open next img")
self.saveFile()
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(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.canvas.image_editing_status != Canvas.LABEL:
filenameWithoutExtension = os.path.splitext(self.filePath)[0]
dest_path = os.path.join(self.currentPath(),
filenameWithoutExtension + '.png')
b, g, r = cv2.split(self.canvas.image_np)
gray = cv2.cvtColor(self.canvas.image_np, cv2.COLOR_BGR2GRAY)
_, a = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)
a_large_canvas = np.zeros((a.shape[0] + 200, a.shape[1] + 200),
np.uint8)
a_large_canvas[100:100 + a.shape[0], 100:100 + a.shape[1]] = a
a_large_canvas = cv2.erode(a_large_canvas, np.ones((3, 3),
np.uint8), 1)
a = a_large_canvas[100:100 + a.shape[0], 100:100 + a.shape[1]]
image = cv2.merge([b, g, r, a])
_, cnts, hier = cv2.findContours(a, cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
cnts_lens = map(lambda x: len(x), cnts)
max_cnt_id = np.argmax(np.array(cnts_lens), axis=0)
bndbox = cv2.boundingRect(cnts[max_cnt_id])
image = image[bndbox[1]:bndbox[1] + bndbox[3],
bndbox[0]:bndbox[0] + bndbox[2]]
cv2.imwrite(dest_path, image)
elif self.defaultSaveDir is not None and len(self.defaultSaveDir):
if self.filePath:
imgFileName = os.path.basename(self.filePath)
prefix = os.path.splitext(imgFileName)[0]
savedFileName = prefix + XML_EXT
savedPath = os.path.join(self.defaultSaveDir, savedFileName)
self._saveFile(savedPath)
else:
imgFileDir = os.path.dirname(self.filePath)
imgFileName = os.path.basename(self.filePath)
prefix = os.path.splitext(imgFileName)[0]
savedFileName = prefix + XML_EXT
savedPath = os.path.join(imgFileDir, savedFileName)
dest_path = savedPath if self.labelFile else self.saveFileDialog()
self._saveFile(dest_path)
def saveFiles(self):
if self.ui.fileList.count() < 1:
return
for i in range(self.ui.fileList.count()):
print('Saving file', i)
item = self.ui.fileList.item(i)
filePath = item.text()
filedir, filename = os.path.split(filePath)
filenamesplit = os.path.splitext(filename)
savedFileName = filenamesplit[0] + XML_EXT
savedPath = os.path.join(filedir, savedFileName)
self.loadFile(filePath)
if not self._saveFile(savedPath):
break
def saveFileAs(self, _value=False):
if self.canvas.image_editing_status != Canvas.LABEL:
dest_path = self.saveImageDialog()
b, g, r = cv2.split(self.canvas.image_np)
gray = cv2.cvtColor(self.canvas.image_np, cv2.COLOR_BGR2GRAY)
_, a = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY)
a = cv2.erode(a, np.ones((5, 5), np.uint8), 1)
image = cv2.merge([b, g, r, a])
cv2.imwrite(dest_path, image)
else:
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_():
return dlg.selectedFiles()[0]
return ''
def saveImageDialog(self):
caption = '%s - Choose File' % __appname__
filters = 'File (*%s)' % ('.png')
openDialogPath = self.currentPath()
dlg = QFileDialog(self, caption, openDialogPath, filters)
dlg.setDefaultSuffix('.png')
dlg.setAcceptMode(QFileDialog.AcceptSave)
filenameWithoutExtension = os.path.splitext(self.filePath)[0]
dlg.selectFile(filenameWithoutExtension)
dlg.setOption(QFileDialog.DontUseNativeDialog, False)
if dlg.exec_():
return dlg.selectedFiles()[0]
return ''
def checkShapes(self):
contour_label = 'outer-contour'
contours = [
shape for shape in self.canvas.shapes
if shape.label == contour_label
]
others = [
shape for shape in self.canvas.shapes
if shape.label != contour_label
]
violations = [[c1, c2] for c1 in contours for c2 in contours
if c1.containsShape(c2)]
if len(violations) > 0:
self.canvas.selectShapes(violations[0], None, None)
return 'Outer-contours should not contain another outer-contour\n外轮廓内部不能包含其它外轮廓'
violations = [[c, o] for c in contours for o in others
if c.intersectsWith(o)]
if len(violations) > 0:
self.canvas.selectShapes(violations[0], None, None)
return 'Outer-contours should not intersect with non-outer-contours\n外轮廓与非外轮廓不能交叉'
violations = [c for c in contours if c.containsNone(others)]
if len(violations) > 0:
self.canvas.selectShapes(violations, None, None)
return 'Outer-contours should contain at least 1 non-outer-contour\n外轮廓内部至少包含1个非外轮廓'
polygons = [
shape for shape in others if shape.d_type == shape.S_POLYGON
]
violations = [[p1, p2] for p1 in polygons for p2 in polygons
if p1 is not p2 and p1.intersectsWith(p2)]
if len(violations) > 0:
self.canvas.selectShapes(violations[0], None, None)
return 'Non-outer-contour-polygons should not intersect with each other\n非外轮廓多边形不能互相交叉'
return None
def _saveFile(self, annotationFilePath):
#self.canvas.overrideCursor(WAIT_CURSOR)
error_info = self.checkShapes()
if error_info is not None:
msg = QMessageBox()
msg.setIcon(QMessageBox.Warning)
msg.setText("Save File Failed! Error:\n" + error_info)
msg.exec_()
return False
if annotationFilePath and self.saveLabels(annotationFilePath):
self.setClean()
self.statusBar().showMessage('Saved to %s' % annotationFilePath)
self.statusBar().show()
return True
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 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 deleteSelectedShape(self):
self.removeLabel(self.canvas.deleteSelected())
self.setDirty()
#if self.noShapes():
# for action in self.actions.onShapesPresent:
# action.setEnabled(False)
def deleteSelectedLabel(self):
if self.ui.labelList.currentItem() is None:
return
try:
shapes = self.itemToShapes[self.ui.labelList.currentItem()]
for shape in shapes:
del self.shapeToItem[shape]
self.canvas.deleteShape(shape)
except KeyError:
self.ui.labelList.takeItem(self.ui.labelList.currentRow())
return
del self.itemToShapes[self.ui.labelList.currentItem()]
self.ui.labelList.takeItem(self.ui.labelList.currentRow())
self.setDirty()
def copyShape(self):
self.canvas.endMove(copy=True)
self.addShapeList(self.canvas.selectedShape)
self.setDirty()
def moveShape(self):
self.canvas.endMove(copy=False)
self.setDirty()
def loadPredefinedClasses(self, predefClassesFile):
label_data = open(predefClassesFile)
data = json.load(label_data)
label_configs = data["labels"]
self.predefined_labels_num = len(label_configs)
if self.predefined_labels_num > len(self.canvas.LABEL_COLORS):
QColor_list = self.canvas.random_colors(
(self.predefined_labels_num - len(self.canvas.LABEL_COLORS)))
self.canvas.LABEL_COLORS.extend(QColor_list)
idx = 0
for config in label_configs:
if config["enable"]:
self.labelInfoDict[config["name"]] = (
config["isPathClosed"], config["type"],
self.canvas.LABEL_COLORS[idx])
idx = idx + 1
if len(self.labelInfoDict) > 0:
self.ui.current_label.setText(list(self.labelInfoDict.keys())[0])
print("predefined labels:", self.labelInfoDict)
def loadPascalXMLByFilename(self, xmlPath):
if self.filePath is None:
return
if os.path.isfile(xmlPath) is False:
return
tVocParseReader = PascalVocReader(xmlPath)
shapes = tVocParseReader.shapes
self.loadLabels(shapes)
self.canvas.verified = tVocParseReader.verified
