class MainWindow(QMainWindow, WindowMixin): FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = 0, 1, 2 def __init__(self, filename=None, output=None): super(MainWindow, self).__init__() self.setWindowTitle(__appname__) # Whether we need to save or not. self.dirty = False self._noSelectionSlot = False self._beginner = True self.screencastViewer = "firefox" self.screencast = "screencast.ogv" # Main widgets and related state. self.labelDialog = LabelDialog(parent=self) self.labelList = QListWidget() self.itemsToShapes = [] 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 = QVBoxLayout() listLayout.setContentsMargins(0, 0, 0, 0) listLayout.addWidget(self.labelList) self.editButton = QToolButton() self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) self.labelListContainer = QWidget() self.labelListContainer.setLayout(listLayout) listLayout.addWidget(self.editButton)#, 0, Qt.AlignCenter) listLayout.addWidget(self.labelList) #define table view self.table_view = QTableView(self) self.table_view.setSelectionBehavior(QAbstractItemView.SelectRows) self.table_view.show() self.table_view.doubleClicked.connect(self.openCurrentData) listLayout.addWidget(self.table_view) self.dock = QDockWidget('Polygon Labels', self) self.dock.setObjectName('Labels') self.dock.setWidget(self.labelListContainer) #self.dock.setWidget(self.infoLayout) 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.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.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', 'Quit application') open = action('&Open', self.openFile, 'Ctrl+O', 'open', 'Open image or label file') save = action('&Save', self.saveFile, 'Ctrl+S', 'save', 'Save labels to file', enabled=False) saveAs = action('&Save As', self.saveFileAs, 'Ctrl+Shift+S', 'save-as', 'Save labels to a different file', enabled=False) close = action('&Close', self.closeFile, 'Ctrl+W', 'close', 'Close current file') color1 = action('Polygon &Line Color', self.chooseColor1, 'Ctrl+L', 'color_line', 'Choose polygon line color') color2 = action('Polygon &Fill Color', self.chooseColor2, 'Ctrl+Shift+L', 'color', 'Choose polygon fill color') createMode = action('Create\nPolygo&ns', self.setCreateMode, 'Ctrl+N', 'new', 'Start drawing polygons', enabled=False) editMode = action('&Edit\nPolygons', self.setEditMode, 'Ctrl+J', 'edit', 'Move and edit polygons', enabled=False) create = action('Create\nPolygo&n', self.createShape, 'Ctrl+N', 'new', 'Draw a new polygon', enabled=False) delete = action('Delete\nPolygon', self.deleteSelectedShape, 'Delete', 'delete', 'Delete', enabled=False) copy = action('&Duplicate\nPolygon', self.copySelectedShape, 'Ctrl+D', 'copy', 'Create a duplicate of the selected polygon', enabled=False) advancedMode = action('&Advanced Mode', self.toggleAdvancedMode, 'Ctrl+Shift+A', 'expert', 'Switch to advanced mode', checkable=True) hideAll = action('&Hide\nPolygons', partial(self.togglePolygons, False), 'Ctrl+H', 'hide', 'Hide all polygons', enabled=False) showAll = action('&Show\nPolygons', partial(self.togglePolygons, True), 'Ctrl+A', 'hide', 'Show all polygons', enabled=False) help = action('&Tutorial', self.tutorial, 'Ctrl+T', 'help', 'Show screencast of introductory tutorial') zoom = QWidgetAction(self) zoom.setDefaultWidget(self.zoomWidget) self.zoomWidget.setWhatsThis( "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', 'Increase zoom level', enabled=False) zoomOut = action('&Zoom Out', partial(self.addZoom, -10), 'Ctrl+-', 'zoom-out', 'Decrease zoom level', enabled=False) zoomOrg = action('&Original size', partial(self.setZoom, 100), 'Ctrl+=', 'zoom', 'Zoom to original size', enabled=False) fitWindow = action('&Fit Window', self.setFitWindow, 'Ctrl+F', 'fit-window', 'Zoom follows window size', checkable=True, enabled=False) fitWidth = action('Fit &Width', self.setFitWidth, 'Ctrl+Shift+F', 'fit-width', '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', 'Modify the label of the selected polygon', enabled=False) self.editButton.setDefaultAction(edit) shapeLineColor = action('Shape &Line Color', self.chshapeLineColor, icon='color_line', tip='Change the line color for this specific shape', enabled=False) shapeFillColor = action('Shape &Fill Color', self.chshapeFillColor, icon='color', tip='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) # Store actions for further handling. self.actions = struct(save=save, saveAs=saveAs, open=open, close=close, lineColor=color1, fillColor=color2, 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,save,saveAs,close,quit), beginner=(), advanced=(), editMenu=(edit, copy, delete, None, color1, color2), 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) addActions(self.menus.file, (open, self.menus.recentFiles, save, saveAs, close, None, quit)) addActions(self.menus.help, (help,)) addActions(self.menus.view, ( 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, save, None, create, copy, delete, None, zoomIn, zoom, zoomOut, fitWindow, fitWidth) self.actions.advanced = ( open, save, None, createMode, editMode, None, hideAll, showAll) self.statusBar().showMessage('%s started.' % __appname__) self.statusBar().show() # Application state. self.image = QImage() self.filename = filename self.labeling_once = output is not None self.output = output self.recentFiles = [] self.maxRecent = 7 self.lineColor = None self.fillColor = None self.zoom_level = 100 self.fit_window = False # XXX: Could be completely declarative. # Restore application settings. self.settings = {} self.recentFiles = self.settings.get('recentFiles', []) size = self.settings.get('window/size', QSize(1000, 800)) position = self.settings.get('window/position', QPoint(0, 0)) self.resize(size) self.move(position) # or simply: #self.restoreGeometry(settings['window/geometry'] self.restoreState(self.settings.get('window/state', QByteArray())) self.lineColor = QColor(self.settings.get('line/color', Shape.line_color)) self.fillColor = QColor(self.settings.get('fill/color', Shape.fill_color)) Shape.line_color = self.lineColor Shape.fill_color = self.fillColor if self.settings.get('advanced', QVariant()): 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. self.queueEvent(partial(self.loadFile, self.filename)) # Callbacks: self.zoomWidget.valueChanged.connect(self.paintCanvas) self.populateModeActions() #self.firstStart = True #if self.firstStart: # QWhatsThis.enterWhatsThisMode() ## Support Functions ## def openCurrentData(self,e): row_idx, key = self.getCurrentData() row_data = self.mymodel.getRow(row_idx) try: if row_data[1].lower()=="yes": self.filename=self.labelFileDict[key] else: self.filename=self.imageFileDict[key] self.loadFile(self.filename) except KeyError as e: print(e, " is not a dictionary key") def getCurrentData(self): index = self.table_view.selectionModel().currentIndex() if index.column(): return index.row(), self.mymodel.getRow(index.row())[0] else: return index.row(), index.data() def updateTable(self, table_data): header = ["filename", "is labeled"] self.mymodel=MyTableModel(table_data, header, self) self.table_view.setModel(self.mymodel) self.table_view.resizeColumnsToContents() self.table_view.update() 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 = [] self.labelList.clear() self.filename = 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, filename): if filename in self.recentFiles: self.recentFiles.remove(filename) elif len(self.recentFiles) >= self.maxRecent: self.recentFiles.pop() self.recentFiles.insert(0, filename) def beginner(self): return self._beginner def advanced(self): return not self.beginner() ## Callbacks ## 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): """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. 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) def updateFileMenu(self): current = self.filename def exists(filename): return os.path.exists(str(filename)) menu = self.menus.recentFiles menu.clear() files = [f for f in self.recentFiles if f != current 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() # React to canvas signals. def shapeSelectionChanged(self, selected=False): if self._noSelectionSlot: self._noSelectionSlot = False else: shape = self.canvas.selectedShape if shape: for item, shape_ in self.itemsToShapes: if shape_ == shape: break item.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 = QListWidgetItem(shape.label) item.setFlags(item.flags() | Qt.ItemIsUserCheckable) item.setCheckState(Qt.Checked) self.itemsToShapes.append((item, shape)) self.labelList.addItem(item) for action in self.actions.onShapesPresent: action.setEnabled(True) def remLabel(self, shape): for index, (item, shape_) in enumerate(self.itemsToShapes): if shape_ == shape: break self.itemsToShapes.pop(index) self.labelList.takeItem(self.labelList.row(item)) def loadLabels(self, shapes): s = [] for label, points, line_color, fill_color in shapes: shape = Shape(label=label) for x, y in points: shape.addPoint(QPointF(x, y)) 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, filename): lf = LabelFile() def format_shape(s): return dict(label=str(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]) shapes = [format_shape(shape) for shape in self.canvas.shapes] try: lf.save(filename, shapes, str(self.filename), self.imageData, self.lineColor.getRgb(), self.fillColor.getRgb()) self.labelFile = lf self.filename = filename return True except LabelFileError as e: self.errorMessage('Error saving label data', '<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 for item_, shape in self.itemsToShapes: if item_ == item: break self.canvas.selectShape(shape) def labelItemChanged(self, item): for item_, shape in self.itemsToShapes: if item_ == item: break label = str(item.text()) if label != shape.label: shape.label = str(item.text()) 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. """ text = self.labelDialog.popUp() if text is not None: self.addLabel(self.canvas.setLastLabel(text)) if self.beginner(): # Switch to edit mode. self.canvas.setEditing(True) self.actions.create.setEnabled(True) else: self.actions.editMode.setEnabled(True) self.setDirty() else: self.canvas.undoLastLine() def scrollRequest(self, delta, orientation): units = - delta * 0.1 # natural scroll 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): units = delta * 0.1 self.addZoom(units) 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: item.setCheckState(Qt.Checked if value else Qt.Unchecked) def loadFile(self, filename=None): """Load the specified file, or the last opened file if None.""" self.resetState() self.canvas.setEnabled(False) if filename is None: filename = self.settings.get('filename', '') filename = str(filename) if QFile.exists(filename): if LabelFile.isLabelFile(filename): try: self.labelFile = LabelFile(filename) # FIXME: PyQt4 installed via Anaconda fails to load JPEG # and JSON encoded images. # https://github.com/ContinuumIO/anaconda-issues/issues/131 if QImage.fromData(self.labelFile.imageData).isNull(): raise LabelFileError( 'Failed loading image data from label file.\n' 'Maybe this is a known issue of PyQt4 built on' ' Anaconda, and may be fixed by installing PyQt5.') except LabelFileError as e: self.errorMessage('Error opening file', ("<p><b>%s</b></p>" "<p>Make sure <i>%s</i> is a valid label file.")\ % (e, filename)) self.status("Error reading %s" % filename) return False self.imageData = self.labelFile.imageData self.lineColor = QColor(*self.labelFile.lineColor) self.fillColor = QColor(*self.labelFile.fillColor) else: # Load image: # read data first and store for saving into label file. self.imageData = read(filename, None) self.labelFile = None image = QImage.fromData(self.imageData) if image.isNull(): formats = ['*.{}'.format(fmt.data().decode()) for fmt in QImageReader.supportedImageFormats()] self.errorMessage( 'Error opening file', '<p>Make sure <i>{0}</i> is a valid image file.<br/>' 'Supported image formats: {1}</p>' .format(filename, ','.join(formats))) self.status("Error reading %s" % filename) return False self.status("Loaded %s" % os.path.basename(str(filename))) self.image = image self.filename = filename 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.filename) self.toggleActions(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() s = self.settings s['filename'] = self.filename if self.filename else '' s['window/size'] = self.size() s['window/position'] = self.pos() s['window/state'] = self.saveState() s['line/color'] = self.lineColor s['fill/color'] = self.fillColor s['recentFiles'] = self.recentFiles s['advanced'] = not self._beginner # ask the use for where to save the labels #s['window/geometry'] = self.saveGeometry() ## User Dialogs ## def loadRecent(self, filename): if self.mayContinue(): self.loadFile(filename) def openFile(self, _value=False): if not self.mayContinue(): return path = os.path.dirname(str(self.filename))\ if self.filename else '.' formats = ['*.{}'.format(fmt.data().decode()) 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 PYQT5: filename, _ = filename filename = str(filename) self.filename=filename self.makeImageAndLabelFileDict(os.path.split(filename)[0]) if hasattr(self, 'imageFileDict'): self.updateTable(self.makeTableViewData()) if filename: self.loadFile(filename) def makeImageAndLabelFileDict(self, path): try: self.imageFileDict={} self.labelFileDict={} for fname in os.listdir(path): if fname.endswith(".png"): self.imageFileDict[fname.split('.')[0]]=os.path.join(path,fname) elif fname.endswith(".json"): self.labelFileDict[fname.split('.')[0]]=os.path.join(path,fname) except FileNotFoundError: msg=QMessageBox() msg.setIcon(QMessageBox.Critical) msg.setText("cannot find given path") def makeTableViewData(self): table_data=[] if len(self.imageFileDict.keys())!=0: for key in self.imageFileDict.keys(): if key in self.labelFileDict.keys(): table_data.append([key,'Yes']) else: table_data.append([key, 'no']) else: for key in self.labelFileDict.keys(): table_data.append([key,'Yes']) return table_data def saveFile(self, _value=False): assert not self.image.isNull(), "cannot save empty image" if self.hasLabels(): if self.labelFile: self._saveFile(self.filename) elif self.output: self._saveFile(self.output) else: self._saveFile(self.saveFileDialog()) image = self.canvas.getLabelImage() if image: path, fname = os.path.split(self.filename) if not os.path.exists(os.path.join(path,"label")): os.mkdir(os.path.join(path,"label")) fname=os.path.join("label","label_"+fname.replace("json","png")) image.save(os.path.join(path,fname),"png") src_filename=os.path.join("label",os.path.split(self.filename)[1]) self.image.save(os.path.join(path,src_filename.replace("json","png")),"png") #table view update row_idx, key = self.getCurrentData() self.labelFileDict[key]=self.filename.replace("png", "json") self.mymodel.setLabeled(row_idx) self.table_view.update() def saveFileAs(self, _value=False): assert not self.image.isNull(), "cannot save empty image" if self.hasLabels(): self._saveFile(self.saveFileDialog()) def saveFileDialog(self): caption = '%s - Choose File' % __appname__ filters = 'Label files (*%s)' % LabelFile.suffix dlg = QFileDialog(self, caption, self.currentPath(), filters) dlg.setDefaultSuffix(LabelFile.suffix[1:]) dlg.setAcceptMode(QFileDialog.AcceptSave) dlg.setOption(QFileDialog.DontConfirmOverwrite, False) dlg.setOption(QFileDialog.DontUseNativeDialog, False) basename = os.path.splitext(self.filename)[0] default_labelfile_name = os.path.join(self.currentPath(), basename + LabelFile.suffix) filename = dlg.getSaveFileName( self, 'Choose File', default_labelfile_name, 'Label files (*%s)' % LabelFile.suffix) if PYQT5: filename, _ = filename filename = str(filename) return filename def _saveFile(self, filename): if filename and self.saveLabels(filename): self.addRecentFile(filename) self.setClean() if self.labeling_once: self.close() 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) # Message Dialogs. # def hasLabels(self): if not self.itemsToShapes: self.errorMessage('No objects labeled', 'You must label at least one object to save the file.') return False return True def mayContinue(self): return not (self.dirty and not self.discardChangesDialog()) def discardChangesDialog(self): yes, no = QMessageBox.Yes, QMessageBox.No msg = 'You have unsaved changes, proceed anyway?' return yes == QMessageBox.warning(self, '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(str(self.filename)) if self.filename else '.' def chooseColor1(self): color = self.colorDialog.getColor(self.lineColor, 'Choose line color', default=DEFAULT_LINE_COLOR) if color: self.lineColor = color # Change the color for all shape lines: Shape.line_color = self.lineColor self.canvas.update() self.setDirty() def chooseColor2(self): color = self.colorDialog.getColor(self.fillColor, '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): yes, no = QMessageBox.Yes, QMessageBox.No msg = 'You are about to permanently delete this polygon, proceed anyway?' if yes == QMessageBox.warning(self, 'Attention', msg, yes|no): 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, '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, '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()
class MainWindow(QMainWindow, WindowMixin): FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3)) def __init__(self, defaultFilename=None): super(MainWindow, self).__init__() self.setWindowTitle(__appname__) # Save as Pascal voc xml self.defaultSaveDir = None self.usingPascalVocFormat = True if self.usingPascalVocFormat: LabelFile.suffix = '.xml' # 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 # Enble auto saving if pressing next self.autoSaving = True self._noSelectionSlot = False self._beginner = True self.screencastViewer = "firefox" self.screencast = "https://youtu.be/p0nR2YsCY_U" self.loadPredefinedClasses() # Main widgets and related state. self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist) self.labelList = QListWidget() self.itemsToShapes = {} self.shapesToItems = {} 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 = QVBoxLayout() listLayout.setContentsMargins(0, 0, 0, 0) listLayout.addWidget(self.labelList) self.editButton = QToolButton() self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) self.labelListContainer = QWidget() self.labelListContainer.setLayout(listLayout) listLayout.addWidget(self.editButton) #, 0, Qt.AlignCenter) listLayout.addWidget(self.labelList) self.dock = QDockWidget(u'Box Labels', self) self.dock.setObjectName(u'Labels') self.dock.setWidget(self.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) self.fileListContainer = QWidget() self.fileListContainer.setLayout(filelistLayout) self.filedock = QDockWidget(u'File List', self) self.filedock.setObjectName(u'Files') self.filedock.setWidget(self.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.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.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.openDir, 'Ctrl+u', 'open', u'Open Dir') changeSavedir = action('&Change default saved Annotation dir', self.changeSavedir, 'Ctrl+r', 'open', u'Change default saved Annotation dir') openAnnotation = action('&Open Annotation', self.openAnnotation, 'Ctrl+q', 'openAnnotation', 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') save = action('&Save', self.saveFile, 'Ctrl+S', 'save', u'Save labels to file', enabled=False) 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') color1 = action('Box &Line Color', self.chooseColor1, 'Ctrl+L', 'color_line', u'Choose Box line color') color2 = action('Box &Fill Color', self.chooseColor2, 'Ctrl+Shift+L', 'color', u'Choose Box fill color') createMode = action('Create\nRectBox', self.setCreateMode, 'Ctrl+N', '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.tutorial, 'Ctrl+T', 'help', u'Show demos') 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) # Store actions for further handling. self.actions = struct( save=save, saveAs=saveAs, open=open, close=close, lineColor=color1, fillColor=color2, 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, quit), beginner=(), advanced=(), editMenu=(edit, copy, delete, None, color1, color2), 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) addActions(self.menus.file, (open, opendir, changeSavedir, openAnnotation, self.menus.recentFiles, save, saveAs, close, None, quit)) addActions(self.menus.help, (help, )) addActions(self.menus.view, (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, openNextImg, openPrevImg, save, None, create, copy, delete, None, zoomIn, zoom, zoomOut, fitWindow, fitWidth) self.actions.advanced = (open, save, None, createMode, editMode, None, hideAll, showAll) self.statusBar().showMessage('%s started.' % __appname__) self.statusBar().show() # Application state. self.image = QImage() self.filePath = u(defaultFilename) self.recentFiles = [] self.maxRecent = 7 self.lineColor = None self.fillColor = None self.zoom_level = 100 self.fit_window = False # XXX: Could be completely declarative. # Restore application settings. if have_qstring(): types = { 'filename': QString, 'recentFiles': QStringList, 'window/size': QSize, 'window/position': QPoint, 'window/geometry': QByteArray, 'line/color': QColor, 'fill/color': QColor, 'advanced': bool, # Docks and toolbars: 'window/state': QByteArray, 'savedir': QString, 'lastOpenDir': QString, } else: types = { 'filename': str, 'recentFiles': list, 'window/size': QSize, 'window/position': QPoint, 'window/geometry': QByteArray, 'line/color': QColor, 'fill/color': QColor, 'advanced': bool, # Docks and toolbars: 'window/state': QByteArray, 'savedir': str, 'lastOpenDir': str, } self.settings = settings = Settings(types) self.recentFiles = list(settings.get('recentFiles', [])) size = settings.get('window/size', QSize(600, 500)) position = settings.get('window/position', QPoint(0, 0)) self.resize(size) self.move(position) saveDir = u(settings.get('savedir', None)) self.lastOpenDir = u(settings.get('lastOpenDir', None)) if os.path.exists(saveDir): self.defaultSaveDir = saveDir self.statusBar().showMessage( '%s started. Annotation will be saved to %s' % (__appname__, self.defaultSaveDir)) self.statusBar().show() # or simply: #self.restoreGeometry(settings['window/geometry'] self.restoreState(settings.get('window/state', QByteArray())) self.lineColor = QColor(settings.get('line/color', Shape.line_color)) self.fillColor = QColor(settings.get('fill/color', Shape.fill_color)) Shape.line_color = self.lineColor Shape.fill_color = self.fillColor def xbool(x): if isinstance(x, QVariant): return x.toBool() return bool(x) if xbool(settings.get('advanced', 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. self.queueEvent(partial(self.loadFile, self.filePath)) # Callbacks: self.zoomWidget.valueChanged.connect(self.paintCanvas) self.populateModeActions() ## Support Functions ## 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() 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() ## Callbacks ## 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): """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) 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() # Tzutalin 20160906 : Add file list and dock to move faster def fileitemDoubleClicked(self, item=None): currIndex = self.mImgList.index(u(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): 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) for action in self.actions.onShapesPresent: action.setEnabled(True) def remLabel(self, shape): 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 in shapes: shape = Shape(label=label) for x, y in points: shape.addPoint(QPointF(x, y)) 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, filePath): lf = LabelFile() 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]) shapes = [format_shape(shape) for shape in self.canvas.shapes] # Can add differrent annotation formats here try: unicodeFilePath = u(filePath) if self.usingPascalVocFormat is True: lf.savePascalVocFormat(unicodeFilePath, shapes, unicodeFilePath, self.imageData, self.lineColor.getRgb(), self.fillColor.getRgb()) else: lf.save(unicodeFilePath, shapes, unicodeFilePath, self.imageData, self.lineColor.getRgb(), self.fillColor.getRgb()) self.labelFile = lf self.filePath = unicodeFilePath 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]) def labelItemChanged(self, item): shape = self.itemsToShapes[item] label = item.text() if label != 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): """Pop-up and give focus to the label editor. position MUST be in global coordinates. """ if len(self.labelHist) > 0: self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist) # text = self.labelDialog.popUp() text = QString(u'pockey') # label only pockey if text is not None: self.addLabel(self.canvas.setLastLabel(text)) 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): units = delta / (8 * 15) scale = 10 self.addZoom(scale * units) 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('filename') unicodeFilePath = u(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) else: # Load image: # read data first and store for saving into label file. 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) ## Label xml file and show bound box according to its filename if self.usingPascalVocFormat is True and \ 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) 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() s = self.settings # If it loads images from dir, don't load it at the begining if self.dirname is None: s['filename'] = self.filePath if self.filePath else '' else: s['filename'] = '' s['window/size'] = self.size() s['window/position'] = self.pos() s['window/state'] = self.saveState() s['line/color'] = self.lineColor s['fill/color'] = self.fillColor s['recentFiles'] = self.recentFiles s['advanced'] = not self._beginner if self.defaultSaveDir is not None and len(self.defaultSaveDir) > 1: s['savedir'] = str(self.defaultSaveDir) else: s['savedir'] = "" if self.lastOpenDir is not None and len(self.lastOpenDir) > 1: s['lastOpenDir'] = self.lastOpenDir else: s['lastOpenDir'] = "" ## User Dialogs ## 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 = u(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 = str(self.defaultSaveDir) else: path = '.' dirpath = str( 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(u(self.filePath))\ if self.filePath else '.' if self.usingPascalVocFormat: formats = ['*.%s' % str(fmt).lower()\ for fmt in QImageReader.supportedImageFormats()] filters = "Open Annotation XML file (%s)" % \ ' '.join(formats + ['*.xml']) filename = str( QFileDialog.getOpenFileName( self, '%s - Choose a xml file' % __appname__, path, filters)) 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 = u( 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 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 if self.autoSaving is True and self.defaultSaveDir is not None: if self.dirty is True and self.hasLabels(): 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(str(self.filePath))\ if self.filePath else '.' formats = ['*.%s' % str(fmt).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: self.loadFile(filename) def saveFile(self, _value=False): assert not self.image.isNull(), "cannot save empty image" if self.hasLabels(): if self.defaultSaveDir is not None and len(str( self.defaultSaveDir)): print('handle the image:' + self.filePath) imgFileName = os.path.basename(self.filePath) savedFileName = os.path.splitext( imgFileName)[0] + LabelFile.suffix savedPath = os.path.join(str(self.defaultSaveDir), savedFileName) self._saveFile(savedPath) else: self._saveFile(self.filePath if self.labelFile\ else self.saveFileDialog()) def saveFileAs(self, _value=False): assert not self.image.isNull(), "cannot save empty image" if self.hasLabels(): 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 _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) # Message Dialogs. # def hasLabels(self): if not self.itemsToShapes: self.errorMessage( u'No objects labeled', u'You must label at least one object to save the file.') return False return True 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 # Change the color for all shape lines: 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): yes, no = QMessageBox.Yes, QMessageBox.No msg = u'You are about to permanently delete this Box, proceed anyway?' if yes == QMessageBox.warning(self, u'Attention', msg, yes | no): 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): predefined_classes_path = os.path.join('data', 'predefined_classes.txt') if os.path.exists(predefined_classes_path) is True: with codecs.open(predefined_classes_path, '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)
class MainWindow(QMainWindow, WindowMixin): FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3)) def __init__(self, filename=None): super(MainWindow, self).__init__() self.setWindowTitle(__appname__) # Save as Pascal voc xml self.defaultSaveDir = None self.usingPascalVocFormat = True if self.usingPascalVocFormat: LabelFile.suffix = '.xml' # 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 # Enble auto saving if pressing next self.autoSaving = True self._noSelectionSlot = False self._beginner = True self.screencastViewer = "firefox" self.screencast = "https://youtu.be/p0nR2YsCY_U" self.loadPredefinedClasses() # Main widgets and related state. self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist) self.labelList = QListWidget() self.itemsToShapes = {} self.shapesToItems = {} 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 = QVBoxLayout() listLayout.setContentsMargins(0, 0, 0, 0) listLayout.addWidget(self.labelList) self.editButton = QToolButton() self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon) self.labelListContainer = QWidget() self.labelListContainer.setLayout(listLayout) listLayout.addWidget(self.editButton)#, 0, Qt.AlignCenter) listLayout.addWidget(self.labelList) self.dock = QDockWidget(u'Box Labels', self) self.dock.setObjectName(u'Labels') self.dock.setWidget(self.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) self.fileListContainer = QWidget() self.fileListContainer.setLayout(filelistLayout) self.filedock = QDockWidget(u'File List', self) self.filedock.setObjectName(u'Files') self.filedock.setWidget(self.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.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.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.openDir, 'Ctrl+u', 'open', u'Open Dir') changeSavedir = action('&Change default saved Annotation dir', self.changeSavedir, 'Ctrl+r', 'open', u'Change default saved Annotation dir') openAnnotation = action('&Open Annotation', self.openAnnotation, 'Ctrl+q', 'openAnnotation', 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') save = action('&Save', self.saveFile, 'Ctrl+S', 'save', u'Save labels to file', enabled=False) 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') color1 = action('Box &Line Color', self.chooseColor1, 'Ctrl+L', 'color_line', u'Choose Box line color') color2 = action('Box &Fill Color', self.chooseColor2, 'Ctrl+Shift+L', 'color', u'Choose Box fill color') createMode = action('Create\nRectBox', self.setCreateMode, 'Ctrl+N', '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.tutorial, 'Ctrl+T', 'help', u'Show demos') 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) # Store actions for further handling. self.actions = struct(save=save, saveAs=saveAs, open=open, close=close, lineColor=color1, fillColor=color2, 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,quit), beginner=(), advanced=(), editMenu=(edit, copy, delete, None, color1, color2), 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) addActions(self.menus.file, (open, opendir,changeSavedir, openAnnotation, self.menus.recentFiles, save, saveAs, close, None, quit)) addActions(self.menus.help, (help,)) addActions(self.menus.view, ( 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, openNextImg, openPrevImg, save, None, create, copy, delete, None, zoomIn, zoom, zoomOut, fitWindow, fitWidth) self.actions.advanced = ( open, save, None, createMode, editMode, None, hideAll, showAll) self.statusBar().showMessage('%s started.' % __appname__) self.statusBar().show() # Application state. self.image = QImage() self.filename = filename self.recentFiles = [] self.maxRecent = 7 self.lineColor = None self.fillColor = None self.zoom_level = 100 self.fit_window = False # XXX: Could be completely declarative. # Restore application settings. if have_qstring(): types = { 'filename': QString, 'recentFiles': QStringList, 'window/size': QSize, 'window/position': QPoint, 'window/geometry': QByteArray, 'line/color': QColor, 'fill/color': QColor, 'advanced': bool, # Docks and toolbars: 'window/state': QByteArray, 'savedir': QString, 'lastOpenDir': QString, } else: types = { 'filename': str, 'recentFiles': list, 'window/size': QSize, 'window/position': QPoint, 'window/geometry': QByteArray, 'line/color': QColor, 'fill/color': QColor, 'advanced': bool, # Docks and toolbars: 'window/state': QByteArray, 'savedir': str, 'lastOpenDir': str, } self.settings = settings = Settings(types) self.recentFiles = list(settings.get('recentFiles', [])) size = settings.get('window/size', QSize(600, 500)) position = settings.get('window/position', QPoint(0, 0)) self.resize(size) self.move(position) saveDir = settings.get('savedir', None) self.lastOpenDir = settings.get('lastOpenDir', None) if os.path.exists(str(saveDir)): self.defaultSaveDir = str(saveDir) self.statusBar().showMessage('%s started. Annotation will be saved to %s' %(__appname__, self.defaultSaveDir)) self.statusBar().show() # or simply: #self.restoreGeometry(settings['window/geometry'] self.restoreState(settings.get('window/state', QByteArray())) self.lineColor = QColor(settings.get('line/color', Shape.line_color)) self.fillColor = QColor(settings.get('fill/color', Shape.fill_color)) Shape.line_color = self.lineColor Shape.fill_color = self.fillColor def xbool(x): if isinstance(x, QVariant): return x.toBool() return bool(x) if xbool(settings.get('advanced', 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. self.queueEvent(partial(self.loadFile, self.filename)) # Callbacks: self.zoomWidget.valueChanged.connect(self.paintCanvas) self.populateModeActions() ## Support Functions ## 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.filename = 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, filename): if filename in self.recentFiles: self.recentFiles.remove(filename) elif len(self.recentFiles) >= self.maxRecent: self.recentFiles.pop() self.recentFiles.insert(0, filename) def beginner(self): return self._beginner def advanced(self): return not self.beginner() ## Callbacks ## 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): """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) def updateFileMenu(self): current = self.filename def exists(filename): return os.path.exists(filename) menu = self.menus.recentFiles menu.clear() files = [f for f in self.recentFiles if f != current 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() # Tzutalin 20160906 : Add file list and dock to move faster def fileitemDoubleClicked(self, item=None): 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): 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) for action in self.actions.onShapesPresent: action.setEnabled(True) def remLabel(self, shape): 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 in shapes: shape = Shape(label=label) for x, y in points: shape.addPoint(QPointF(x, y)) 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, filename): lf = LabelFile() 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]) shapes = [format_shape(shape) for shape in self.canvas.shapes] # Can add differrent annotation formats here try: if self.usingPascalVocFormat is True: print('savePascalVocFormat save to:' + filename) lf.savePascalVocFormat(filename, shapes, str(self.filename), self.imageData, self.lineColor.getRgb(), self.fillColor.getRgb()) else: lf.save(filename, shapes, str(self.filename), self.imageData, self.lineColor.getRgb(), self.fillColor.getRgb()) self.labelFile = lf self.filename = filename 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]) def labelItemChanged(self, item): shape = self.itemsToShapes[item] label = item.text() if label != 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): """Pop-up and give focus to the label editor. position MUST be in global coordinates. """ if len(self.labelHist) > 0: self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist) text = self.labelDialog.popUp() if text is not None: self.addLabel(self.canvas.setLastLabel(text)) 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): units = delta / (8 * 15) scale = 10 self.addZoom(scale * units) 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, filename=None): """Load the specified file, or the last opened file if None.""" self.resetState() self.canvas.setEnabled(False) if filename is None: filename = self.settings.get('filename') # Tzutalin 20160906 : Add file list and dock to move faster # Highlight the file item if filename and self.fileListWidget.count() > 0: index = self.mImgList.index(filename) fileWidgetItem = self.fileListWidget.item(index) fileWidgetItem.setSelected(True) if filename and QFile.exists(filename): if LabelFile.isLabelFile(filename): try: self.labelFile = LabelFile(filename) 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, filename)) self.status("Error reading %s" % filename) return False self.imageData = self.labelFile.imageData self.lineColor = QColor(*self.labelFile.lineColor) self.fillColor = QColor(*self.labelFile.fillColor) else: # Load image: # read data first and store for saving into label file. self.imageData = read(filename, 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." % filename) self.status("Error reading %s" % filename) return False self.status("Loaded %s" % os.path.basename(str(filename))) self.image = image self.filename = filename 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.filename) self.toggleActions(True) ## Label xml file and show bound box according to its filename if self.usingPascalVocFormat is True and \ self.defaultSaveDir is not None: basename = os.path.basename(os.path.splitext(self.filename)[0]) + '.xml' xmlPath = os.path.join(self.defaultSaveDir, basename) self.loadPascalXMLByFilename(xmlPath) 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() s = self.settings # If it loads images from dir, don't load it at the begining if self.dirname is None: s['filename'] = self.filename if self.filename else '' else: s['filename'] = '' s['window/size'] = self.size() s['window/position'] = self.pos() s['window/state'] = self.saveState() s['line/color'] = self.lineColor s['fill/color'] = self.fillColor s['recentFiles'] = self.recentFiles s['advanced'] = not self._beginner if self.defaultSaveDir is not None and len(self.defaultSaveDir) > 1: s['savedir'] = str(self.defaultSaveDir) else: s['savedir'] = "" if self.lastOpenDir is not None and len(self.lastOpenDir) > 1: s['lastOpenDir'] = str(self.lastOpenDir) else: s['lastOpenDir'] = "" #ask the use for where to save the labels #s['window/geometry'] = self.saveGeometry() ## User Dialogs ## 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 = u(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 = str(self.defaultSaveDir) else: path = '.' dirpath = str(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.filename is None: return path = os.path.dirname(str(self.filename))\ if self.filename else '.' if self.usingPascalVocFormat: formats = ['*.%s' % str(fmt).lower()\ for fmt in QImageReader.supportedImageFormats()] filters = "Open Annotation XML file (%s)" % \ ' '.join(formats + ['*.xml']) filename = str(QFileDialog.getOpenFileName(self, '%s - Choose a xml file' % __appname__, path, filters)) self.loadPascalXMLByFilename(filename) def openDir(self, _value=False): if not self.mayContinue(): return path = os.path.dirname(self.filename)\ if self.filename else '.' if self.lastOpenDir is not None and len(self.lastOpenDir) > 1: path = self.lastOpenDir dirpath = str(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.mImgList = self.scanAllImages(dirpath) self.openNextImg() for imgPath in self.mImgList: item = QListWidgetItem(imgPath) self.fileListWidget.addItem(item) def openPrevImg(self, _value=False): if not self.mayContinue(): return if len(self.mImgList) <= 0: return if self.filename is None: return currIndex = self.mImgList.index(self.filename) 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 if self.autoSaving is True and self.defaultSaveDir is not None: if self.dirty is True and self.hasLabels(): self.saveFile() if not self.mayContinue(): return if len(self.mImgList) <= 0: return filename = None if self.filename is None: filename = self.mImgList[0] else: currIndex = self.mImgList.index(self.filename) 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(str(self.filename))\ if self.filename else '.' formats = ['*.%s' % str(fmt).lower()\ for fmt in QImageReader.supportedImageFormats()] filters = "Image & Label files (%s)" % \ ' '.join(formats + ['*%s' % LabelFile.suffix]) filename = str(QFileDialog.getOpenFileName(self, '%s - Choose Image or Label file' % __appname__, path, filters)) if filename: self.loadFile(filename) def saveFile(self, _value=False): assert not self.image.isNull(), "cannot save empty image" if self.hasLabels(): if self.defaultSaveDir is not None and len(str(self.defaultSaveDir)): print('handle the image:' + self.filename) imgFileName = os.path.basename(self.filename) savedFileName = os.path.splitext(imgFileName)[0] + LabelFile.suffix savedPath = os.path.join(str(self.defaultSaveDir), savedFileName) self._saveFile(savedPath) else: self._saveFile(self.filename if self.labelFile\ else self.saveFileDialog()) def saveFileAs(self, _value=False): assert not self.image.isNull(), "cannot save empty image" if self.hasLabels(): 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.filename)[0] dlg.selectFile(filenameWithoutExtension) dlg.setOption(QFileDialog.DontUseNativeDialog, False) if dlg.exec_(): return dlg.selectedFiles()[0] return '' def _saveFile(self, filename): if filename and self.saveLabels(filename): self.addRecentFile(filename) self.setClean() self.statusBar().showMessage('Saved to %s' % filename) 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) # Message Dialogs. # def hasLabels(self): if not self.itemsToShapes: self.errorMessage(u'No objects labeled', u'You must label at least one object to save the file.') return False return True 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(str(self.filename)) if self.filename else '.' def chooseColor1(self): color = self.colorDialog.getColor(self.lineColor, u'Choose line color', default=DEFAULT_LINE_COLOR) if color: self.lineColor = color # Change the color for all shape lines: 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): yes, no = QMessageBox.Yes, QMessageBox.No msg = u'You are about to permanently delete this Box, proceed anyway?' if yes == QMessageBox.warning(self, u'Attention', msg, yes|no): 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): predefined_classes_path = os.path.join('data', 'predefined_classes.txt') if os.path.exists(predefined_classes_path) is True: with codecs.open(predefined_classes_path, '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.filename is None: return if os.path.isfile(xmlPath) is False: return tVocParseReader = PascalVocReader(xmlPath) shapes = tVocParseReader.getShapes() self.loadLabels(shapes)