class MainWindow(QMainWindow, WindowMixin):
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
def __init__(self, defaultFilename=None, defaaultPrefdefClassFile=None, defaultSaveDir=None):
super(MainWindow, self).__init__()
self.setWindowTitle(__appname__)
#Load setting in the main thread
self.settings = Settings()
self.settings.load()
settings = self.settings
#load string bundle for i18n
self.stringBundle = StringBundle.getBundle()
getStr = lambda strId: self.stringBundle.getString(strId)
#save as pascal voc xml
self.defaultSaveDir = defaultSaveDir
self.labelFileFormat = settings.get(SETTING_LABEL_FILE_FORMAT, LabelFileFormat.PASCAL_VOC)
#for loading all image under a directory
self.mImgList = []
self.dirname = None
self.labelHist = []
self.lastOpenDir = None
#NEED TO DO
#whether we need to save or not
self.dirty = False
self._noSelectionSlot = False
self._beginner = True
self.screencastViewer = self.getAvailableScreencastViewer()
self.screencast = "https://youtu.be/p0nR2YsCY_U"
#Load predefined classes to the list
self.loadPredefinedClasses(defaaultPrefdefClassFile)
#main widgets and related state
self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
self.itemsToShapes = {}
self.shapesToItems = {}
self.prevLabelText = ''
listLayout = QVBoxLayout()
listLayout.setContentsMargins(0,0,0,0)
self.useDefaultLableCheckbox = QCheckBox(getStr('useDefaultLabel'))
self.useDefaultLableCheckbox.setChecked(False)
self.defaultLabelTextLine = QLineEdit()
useDefaultLabelQHBoxLayout = QHBoxLayout()
useDefaultLabelQHBoxLayout.addWidget(self.useDefaultLableCheckbox)
useDefaultLabelQHBoxLayout.addWidget(self.defaultLabelTextLine)
useDefaultLabelContainer = QWidget()
useDefaultLabelContainer.setLayout((useDefaultLabelQHBoxLayout))
self.diffcButton = QCheckBox(getStr('useDifficult'))
self.diffcButton.setChecked(False)
self.diffcButton.stateChanged.connect(self.btnstate)
self.editButton = QToolButton()
self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
listLayout.addWidget(self.editButton)
listLayout.addWidget(self.diffcButton)
listLayout.addWidget(useDefaultLabelContainer)
self.comboBox = ComboBox(self)
listLayout.addWidget(self.comboBox)
# Create and add a widget for showing current label items
self.labelList = QListWidget()
labelListContainer = QWidget()
labelListContainer.setLayout(listLayout)
self.labelList.itemActivated.connect(self.labelSelectionChanged)
self.labelList.itemSelectionChanged.connect(self.labelSelectionChanged)
self.labelList.itemDoubleClicked.connect(self.editLabel)
# Connect to itemChanged to detect checkbox changes.
self.labelList.itemChanged.connect(self.labelItemChanged)
listLayout.addWidget(self.labelList)
self.dock = QDockWidget(getStr('boxLabelText'), self)
self.dock.setObjectName(getStr('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(getStr('fileList'), self)
self.filedock.setObjectName(getStr('files'))
self.filedock.setWidget(fileListContainer)
self.zoomWidget = ZoomWidget()
self.colorDialog = ColorDialog(parent=self)
self.canvas = Canvas(parent=self)
self.canvas.zoomRequest.connect(self.zoomRequest)
self.canvas.setDrawingShapeToSquare(settings.get(SETTING_DRAW_SQUARE, False))
scroll = QScrollArea()
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(True)
self.scrollBars = {
Qt.Vertical: scroll.verticalScrollBar(),
Qt.Horizontal: scroll.horizontalScrollBar()
}
self.scrollArea = scroll
self.canvas.scrollRequest.connect(self.scrollRequest)
self.canvas.newShape.connect(self.newShape)
self.canvas.shapeMoved.connect(self.setDirty)
self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
self.setCentralWidget(scroll)
self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
self.filedock.setFeatures(QDockWidget.DockWidgetFloatable)
self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
def labelSelectionChanged(self):
item = self.currentItem()
if item and self.canvas.editing():
self._noSelectionSlot = True
self.canvas.selectShape(self.itemsToShapes[item])
shape = self.itemsToShapes[item]
# Add Chris
self.diffcButton.setChecked(shape.difficult)
def editLabel(self):
if not self.canvas.editing():
return
item = self.currentItem()
if not item:
return
text = self.labelDialog.popUp(item.text())
if text is not None:
item.setText(text)
item.setBackground(generateColorByText(text))
self.setDirty()
self.updateComboBox()
def labelItemChanged(self, item):
shape = self.itemsToShapes[item]
label = item.text()
if label != shape.label:
shape.label = item.text()
shape.line_color = generateColorByText(shape.label)
self.setDirty()
else: # User probably changed item visibility
self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
def zoomRequest(self, delta):
# get the current scrollbar positions
# calculate the percentages ~ coordinates
h_bar = self.scrollBars[Qt.Horizontal]
v_bar = self.scrollBars[Qt.Vertical]
# get the current maximum, to know the difference after zooming
h_bar_max = h_bar.maximum()
v_bar_max = v_bar.maximum()
# get the cursor position and canvas size
# calculate the desired movement from 0 to 1
# where 0 = move left
# 1 = move right
# up and down analogous
cursor = QCursor()
pos = cursor.pos()
relative_pos = QWidget.mapFromGlobal(self, pos)
cursor_x = relative_pos.x()
cursor_y = relative_pos.y()
w = self.scrollArea.width()
h = self.scrollArea.height()
# the scaling from 0 to 1 has some padding
# you don't have to hit the very leftmost pixel for a maximum-left movement
margin = 0.1
move_x = (cursor_x - margin * w) / (w - 2 * margin * w)
move_y = (cursor_y - margin * h) / (h - 2 * margin * h)
# clamp the values from 0 to 1
move_x = min(max(move_x, 0), 1)
move_y = min(max(move_y, 0), 1)
# zoom in
units = delta / (8 * 15)
scale = 10
self.addZoom(scale * units)
# get the difference in scrollbar values
# this is how far we can move
d_h_bar_max = h_bar.maximum() - h_bar_max
d_v_bar_max = v_bar.maximum() - v_bar_max
# get the new scrollbar values
new_h_bar_value = h_bar.value() + move_x * d_h_bar_max
new_v_bar_value = v_bar.value() + move_y * d_v_bar_max
h_bar.setValue(new_h_bar_value)
v_bar.setValue(new_v_bar_value)
def scrollRequest(self, delta, orientation):
units = - delta / (8 * 15)
bar = self.scrollBars[orientation]
bar.setValue(bar.value() + bar.singleStep() * units)
def newShape(self):
"""Pop-up and give focus to the label editor.
position MUST be in global coordinates.
"""
if not self.useDefaultLabelCheckbox.isChecked() or not self.defaultLabelTextLine.text():
if len(self.labelHist) > 0:
self.labelDialog = LabelDialog(
parent=self, listItem=self.labelHist)
# Sync single class mode from PR#106
if self.singleClassMode.isChecked() and self.lastLabel:
text = self.lastLabel
else:
text = self.labelDialog.popUp(text=self.prevLabelText)
self.lastLabel = text
else:
text = self.defaultLabelTextLine.text()
def setDirty(self):
self.dirty = True
self.actions.save.setEnabled(True)
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 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 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 btnstate(self, item= None):
""" Function to handle difficult examples
Update on each object """
if not self.canvas.editing():
return
item = self.currentItem()
if not item: # If not selected Item, take the first one
item = self.labelList.item(self.labelList.count()-1)
difficult = self.diffcButton.isChecked()
try:
shape = self.itemsToShapes[item]
except:
pass
# Checked and Update
try:
if difficult != shape.difficult:
shape.difficult = difficult
self.setDirty()
else: # User probably changed item visibility
self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
except:
pass
def loadFile(self, filePath=None):
self.resetState()
self.canvas.setEnabled(False)
if filePath is None:
filePath = self.settings.get(SETTING_FILENAME)
filePath = ustr(filePath)
unicodeFilePath = ustr(filePath)
unicodeFilePath = os.path.abspath(unicodeFilePath)
if unicodeFilePath and self.fileListWidget.count() > 0:
if unicodeFilePath in self.mImgList:
index = self.mImgList.index(unicodeFilePath)
fileWidgetItem = self.fileListWidget.item(index)
fileWidgetItem.setSelected(True)
else:
self.fileListWidget.clear()
self.mImgList.clear()
if unicodeFilePath and os.path.exists(unicodeFilePath):
if LabelFile.isLabelFile(unicodeFilePath):
try:
self.labelFile = LabelFile(unicodeFilePath)
except LabelFileError as e:
self.errorMessage(u'Error opening file',
(u"<p><b>%s</b></p>"
u"<p>Make sure <i>%s</i> is a valid label file.")
% (e, unicodeFilePath))
self.status("Error reading %s" % unicodeFilePath)
return False
self.imageData = self.labelFile.imageData
self.lineColor = QColor(*self.labelFile.lineColor)
self.fillColor = QColor(*self.labelFile.fillColor)
self.canvas.verified = self.labelFile.verified
else:
# Load image:
# read data first and store for saving into label file.
self.imageData = read(unicodeFilePath, None)
self.labelFile = None
self.canvas.verified = False
if isinstance(self.imageData, QImage):
image = self.imageData
else:
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)
self.showBoundingBoxFromAnnotationFile(filePath)
self.setWindowTitle(__appname__ + ' ' + filePath)
# Default : select last item if there is at least one item
if self.labelList.count():
self.labelList.setCurrentItem(self.labelList.item(self.labelList.count()-1))
self.labelList.item(self.labelList.count()-1).setSelected(True)
self.canvas.setFocus(True)
return True
return False
def getAvailableScreencastViewer(self):
osName = platform.system()
if osName == 'Windows':
return ['C:\\Program Files\\Internet Explorer\\iexplore.exe']
elif osName == 'Linux':
return ['xdg-open']
elif osName == 'Darwin':
return ['open']
def loadPredefinedClasses(self, predefClassesFile):
if os.path.exists(predefClassesFile) is True:
with codecs.open(predefClassesFile, 'r', 'utf8') as f:
for line in f:
line = line.strip()
if self.labelHist is None:
self.labelHist = [line]
else:
self.labelHist.append(line)
def comboSelectionChanged(self, index):
text = self.comboBox.cb.itemText(index)
for i in range(self.labelList.count()):
if text == "":
self.labelList.item(i).setCheckState(2)
elif text != self.labelList.item(i).text():
self.labelList.item(i).setCheckState(0)
else:
self.labelList.item(i).setCheckState(2)
class MainWindow(QMainWindow, WindowMixin):
_, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
def __init__(self, defaultFilename=None, defaultPrefdefClassFile=None):
super(MainWindow, self).__init__()
self.setWindowTitle(__appname__)
# Load setting in the main thread
self.settings = Settings()
self.settings.load()
settings = self.settings
self.defaultSaveDir = None
# Whether we need to save or not.
self.dirty = False
# For loading all image under a directory
self.mImgList = []
self.dirname = None
# Main widgets and related state.
# self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
# File and Notes
self.noteArea = QPlainTextEdit()
self.fileListWidget = QListWidget()
self.fileListWidget.itemDoubleClicked.connect(
self.fileitemDoubleClicked)
self.fileNoteLayout = QVBoxLayout()
self.fileNoteLayout.addWidget(self.fileListWidget)
self.fileNoteLayout.addWidget(self.noteArea)
self.fileNoteArea = QWidget()
self.fileNoteArea.setLayout(self.fileNoteLayout)
# File and Notes to Dock
self.fileNoteDock = QDockWidget(u'Files and Notes', self)
self.fileNoteDock.setObjectName(u'file_note')
self.fileNoteDock.setWidget(self.fileNoteArea)
self.addDockWidget(Qt.LeftDockWidgetArea, self.fileNoteDock)
#First Search
self.firstText = QLineEdit()
self.firsrSearchResultArea = QListWidget()
self.firsrSearchResultArea.itemDoubleClicked.connect(
self.fileitemDoubleClicked)
self.firstSearchLayout = QVBoxLayout()
self.firstSearchLayout.addWidget(self.firstText)
self.firstSearchLayout.addWidget(self.firsrSearchResultArea)
self.firstSearchArea = QWidget()
self.firstSearchArea.setLayout(self.firstSearchLayout)
#Locode Search
self.locodeText = QLineEdit()
self.locodeDataLayout = QFormLayout()
self.storeKw = QLineEdit()
self.locodeDataLayout.addRow('Store KeyWord', self.storeKw)
self.mallKw = QLineEdit()
self.locodeDataLayout.addRow('Mall KeyWord', self.mallKw)
self.gstNo = QLineEdit()
self.locodeDataLayout.addRow('GST NO', self.gstNo)
self.locodeSearchLayout = QVBoxLayout()
self.locodeDataArea = QWidget()
self.locodeDataArea.setLayout(self.locodeDataLayout)
self.locodeSearchLayout.addWidget(self.locodeText)
self.locodeSearchLayout.addWidget(self.locodeDataArea)
self.locodeSearchArea = QWidget()
self.locodeSearchArea.setLayout(self.locodeSearchLayout)
# FirstchSearch and LocodeSearch to Doct
self.inputAreaLayout = QHBoxLayout()
self.inputAreaLayout.addWidget(self.firstSearchArea)
self.inputAreaLayout.addWidget(self.locodeSearchArea)
self.inputArea = QWidget()
self.inputArea.setLayout(self.inputAreaLayout)
self.inputDock = QDockWidget(u'Input Area', self)
self.inputDock.setObjectName(u'input_area')
self.inputDock.setWidget(self.inputArea)
self.addDockWidget(Qt.RightDockWidgetArea, self.inputDock)
#Canvas
self.zoomWidget = ZoomWidget()
self.canvas = Canvas(parent=self)
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.setCentralWidget(scroll)
# Actions
action = partial(newAction, self)
quit = action('&Quit', self.close, 'Ctrl+Q', 'quit',
u'Quit application')
opendir = action('&Open Dir', self.openDirDialog, 'Ctrl+o', 'open',
u'Open Dir')
changeSavedir = action('&Change Save Dir', self.changeSavedirDialog,
'Ctrl+r', 'open',
u'Change default saved Annotation dir')
openNextImg = action('&Next Image', self.openNextImg, 'd', 'next',
u'Open Next')
openPrevImg = action('&Prev Image', self.openPrevImg, 'a', 'prev',
u'Open Prev')
saveAs = action('&Save As',
self.saveFileAs,
'Ctrl+s',
'save-as',
u'Save labels to a different file',
enabled=False)
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)
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, fitWidth)
self.zoomMode = self.MANUAL_ZOOM
self.scalers = {
self.FIT_WIDTH:
self.scaleFitWidth,
# Set to one to scale to 100% when loading files.
self.MANUAL_ZOOM:
lambda: 1,
}
# Store actions for further handling.
self.actions = struct(saveAs=saveAs,
zoom=zoom,
fitWidth=fitWidth,
zoomActions=zoomActions,
fileMenuActions=(open, saveAs, quit))
self.menus = struct(file=self.menu('&File'),
edit=self.menu('&Edit'),
view=self.menu('&View'))
addActions(self.menus.file, (opendir, changeSavedir, saveAs, quit))
addActions(self.menus.view, (fitWidth, openPrevImg, openNextImg))
self.statusBar().showMessage('%s started.' % __appname__)
self.statusBar().show()
# Application state.
self.image = QImage()
self.filePath = ustr(defaultFilename)
self.recentFiles = []
self.maxRecent = 7
self.lineColor = None
self.fillColor = None
self.zoom_level = 100
self.fit_window = False
# Add Chris
self.difficult = False
size = settings.get(SETTING_WIN_SIZE, QSize(600, 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()))
print self.filePath
# Since loading the file may take some time, make sure it runs in the background.
if self.filePath and os.path.isdir(self.filePath):
self.queueEvent(partial(self.importDirImages, self.filePath or ""))
elif self.filePath:
self.queueEvent(partial(self.loadFile, self.filePath or ""))
# Callbacks:
self.zoomWidget.valueChanged.connect(self.paintCanvas)
# Display cursor coordinates at the right of status bar
self.labelCoordinates = QLabel('')
self.statusBar().addPermanentWidget(self.labelCoordinates)
# Open Dir if deafult file
if self.filePath and os.path.isdir(self.filePath):
self.openDirDialog(dirpath=self.filePath)
def setDirty(self):
self.dirty = True
self.actions.save.setEnabled(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)
def queueEvent(self, function):
QTimer.singleShot(0, function)
def status(self, message, delay=5000):
self.statusBar().showMessage(message, delay)
def resetState(self):
self.filePath = None
self.imageData = None
self.canvas.resetState()
self.labelCoordinates.clear()
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)
## Callbacks ##
def showTutorialDialog(self):
subprocess.Popen([self.screencastViewer, self.screencast])
def showInfoDialog(self):
msg = u'Name:{0} \nApp Version:{1} \n{2} '.format(
__appname__, __version__, sys.version_info)
QMessageBox.information(self, u'Information', msg)
def createShape(self):
assert self.beginner()
self.canvas.setEditing(False)
self.actions.create.setEnabled(False)
def toggleDrawingSensitive(self, drawing=True):
"""In the middle of drawing, toggling between modes should be disabled."""
self.actions.editMode.setEnabled(not drawing)
if not drawing and self.beginner():
# Cancel creation.
print('Cancel creation.')
self.canvas.setEditing(True)
self.canvas.restoreCursor()
self.actions.create.setEnabled(True)
def popLabelListMenu(self, point):
self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
# Tzutalin 20160906 : Add file list and dock to move faster
def fileitemDoubleClicked(self, item=None):
currIndex = self.mImgList.index(ustr(item.text()))
if currIndex < len(self.mImgList):
filename = self.mImgList[currIndex]
if filename:
self.loadFile(filename)
# Add chris
def btnstate(self, item=None):
""" Function to handle difficult examples
Update on each object """
if not self.canvas.editing():
return
item = self.currentItem()
if not item: # If not selected Item, take the first one
item = self.labelList.item(self.labelList.count() - 1)
difficult = self.diffcButton.isChecked()
try:
shape = self.itemsToShapes[item]
except:
pass
# Checked and Update
try:
if difficult != shape.difficult:
shape.difficult = difficult
self.setDirty()
else: # User probably changed item visibility
self.canvas.setShapeVisible(shape,
item.checkState() == Qt.Checked)
except:
pass
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.zoomMode = self.MANUAL_ZOOM
self.zoomWidget.setValue(value)
def addZoom(self, increment=10):
self.setZoom(self.zoomWidget.value() + increment)
def zoomRequest(self, delta):
# get the current scrollbar positions
# calculate the percentages ~ coordinates
h_bar = self.scrollBars[Qt.Horizontal]
v_bar = self.scrollBars[Qt.Vertical]
# get the current maximum, to know the difference after zooming
h_bar_max = h_bar.maximum()
v_bar_max = v_bar.maximum()
# get the cursor position and canvas size
# calculate the desired movement from 0 to 1
# where 0 = move left
# 1 = move right
# up and down analogous
cursor = QCursor()
pos = cursor.pos()
relative_pos = QWidget.mapFromGlobal(self, pos)
cursor_x = relative_pos.x()
cursor_y = relative_pos.y()
w = self.scrollArea.width()
h = self.scrollArea.height()
# the scaling from 0 to 1 has some padding
# you don't have to hit the very leftmost pixel for a maximum-left movement
margin = 0.1
move_x = (cursor_x - margin * w) / (w - 2 * margin * w)
move_y = (cursor_y - margin * h) / (h - 2 * margin * h)
# clamp the values from 0 to 1
move_x = min(max(move_x, 0), 1)
move_y = min(max(move_y, 0), 1)
# zoom in
units = delta / (8 * 15)
scale = 10
self.addZoom(scale * units)
# get the difference in scrollbar values
# this is how far we can move
d_h_bar_max = h_bar.maximum() - h_bar_max
d_v_bar_max = v_bar.maximum() - v_bar_max
# get the new scrollbar values
new_h_bar_value = h_bar.value() + move_x * d_h_bar_max
new_v_bar_value = v_bar.value() + move_y * d_v_bar_max
h_bar.setValue(new_h_bar_value)
v_bar.setValue(new_v_bar_value)
def setFitWidth(self, value=True):
if value:
self.actions.fitWindow.setChecked(False)
self.zoomMode = self.FIT_WIDTH if value else self.MANUAL_ZOOM
self.adjustScale()
def togglePolygons(self, value):
for item, shape in self.itemsToShapes.items():
item.setCheckState(Qt.Checked if value else Qt.Unchecked)
def loadFile(self, filePath=None):
"""Load the specified file, or the last opened file if None."""
self.resetState()
self.canvas.setEnabled(False)
if filePath is None:
filePath = self.settings.get(SETTING_FILENAME)
# Make sure that filePath is a regular python string, rather than QString
filePath = str(filePath)
unicodeFilePath = ustr(filePath)
# Tzutalin 20160906 : Add file list and dock to move faster
# Highlight the file item
if unicodeFilePath and self.fileListWidget.count() > 0:
index = self.mImgList.index(unicodeFilePath)
fileWidgetItem = self.fileListWidget.item(index)
fileWidgetItem.setSelected(True)
if unicodeFilePath and os.path.exists(unicodeFilePath):
# Load image:
# read data first and store for saving into label file.
self.imageData = read(unicodeFilePath, 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))
self.setClean()
self.canvas.setEnabled(True)
self.adjustScale(initial=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])
# xmlPath = os.path.join(self.defaultSaveDir, basename + XML_EXT)
# txtPath = os.path.join(self.defaultSaveDir, basename + TXT_EXT)
#
# """Annotation file priority:
# PascalXML > YOLO
# """
# if os.path.isfile(xmlPath):
# self.loadPascalXMLByFilename(xmlPath)
# elif os.path.isfile(txtPath):
# self.loadYOLOTXTByFilename(txtPath)
# else:
# xmlPath = os.path.splitext(filePath)[0] + XML_EXT
# txtPath = os.path.splitext(filePath)[0] + TXT_EXT
# if os.path.isfile(xmlPath):
# self.loadPascalXMLByFilename(xmlPath)
# elif os.path.isfile(txtPath):
# self.loadYOLOTXTByFilename(txtPath)
self.setWindowTitle(__appname__ + ' ' + filePath)
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.zoomMode]()
self.zoomWidget.setValue(int(100 * value))
def scaleFitWindow(self):
"""Figure out the size of the pixmap in order to fit the main widget."""
e = 2.0 # So that no scrollbars are generated.
w1 = self.centralWidget().width() - e
h1 = self.centralWidget().height() - e
a1 = w1 / h1
# Calculate a new scale value based on the pixmap's aspect ratio.
w2 = self.canvas.pixmap.width() - 0.0
h2 = self.canvas.pixmap.height() - 0.0
a2 = w2 / h2
return w1 / w2 if a2 >= a1 else h1 / h2
def scaleFitWidth(self):
# The epsilon does not seem to work too well here.
w = self.centralWidget().width() - 2.0
return w / self.canvas.pixmap.width()
def closeEvent(self, event):
if not self.mayContinue():
event.ignore()
settings = self.settings
# If it loads images from dir, don't load it at the begining
if self.dirname is None:
settings[SETTING_FILENAME] = self.filePath if self.filePath else ''
else:
settings[SETTING_FILENAME] = ''
settings[SETTING_WIN_SIZE] = self.size()
settings[SETTING_WIN_POSE] = self.pos()
settings[SETTING_WIN_STATE] = self.saveState()
if self.defaultSaveDir and os.path.exists(self.defaultSaveDir):
settings[SETTING_SAVE_DIR] = ustr(self.defaultSaveDir)
else:
settings[SETTING_SAVE_DIR] = ""
if self.lastOpenDir and os.path.exists(self.lastOpenDir):
settings[SETTING_LAST_OPEN_DIR] = self.lastOpenDir
else:
settings[SETTING_LAST_OPEN_DIR] = ""
settings.save()
## User Dialogs ##
def loadRecent(self, filename):
if self.mayContinue():
self.loadFile(filename)
def scanAllImages(self, folderPath):
extensions = [
'.%s' % fmt.data().decode("ascii").lower()
for fmt in QImageReader.supportedImageFormats()
]
images = []
for root, dirs, files in os.walk(folderPath):
for file in files:
if file.lower().endswith(tuple(extensions)):
relativePath = os.path.join(root, file)
path = ustr(os.path.abspath(relativePath))
images.append(path)
images.sort(key=lambda x: x.lower())
return images
def changeSavedirDialog(self, _value=False):
if self.defaultSaveDir is not None:
path = ustr(self.defaultSaveDir)
else:
path = '.'
dirpath = ustr(
QFileDialog.getExistingDirectory(
self, '%s - Save annotations to the directory' % __appname__,
path, QFileDialog.ShowDirsOnly
| QFileDialog.DontResolveSymlinks))
if dirpath is not None and len(dirpath) > 1:
self.defaultSaveDir = dirpath
self.statusBar().showMessage(
'%s . Annotation will be saved to %s' %
('Change saved folder', self.defaultSaveDir))
self.statusBar().show()
def openAnnotationDialog(self, _value=False):
if self.filePath is None:
self.statusBar().showMessage('Please select image first')
self.statusBar().show()
return
path = os.path.dirname(ustr(self.filePath))\
if self.filePath else '.'
if self.usingPascalVocFormat:
filters = "Open Annotation XML file (%s)" % ' '.join(['*.xml'])
filename = ustr(
QFileDialog.getOpenFileName(
self, '%s - Choose a xml file' % __appname__, path,
filters))
if filename:
if isinstance(filename, (tuple, list)):
filename = filename[0]
self.loadPascalXMLByFilename(filename)
def openDirDialog(self, _value=False, dirpath=None):
if not self.mayContinue():
return
defaultOpenDirPath = dirpath if dirpath else '.'
if self.lastOpenDir and os.path.exists(self.lastOpenDir):
defaultOpenDirPath = self.lastOpenDir
else:
defaultOpenDirPath = os.path.dirname(
self.filePath) if self.filePath else '.'
targetDirPath = ustr(
QFileDialog.getExistingDirectory(
self, '%s - Open Directory' % __appname__, defaultOpenDirPath,
QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks))
self.importDirImages(targetDirPath)
def importDirImages(self, dirpath):
if not self.mayContinue() or not dirpath:
return
self.lastOpenDir = dirpath
self.dirname = dirpath
self.filePath = None
self.fileListWidget.clear()
self.mImgList = self.scanAllImages(dirpath)
self.openNextImg()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileListWidget.addItem(item)
def verifyImg(self, _value=False):
# Proceding next image without dialog if having any label
if self.filePath is not None:
try:
self.labelFile.toggleVerify()
except AttributeError:
# If the labelling file does not exist yet, create if and
# re-save it with the verified attribute.
self.saveFile()
self.labelFile.toggleVerify()
self.canvas.verified = self.labelFile.verified
self.paintCanvas()
self.saveFile()
def openPrevImg(self, _value=False):
# Proceding prev image without dialog if having any label
if self.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):
# Proceding prev image without dialog if having any label
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]
self.loadFile(filename)
def saveFile(self, _value=False):
if self.defaultSaveDir is not None and len(ustr(self.defaultSaveDir)):
if self.filePath:
imgFileName = os.path.basename(self.filePath)
savedFileName = os.path.splitext(imgFileName)[0]
savedPath = os.path.join(ustr(self.defaultSaveDir),
savedFileName)
self._saveFile(savedPath)
else:
imgFileDir = os.path.dirname(self.filePath)
imgFileName = os.path.basename(self.filePath)
savedFileName = os.path.splitext(imgFileName)[0]
savedPath = os.path.join(imgFileDir, savedFileName)
self._saveFile(
savedPath if self.labelFile else self.saveFileDialog())
def saveFileAs(self, _value=False):
assert not self.image.isNull(), "cannot save empty image"
self._saveFile(self.saveFileDialog())
def saveFileDialog(self):
caption = '%s - Choose File' % __appname__
filters = 'File (*%s)' % LabelFile.suffix
openDialogPath = self.currentPath()
dlg = QFileDialog(self, caption, openDialogPath, filters)
dlg.setDefaultSuffix(LabelFile.suffix[1:])
dlg.setAcceptMode(QFileDialog.AcceptSave)
filenameWithoutExtension = os.path.splitext(self.filePath)[0]
dlg.selectFile(filenameWithoutExtension)
dlg.setOption(QFileDialog.DontUseNativeDialog, False)
if dlg.exec_():
return dlg.selectedFiles()[0]
return ''
def closeFile(self, _value=False):
if not self.mayContinue():
return
self.resetState()
self.setClean()
self.toggleActions(False)
self.canvas.setEnabled(False)
self.actions.saveAs.setEnabled(False)
def resetAll(self):
self.settings.reset()
self.close()
proc = QProcess()
proc.startDetached(os.path.abspath(__file__))
def mayContinue(self):
return not (self.dirty and not self.discardChangesDialog())
def discardChangesDialog(self):
yes, no = QMessageBox.Yes, QMessageBox.No
msg = u'You have unsaved changes, proceed anyway?'
return yes == QMessageBox.warning(self, u'Attention', msg, yes | no)
def errorMessage(self, title, message):
return QMessageBox.critical(self, title,
'<p><b>%s</b></p>%s' % (title, message))
def currentPath(self):
return os.path.dirname(self.filePath) if self.filePath else '.'
def chooseColor1(self):
color = self.colorDialog.getColor(self.lineColor,
u'Choose line color',
default=DEFAULT_LINE_COLOR)
if color:
self.lineColor = color
Shape.line_color = color
self.canvas.setDrawingColor(color)
self.canvas.update()
self.setDirty()
def deleteSelectedShape(self):
self.remLabel(self.canvas.deleteSelected())
self.setDirty()
if self.noShapes():
for action in self.actions.onShapesPresent:
action.setEnabled(False)
def chshapeLineColor(self):
color = self.colorDialog.getColor(self.lineColor,
u'Choose line color',
default=DEFAULT_LINE_COLOR)
if color:
self.canvas.selectedShape.line_color = color
self.canvas.update()
self.setDirty()
def chshapeFillColor(self):
color = self.colorDialog.getColor(self.fillColor,
u'Choose fill color',
default=DEFAULT_FILL_COLOR)
if color:
self.canvas.selectedShape.fill_color = color
self.canvas.update()
self.setDirty()
def copyShape(self):
self.canvas.endMove(copy=True)
self.addLabel(self.canvas.selectedShape)
self.setDirty()
def moveShape(self):
self.canvas.endMove(copy=False)
self.setDirty()
def loadPredefinedClasses(self, predefClassesFile):
if os.path.exists(predefClassesFile) is True:
with codecs.open(predefClassesFile, 'r', 'utf8') as f:
for line in f:
line = line.strip()
if self.labelHist is None:
self.labelHist = [line]
else:
self.labelHist.append(line)
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, Ui_MainWindow):
def __init__(self,
defaultFilename=None,
defaultPrefdefClassFile=None,
defaultSaveDir=None):
super(MainWindow, self).__init__()
self.setupUi(self)
# For loading all image under a directory
self.mImgList = []
self.dirname = None
self.labelHist = []
self.lastOpenDir = None
self.filePath = ustr(defaultFilename)
# Load setting in the main thread
self.settings = Settings()
self.settings.load()
settings = self.settings
# Whether we need to save or not.
self.dirty = False
self.statusBar().showMessage('%s started.' % __appname__)
self.statusBar().show()
# Load string bundle for i18n
self.stringBundle = StringBundle.getBundle()
getStr = lambda strId: self.stringBundle.getString(strId) # 作用？
# Save as Pascal voc xml
self.defaultSaveDir = None # defaultSaveDir
self.labelFileFormat = settings.get(SETTING_LABEL_FILE_FORMAT,
LabelFileFormat.PASCAL_VOC)
# Auto saving : Enable auto saving if pressing next
self.autoSaving = QAction(getStr('autoSaveMode'), self)
self.autoSaving.setCheckable(True)
self.autoSaving.setChecked(settings.get(SETTING_AUTO_SAVE,
False)) # 未加入setting
self.canvas = Canvas(parent=self)
# self.canvas.zoomRequest.connect(self.zoomRequest)
self.canvas.setDrawingShapeToSquare(
settings.get(SETTING_DRAW_SQUARE, False))
self.listView.itemDoubleClicked.connect(
self.fileitemDoubleClicked) # 双击
self.open.triggered.connect(self.openFile)
# Actions
action = partial(newAction, self)
quit = action(getStr('quit'), self.close, 'Ctrl+Q', 'quit',
getStr('quitApp'))
open = action(getStr('openFile'), self.openFile, 'Ctrl+O', 'open',
getStr('openFileDetail'))
opendir = action(getStr('openDir'), self.openDirDialog, 'Ctrl+u',
'open', getStr('openDir'))
copyPrevBounding = action(getStr('copyPrevBounding'),
self.copyPreviousBoundingBoxes, 'Ctrl+v',
'paste', getStr('copyPrevBounding'))
changeSavedir = action(getStr('changeSaveDir'),
self.changeSavedirDialog, 'Ctrl+r', 'open',
getStr('changeSavedAnnotationDir'))
# openAnnotation = action(getStr('openAnnotation'), self.openAnnotationDialog,
# 'Ctrl+Shift+O', 'open', getStr('openAnnotationDetail'))
openNextImg = action(getStr('nextImg'), self.openNextImg, 'd', 'next',
getStr('nextImgDetail'))
openPrevImg = action(getStr('prevImg'), self.openPrevImg, 'a', 'prev',
getStr('prevImgDetail'))
# verify = action(getStr('verifyImg'), self.verifyImg,
# 'space', 'verify', getStr('verifyImgDetail'))
save = action(getStr('save'),
self.saveFile,
'Ctrl+S',
'save',
getStr('saveDetail'),
enabled=False)
isUsingPascalVoc = self.labelFileFormat == LabelFileFormat.PASCAL_VOC
# save_format = action('&PascalVOC' if isUsingPascalVoc else '&YOLO',
# self.change_format, 'Ctrl+',
# 'format_voc' if isUsingPascalVoc else 'format_yolo',
# getStr('changeSaveFormat'), enabled=True)
saveAs = action(getStr('saveAs'),
self.saveFileAs,
'Ctrl+Shift+S',
'save-as',
getStr('saveAsDetail'),
enabled=False)
close = action(getStr('closeCur'), self.closeFile, 'Ctrl+W', 'close',
getStr('closeCurDetail'))
deleteImg = action(getStr('deleteImg'), self.deleteImg, 'Ctrl+D',
'close', getStr('deleteImgDetail'))
resetAll = action(getStr('resetAll'), self.resetAll, None, 'resetall',
getStr('resetAllDetail'))
# color1 = action(getStr('boxLineColor'), self.chooseColor1,
# 'Ctrl+L', 'color_line', getStr('boxLineColorDetail'))
# createMode = action(getStr('crtBox'), self.setCreateMode,
# 'w', 'new', getStr('crtBoxDetail'), enabled=False)
# editMode = action('&Edit\nRectBox', self.setEditMode,
# 'Ctrl+J', 'edit', u'Move and edit Boxs', enabled=False)
#
# create = action(getStr('crtBox'), self.createShape,
# 'w', 'new', getStr('crtBoxDetail'), enabled=False)
# delete = action(getStr('delBox'), self.deleteSelectedShape,
# 'Delete', 'delete', getStr('delBoxDetail'), enabled=False)
# copy = action(getStr('dupBox'), self.copySelectedShape,
# 'Ctrl+D', 'copy', getStr('dupBoxDetail'),
# enabled=False)
# advancedMode = action(getStr('advancedMode'), self.toggleAdvancedMode,
# 'Ctrl+Shift+A', 'expert', getStr('advancedModeDetail'),
# checkable=True)
# hideAll = action('&Hide\nRectBox', partial(self.togglePolygons, False),
# 'Ctrl+H', 'hide', getStr('hideAllBoxDetail'),
# enabled=False)
# showAll = action('&Show\nRectBox', partial(self.togglePolygons, True),
# 'Ctrl+A', 'hide', getStr('showAllBoxDetail'),
# enabled=False)
# help = action(getStr('tutorial'), self.showTutorialDialog, None, 'help', getStr('tutorialDetail'))
# showInfo = action(getStr('info'), self.showInfoDialog, None, 'help', getStr('info'))
self.zoomWidget = ZoomWidget()
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(getStr('zoomin'), partial(self.addZoom, 10),
# 'Ctrl++', 'zoom-in', getStr('zoominDetail'), enabled=False)
# zoomOut = action(getStr('zoomout'), partial(self.addZoom, -10),
# 'Ctrl+-', 'zoom-out', getStr('zoomoutDetail'), enabled=False)
# zoomOrg = action(getStr('originalsize'), partial(self.setZoom, 100),
# 'Ctrl+=', 'zoom', getStr('originalsizeDetail'), enabled=False)
# fitWindow = action(getStr('fitWin'), self.setFitWindow,
# 'Ctrl+F', 'fit-window', getStr('fitWinDetail'),
# checkable=True, enabled=False)
# fitWidth = action(getStr('fitWidth'), self.setFitWidth,
# 'Ctrl+Shift+F', 'fit-width', getStr('fitWidthDetail'),
# 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(getStr('editLabel'), self.editLabel,
# 'Ctrl+E', 'edit', getStr('editLabelDetail'),
# enabled=False)
# self.editButton.setDefaultAction(edit)
# shapeLineColor = action(getStr('shapeLineColor'), self.chshapeLineColor,
# icon='color_line', tip=getStr('shapeLineColorDetail'),
# enabled=False)
# shapeFillColor = action(getStr('shapeFillColor'), self.chshapeFillColor,
# icon='color', tip=getStr('shapeFillColorDetail'),
# enabled=False)
# 文件夹
# Store actions for further handling.
self.actions = struct(
save=save,
saveAs=saveAs,
open=open,
close=close, #save_format=save_format,
resetAll=resetAll,
deleteImg=deleteImg,
#lineColor=color1, create=create, delete=delete, edit=edit, copy=copy,
#createMode=createMode, editMode=editMode, advancedMode=advancedMode,
# shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
# zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
# fitWindow=fitWindow, fitWidth=fitWidth,
# zoomActions=zoomActions,
fileMenuActions=(open, opendir, save, saveAs, close, resetAll,
quit),
beginner=(),
advanced=(),
# editMenu=(edit, copy, delete,
# None, color1, self.drawSquaresOption),
# beginnerContext=(create, edit, copy, delete),
# advancedContext=(createMode, editMode, edit, copy,
# delete, shapeLineColor, shapeFillColor),
onLoadActive=(close), #, create, createMode, editMode),
onShapesPresent=(saveAs)) #, hideAll, showAll))
# Custom context menu for the canvas widget:
# addActions(self.canvas.menus[0], self.actions.beginnerContext)
# addActions(self.canvas.menus[1], (
# action('&Copy here', self.copyShape),
# action('&Move here', self.moveShape)))
# self.tools = self.toolbar('Tools')
self.actions.beginner = (
open,
opendir,
changeSavedir,
openNextImg,
openPrevImg,
save,
None, #create, copy, verify, save_format,
#delete, None,
) # zoomIn, zoom, zoomOut, fitWindow, fitWidth)
self.actions.advanced = (
open,
opendir,
changeSavedir,
openNextImg,
openPrevImg,
save,
None, #save_format,
#createMode, editMode, None,
) # hideAll, showAll)
scroll = QScrollArea()
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(True)
self.scrollBars = {
Qt.Vertical: scroll.verticalScrollBar(),
Qt.Horizontal: scroll.horizontalScrollBar()
}
# 图片相关
# Tzutalin 20160906 : Add file list and dock to move faster
def fileitemDoubleClicked(self, item=None):
currIndex = self.mImgList.index(ustr(item.text()))
if currIndex < len(self.mImgList):
filename = self.mImgList[currIndex]
if filename:
self.loadFile(filename)
def openDirDialog(self, _value=False, dirpath=None, silent=False):
if not self.mayContinue():
return
defaultOpenDirPath = dirpath if dirpath else '.'
if self.lastOpenDir and os.path.exists(self.lastOpenDir):
defaultOpenDirPath = self.lastOpenDir
else:
defaultOpenDirPath = os.path.dirname(
self.filePath) if self.filePath else '.'
if silent != True:
targetDirPath = ustr(
QFileDialog.getExistingDirectory(
self, '%s - Open Directory' % __appname__,
defaultOpenDirPath, QFileDialog.ShowDirsOnly
| QFileDialog.DontResolveSymlinks))
else:
targetDirPath = ustr(defaultOpenDirPath)
self.lastOpenDir = targetDirPath
self.importDirImages(targetDirPath)
def importDirImages(self, dirpath):
if not self.mayContinue() or not dirpath:
return
self.lastOpenDir = dirpath
self.dirname = dirpath
self.filePath = None
self.listView.clear() # 1
self.mImgList = self.scanAllImages(dirpath)
self.openNextImg()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.listView.addItem(item) # 1
def openPrevImg(self, _value=False):
# Proceding prev image without dialog if having any label
if self.autoSaving.isChecked():
if self.defaultSaveDir is not None:
if self.dirty is True:
self.saveFile()
else:
self.changeSavedirDialog()
return
if not self.mayContinue():
return
if len(self.mImgList) <= 0:
return
if self.filePath is None:
return
currIndex = self.mImgList.index(self.filePath)
if currIndex - 1 >= 0:
filename = self.mImgList[currIndex - 1]
if filename:
self.loadFile(filename)
def openNextImg(self, _value=False):
# Proceding prev image without dialog if having any label
if self.autoSaving.isChecked():
if self.defaultSaveDir is not None:
if self.dirty is True:
self.saveFile()
else:
self.changeSavedirDialog()
return
if not self.mayContinue():
return
if len(self.mImgList) <= 0:
return
filename = None
if self.filePath is None:
filename = self.mImgList[0]
else:
currIndex = self.mImgList.index(self.filePath)
if currIndex + 1 < len(self.mImgList):
filename = self.mImgList[currIndex + 1]
if filename:
self.loadFile(filename)
def changeSavedirDialog(self, _value=False):
if self.defaultSaveDir is not None:
path = ustr(self.defaultSaveDir)
else:
path = '.'
dirpath = ustr(
QFileDialog.getExistingDirectory(
self, '%s - Save annotations to the directory' % __appname__,
path, QFileDialog.ShowDirsOnly
| QFileDialog.DontResolveSymlinks))
if dirpath is not None and len(dirpath) > 1:
self.defaultSaveDir = dirpath
self.statusBar().showMessage(
'%s . Annotation will be saved to %s' %
('Change saved folder', self.defaultSaveDir))
self.statusBar().show()
def resetState(self):
self.itemsToShapes.clear()
self.shapesToItems.clear()
self.labelList.clear()
self.filePath = None
self.imageData = None
self.labelFile = None
self.canvas.resetState()
self.labelCoordinates.clear()
self.comboBox.cb.clear()
# def openfile(self):
# openfile_name = QFileDialog.getOpenFileName(self, '选择文件', '', 'Excel files(*.xlsx , *.xls)')
#
# def openImage(self):
# imagePath, _ = QFileDialog.getOpenFileName()
# pixmap = QPixmap(imagePath)
# self.label.setPixmap(pixmap)
# self.resize(pixmap.size())
# self.adjustSize()
def openFile(self, _value=False):
if not self.mayContinue():
print(1)
return
path = os.path.dirname(ustr(self.filePath)) if self.filePath else '.'
print(2)
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 loadFile(self, filePath=None):
"""Load the specified file, or the last opened file if None."""
self.resetState()
self.canvas.setEnabled(False)
if filePath is None:
filePath = self.settings.get(SETTING_FILENAME)
# Make sure that filePath is a regular python string, rather than QString
filePath = ustr(filePath)
# Fix bug: An index error after select a directory when open a new file.
unicodeFilePath = ustr(filePath)
unicodeFilePath = os.path.abspath(unicodeFilePath)
# Tzutalin 20160906 : Add file list and dock to move faster
# Highlight the file item
if unicodeFilePath and self.listView.count() > 0:
if unicodeFilePath in self.mImgList:
index = self.mImgList.index(unicodeFilePath)
fileWidgetItem = self.listView.addItem(index) # item 没找到
fileWidgetItem.setSelected(True)
else:
self.listView.clear()
self.mImgList.clear()
if unicodeFilePath and os.path.exists(unicodeFilePath):
if LabelFile.isLabelFile(unicodeFilePath):
try:
self.labelFile = LabelFile(unicodeFilePath)
except LabelFileError as e:
self.errorMessage(
u'Error opening file',
(u"<p><b>%s</b></p>"
u"<p>Make sure <i>%s</i> is a valid label file.") %
(e, unicodeFilePath))
self.status("Error reading %s" % unicodeFilePath)
return False
self.imageData = self.labelFile.imageData
self.lineColor = QColor(*self.labelFile.lineColor)
self.fillColor = QColor(*self.labelFile.fillColor)
self.canvas.verified = self.labelFile.verified
else:
# Load image:
# read data first and store for saving into label file.
self.imageData = read(unicodeFilePath, None)
self.labelFile = None
self.canvas.verified = False
image = QImage.fromData(self.imageData)
if image.isNull():
self.errorMessage(
u'Error opening file',
u"<p>Make sure <i>%s</i> is a valid image file." %
unicodeFilePath)
self.status("Error reading %s" % unicodeFilePath)
return False
self.status("Loaded %s" % os.path.basename(unicodeFilePath))
self.image = image
self.filePath = unicodeFilePath
self.canvas.loadPixmap(QPixmap.fromImage(image))
if self.labelFile:
self.loadLabels(self.labelFile.shapes)
self.setClean()
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.addRecentFile(self.filePath)
self.toggleActions(True)
self.showBoundingBoxFromAnnotationFile(filePath)
self.setWindowTitle(__appname__ + ' ' + filePath)
# Default : select last item if there is at least one item
if self.labelList.count():
self.labelList.setCurrentItem(
self.labelList.item(self.labelList.count() - 1))
self.labelList.item(self.labelList.count() -
1).setSelected(True)
self.canvas.setFocus(True)
return True
return False
def saveFile(self, _value=False):
if self.defaultSaveDir is not None and len(ustr(self.defaultSaveDir)):
if self.filePath:
imgFileName = os.path.basename(self.filePath)
savedFileName = os.path.splitext(imgFileName)[0]
savedPath = os.path.join(ustr(self.defaultSaveDir),
savedFileName)
self._saveFile(savedPath)
else:
imgFileDir = os.path.dirname(self.filePath)
imgFileName = os.path.basename(self.filePath)
savedFileName = os.path.splitext(imgFileName)[0]
savedPath = os.path.join(imgFileDir, savedFileName)
self._saveFile(savedPath if self.labelFile else self.
saveFileDialog(removeExt=False))
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)
self.actions.saveAs.setEnabled(False)
def deleteImg(self):
deletePath = self.filePath
if deletePath is not None:
self.openNextImg()
os.remove(deletePath)
self.importDirImages(self.lastOpenDir)
def resetAll(self):
self.settings.reset()
self.close()
proc = QProcess()
proc.startDetached(os.path.abspath(__file__))
def saveFileDialog(self, removeExt=True):
caption = '%s - Choose File' % __appname__
filters = 'File (*%s)' % LabelFile.suffix
openDialogPath = self.currentPath()
dlg = QFileDialog(self, caption, openDialogPath, filters)
dlg.setDefaultSuffix(LabelFile.suffix[1:])
dlg.setAcceptMode(QFileDialog.AcceptSave)
filenameWithoutExtension = os.path.splitext(self.filePath)[0]
dlg.selectFile(filenameWithoutExtension)
dlg.setOption(QFileDialog.DontUseNativeDialog, False)
if dlg.exec_():
fullFilePath = ustr(dlg.selectedFiles()[0])
if removeExt:
return os.path.splitext(fullFilePath)[
0] # Return file path without the extension.
else:
return fullFilePath
return ''
def currentPath(self):
return os.path.dirname(self.filePath) if self.filePath else '.'
def scanAllImages(self, folderPath):
extensions = [
'.%s' % fmt.data().decode("ascii").lower()
for fmt in QImageReader.supportedImageFormats()
]
images = []
for root, dirs, files in os.walk(folderPath):
for file in files:
if file.lower().endswith(tuple(extensions)):
relativePath = os.path.join(root, file)
path = ustr(os.path.abspath(relativePath))
images.append(path)
natural_sort(images, key=lambda x: x.lower())
return images
def mayContinue(self):
if not self.dirty:
return True
else:
discardChanges = self.discardChangesDialog()
if discardChanges == QMessageBox.No:
return True
elif discardChanges == QMessageBox.Yes:
self.saveFile()
return True
else:
return False
def discardChangesDialog(self):
yes, no, cancel = QMessageBox.Yes, QMessageBox.No, QMessageBox.Cancel
msg = u'You have unsaved changes, would you like to save them and proceed?\nClick "No" to undo all changes.'
return QMessageBox.warning(self, u'Attention', msg, yes | no | cancel)
def zoomRequest(self, delta):
# get the current scrollbar positions
# calculate the percentages ~ coordinates
h_bar = self.scrollBars[Qt.Horizontal]
v_bar = self.scrollBars[Qt.Vertical]
# get the current maximum, to know the difference after zooming
h_bar_max = h_bar.maximum()
v_bar_max = v_bar.maximum()
# get the cursor position and canvas size
# calculate the desired movement from 0 to 1
# where 0 = move left
# 1 = move right
# up and down analogous
cursor = QCursor()
pos = cursor.pos()
relative_pos = QWidget.mapFromGlobal(self, pos)
cursor_x = relative_pos.x()
cursor_y = relative_pos.y()
w = self.scrollArea.width()
h = self.scrollArea.height()
# the scaling from 0 to 1 has some padding
# you don't have to hit the very leftmost pixel for a maximum-left movement
margin = 0.1
move_x = (cursor_x - margin * w) / (w - 2 * margin * w)
move_y = (cursor_y - margin * h) / (h - 2 * margin * h)
# clamp the values from 0 to 1
move_x = min(max(move_x, 0), 1)
move_y = min(max(move_y, 0), 1)
# zoom in
units = delta / (8 * 15)
scale = 10
self.addZoom(scale * units)
# get the difference in scrollbar values
# this is how far we can move
d_h_bar_max = h_bar.maximum() - h_bar_max
d_v_bar_max = v_bar.maximum() - v_bar_max
# get the new scrollbar values
new_h_bar_value = h_bar.value() + move_x * d_h_bar_max
new_v_bar_value = v_bar.value() + move_y * d_v_bar_max
h_bar.setValue(new_h_bar_value)
v_bar.setValue(new_v_bar_value)
def addZoom(self, increment=10):
self.setZoom(self.zoomWidget.value() + increment)
def copyPreviousBoundingBoxes(self):
currIndex = self.mImgList.index(self.filePath)
if currIndex - 1 >= 0:
prevFilePath = self.mImgList[currIndex - 1]
self.showBoundingBoxFromAnnotationFile(prevFilePath)
self.saveFile()
def copySelectedShape(self):
self.addLabel(self.canvas.copySelectedShape())
# fix copy and delete
self.shapeSelectionChanged(True)
class MainWindow(QMainWindow, WindowMixin):
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
def __init__(self, defaultFilename=None, defaultPrefdefClassFile=None, defaultSaveDir=None):
super(MainWindow, self).__init__()
self.setWindowTitle(__appname__)
self.settings = Settings()
self.settings.load()
settings = self.settings
self.stringBundle = StringBundle.getBundle()
getStr = lambda strId: self.stringBundle.getString(strId)
self.defaultSaveDir = defaultSaveDir
self.mImgList = []
self.dirname = None
self.labelHist = []
self.lastOpenDir = None
self.dirty = False
self.countEnable = 0
self._noSelectionSlot = False
self._beginner = True
self.loadPredefinedClasses(defaultPrefdefClassFile)
self.defaultPrefdefClassFile = defaultPrefdefClassFile
self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
self.itemsToShapes = {}
self.shapesToItems = {}
self.prevLabelText = ''
listLayout = QVBoxLayout()
listLayout.setContentsMargins(0, 0, 0, 0)
self.defaultLabelTextLine = QLineEdit()
self.labelList = QListWidget()
self.logContainer = QListWidget()
labelListContainer = QWidget()
labelListContainer.setLayout(listLayout)
self.labelList.itemClicked.connect(self.labelSelectionChanged)
listLayout.addWidget(self.labelList)
self.logContainer.itemChanged.connect(self.pointListChanged)
self.logContainer.itemClicked.connect(self.pointSelectionChange)
self.dock = QDockWidget("Log", self)
self.dock.setObjectName(getStr('labels'))
self.dock.setWidget(self.logContainer)
self.logDock = QDockWidget(getStr('boxLabelText'), self)
self.logDock.setObjectName("Log")
self.logDock.setWidget(labelListContainer)
self.zoomWidget = ZoomWidget()
self.canvas = Canvas(parent=self)
self.canvas.zoomRequest.connect(self.zoomRequest)
self.canvas.setDrawingShapeToSquare(settings.get(SETTING_DRAW_SQUARE, False))
scroll = QScrollArea()
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(True)
self.scrollBars = {
Qt.Vertical: scroll.verticalScrollBar(),
Qt.Horizontal: scroll.horizontalScrollBar()
}
self.scrollArea = scroll
self.canvas.scrollRequest.connect(self.scrollRequest)
self.canvas.newShape.connect(self.newShape)
self.canvas.shapeMoved.connect(self.setDirty)
self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
self.setCentralWidget(scroll)
self.addDockWidget(Qt.RightDockWidgetArea, self.logDock)
self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
self.logDock.setFeatures(self.logDock.features() ^ self.dockFeatures)
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
action = partial(newAction, self)
quit = action(getStr('quit'), self.close,
'', 'quit', getStr('quitApp'))
open = action(getStr('openFile'), self.openFile,
'', 'open', getStr('openFileDetail'))
save = action(getStr('save'), self.saveFile,
'', 'save', getStr('saveDetail'), enabled=False)
count = action('Count', self.countCell,
'', 'color_line', 'count', enabled=False)
self.partition = action("Auto Partition", self.createFixedPartition,
'', 'verify', "Partition", enabled=False)
self.selectNone = action('Select None', self.selectFun,
'', 'done', 'select', enabled=False)
openNextImg = action("Next Image", self.openNextImg,
'', 'next', "next image", enabled=False)
create = action("Create RectBox", self.createShape,
'', 'new', getStr('crtBoxDetail'), enabled=False)
delete = action("Delete RectBox", self.deleteSelectedShape,
'Delete', 'delete', getStr('delBoxDetail'), enabled=False)
copy = action("Duplicate RectBox", self.copySelectedShape,
'', 'copy', getStr('dupBoxDetail'),
enabled=False)
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)
help = action('Help',self.helpDialog,None,'help')
about = action('About',self.aboutDialog,None,'help')
zoomIn = action(getStr('zoomin'), partial(self.addZoom, 10),
'', 'zoom-in', getStr('zoominDetail'), enabled=False)
zoomOut = action(getStr('zoomout'), partial(self.addZoom, -10),
'', 'zoom-out', getStr('zoomoutDetail'), enabled=False)
zoomOrg = action(getStr('originalsize'), partial(self.setZoom, 100),
'', 'zoom', getStr('originalsizeDetail'), enabled=False)
fitWindow = action(getStr('fitWin'), self.setFitWindow,
'', 'fit-window', getStr('fitWinDetail'),
checkable=True, enabled=False)
fitWidth = action(getStr('fitWidth'), self.setFitWidth,
'', 'fit-width', getStr('fitWidthDetail'),
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,
}
labelMenu = QMenu()
addActions(labelMenu, (delete,))
self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
self.labelList.customContextMenuRequested.connect(
self.popLabelListMenu)
self.sporeCountOption = QAction('Spore count', self)
self.sporeCountOption.setCheckable(True)
self.sporeCountOption.setDisabled(True)
self.sporeCountOption.triggered.connect(self.sporeCountFun)
self.rbcCountOption = QAction('RBC count', self)
self.rbcCountOption.setCheckable(True)
self.rbcCountOption.setDisabled(True)
self.rbcCountOption.triggered.connect(self.rbcCountFun)
self.wbcCountOption = QAction('WBC count', self)
self.wbcCountOption.setCheckable(True)
self.wbcCountOption.setDisabled(True)
self.wbcCountOption.triggered.connect(self.wbcCountFun)
self.bacteriaCountOption = QAction('Bacteria count', self)
self.bacteriaCountOption.setCheckable(True)
self.bacteriaCountOption.setDisabled(True)
self.bacteriaCountOption.triggered.connect(self.bacteriaCountFun)
self.labelImage = QAction('Label Image', self)
self.labelImage.setDisabled(True)
self.labelImage.triggered.connect(self.labelImageFun)
self.modifyLabel = QAction('Modify Label', self)
self.modifyLabel.setDisabled(False)
self.modifyLabel.triggered.connect(self.modifyLabelFun)
self.createModel = QAction('Create Model', self)
self.createModel.setDisabled(True)
self.createModel.triggered.connect(self.createModelFun)
self.useModel = QAction('Use Existing Model', self)
self.useModel.setDisabled(True)
self.useModel.triggered.connect(self.useModelFun)
self.actions = struct(save=save, open=open, count=(count,), selectNone=self.selectNone, partition=self.partition, create=create, delete=delete, copy=copy,
zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
fitWindow=fitWindow, fitWidth=fitWidth,
zoomActions=zoomActions,
beginner=(), advanced=(),
menuMenu=(self.sporeCountOption, self.rbcCountOption,
self.wbcCountOption,
self.bacteriaCountOption),
toolMenu=(self.labelImage, self.modifyLabel),
modelMenu=(self.createModel, self.useModel),
beginnerContext=(create, copy, delete),
advancedContext=(copy,
delete),
onLoadActive=(create,),
onNullImage=(self.labelImage, self.sporeCountOption, self.rbcCountOption, self.bacteriaCountOption, self.createModel, self.useModel),
onLabelSelect=(openNextImg, ),
resetDeActive=(count, copy, delete, save, create, openNextImg, self.selectNone, self.partition))
self.menus = struct(
file=self.menu('&File'),
menu=self.menu('&Menu'),
tool=self.menu('&Tools'),
model=self.menu('&Models'),
help=self.menu('&Help'))
self.lastLabel = None
addActions(self.menus.file,
(open, save, quit))
addActions(self.menus.help,(help,about))
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.beginner = (
open, count, self.partition, self.selectNone, save, None, openNextImg, create, copy, delete, None,
zoomIn, zoom, zoomOut, fitWindow, fitWidth)
self.actions.advanced = (open, save, count, self.partition, self.selectNone, None)
self.statusBar().showMessage('%s started.' % __appname__)
self.statusBar().show()
self.image = QImage()
self.filePath = defaultFilename
self.recentFiles = []
self.maxRecent = 7
self.lineColor = None
self.fillColor = None
self.zoom_level = 100
self.fit_window = False
self.difficult = False
if settings.get(SETTING_RECENT_FILES):
if have_qstring():
recentFileQStringList = settings.get(SETTING_RECENT_FILES)
self.recentFiles = [i for i in recentFileQStringList]
else:
self.recentFiles = recentFileQStringList = settings.get(SETTING_RECENT_FILES)
size = settings.get(SETTING_WIN_SIZE, QSize(600, 500))
position = QPoint(0, 0)
saved_position = settings.get(SETTING_WIN_POSE, position)
# Fix the multiple monitors issue
for i in range(QApplication.desktop().screenCount()):
if QApplication.desktop().availableGeometry(i).contains(saved_position):
position = saved_position
break
self.resize(size)
self.move(position)
saveDir = settings.get(SETTING_SAVE_DIR, None)
self.lastOpenDir = settings.get(SETTING_LAST_OPEN_DIR, None)
if self.defaultSaveDir is None and saveDir is not None and os.path.exists(saveDir):
self.defaultSaveDir = saveDir
self.statusBar().showMessage('%s started. Annotation will be saved to %s' %
(__appname__, self.defaultSaveDir))
self.statusBar().show()
self.restoreState(settings.get(SETTING_WIN_STATE, QByteArray()))
Shape.line_color = self.lineColor = QColor(settings.get(SETTING_LINE_COLOR, DEFAULT_LINE_COLOR))
Shape.fill_color = self.fillColor = QColor(settings.get(SETTING_FILL_COLOR, DEFAULT_FILL_COLOR))
self.canvas.setDrawingColor(self.lineColor)
Shape.difficult = self.difficult
self.updateFileMenu()
self.zoomWidget.valueChanged.connect(self.paintCanvas)
self.populateModeActions()
self.labelCoordinates = QLabel('')
self.statusBar().addPermanentWidget(self.labelCoordinates)
def toggaleCount(self, menuObj):
if self.countEnable == 1:
for count in self.actions.count:
count.setEnabled(True)
for menu in self.actions.menuMenu:
menu.setChecked(False)
menuObj.setChecked(True)
def helpDialog(self):
webbrowser.open("resources\\pdf\\help.pdf")
def aboutDialog(self):
webbrowser.open('resources\\pdf\\about.pdf')
def pointListChanged(self, item):
tag = int(item.data(0).split(" ")[0].split("(")[1].split(")")[0])-1
if tag in self.selectedPointList:
item.setCheckState(Qt.Unchecked)
self.selectedPointList.remove(tag)
else:
self.selectedPointList.append(tag)
item.setCheckState(Qt.Checked)
def createFixedPartition(self):
self.pointList = []
if (self.partition.iconText() == "Fixed Partition"):
self.partition.setIconText("Auto Partition")
self.imageObj.findHoughLines()
self.imageObj.sort_lines()
else:
self.imageObj.fixedIntersectionPoint()
self.partition.setIconText("Fixed Partition")
for j in range(len(self.imageObj.intersection_point) - 1):
for i in range(len(self.imageObj.intersection_point[0]) - 1):
self.pointList.append((self.imageObj.intersection_point[j + 1][i],
self.imageObj.intersection_point[j][i],
self.imageObj.intersection_point[j][i + 1],
self.imageObj.intersection_point[j + 1][i + 1],
self.imageObj.find_area(
self.imageObj.intersection_point[j][i],
self.imageObj.intersection_point[j + 1][i + 1])))
n = len(self.pointList)
for i in range(n):
for j in range(0, n - i - 1):
if self.pointList[j][4] < self.pointList[j + 1][4]:
self.pointList[j], self.pointList[j + 1] = self.pointList[j + 1], self.pointList[j]
count = 1
self.selectedPointList = []
self.logContainer.clear()
for i in self.pointList:
item = QListWidgetItem("(" + str(count) + ") " + str(i[4]))
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
if count <= 16:
item.setCheckState(Qt.Checked)
self.selectedPointList.append(count - 1)
else:
item.setCheckState(Qt.Unchecked)
self.logContainer.addItem(item)
count = count + 1
def selectFun(self):
if (self.selectNone.iconText() == "Select None"):
for i in range(len(self.pointList)):
item = self.logContainer.item(i)
tag = int(item.data(0).split(" ")[0].split("(")[1].split(")")[0]) - 1
item.setCheckState(Qt.Unchecked)
self.selectNone.setIconText("Select All")
else:
for i in range(len(self.pointList)):
item = self.logContainer.item(i)
tag = int(item.data(0).split(" ")[0].split("(")[1].split(")")[0]) - 1
item.setCheckState(Qt.Checked)
self.selectNone.setIconText("Select None")
def openNextImg(self):
self.actions.create.setEnabled(True)
self.actions.save.setEnabled(False)
self.logContainer.clear()
self.labelList.clear()
rootPath = "models/research/object_detection/"
if not self.labelFlag:
self.actions.selectNone.setEnabled(False)
self.actions.partition.setEnabled(False)
filelist = [f for f in os.listdir(rootPath+"images") if f.endswith(".jpg")]
for f in filelist:
os.remove(os.path.join(rootPath+"images", f))
filelist = [f for f in os.listdir(rootPath+"images") if f.endswith(".xml")]
for f in filelist:
os.remove(os.path.join(rootPath+"images", f))
for i in self.selectedPointList:
A, B, C, D = self.pointList[i][0], self.pointList[i][1], self.pointList[i][2], self.pointList[i][3]
cv2.imwrite(rootPath+"images/img_"+str(i)+".jpg", self.imageObj.frame[A[1]:D[1], B[0]:A[0]])
self.labelFlag = True
tempImage = read(rootPath+"images/img_"+str(self.selectedPointList[self.currentIndexOfSelectedPoint])+".jpg", None)
self.canvas.verified = False
Image = QImage.fromData(tempImage)
self.canvas.loadPixmap(QPixmap.fromImage(Image))
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.canvas.setFocus(True)
if self.currentIndexOfSelectedPoint + 1 == len(self.selectedPointList):
self.currentIndexOfSelectedPoint = 0
else:
self.currentIndexOfSelectedPoint = self.currentIndexOfSelectedPoint + 1
def pointSelectionChange(self, item):
tag = int(item.data(0).split(" ")[0].split("(")[1].split(")")[0])-1
A, B, C, D = self.pointList[tag][0], self.pointList[tag][1], self.pointList[tag][2], self.pointList[tag][3]
tempImage = np.copy(self.imageObj.ansFrame)
cv2.line(tempImage, A, B, (0, 0, 255), 10)
cv2.line(tempImage, B, C, (0, 0, 255), 10)
cv2.line(tempImage, C, D, (0, 0, 255), 10)
cv2.line(tempImage, D, A, (0, 0, 255), 10)
cv2.imwrite("images/current.jpg", tempImage)
tempImage = read("images/current.jpg", None)
self.canvas.verified = False
Image = QImage.fromData(tempImage)
self.canvas.loadPixmap(QPixmap.fromImage(Image))
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.canvas.setFocus(True)
# self.canvas.verified = False
# Image = QImage(tempImage, tempImage.shape[1], tempImage.shape[0],
# QImage.Format_RGB888)
# self.canvas.loadPixmap(QPixmap.fromImage(Image))
# self.canvas.setEnabled(True)
# self.adjustScale(initial=True)
# self.paintCanvas()
# self.canvas.setFocus(True)
def countCell(self):
self.sporeCountObj.image.ansFrame = np.copy(self.sporeCountObj.image.frame)
sporeCountList = {}
for i in self.selectedPointList:
sporeCountList[i] = self.sporeCountObj.find_spore(self.pointList[i][0], self.pointList[i][1],
self.pointList[i][2], self.pointList[i][3])
print(sporeCountList)
self.canvas.verified = False
tempImage = np.copy(self.sporeCountObj.image.ansFrame)
cv2.imwrite("images/current.jpg", tempImage)
tempImage = read("images/current.jpg", None)
self.canvas.verified = False
Image = QImage.fromData(tempImage)
self.canvas.loadPixmap(QPixmap.fromImage(Image))
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.canvas.setFocus(True)
# plt.show()
# Image = QImage(tempImage, tempImage.shape[1], tempImage.shape[0],
# QImage.Format_RGB888)
# self.canvas.loadPixmap(QPixmap.fromImage(Image))
# self.canvas.setEnabled(True)
# self.adjustScale(initial=True)
# self.paintCanvas()
# self.canvas.setFocus(True)
# self.logContainer.clear()
count = 0
sporeCount = 0
self.logContainer.clear()
for i in self.pointList:
if count in self.selectedPointList:
item = QListWidgetItem("(" + str(count + 1) + ") " + str(i[4]) + " [" + str(sporeCountList[count]) + "]")
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
self.logContainer.addItem(item)
sporeCount = sporeCount + sporeCountList[count]
else:
item = QListWidgetItem("("+str(count+1)+") "+str(i[4]))
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Unchecked)
self.logContainer.addItem(item)
count = count + 1
self.dock.setWindowTitle("Total Count : "+str(sporeCount))
def sporeCountFun(self):
self.setClean()
self.sporeCountObj = Spore("inference_graph/spore.pb", self.imageObj)
self.sporeCountObj.loadModel()
self.createFixedPartition()
self.toggaleCount(self.sporeCountOption)
self.actions.selectNone.setEnabled(True)
self.actions.partition.setEnabled(True)
def rbcCountFun(self):
self.setClean()
self.sporeCountObj = Spore("inference_graph/rbc_second.pb", self.imageObj)
self.sporeCountObj.loadModel()
self.createFixedPartition()
self.toggaleCount(self.rbcCountOption)
self.actions.selectNone.setEnabled(True)
self.actions.partition.setEnabled(True)
def wbcCountFun(self):
print("WBC option.")
self.toggaleCount(self.wbcCountOption)
def bacteriaCountFun(self):
self.setClean()
self.sporeCountObj = Spore("inference_graph/bacteria.pb", self.imageObj)
self.sporeCountObj.loadModel()
self.createFixedPartition()
self.toggaleCount(self.bacteriaCountOption)
self.actions.selectNone.setEnabled(True)
self.actions.partition.setEnabled(True)
def labelImageFun(self):
self.setClean()
self.createFixedPartition()
self.actions.selectNone.setEnabled(True)
self.actions.partition.setEnabled(True)
self.currentIndexOfSelectedPoint = 0
self.labelFlag = False
self.shapeFlag = False
rootPath = "models/research/object_detection/images/"
filelist = os.listdir(rootPath+"train/")
for file in filelist:
os.remove(rootPath+"train/"+file)
filelist = os.listdir(rootPath+"test/")
for file in filelist:
os.remove(rootPath+"test/"+file)
if self.countEnable == 1:
for action in self.actions.onLabelSelect:
action.setEnabled(True)
def modifyLabelFun(self):
self.OPEN_window = QtWidgets.QWidget()
self.OPENModifyLabelui = Ui_modifyLabel()
self.OPENModifyLabelui.setupUi(self.OPEN_window)
self.OPEN_window.show()
self.selectedLabel = ""
self.tempLabelHist = []
for label in self.labelHist:
self.OPENModifyLabelui.labelListWidget.addItem(label)
self.tempLabelHist.append(label)
self.OPENModifyLabelui.labelListWidget.clicked.connect(self.clickOnLabel)
self.OPENModifyLabelui.deleteButton.clicked.connect(self.deleteLabel)
self.OPENModifyLabelui.addButton.clicked.connect(self.addLabelFun)
self.OPENModifyLabelui.saveButton.clicked.connect(self.saveLabelHist)
def createModelFun(self):
self.OPENCreateModelWindow = QtWidgets.QWidget()
self.OPENCreateModelui = Ui_crateModel()
self.OPENCreateModelui.setupUi(self.OPENCreateModelWindow)
self.OPENCreateModelWindow.show()
pwd = os.popen('cd').read()[:-1].replace('\\', '/')
pre = "<!DOCTYPE HTML PUBLIC \"-//W3C//DTD HTML 4.0//EN\" \"http://www.w3.org/TR/REC-html40/strict.dtd\"><html><head><meta name=\"qrichtext\" content=\"1\" /><style type=\"text/css\">p, li { white-space: pre-wrap; }</style></head><body style=\" font-family:'MS Shell Dlg 2'; font-size:7.8pt; font-weight:400; font-style:normal;\"><p style=\" margin-top:12px; margin-bottom:12px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px; background-color:#ffffff;\"><span style=\" font-family:'Consolas'; font-size:9.8pt; font-weight:600; color:#008080;\">"
post = "</span></p></body></html>"
cmd1 = "set PYTHONPATH="
cmd1 += pwd + "/models/research/slim"
cmd2 = "python models/research/object_detection/train.py --logtostderr --train_dir=models/research/object_detection/training/ --pipeline_config_path=models/research/object_detection/training/faster_rcnn_inception_v2_pets.config"
cmd = pre + cmd1 + ' && ' + cmd2 + post
self.setClean()
self.OPENCreateModelui.pathBrowser.setText(cmd)
self.OPENCreateModelui.nextButton.clicked.connect(self.openCmd)
self.OPENCreateModelui.createButton.clicked.connect(self.createPbModelFun)
def useModelFun(self):
self.OPENSelectModelWindow = QtWidgets.QWidget()
self.OPENSelectModelui = Ui_selectModel()
self.OPENSelectModelui.setupUi(self.OPENSelectModelWindow)
self.OPENSelectModelWindow.show()
self.selectedModelName = ""
modelNameList = [f for f in os.listdir("inference_graph") if f.endswith(".pb")]
for modelName in modelNameList:
self.OPENSelectModelui.modelNameListWidget.addItem(modelName)
self.OPENSelectModelui.modelNameListWidget.clicked.connect(self.clickOnModelName)
self.OPENSelectModelui.doneButton.clicked.connect(self.openModel)
def clickOnModelName(self):
self.selectedModelName = self.OPENSelectModelui.modelNameListWidget.currentItem().text()
def openModel(self):
if self.selectedModelName:
self.setClean()
self.sporeCountObj = Spore("inference_graph/"+self.selectedModelName, self.imageObj)
self.sporeCountObj.loadModel()
self.createFixedPartition()
for count in self.actions.count:
count.setEnabled(True)
self.actions.selectNone.setEnabled(True)
self.actions.partition.setEnabled(True)
self.OPENSelectModelWindow.close()
else:
error_dialog = QtWidgets.QErrorMessage()
error_dialog.showMessage('Oops! Please Select model')
error_dialog.exec_()
def openCmd(self):
if self.checkImageFileLable():
if self.checkValidModelName():
rootPath = "models/research/object_detection/images/"
filelist = [f for f in os.listdir(rootPath) if f.endswith(".jpg")]
testFileList = random.sample(filelist, k=math.ceil(len(filelist) * 0.3))
for file in filelist:
if file in testFileList:
shutil.move(rootPath + file, rootPath+"test/" + file)
shutil.move(rootPath + file[:-4] + ".xml", rootPath + "test/" + file[:-4] + ".xml")
else:
shutil.move(rootPath + file[:-4] + ".xml", rootPath + "train/" + file[:-4] + ".xml")
shutil.move(rootPath + file, rootPath + "train/" + file)
MyOut = subprocess.Popen(
'python models/research/object_detection/xml_to_csv.py',
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT, shell=True)
stdout = MyOut.communicate()[0]
print(stdout)
filelist = os.listdir("models/research/object_detection/training")
if "labelmap.pbtxt" in filelist:
filelist.remove("labelmap.pbtxt")
if "pipeline.config" in filelist:
filelist.remove("pipeline.config")
if "faster_rcnn_inception_v2_pets.config" in filelist:
filelist.remove("faster_rcnn_inception_v2_pets.config")
for f in filelist:
os.remove(os.path.join("models/research/object_detection/training", f))
MyOut = subprocess.Popen('python models/research/object_detection/generate_tfrecord.py --csv_input=models/research/object_detection/images/train_labels.csv --image_dir=models/research/object_detection/images/train --output_path=models/research/object_detection/train.record',
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT, shell=True)
stdout = MyOut.communicate()[0]
print(stdout)
MyOut = subprocess.Popen(
'python models/research/object_detection/generate_tfrecord.py --csv_input=models/research/object_detection/images/test_labels.csv --image_dir=models/research/object_detection/images/test --output_path=models/research/object_detection/test.record',
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT, shell=True)
stdout = MyOut.communicate()[0]
print(stdout)
os.popen('start')
self.OPENCreateModelui.nextButton.setEnabled(False)
self.OPENCreateModelui.createButton.setEnabled(True)
else:
error_dialog = QtWidgets.QErrorMessage()
error_dialog.showMessage('Oops! Please Enter Model Name')
error_dialog.exec_()
else:
error_dialog = QtWidgets.QErrorMessage()
error_dialog.showMessage('Oops! Please Make Perfect Labeling, Some XML File Missing')
error_dialog.exec_()
def checkValidModelName(self):
modelName = self.OPENCreateModelui.modelNameLineEdit.text()+".pb"
filelist = [f for f in os.listdir("inference_graph") if f.endswith(".pb")]
if modelName in filelist:
return False
return True
def createPbModelFun(self):
if self.checkValidModelName():
rootPath = "models/research/object_detection/"
filelist = [f for f in os.listdir(rootPath+"training") if f.endswith(".index")]
if filelist:
checkPoint = -1
for fileName in filelist:
if checkPoint < int(fileName[:-6].split('-')[1]):
checkPoint = int(fileName[:-6].split('-')[1])
if os.path.exists(rootPath+'inference_graph'):
shutil.rmtree(rootPath+'inference_graph')
pwd = os.popen('cd').read()[:-1].replace('\\', '/')
cmd1 = "set PYTHONPATH="
cmd1 += pwd + "/models/research/slim"
cmd2 = 'python '+rootPath+'export_inference_graph.py --input_type image_tensor --pipeline_config_path '+rootPath+'training/faster_rcnn_inception_v2_pets.config --trained_checkpoint_prefix '+rootPath+'training/model.ckpt-'+str(checkPoint)+' - -output_directory '+rootPath+'inference_graph'
MyOut = subprocess.Popen( cmd1 +" && " + cmd2,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT, shell=True)
stdout = MyOut.communicate()[0]
print(stdout)
shutil.move(rootPath + "inference_graph/frozen_inference_graph.pb", "inference_graph/" + self.OPENCreateModelui.modelNameLineEdit.text()+".pb")
self.OPENCreateModelWindow.close()
else:
error_dialog = QtWidgets.QErrorMessage()
error_dialog.showMessage('Oops! Please try again, Something was wrong.')
error_dialog.exec_()
else:
error_dialog = QtWidgets.QErrorMessage()
error_dialog.showMessage('Oops! Please Enter Model Name.')
error_dialog.exec_()
def checkImageFileLable(self):
rootPath = "models/research/object_detection/images/"
filelist = [f for f in os.listdir(rootPath) if f.endswith(".jpg")]
for file in filelist:
if not os.path.exists(rootPath+file[:-4]+".xml"):
return False
return True
def clickOnLabel(self):
self.selectedLabel = self.OPENModifyLabelui.labelListWidget.currentItem().text()
self.OPENModifyLabelui.labelNameLineEdit.setText(self.selectedLabel)
def deleteLabel(self):
if self.selectedLabel in self.tempLabelHist:
temp = []
for label in self.tempLabelHist:
if self.selectedLabel != label:
temp.append(label)
self.tempLabelHist = temp
self.OPENModifyLabelui.labelListWidget.clear()
for label in self.tempLabelHist:
self.OPENModifyLabelui.labelListWidget.addItem(label)
def addLabelFun(self):
label = self.OPENModifyLabelui.labelNameLineEdit.text()
if label not in self.tempLabelHist:
self.tempLabelHist.append(label)
self.OPENModifyLabelui.labelListWidget.addItem(label)
def saveLabelHist(self):
with open(self.defaultPrefdefClassFile, 'w') as filehandle:
self.labelHist.clear()
for label in self.tempLabelHist:
filehandle.write(label+"\n")
self.labelHist.append(label)
self.OPEN_window.close()
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.menu.clear()
addActions(self.menus.menu, self.actions.menuMenu)
addActions(self.menus.tool, self.actions.toolMenu)
addActions(self.menus.model, self.actions.modelMenu)
def setDirty(self):
self.dirty = True
self.actions.save.setEnabled(True)
def setClean(self):
self.dirty = False
self.logContainer.clear()
for menu in self.actions.menuMenu:
menu.setChecked(False)
for menu in self.actions.resetDeActive:
menu.setEnabled(False)
for count in self.actions.onNullImage:
count.setDisabled(False)
self.dock.setWindowTitle("Log")
def toggleActions(self):
for z in self.actions.zoomActions:
z.setEnabled(True)
for action in self.actions.onLoadActive:
action.setEnabled(False)
self.countEnable = 1
def status(self, message, delay=5000):
self.statusBar().showMessage(message, delay)
def resetState(self):
self.itemsToShapes.clear()
self.shapesToItems.clear()
self.labelList.clear()
self.filePath = None
self.imageData = None
self.labelFile = None
self.canvas.resetState()
self.labelCoordinates.clear()
def currentItem(self):
items = self.labelList.selectedItems()
if items:
return items[0]
return None
def addRecentFile(self, filePath):
if filePath in self.recentFiles:
self.recentFiles.remove(filePath)
elif len(self.recentFiles) >= self.maxRecent:
self.recentFiles.pop()
self.recentFiles.insert(0, filePath)
def beginner(self):
return self._beginner
def createShape(self):
assert self.beginner()
self.canvas.setEditing(False)
self.actions.create.setEnabled(False)
def toggleDrawingSensitive(self, drawing=True):
if not drawing and self.beginner():
print('Cancel creation.')
self.canvas.setEditing(True)
self.canvas.restoreCursor()
self.actions.create.setEnabled(True)
def updateFileMenu(self):
currFilePath = self.filePath
def exists(filename):
return os.path.exists(filename)
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))
def popLabelListMenu(self, point):
self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
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)
def addLabel(self, shape):
item = HashableQListWidgetItem(shape.label)
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
item.setBackground(generateColorByText(shape.label))
self.itemsToShapes[item] = shape
self.shapesToItems[shape] = item
self.labelList.addItem(item)
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 saveFile(self):
annotationFilePath = "models/research/object_detection/images/img_"+str(self.selectedPointList[self.currentIndexOfSelectedPoint-1])+".xml"
filePath = "models/research/object_detection/images/img_"+str(self.selectedPointList[self.currentIndexOfSelectedPoint-1])+".jpg"
if self.labelFile is None:
self.labelFile = LabelFile()
self.labelFile.verified = self.canvas.verified
def format_shape(s):
return dict(label=s.label,
line_color=s.line_color.getRgb(),
fill_color=s.fill_color.getRgb(),
points=[(p.x(), p.y()) for p in s.points],
difficult = s.difficult)
self.currentTextLable = self.canvas.shapes[0].label
# data = "item {\n\tid: 1\n\tname: '"+self.currentTextLable+"'\n}"
data = "item {\n\tid: 1\n\tname: 'cell'\n}"
file = open("models/research/object_detection/training/labelmap.pbtxt", "w")
file.write(data)
file.close()
shapes = [format_shape(shape) for shape in self.canvas.shapes]
try:
imageFilePath = os.popen('cd').read()[:-1]
self.labelFile.savePascalVocFormat(annotationFilePath, shapes, imageFilePath+"/"+filePath)
return True
except:
self.errorMessage(u'Error', u'Error in saving label data')
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 newShape(self):
if not self.shapeFlag:
if len(self.labelHist) > 0:
self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
self.shapeFlag = True
else:
self.labelDialog = LabelDialog(parent=self, listItem=[self.lastLabel])
text = self.labelDialog.popUp(text=self.prevLabelText)
self.lastLabel = text
if text is not None:
self.prevLabelText = text
generate_color = generateColorByText(text)
shape = self.canvas.setLastLabel(text, generate_color, generate_color)
self.addLabel(shape)
if self.beginner(): # Switch to edit mode.
self.canvas.setEditing(True)
self.actions.create.setEnabled(True)
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 loadFile(self, filePath=None):
self.resetState()
self.canvas.setEnabled(False)
if filePath is None:
filePath = self.settings.get(SETTING_FILENAME)
absFilePath = os.path.abspath(filePath)
if absFilePath and os.path.exists(absFilePath):
self.imageData = read(absFilePath, None)
self.labelFile = None
self.canvas.verified = False
image = QImage.fromData(self.imageData)
if image.isNull():
self.setClean()
for item in self.actions.onNullImage:
item.setDisabled(True)
return False
self.status("Loaded %s" % os.path.basename(absFilePath))
self.image = image
self.filePath = absFilePath
self.canvas.loadPixmap(QPixmap.fromImage(image))
self.setClean()
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.addRecentFile(self.filePath)
self.toggleActions()
self.setWindowTitle(__appname__ + ' ' + filePath)
self.canvas.setFocus(True)
# print(filePath)
self.imageObj = ImageOperation(filePath)
self.imageObj.loadImage()
self.imageObj.cannyConvert()
return True
return False
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 loadRecent(self, filename):
if self.mayContinue():
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 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.remLabel(self.canvas.deleteSelected())
self.setDirty()
def copyShape(self):
self.canvas.endMove(copy=True)
self.addLabel(self.canvas.selectedShape)
self.setDirty()
def moveShape(self):
self.canvas.endMove(copy=False)
self.setDirty()
def loadPredefinedClasses(self, predefClassesFile):
if os.path.exists(predefClassesFile) is True:
with codecs.open(predefClassesFile, 'r', 'utf8') as f:
for line in f:
line = line.strip()
if self.labelHist is None:
self.labelHist = [line]
else:
self.labelHist.append(line)
class MainWindow(QMainWindow, WindowMixin):
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
def __init__(self, defaultFilename=None, defaultPrefdefClassFile=None, defaultSaveDir=None):
super(MainWindow, self).__init__()
self.setWindowTitle(__appname__)
self.settings = Settings()
self.settings.load()
settings = self.settings
self.stringBundle = StringBundle.getBundle()
getStr = lambda strId: self.stringBundle.getString(strId)
self.defaultSaveDir = defaultSaveDir
self.mImgList = []
self.dirname = None
self.labelHist = []
self.lastOpenDir = None
self.dirty = False
self.countEnable = 0
self._noSelectionSlot = False
self._beginner = True
self.loadPredefinedClasses(defaultPrefdefClassFile)
self.defaultPrefdefClassFile = defaultPrefdefClassFile
self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
self.itemsToShapes = {}
self.shapesToItems = {}
self.prevLabelText = ''
listLayout = QVBoxLayout()
listLayout.setContentsMargins(0, 0, 0, 0)
self.defaultLabelTextLine = QLineEdit()
self.labelList = QListWidget()
self.logContainer = QListWidget()
labelListContainer = QWidget()
labelListContainer.setLayout(listLayout)
self.labelList.itemClicked.connect(self.labelSelectionChanged)
listLayout.addWidget(self.labelList)
self.logContainer.itemChanged.connect(self.pointListChanged)
self.logContainer.itemClicked.connect(self.pointSelectionChange)
self.dock = QDockWidget("Log", self)
self.dock.setObjectName(getStr('labels'))
self.dock.setWidget(self.logContainer)
self.logDock = QDockWidget(getStr('boxLabelText'), self)
self.logDock.setObjectName("Log")
self.logDock.setWidget(labelListContainer)
self.zoomWidget = ZoomWidget()
self.canvas = Canvas(parent=self)
self.canvas.zoomRequest.connect(self.zoomRequest)
self.canvas.setDrawingShapeToSquare(settings.get(SETTING_DRAW_SQUARE, False))
scroll = QScrollArea()
scroll.setWidget(self.canvas)
scroll.setWidgetResizable(True)
self.scrollBars = {
Qt.Vertical: scroll.verticalScrollBar(),
Qt.Horizontal: scroll.horizontalScrollBar()
}
self.scrollArea = scroll
self.canvas.scrollRequest.connect(self.scrollRequest)
self.canvas.newShape.connect(self.newShape)
self.canvas.shapeMoved.connect(self.setDirty)
self.canvas.selectionChanged.connect(self.shapeSelectionChanged)
self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
self.setCentralWidget(scroll)
self.addDockWidget(Qt.RightDockWidgetArea, self.logDock)
self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable
self.logDock.setFeatures(self.logDock.features() ^ self.dockFeatures)
self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
action = partial(newAction, self)
quit = action(getStr('quit'), self.close,
'Ctrl+Q', 'quit', getStr('quitApp'))
open = action(getStr('openFile'), self.openFile,
'Ctrl+O', 'open', getStr('openFileDetail'))
save = action(getStr('save'), self.saveFile,
'Ctrl+S', 'save', getStr('saveDetail'), enabled=False)
count = action('Count', self.countCell,
'Ctrl+J', 'color_line', 'count', enabled=False)
self.selectNone = action('Select None', self.selectFun,
'Ctrl+A','done', 'select', enabled=False)
openNextImg = action("Next Image", self.openNextImg,
'd', 'next', "next image", enabled=False)
create = action(getStr('crtBox'), self.createShape,
'w', 'new', getStr('crtBoxDetail'), enabled=False)
delete = action(getStr('delBox'), self.deleteSelectedShape,
'Delete', 'delete', getStr('delBoxDetail'), enabled=False)
copy = action(getStr('dupBox'), self.copySelectedShape,
'Ctrl+D', 'copy', getStr('dupBoxDetail'),
enabled=False)
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)
help = action('Help',self.helpDialog,None,'help')
about = action('About',self.aboutDialog,None,'help')
zoomIn = action(getStr('zoomin'), partial(self.addZoom, 10),
'Ctrl++', 'zoom-in', getStr('zoominDetail'), enabled=False)
zoomOut = action(getStr('zoomout'), partial(self.addZoom, -10),
'Ctrl+-', 'zoom-out', getStr('zoomoutDetail'), enabled=False)
zoomOrg = action(getStr('originalsize'), partial(self.setZoom, 100),
'Ctrl+=', 'zoom', getStr('originalsizeDetail'), enabled=False)
fitWindow = action(getStr('fitWin'), self.setFitWindow,
'Ctrl+F', 'fit-window', getStr('fitWinDetail'),
checkable=True, enabled=False)
fitWidth = action(getStr('fitWidth'), self.setFitWidth,
'Ctrl+Shift+F', 'fit-width', getStr('fitWidthDetail'),
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,
}
labelMenu = QMenu()
addActions(labelMenu, (delete,))
self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
self.labelList.customContextMenuRequested.connect(
self.popLabelListMenu)
self.sporeCountOption = QAction('Spore count', self)
self.sporeCountOption.setCheckable(True)
self.sporeCountOption.setDisabled(True)
self.sporeCountOption.triggered.connect(self.sporeCountFun)
self.rbcCountOption = QAction('RBC count', self)
self.rbcCountOption.setCheckable(True)
self.rbcCountOption.setDisabled(True)
self.rbcCountOption.triggered.connect(self.rbcCountFun)
self.wbcCountOption = QAction('WBC count', self)
self.wbcCountOption.setCheckable(True)
self.wbcCountOption.setDisabled(True)
self.wbcCountOption.triggered.connect(self.wbcCountFun)
self.plateletsCountOption = QAction('Platelets count', self)
self.plateletsCountOption.setCheckable(True)
self.plateletsCountOption.setDisabled(True)
self.plateletsCountOption.triggered.connect(self.plateletsCountFun)
self.bacteriaCountOption = QAction('Bacteria count', self)
self.bacteriaCountOption.setCheckable(True)
self.bacteriaCountOption.setDisabled(True)
self.bacteriaCountOption.triggered.connect(self.bacteriaCountFun)
self.labelImage = QAction('Label Image', self)
self.labelImage.setDisabled(True)
self.labelImage.triggered.connect(self.labelImageFun)
self.modifyLabel = QAction('Modify Label', self)
self.modifyLabel.setDisabled(False)
self.modifyLabel.triggered.connect(self.modifyLabelFun)
self.createModel = QAction('Create Model', self)
self.createModel.setDisabled(False)
self.createModel.triggered.connect(self.createModelFun)
self.actions = struct(save=save, open=open, count=(count,), selectNone=self.selectNone, create=create, delete=delete, copy=copy,
zoom=zoom, zoomIn=zoomIn, zoomOut=zoomOut, zoomOrg=zoomOrg,
fitWindow=fitWindow, fitWidth=fitWidth,
zoomActions=zoomActions,
fileMenuActions=(
open, self.selectNone, save, quit),
beginner=(), advanced=(),
menuMenu=(self.sporeCountOption, self.rbcCountOption,
self.wbcCountOption, self.plateletsCountOption,
self.bacteriaCountOption),
toolMenu=(self.labelImage, self.modifyLabel),
modelMenu=(self.createModel,),
beginnerContext=(create, copy, delete),
advancedContext=(copy,
delete),
onLoadActive=(create,),
onLabelSelect=(openNextImg, ),
resetDeActive=(count, copy, delete, create, openNextImg))
self.menus = struct(
file=self.menu('&File'),
menu=self.menu('&Menu'),
tool=self.menu('&Tools'),
model=self.menu('&Models'),
help=self.menu('&Help'))
self.lastLabel = None
addActions(self.menus.file,
(open, save, quit))
addActions(self.menus.help,(help,about))
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.beginner = (
open, count, self.selectNone, save, None, openNextImg, create, copy, delete, None,
zoomIn, zoom, zoomOut, fitWindow, fitWidth)
self.actions.advanced = (open, save, count, self.selectNone, None)
self.statusBar().showMessage('%s started.' % __appname__)
self.statusBar().show()
self.image = QImage()
self.filePath = defaultFilename
self.recentFiles = []
self.maxRecent = 7
self.lineColor = None
self.fillColor = None
self.zoom_level = 100
self.fit_window = False
self.difficult = False
if settings.get(SETTING_RECENT_FILES):
if have_qstring():
recentFileQStringList = settings.get(SETTING_RECENT_FILES)
self.recentFiles = [i for i in recentFileQStringList]
else:
self.recentFiles = recentFileQStringList = settings.get(SETTING_RECENT_FILES)
size = settings.get(SETTING_WIN_SIZE, QSize(600, 500))
position = QPoint(0, 0)
saved_position = settings.get(SETTING_WIN_POSE, position)
# Fix the multiple monitors issue
for i in range(QApplication.desktop().screenCount()):
if QApplication.desktop().availableGeometry(i).contains(saved_position):
position = saved_position
break
self.resize(size)
self.move(position)
saveDir = settings.get(SETTING_SAVE_DIR, None)
self.lastOpenDir = settings.get(SETTING_LAST_OPEN_DIR, None)
if self.defaultSaveDir is None and saveDir is not None and os.path.exists(saveDir):
self.defaultSaveDir = saveDir
self.statusBar().showMessage('%s started. Annotation will be saved to %s' %
(__appname__, self.defaultSaveDir))
self.statusBar().show()
self.restoreState(settings.get(SETTING_WIN_STATE, QByteArray()))
Shape.line_color = self.lineColor = QColor(settings.get(SETTING_LINE_COLOR, DEFAULT_LINE_COLOR))
Shape.fill_color = self.fillColor = QColor(settings.get(SETTING_FILL_COLOR, DEFAULT_FILL_COLOR))
self.canvas.setDrawingColor(self.lineColor)
Shape.difficult = self.difficult
self.updateFileMenu()
self.zoomWidget.valueChanged.connect(self.paintCanvas)
self.populateModeActions()
self.labelCoordinates = QLabel('')
self.statusBar().addPermanentWidget(self.labelCoordinates)
def toggaleCount(self, menuObj):
if self.countEnable == 1:
for count in self.actions.count:
count.setEnabled(True)
for menu in self.actions.menuMenu:
menu.setChecked(False)
menuObj.setChecked(True)
def helpDialog(self):
webbrowser.open("resources\\pdf\\help.pdf")
def aboutDialog(self):
webbrowser.open('resources\\pdf\\about.pdf')
def pointListChanged(self, item):
tag = int(item.data(0).split(" ")[0].split("(")[1].split(")")[0])-1
if tag in self.selectedPointList:
item.setCheckState(Qt.Unchecked)
self.selectedPointList.remove(tag)
else:
self.selectedPointList.append(tag)
item.setCheckState(Qt.Checked)
def selectFun(self):
if (self.selectNone.iconText() == "Select None"):
for i in range(len(self.pointList)):
item = self.logContainer.item(i)
tag = int(item.data(0).split(" ")[0].split("(")[1].split(")")[0]) - 1
item.setCheckState(Qt.Unchecked)
self.selectNone.setIconText("Select All")
else:
for i in range(len(self.pointList)):
item = self.logContainer.item(i)
tag = int(item.data(0).split(" ")[0].split("(")[1].split(")")[0]) - 1
item.setCheckState(Qt.Checked)
self.selectNone.setIconText("Select None")
def openNextImg(self):
self.actions.create.setEnabled(True)
self.actions.save.setEnabled(False)
self.logContainer.clear()
self.labelList.clear()
if not self.labelFlag:
self.actions.selectNone.setEnabled(False)
filelist = [f for f in os.listdir("images") if f.endswith(".jpg")]
for f in filelist:
os.remove(os.path.join("images", f))
filelist = [f for f in os.listdir("images") if f.endswith(".xml")]
for f in filelist:
os.remove(os.path.join("images", f))
for i in self.selectedPointList:
A, B, C, D = self.pointList[i][0], self.pointList[i][1], self.pointList[i][2], self.pointList[i][3]
cv2.imwrite("images/img_"+str(i)+".jpg", self.imageObj.frame[A[1]:D[1], B[0]:A[0]])
self.labelFlag = True
tempImage = read("images/img_"+str(self.selectedPointList[self.currentIndexOfSelectedPoint])+".jpg", None)
self.canvas.verified = False
Image = QImage.fromData(tempImage)
self.canvas.loadPixmap(QPixmap.fromImage(Image))
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.canvas.setFocus(True)
if self.currentIndexOfSelectedPoint + 1 == len(self.selectedPointList):
self.currentIndexOfSelectedPoint = 0
else:
self.currentIndexOfSelectedPoint = self.currentIndexOfSelectedPoint + 1
def pointSelectionChange(self, item):
tag = int(item.data(0).split(" ")[0].split("(")[1].split(")")[0])-1
A, B, C, D = self.pointList[tag][0], self.pointList[tag][1], self.pointList[tag][2], self.pointList[tag][3]
tempImage = np.copy(self.imageObj.ansFrame)
cv2.line(tempImage, A, B, (0, 0, 255), 5)
cv2.line(tempImage, B, C, (0, 0, 255), 5)
cv2.line(tempImage, C, D, (0, 0, 255), 5)
cv2.line(tempImage, D, A, (0, 0, 255), 5)
self.canvas.verified = False
Image = QImage(tempImage, tempImage.shape[1], tempImage.shape[0],
QImage.Format_RGB888)
self.canvas.loadPixmap(QPixmap.fromImage(Image))
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.canvas.setFocus(True)
def countCell(self):
self.sporeCountObj.image.ansFrame = np.copy(self.sporeCountObj.image.frame)
sporeCountList = {}
for i in self.selectedPointList:
sporeCountList[i] = self.sporeCountObj.find_spore(self.pointList[i][0], self.pointList[i][1],
self.pointList[i][2], self.pointList[i][3])
print(sporeCountList)
self.canvas.verified = False
tempImage = np.copy(self.sporeCountObj.image.ansFrame)
Image = QImage(tempImage, tempImage.shape[1], tempImage.shape[0],
QImage.Format_RGB888)
self.canvas.loadPixmap(QPixmap.fromImage(Image))
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.canvas.setFocus(True)
self.logContainer.clear()
count = 0
sporeCount = 0
for i in self.pointList:
if count in self.selectedPointList:
item = QListWidgetItem("(" + str(count + 1) + ") " + str(i[4]) + " [" + str(sporeCountList[count]) + "]")
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
self.logContainer.addItem(item)
sporeCount = sporeCount + sporeCountList[count]
else:
item = QListWidgetItem("("+str(count+1)+") "+str(i[4]))
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Unchecked)
self.logContainer.addItem(item)
count = count + 1
self.dock.setWindowTitle("Total Count : "+str(sporeCount))
def sporeCountFun(self):
self.setClean()
self.sporeCountObj = Spore("inference_graph/frozen_inference_graph_.pb", self.imageObj)
self.sporeCountObj.loadModel()
self.pointList=[]
for j in range(len(self.sporeCountObj.image.intersection_point) - 1):
for i in range(len(self.sporeCountObj.image.intersection_point[0]) - 1):
self.pointList.append((self.sporeCountObj.image.intersection_point[j + 1][i],
self.sporeCountObj.image.intersection_point[j][i],
self.sporeCountObj.image.intersection_point[j][i + 1],
self.sporeCountObj.image.intersection_point[j + 1][i + 1],
self.sporeCountObj.image.find_area(self.sporeCountObj.image.intersection_point[j][i],
self.sporeCountObj.image.intersection_point[j + 1][i + 1])))
n = len(self.pointList)
for i in range(n):
for j in range(0, n - i - 1):
if self.pointList[j][4] < self.pointList[j + 1][4]:
self.pointList[j], self.pointList[j + 1] = self.pointList[j + 1], self.pointList[j]
count = 1
self.selectedPointList = []
self.logContainer.clear()
for i in self.pointList:
item = QListWidgetItem("("+str(count)+") "+str(i[4]))
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
if count<=16:
item.setCheckState(Qt.Checked)
self.selectedPointList.append(count-1)
else:
item.setCheckState(Qt.Unchecked)
self.logContainer.addItem(item)
count = count + 1
self.toggaleCount(self.sporeCountOption)
self.actions.selectNone.setEnabled(True)
def rbcCountFun(self):
print("RBC option.")
self.toggaleCount(self.rbcCountOption)
def wbcCountFun(self):
print("WBC option.")
self.toggaleCount(self.wbcCountOption)
def plateletsCountFun(self):
print("Platelets option.")
self.toggaleCount(self.plateletsCountOption)
def bacteriaCountFun(self):
print("Bacteria option.")
self.toggaleCount(self.bacteriaCountOption)
def labelImageFun(self):
self.setClean()
self.pointList = []
for j in range(len(self.imageObj.intersection_point) - 1):
for i in range(len(self.imageObj.intersection_point[0]) - 1):
self.pointList.append((self.imageObj.intersection_point[j + 1][i],
self.imageObj.intersection_point[j][i],
self.imageObj.intersection_point[j][i + 1],
self.imageObj.intersection_point[j + 1][i + 1],
self.imageObj.find_area(self.imageObj.intersection_point[j][i],
self.imageObj.intersection_point[j + 1][
i + 1])))
n = len(self.pointList)
for i in range(n):
for j in range(0, n - i - 1):
if self.pointList[j][4] < self.pointList[j + 1][4]:
self.pointList[j], self.pointList[j + 1] = self.pointList[j + 1], self.pointList[j]
count = 1
self.selectedPointList = []
self.logContainer.clear()
for i in self.pointList:
item = QListWidgetItem("(" + str(count) + ") " + str(i[4]))
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
if count <= 16:
item.setCheckState(Qt.Checked)
self.selectedPointList.append(count - 1)
else:
item.setCheckState(Qt.Unchecked)
self.logContainer.addItem(item)
count = count + 1
self.actions.selectNone.setEnabled(True)
self.currentIndexOfSelectedPoint = 0
self.labelFlag = False
self.shapeFlag = False
if self.countEnable == 1:
for action in self.actions.onLabelSelect:
action.setEnabled(True)
def modifyLabelFun(self):
self.OPEN_window = QtWidgets.QWidget()
self.OPENModifyLabelui = Ui_modifyLabel()
self.OPENModifyLabelui.setupUi(self.OPEN_window)
self.OPEN_window.show()
self.selectedLabel = ""
self.tempLabelHist = []
for label in self.labelHist:
self.OPENModifyLabelui.labelListWidget.addItem(label)
self.tempLabelHist.append(label)
self.OPENModifyLabelui.labelListWidget.clicked.connect(self.clickOnLabel)
self.OPENModifyLabelui.deleteButton.clicked.connect(self.deleteLabel)
self.OPENModifyLabelui.addButton.clicked.connect(self.addLabelFun)
self.OPENModifyLabelui.saveButton.clicked.connect(self.saveLabelHist)
def createModelFun(self):
self.OPENCreateModelWindow = QtWidgets.QWidget()
self.OPENCreateModelui = Ui_crateModel()
self.OPENCreateModelui.setupUi(self.OPENCreateModelWindow)
self.OPENCreateModelWindow.show()
self.selectedImageFilePath = []
self.selectedXmlFilePath = []
self.OPENCreateModelui.imageAddButton.clicked.connect(self.addImageLabel)
self.OPENCreateModelui.imageDeleteButton.clicked.connect(self.deleteImageLabel)
self.OPENCreateModelui.xmlAddButton.clicked.connect(self.addXmlLabel)
self.OPENCreateModelui.xmlDeleteButton.clicked.connect(self.deleteXmlLabel)
self.OPENCreateModelui.cancleButton.clicked.connect(self.closeCreateModelUi)
self.OPENCreateModelui.nextButton.clicked.connect(self.nextCreateModelUi)
self.OPENCreateModelui.imageListWidget.clicked.connect(self.clickOnImagePath)
self.OPENCreateModelui.xmlListWidget.clicked.connect(self.clickOnXmlPath)
def addImageLabel(self):
formats = ['*.jpeg', '*.jpg', '*.png', '*.bmp']
filters = "Image files (%s)" % ' '.join(formats)
filename = QFileDialog.getOpenFileName(self, '%s - Choose Image or Label file' % __appname__, '.', filters)
if filename[0]:
self.selectedImageFilePath.append(filename[0])
self.OPENCreateModelui.imageListWidget.addItem(filename[0])
def addXmlLabel(self):
formats = ['*.xml']
filters = "Image files (%s)" % ' '.join(formats)
filename = QFileDialog.getOpenFileName(self, '%s - Choose Image or Label file' % __appname__, '.', filters)
if filename[0]:
self.selectedXmlFilePath.append(filename[0])
self.OPENCreateModelui.xmlListWidget.addItem(filename[0])
def closeCreateModelUi(self):
self.setEnabled(True)
self.OPENCreateModelWindow.close()
def clickOnImagePath(self):
self.selectedImage = self.OPENCreateModelui.imageListWidget.currentItem().text()
def clickOnXmlPath(self):
self.selectedXml = self.OPENCreateModelui.xmlListWidget.currentItem().text()
def clickOnLabel(self):
self.selectedLabel = self.OPENModifyLabelui.labelListWidget.currentItem().text()
self.OPENModifyLabelui.labelNameLineEdit.setText(self.selectedLabel)
def deleteLabel(self):
if self.selectedLabel in self.tempLabelHist:
temp = []
for label in self.tempLabelHist:
if self.selectedLabel != label:
temp.append(label)
self.tempLabelHist = temp
self.OPENModifyLabelui.labelListWidget.clear()
for label in self.tempLabelHist:
self.OPENModifyLabelui.labelListWidget.addItem(label)
def deleteImageLabel(self):
if self.selectedImage in self.selectedImageFilePath:
temp = []
for path in self.selectedImageFilePath:
if self.selectedImage != path:
temp.append(path)
self.selectedImageFilePath = temp
self.OPENCreateModelui.imageListWidget.clear()
for label in self.selectedImageFilePath:
self.OPENCreateModelui.imageListWidget.addItem(label)
def deleteXmlLabel(self):
if self.selectedXml in self.selectedXmlFilePath:
temp = []
for path in self.selectedXmlFilePath:
if self.selectedXml != path:
temp.append(path)
self.selectedXmlFilePath = temp
self.OPENCreateModelui.xmlListWidget.clear()
for label in self.selectedXmlFilePath:
self.OPENCreateModelui.xmlListWidget.addItem(label)
def nextCreateModelUi(self):
if not self.OPENCreateModelui.modelNameLineEdit.text():
error_dialog = QtWidgets.QErrorMessage()
error_dialog.showMessage('Oops! Please Enter Model Name')
error_dialog.exec_()
elif len(self.selectedImageFilePath)!=len(self.selectedXmlFilePath) or not len(self.selectedImageFilePath):
error_dialog = QtWidgets.QErrorMessage()
error_dialog.showMessage('Oops! Please select same number of image and xml file.')
error_dialog.exec_()
else:
print("Done")
def addLabelFun(self):
label = self.OPENModifyLabelui.labelNameLineEdit.text()
if label not in self.tempLabelHist:
self.tempLabelHist.append(label)
self.OPENModifyLabelui.labelListWidget.addItem(label)
def saveLabelHist(self):
with open(self.defaultPrefdefClassFile, 'w') as filehandle:
self.labelHist.clear()
for label in self.tempLabelHist:
filehandle.write(label+"\n")
self.labelHist.append(label)
self.OPEN_window.close()
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.menu.clear()
addActions(self.menus.menu, self.actions.menuMenu)
addActions(self.menus.tool, self.actions.toolMenu)
addActions(self.menus.model, self.actions.modelMenu)
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(False)
self.logContainer.clear()
for menu in self.actions.menuMenu:
menu.setChecked(False)
for menu in self.actions.resetDeActive:
menu.setEnabled(False)
self.sporeCountOption.setDisabled(False)
for count in self.actions.toolMenu:
count.setDisabled(False)
def toggleActions(self):
for z in self.actions.zoomActions:
z.setEnabled(True)
for action in self.actions.onLoadActive:
action.setEnabled(False)
self.countEnable = 1
def status(self, message, delay=5000):
self.statusBar().showMessage(message, delay)
def resetState(self):
self.itemsToShapes.clear()
self.shapesToItems.clear()
self.labelList.clear()
self.filePath = None
self.imageData = None
self.labelFile = None
self.canvas.resetState()
self.labelCoordinates.clear()
def currentItem(self):
items = self.labelList.selectedItems()
if items:
return items[0]
return None
def addRecentFile(self, filePath):
if filePath in self.recentFiles:
self.recentFiles.remove(filePath)
elif len(self.recentFiles) >= self.maxRecent:
self.recentFiles.pop()
self.recentFiles.insert(0, filePath)
def beginner(self):
return self._beginner
def createShape(self):
assert self.beginner()
self.canvas.setEditing(False)
self.actions.create.setEnabled(False)
def toggleDrawingSensitive(self, drawing=True):
if not drawing and self.beginner():
print('Cancel creation.')
self.canvas.setEditing(True)
self.canvas.restoreCursor()
self.actions.create.setEnabled(True)
def updateFileMenu(self):
currFilePath = self.filePath
def exists(filename):
return os.path.exists(filename)
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))
def popLabelListMenu(self, point):
self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
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)
def addLabel(self, shape):
item = HashableQListWidgetItem(shape.label)
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
item.setCheckState(Qt.Checked)
item.setBackground(generateColorByText(shape.label))
self.itemsToShapes[item] = shape
self.shapesToItems[shape] = item
self.labelList.addItem(item)
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 saveFile(self):
annotationFilePath = "images\\img_"+str(self.selectedPointList[self.currentIndexOfSelectedPoint-1])+".xml"
filePath = "images\\img_"+str(self.selectedPointList[self.currentIndexOfSelectedPoint-1])+".jpg"
if self.labelFile is None:
self.labelFile = LabelFile()
self.labelFile.verified = self.canvas.verified
def format_shape(s):
return dict(label=s.label,
line_color=s.line_color.getRgb(),
fill_color=s.fill_color.getRgb(),
points=[(p.x(), p.y()) for p in s.points],
difficult = s.difficult)
shapes = [format_shape(shape) for shape in self.canvas.shapes]
try:
imageFilePath = os.popen('cd').read()[:-1]
self.labelFile.savePascalVocFormat(annotationFilePath, shapes, imageFilePath+"\\"+filePath)
return True
except:
self.errorMessage(u'Error', u'Error in saving label data')
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 newShape(self):
if not self.shapeFlag:
if len(self.labelHist) > 0:
self.labelDialog = LabelDialog(parent=self, listItem=self.labelHist)
self.shapeFlag = True
else:
self.labelDialog = LabelDialog(parent=self, listItem=[self.lastLabel])
text = self.labelDialog.popUp(text=self.prevLabelText)
self.lastLabel = text
if text is not None:
self.prevLabelText = text
generate_color = generateColorByText(text)
shape = self.canvas.setLastLabel(text, generate_color, generate_color)
self.addLabel(shape)
if self.beginner(): # Switch to edit mode.
self.canvas.setEditing(True)
self.actions.create.setEnabled(True)
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 loadFile(self, filePath=None):
self.resetState()
self.canvas.setEnabled(False)
if filePath is None:
filePath = self.settings.get(SETTING_FILENAME)
absFilePath = os.path.abspath(filePath)
if absFilePath and os.path.exists(absFilePath):
self.imageData = read(absFilePath, None)
self.labelFile = None
self.canvas.verified = False
image = QImage.fromData(self.imageData)
if image.isNull():
self.setClean()
self.sporeCountOption.setDisabled(True)
self.labelImage.setDisabled(True)
return False
self.status("Loaded %s" % os.path.basename(absFilePath))
self.image = image
self.filePath = absFilePath
self.canvas.loadPixmap(QPixmap.fromImage(image))
self.setClean()
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.addRecentFile(self.filePath)
self.toggleActions()
self.setWindowTitle(__appname__ + ' ' + filePath)
self.canvas.setFocus(True)
# print(filePath)
self.imageObj = ImageOperation(filePath)
self.imageObj.loadImage()
self.imageObj.cannyConvert()
self.imageObj.findHoughLines()
self.imageObj.sort_lines()
return True
return False
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 loadRecent(self, filename):
if self.mayContinue():
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 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.remLabel(self.canvas.deleteSelected())
self.setDirty()
def copyShape(self):
self.canvas.endMove(copy=True)
self.addLabel(self.canvas.selectedShape)
self.setDirty()
def moveShape(self):
self.canvas.endMove(copy=False)
self.setDirty()
def loadPredefinedClasses(self, predefClassesFile):
if os.path.exists(predefClassesFile) is True:
with codecs.open(predefClassesFile, 'r', 'utf8') as f:
for line in f:
line = line.strip()
if self.labelHist is None:
self.labelHist = [line]
else:
self.labelHist.append(line)
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