class labelMaster(QMainWindow, WindowMixin):
FIT_WINDOW, FIT_WIDTH, MANUAL_ZOOM = list(range(3))
def __init__(self):
QWidget.__init__(self)
# self.setFixedSize(800, 600)
# self.descisionDlg = ouput(self)
self.setWindowTitle("Label Master")
self.labelCoordinates = QLabel('')
self.statusBar().addPermanentWidget(self.labelCoordinates)
# intit
self.cvMat = None
self.qImage = None
self.dirty = False
self.saveFolder = None
self.filePath = None
self.lastOpenDir = None
self.mImgList = []
self.prevLabel = "Enter object label"
self.bLoopImplement = True
self.bAutoImplement = False
self.bSaveOutput = False
self.bShowResult = True
self.bSelectionChanged = False
self.pool = Pool(5)
self.output = []
self.bEndThreads = []
self.t0 = 0
# self.stringBundle = StringBundle.getBundle()
# getStr = lambda strId: self.stringBundle.getString(strId)
self.labelDialog = LabelDialog(listItem=listLabel)
self.colorDialog = ColorDialog(parent=self)
self.lineColor = None
self.textColor = None
self.itemsToShapes = {}
self.shapesToItems = {}
self._noSelectionSlot = []
self._beginner = True
self.statusBar().showMessage('%s' % __appname__)
# self.statusBar().show()
#============ menu=========
self.menus = struct(file=self.menu('&File'),
edit=self.menu('&Edit'),
view=self.menu('&View'),
help=self.menu('&Help'),
labelList=QMenu(),
fileLog=QMenu())
#================init UI ==============
action = partial(newAction, self)
#=======toolBar==========
openFile = action("Open", self.open, OPEN_FILE, "res/open.png",
"Open image file")
openDir = action("Open Dir", self.openDirDialog, OPEN_DIR,
"res/openDir.png", "Open image folder")
createShape = action("Create rect box", self.createShape, CREATE,
"res/draw.png", "Start draw rectangle")
createPolygon = action("Create polygon", self.createPolygon, POLYGON,
"res/polygon.png", "Start draw polygon")
copy = action('copy',
self.copySelectedShape,
COPY,
'res/copy.png',
'dupBoxDetail',
enabled=False)
saveFile = action("Save", self.save, SAVE, "res/save.png",
"Save parameters")
editLabel = action("Edit label",
self.editLabel,
EDIT,
"res/edit.png",
"Editting label",
enabled=False)
deleteShape = action("Delete shape",
self.deleteSelectedShape,
DLETE,
"res/delete.png",
"Delete shape selected",
enabled=False)
implement = action("Implement",
self.implement,
IMPLEMENT,
"res/event.png",
"Implement",
enabled=False)
lineColor = action("Line color",
lambda: self.boxColor("line"),
LINE_COLOR,
"res/lineColor.png",
"Box line color",
enabled=True)
textColor = action("Text color",
lambda: self.boxColor("text"),
TEXT_COLOR,
"res/textColor.png",
"Box text color",
enabled=True)
font = action("Font",
self.boxFont,
TEXT_FONT,
"res/font.png",
"Box font text",
enabled=True)
nextImage = action("Next Image",
self.openNextImg,
NEXT,
"res/next.png",
"Next image",
enabled=True)
backImage = action("Back Image",
self.openPrevImg,
BACK,
"res/back.png",
"Back image",
enabled=True)
zoomIn = action('zoomin',
partial(self.addZoom, 10),
ZOOM_IN,
'res/zoom-in.png',
'zoominDetail',
enabled=False)
zoomOut = action('zoomout',
partial(self.addZoom, -10),
ZOOM_OUT,
'res/zoom-out.png',
'zoomoutDetail',
enabled=False)
zoomOrg = action('originalsize',
partial(self.setZoom, 100),
ZOOM_ORG,
'res/zoom.png',
'originalsizeDetail',
enabled=False)
fitWindow = action('fitWin',
self.setFitWindow,
FIT_WINDOW,
'res/fit-window.png',
'fitWinDetail',
checkable=True,
enabled=False)
fitWidth = action('fitWidth',
self.setFitWidth,
FIT_WIDTH,
'res/fit-width.png',
'fitWidthDetail',
checkable=True,
enabled=False)
clearAll = action('clearAll',
self.clearList,
CLEAR_ALL,
'res/clear-all.png',
'Clear all list items',
checkable=True,
enabled=True)
showCamera = action('camera',
self.showCamera,
SHOW_CAMERA,
'res/camera2.png',
'Camera',
checkable=True,
enabled=True)
actionMenuFile = [openFile, openDir, saveFile]
addActions(self.menus.file, actionMenuFile)
actionMenuEdit = [
createShape, implement, editLabel, copy, lineColor, textColor,
font, deleteShape
]
addActions(self.menus.edit, actionMenuEdit)
actionToolbar = (openFile, openDir, saveFile, createShape,
createPolygon, editLabel, implement, nextImage,
backImage)
self.toolbar("Tools", actions=actionToolbar)
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)
# Group zoom controls into a list for easier toggling.
beginner = (implement, editLabel, copy, deleteShape)
self.actions = struct(openFile=openFile,
openDir=openDir,
saveFile=saveFile,
createShape=createShape,
createPolygon=createPolygon,
editLabel=editLabel,
deleteShape=deleteShape,
copy=copy,
implement=implement,
lineColor=lineColor,
textColor=textColor,
font=font,
nextImage=nextImage,
backImage=backImage,
zoomIn=zoomIn,
zoomOut=zoomOut,
zoomOrg=zoomOrg,
fitWidth=fitWidth,
fitWindow=fitWindow,
beginner=beginner,
clearAll=clearAll)
self.zoomWidget.valueChanged.connect(self.paintCanvas)
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,
}
listLayout = QVBoxLayout()
listLayout.setContentsMargins(0, 0, 0, 0)
# Create a widget for using default label
self.useDefaultLabelCheckbox = QCheckBox('useDefaultLabel')
self.useDefaultLabelCheckbox.setChecked(False)
self.defaultLabelTextLine = QLineEdit()
self.autoImplement = QCheckBox('autoImplement')
self.autoImplement.setChecked(False)
self.saveOutput = QCheckBox('saveOutput')
self.saveOutput.setChecked(False)
self.updateOuput = QCheckBox('showOutput')
self.updateOuput.setChecked(True)
self.autoImplement.stateChanged.connect(self.stateChanged)
self.saveOutput.stateChanged.connect(self.stateChanged)
self.updateOuput.stateChanged.connect(self.stateChanged)
useDefaultLabelQHBoxLayout = QVBoxLayout()
useDefaultLabelQHBoxLayout.addWidget(self.useDefaultLabelCheckbox)
useDefaultLabelQHBoxLayout.addWidget(self.defaultLabelTextLine)
useDefaultLabelQHBoxLayout.addWidget(self.saveOutput)
useDefaultLabelQHBoxLayout.addWidget(self.updateOuput)
useDefaultLabelQHBoxLayout.addWidget(self.autoImplement)
useDefaultLabelContainer = QWidget()
useDefaultLabelContainer.setLayout(useDefaultLabelQHBoxLayout)
# Create a widget for edit and diffc buttonp
self.editButton = QToolButton()
self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
self.editButton.setDefaultAction(editLabel)
self.implementButton = QToolButton()
self.implementButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
self.implementButton.setDefaultAction(implement)
buttonLayout = QHBoxLayout()
buttonLayout.setContentsMargins(0, 0, 0, 0)
# Add some of widgets to listLayout
buttonLayout.addWidget(self.editButton)
buttonLayout.addWidget(self.implementButton)
listLayout.addLayout(buttonLayout)
listLayout.addWidget(useDefaultLabelContainer)
# Create and add a dock widget for showing current label items
self.labelList = QListWidget()
addActions(self.menus.labelList, (editLabel, deleteShape))
self.labelList.setContextMenuPolicy(Qt.CustomContextMenu)
self.labelList.customContextMenuRequested.connect(
self.popLabelListMenu)
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)
labelListContainer = QWidget()
labelListContainer.setLayout(listLayout)
self.fileLogWidget = QListWidget()
addActions(self.menus.fileLog, [clearAll])
self.fileLogWidget.setContextMenuPolicy(Qt.CustomContextMenu)
self.fileLogWidget.customContextMenuRequested.connect(
self.popfileLogMenu)
self.fileLogWidget.itemDoubleClicked.connect(
self.fileitemDoubleClicked)
# self.fileLogWidget.itemChanged.connect(self.fileLogRowChanged)
fileLogLayout = QVBoxLayout()
fileLogLayout.setContentsMargins(0, 0, 0, 0)
fileLogLayout.addWidget(self.fileLogWidget)
fileLogContainer = QWidget()
fileLogContainer.setLayout(fileLogLayout)
self.dock = QDockWidget('boxLabelText', self)
self.dock.setObjectName('labels')
self.dock.setWidget(labelListContainer)
self.filedock = QDockWidget('fileLog', self)
self.filedock.setObjectName('files')
self.filedock.setWidget(fileLogContainer)
self.parameterdock = QDockWidget('paramsVision', self)
self.parameterdock.setObjectName('params')
self.paramerterWidget = ParamerterDialog(self)
# self.paramerterWidget.ln_kThresh.textChanged.connect(self.textChanged)
self.parameterdock.setWidget(self.paramerterWidget)
# self.cameradock = QDockWidget('camera', self)
# self.cameradock.setObjectName('cameara')
# self.cameraDlg = cameraDialog(self)
# self.cameradock.setWidget(self.cameraDlg)
self.addDockWidget(Qt.RightDockWidgetArea, self.dock)
self.dockFeatures = QDockWidget.DockWidgetClosable | QDockWidget.DockWidgetFloatable | QDockWidget.DockWidgetMovable
# self.dock.setFeatures(self.dock.features() ^ self.dockFeatures)
self.dock.setFeatures(self.dockFeatures)
self.addDockWidget(Qt.RightDockWidgetArea, self.filedock)
self.filedock.setFeatures(self.dockFeatures)
self.addDockWidget(Qt.RightDockWidgetArea, self.parameterdock)
self.parameterdock.setFeatures(self.dockFeatures)
# self.addDockWidget(Qt.RightDockWidgetArea, self.cameradock)
# self.cameradock.setFeatures(self.dockFeatures)
toggleDock = self.dock.toggleViewAction()
toggleDock.setText('boxLabel')
toggleDock.setShortcut(TOGGLE_LABEL_LIST)
toggleFileDock = self.filedock.toggleViewAction()
toggleFileDock.setText('logFile')
toggleFileDock.setShortcut(TOGGLE_LOG_FILE)
toggleParaDock = self.parameterdock.toggleViewAction()
toggleParaDock.setText('parameter')
toggleParaDock.setShortcut(TOGGLE_PARAMETER)
# toggleCamDock = self.cameradock.toggleViewAction()
# toggleCamDock.setText('camera')
# toggleCamDock.setShortcut(TOGGLE_SHOW_CAMERA)
# self.dock.hide()
self.filedock.hide()
self.parameterdock.hide()
# self.cameradock.hide()
addActions(self.menus.view,
[toggleDock, toggleFileDock, toggleParaDock, showCamera])
# ==============Canvas================
self.canvas = Canvas(self)
# load params
self.loadParams()
# contextMenu
addActions(self.canvas.menus[0], self.actions.beginner)
#=================
qImage = QImage(640, 480, QImage.Format_RGBA8888)
self.canvas.loadPixmap(QPixmap.fromImage(qImage))
scroll = QScrollArea()
scroll.setWidget(self.canvas)
# scroll.setWidget(self.listWidget)
scroll.setWidgetResizable(True)
self.scrollBars = {
Qt.Vertical: scroll.verticalScrollBar(),
Qt.Horizontal: scroll.horizontalScrollBar()
}
self.scrollArea = scroll
#signal Canvas
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.zoomRequest.connect(self.zoomRequest)
self.canvas.drawingPolygon.connect(self.toggleDrawingSensitive)
self.setCentralWidget(scroll)
self.fitState()
#
# ==============================
palette = self.palette()
role = self.backgroundRole()
palette.setColor(role, QColor(200, 232, 232, 255))
self.setPalette(palette)
palette = self.canvas.palette()
role = self.canvas.backgroundRole()
palette.setColor(role, QColor(200, 232, 232, 255))
self.canvas.setPalette(palette)
#==start thead implement
self.myThread(target=self.loopImplement, args=())
#==start thead OpenCam
#======Functions==
def getStrDateTime(self):
return time.strftime("%y%m%d_%H%M%S")
def saveImage(self, image=None):
if self.saveFolder is None:
self.saveFolder = QFileDialog.getExistingDirectory(
self, "Open Directory", os.getcwd(),
QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks)
if os.path.exists(self.saveFolder):
if isinstance(image, np.ndarray):
filepath = os.path.join(self.saveFolder,
self.getStrDateTime() + ".jpg")
print(filepath)
cv2.imwrite(filepath, image)
def logFile(self, obj):
txt = time.strftime("%H:%M:%S ") + str(obj)
self.fileLogWidget.addItem(txt)
def shapeProcess(self, shape):
#
print("prosess : ", shape)
# index = self.canvas.shapes.index(shape)
index, labels, rect = self.formatShape(shape)
self.bEndThreads[index] = False
# result
shape.result["in"]["id"] = index
shape.result["in"]["label"] = labels
shape.result["in"]["rect"] = rect
try:
roi = self.qImage.copy(rect)
cvRoi = qImageToCvMat(roi)
except:
cvRoi = None
return
lbs = labels.split(",")[:-1]
for label in lbs:
if label == OCR:
text = self.getOCR(shape, label, cvRoi)
elif label == BARCODE:
codes = self.getBarCode(shape, label, cvRoi)
# elif label == CROP:
# self.saveImage(cvRoi)
# update text , barcode on canvas
if self.bShowResult and shape.result:
w, h = rect.width(), rect.height()
txt = ""
for key in lbs:
txt += str(shape.result["out"][key]) + "\n"
self.canvas.text[index] = txt
rect.setHeight(len(lbs) * rect.height())
self.canvas.locText[index] = rect.translated(QPoint(0, -h))
self.canvas.update()
self.bEndThreads[index] = True
# if self.bEndThreads == len(self.canvas.shapes)*[True]:
# self.logFile("dt : %.3f"%(time.time() - self.t0))
if shape.result:
print(shape.result)
return shape.result
else:
return None
def getOCR(self, shape, key, img=None):
lang = self.paramerterWidget.cbb_language.currentText()
oem = self.paramerterWidget.cbb_oem.currentIndex()
psm = self.paramerterWidget.cbb_psm.currentIndex()
config = '--oem %d --psm %d' % (oem, psm)
# print(config)
txt = get_text(img, lang, config=config)
if shape:
shape.result["out"][key] = txt
return txt
def getBarCode(self, shape, key, img=None):
if img is None:
return
codes = getBarcode(img)
if shape:
shape.result["out"][key] = codes
return codes
# actions
def read(self, filename, default=None):
try:
with open(filename, "rb") as f:
return f.read()
except:
return default
def createPolygon(self):
self.canvas.setDrawPolygon()
self.actions.createShape.setEnabled(True)
def createShape(self):
self.canvas.setEditing(False)
self.actions.createShape.setEnabled(False)
def remLabel(self, shape):
if shape is None:
# print('rm empty label')
return
item = self.shapesToItems[shape]
# index = self.canvas.shapes.index(shape)
self.labelList.takeItem(self.labelList.row(item))
self.canvas.text = []
self.canvas.locText = []
if self.labelList.count() == 0:
self.actions.implement.setEnabled(False)
self.actions.editLabel.setEnabled(False)
self.actions.copy.setEnabled(False)
self.actions.deleteShape.setEnabled(False)
del self.shapesToItems[shape]
del self.itemsToShapes[item]
# del self.bEndThreads[index]
def boxFont(self):
font, ok = QFontDialog.getFont()
if ok:
self.canvas.fontText = font
def boxColor(self, obj):
color = self.colorDialog.getColor(self.lineColor,
u'Choose %s color' % obj,
default=DEFAULT_LINE_COLOR)
# print(color,type(color))
if color:
if obj == LINE:
self.lineColor = color
Shape.line_color = color
self.canvas.setDrawingColor(color)
elif obj == TEXT:
# pass
self.canvas.drawingTextColor = color
self.canvas.update()
self.setDirty()
def deleteSelectedShape(self):
self.remLabel(self.canvas.deleteSelected())
def implement(self):
if self.qImage is None:
return
shapes = self.canvas.shapes
self.canvas.text = len(shapes) * [None]
self.canvas.locText = len(shapes) * [None]
# self.canvas.drawingTextColor = len(shapes)*[None]
# print(shape.points)
self.t0 = time.time()
self.pool.map(self.shapeProcess, shapes)
self.logFile("dt : %.3f" % (time.time() - self.t0))
# multi thread
# targets = len(shapes)*[self.shapeProcess]
# args = [(shape,) for shape in shapes]
# self.myMultiThread(targets = targets, args=args)
#========
# dt = time.time() - t0
# self.logFile("tactime : %s"%dt)
if self.bSaveOutput:
pass
def open(self):
filters = "All File(*.*);;JPG iamge (*.jpg)"
self.filePath, _ = QFileDialog.getOpenFileName(self, 'Select File', "",
filters)
if self.filePath:
self.mImgList = []
self.loadFile(self.filePath)
def loadFile(self, filePath):
self.filePath = filePath
_, extension = os.path.splitext(filePath)
if filePath != "" and extension in __imageFile__ and os.path.exists(
filePath):
self.setWindowTitle("Label Master " + filePath)
self.resetState()
self.qImage = QImage(filePath)
# self.qImage = self.qImage.convertToFormat(QImage.Format_ARGB32)
self.logFile(getFormatQImage(self.qImage))
# self.cvMat = cv2.imread(self.filePath)
# if isinstance(self.cvMat,np.ndarray):
# qImg = cVMatToQImage(self.cvMat[...,::-1])
self.canvas.loadPixmap(QPixmap.fromImage(self.qImage))
self.fitState()
pass
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 dirpath:
return
self.lastOpenDir = dirpath
self.dirname = dirpath
self.filePath = None
self.fileLogWidget.clear()
self.mImgList = self.scanAllImages(dirpath)
self.openNextImg()
for imgPath in self.mImgList:
item = QListWidgetItem(imgPath)
self.fileLogWidget.addItem(item)
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 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 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 loadParams(self):
try:
params = load_from_json("params/param.json")
shapes = params["shapes"]
for sh in shapes:
index, label, r = sh["id"], sh["label"], sh["rect"]
generate_color = generateColorByText(label)
shape = Shape(label=label,
line_color=generate_color,
paintLabel=True)
# shape.fill = True
shape.fill_color = generate_color
shape.points = 4 * [QPoint()]
shape.points[0] = QPointF(r[0], r[1])
shape.points[1] = QPointF(r[0] + r[2], r[1])
shape.points[2] = QPointF(r[0] + r[2], r[1] + r[3])
shape.points[3] = QPointF(r[0], r[1] + r[3])
shape.close()
self.canvas.shapes.append(shape)
self.addLabel(shape)
except:
pass
def save(self):
a = [self.formatShape(shape) for shape in self.canvas.shapes]
b = [{
"id": i,
"label": lb,
"rect": [r.x(), r.y(), r.width(),
r.height()]
} for i, lb, r in a]
raw_data = {
"k-thresh": self.paramerterWidget.ln_threshold.text(),
"block size": self.paramerterWidget.ln_blockSize.text(),
"shapes": b
}
save_to_json("params/param.json", raw_data)
pass
def setDirty(self):
self.dirty = True
# self.actions.createShape.setEnabled(False)
shape = self.canvas.selectedShape
# self.implement()
def formatShape(self, shape):
if shape is None:
return None
points = shape.points
tl = points[0].toPoint()
br = points[2].toPoint()
idx = self.canvas.shapes.index(shape)
label = shape.label
cvRect = QRect(tl, br)
return idx, label, cvRect
def myThread(self, target, args=()):
thread = threading.Thread(target=target, args=args)
thread.start()
def myMultiThread(self, targets, args):
for tar, arg in zip(targets, args):
thread = threading.Thread(target=tar, args=arg)
thread.start()
# run thread auto implement
def loopImplement(self):
while self.bLoopImplement:
if self.bAutoImplement and self.bEndThreads == len(
self.canvas.shapes) * [True]:
self.implement()
time.sleep(0.2)
# redraw rect
# ======signal slot================
def copySelectedShape(self):
self.addLabel(self.canvas.copySelectedShape())
# fix copy and delete
self.shapeSelectionChanged(True)
def textChanged(self, text):
# self.canvas.text = text
self.canvas.update()
# ================add and edit Label===========================
def stateChanged(self, iState):
self.bAutoImplement = self.autoImplement.isChecked()
if self.bAutoImplement:
self.bEndThreads = len(self.canvas.shapes) * [True]
self.bSaveOutput = self.saveOutput.isChecked()
self.bShowResult = self.updateOuput.isChecked()
self.canvas.update()
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 beginner(self):
return self._beginner
def advanced(self):
return not self.beginner()
def toggleDrawingSensitive(self, drawing=True):
"""In the middle of drawing, toggling between modes should be disabled."""
# self.actions.editLabel.setEnabled(not drawing)
# self.actions.implement.setEnabled(not drawing)
if not drawing and self.beginner():
# Cancel creation.
print('Cancel creation.')
self.canvas.setEditing(True)
self.canvas.restoreCursor()
self.actions.createShape.setEnabled(True)
def popLabelListMenu(self, point):
self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
def popfileLogMenu(self, point):
self.menus.fileLog.exec_(self.fileLogWidget.mapToGlobal(point))
def currentItem(self):
items = self.labelList.selectedItems()
if items:
return items[0]
return None
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.bSelectionChanged = selected
self.actions.implement.setEnabled(selected)
self.actions.editLabel.setEnabled(selected)
self.actions.copy.setEnabled(selected)
self.actions.deleteShape.setEnabled(selected)
# self.actions.shapeLineColor.setEnabled(selected)
# self.actions.shapeFillColor.setEnabled(selected)
def editLabel(self):
if not self.canvas.editing():
return
item = self.currentItem()
if not item:
return
# self.descisionDlg.show()
text = self.labelDialog.popUp(item.text())
if text is not None:
shape = self.itemsToShapes[item]
index, label, rect = self.formatShape(shape)
item.setText(text)
item.setBackground(generateColorByText(text))
self.setDirty()
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]
# self.actions.editLabel.setEnabled(True)
# self.actions.implement.setEnabled(True)
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 addLabel(self, shape):
shape.paintLabel = True
index, label, rect = self.formatShape(shape)
item = HashableQListWidgetItem(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)
self.bEndThreads.append(False)
# =================signal Canvas==========================
def newShape(self):
#Pop-up and give focus to the label editor.
if self.useDefaultLabelCheckbox.isChecked(
) and self.defaultLabelTextLine.text():
label = self.defaultLabelTextLine.text()
else:
label = self.prevLabel
text = self.labelDialog.popUp(text=label)
if text:
generate_color = generateColorByText(text)
shape = self.canvas.setLastLabel(text, generate_color,
generate_color)
self.prevLabel = text
self.canvas.setEditing(True)
self.actions.createShape.setEnabled(True)
self.addLabel(shape)
else:
self.canvas.resetAllLines()
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):
# print(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 setCreateMode(self):
self.canvas.setEditing(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 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 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
if h2 == 0:
h2 = 1
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 clearList(self):
# nRow = self.fileLogWidget.count()
# for i in range(nRow):
# self.fileLogWidget.takeItem(i)
self.fileLogWidget.clear()
def showCamera(self):
self.cameraDlg = cameraDialog(self)
self.cameraDlg.show()
def resetState(self):
# self.canvas.deleteSelected()
# self.labelList.clear()
# self.itemsToShapes.clear()
# self.shapesToItems.clear()
self.canvas.locText = len(self.canvas.shapes) * [None]
self.canvas.text = len(self.canvas.shapes) * [None]
# self.canvas.drawingTextColor = len(self.canvas.shapes)*[None]
self.canvas.resetState()
self.labelCoordinates.clear()
self.canvas.setEnabled(False)
self.canvas.verified = False
def fitState(self):
self.canvas.setEnabled(True)
self.adjustScale(initial=True)
self.paintCanvas()
self.canvas.setFocus(True)
def resizeEvent(self, event):
if self.canvas and self.zoomMode != self.MANUAL_ZOOM:
self.adjustScale()
super(labelMaster, self).resizeEvent(event)
def closeEvent(self, event):
print("Close main")
self.bLoopImplement = False
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_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):
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, 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)