Assign different labels with different colors
This commit is contained in:
tzutalin
parent
713cf4537e
commit
6cf04adfa5
95
labelImg.py
95
labelImg.py
@ -27,7 +27,7 @@ except ImportError:
|
||||
import resources
|
||||
# Add internal libs
|
||||
from libs.constants import *
|
||||
from libs.lib import struct, newAction, newIcon, addActions, fmtShortcut
|
||||
from libs.lib import struct, newAction, newIcon, addActions, fmtShortcut, generateColorByText
|
||||
from libs.settings import Settings
|
||||
from libs.shape import Shape, DEFAULT_LINE_COLOR, DEFAULT_FILL_COLOR
|
||||
from libs.canvas import Canvas
|
||||
@ -209,13 +209,13 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
open = action('&Open', self.openFile,
|
||||
'Ctrl+O', 'open', u'Open image or label file')
|
||||
|
||||
opendir = action('&Open Dir', self.openDir,
|
||||
opendir = action('&Open Dir', self.openDirDialog,
|
||||
'Ctrl+u', 'open', u'Open Dir')
|
||||
|
||||
changeSavedir = action('&Change Save Dir', self.changeSavedir,
|
||||
changeSavedir = action('&Change Save Dir', self.changeSavedirDialog,
|
||||
'Ctrl+r', 'open', u'Change default saved Annotation dir')
|
||||
|
||||
openAnnotation = action('&Open Annotation', self.openAnnotation,
|
||||
openAnnotation = action('&Open Annotation', self.openAnnotationDialog,
|
||||
'Ctrl+Shift+O', 'open', u'Open Annotation')
|
||||
|
||||
openNextImg = action('&Next Image', self.openNextImg,
|
||||
@ -243,7 +243,7 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
'Ctrl+Shift+L', 'color', u'Choose Box fill color')
|
||||
|
||||
createMode = action('Create\nRectBox', self.setCreateMode,
|
||||
'Ctrl+N', 'new', u'Start drawing Boxs', enabled=False)
|
||||
'w', 'new', u'Start drawing Boxs', enabled=False)
|
||||
editMode = action('&Edit\nRectBox', self.setEditMode,
|
||||
'Ctrl+J', 'edit', u'Move and edit Boxs', enabled=False)
|
||||
|
||||
@ -427,13 +427,10 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
(__appname__, self.defaultSaveDir))
|
||||
self.statusBar().show()
|
||||
|
||||
# or simply:
|
||||
# self.restoreGeometry(settings[SETTING_WIN_GEOMETRY]
|
||||
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.line_color = self.lineColor = QColor(settings.get(SETTING_LINE_COLOR, DEFAULT_LINE_COLOR))
|
||||
Shape.fill_color = self.fillColor = QColor(settings.get(SETTING_FILL_COLOR, DEFAULT_FILL_COLOR))
|
||||
self.canvas.setDrawingColor(self.lineColor)
|
||||
# Add chris
|
||||
Shape.difficult = self.difficult
|
||||
|
||||
@ -448,9 +445,12 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
|
||||
# Populate the File menu dynamically.
|
||||
self.updateFileMenu()
|
||||
# Since loading the file may take some time, make sure it runs in the
|
||||
# background.
|
||||
self.queueEvent(partial(self.loadFile, self.filePath or ""))
|
||||
|
||||
# Since loading the file may take some time, make sure it runs in the background.
|
||||
if self.filePath and os.path.isdir(self.filePath):
|
||||
self.queueEvent(partial(self.importDirImages, self.filePath or ""))
|
||||
elif self.filePath:
|
||||
self.queueEvent(partial(self.loadFile, self.filePath or ""))
|
||||
|
||||
# Callbacks:
|
||||
self.zoomWidget.valueChanged.connect(self.paintCanvas)
|
||||
@ -598,10 +598,10 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
def popLabelListMenu(self, point):
|
||||
self.menus.labelList.exec_(self.labelList.mapToGlobal(point))
|
||||
|
||||
def editLabel(self, item=None):
|
||||
def editLabel(self):
|
||||
if not self.canvas.editing():
|
||||
return
|
||||
item = item if item else self.currentItem()
|
||||
item = self.currentItem()
|
||||
text = self.labelDialog.popUp(item.text())
|
||||
if text is not None:
|
||||
item.setText(text)
|
||||
@ -686,12 +686,18 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
shape.difficult = difficult
|
||||
shape.close()
|
||||
s.append(shape)
|
||||
self.addLabel(shape)
|
||||
|
||||
if line_color:
|
||||
shape.line_color = QColor(*line_color)
|
||||
else:
|
||||
shape.line_color = generateColorByText(label)
|
||||
|
||||
if fill_color:
|
||||
shape.fill_color = QColor(*fill_color)
|
||||
else:
|
||||
shape.fill_color = generateColorByText(label)
|
||||
|
||||
self.addLabel(shape)
|
||||
|
||||
self.canvas.loadShapes(s)
|
||||
|
||||
@ -703,10 +709,8 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
|
||||
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,
|
||||
line_color=s.line_color.getRgb(),
|
||||
fill_color=s.fill_color.getRgb(),
|
||||
points=[(p.x(), p.y()) for p in s.points],
|
||||
# add chris
|
||||
difficult = s.difficult)
|
||||
@ -723,8 +727,7 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
self.lineColor.getRgb(), self.fillColor.getRgb())
|
||||
return True
|
||||
except LabelFileError as e:
|
||||
self.errorMessage(u'Error saving label data',
|
||||
u'<b>%s</b>' % e)
|
||||
self.errorMessage(u'Error saving label data', u'<b>%s</b>' % e)
|
||||
return False
|
||||
|
||||
def copySelectedShape(self):
|
||||
@ -774,7 +777,9 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
self.diffcButton.setChecked(False)
|
||||
if text is not None:
|
||||
self.prevLabelText = text
|
||||
self.addLabel(self.canvas.setLastLabel(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)
|
||||
@ -904,6 +909,7 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
# read data first and store for saving into label file.
|
||||
self.imageData = read(unicodeFilePath, None)
|
||||
self.labelFile = None
|
||||
|
||||
image = QImage.fromData(self.imageData)
|
||||
if image.isNull():
|
||||
self.errorMessage(u'Error opening file',
|
||||
@ -996,12 +1002,12 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
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:
|
||||
if self.defaultSaveDir and os.path.exists(self.defaultSaveDir):
|
||||
settings[SETTING_SAVE_DIR] = ustr(self.defaultSaveDir)
|
||||
else:
|
||||
settings[SETTING_SAVE_DIR] = ""
|
||||
|
||||
if self.lastOpenDir is not None and len(self.lastOpenDir) > 1:
|
||||
if self.lastOpenDir and os.path.exists(self.lastOpenDir):
|
||||
settings[SETTING_LAST_OPEN_DIR] = self.lastOpenDir
|
||||
else:
|
||||
settings[SETTING_LAST_OPEN_DIR] = ""
|
||||
@ -1028,7 +1034,7 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
images.sort(key=lambda x: x.lower())
|
||||
return images
|
||||
|
||||
def changeSavedir(self, _value=False):
|
||||
def changeSavedirDialog(self, _value=False):
|
||||
if self.defaultSaveDir is not None:
|
||||
path = ustr(self.defaultSaveDir)
|
||||
else:
|
||||
@ -1045,7 +1051,7 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
('Change saved folder', self.defaultSaveDir))
|
||||
self.statusBar().show()
|
||||
|
||||
def openAnnotation(self, _value=False):
|
||||
def openAnnotationDialog(self, _value=False):
|
||||
if self.filePath is None:
|
||||
self.statusBar().showMessage('Please select image first')
|
||||
self.statusBar().show()
|
||||
@ -1061,23 +1067,26 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
filename = filename[0]
|
||||
self.loadPascalXMLByFilename(filename)
|
||||
|
||||
def openDir(self, _value=False):
|
||||
def openDirDialog(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
|
||||
defaultOpenDirPath = '.'
|
||||
if self.lastOpenDir and os.path.exists(self.lastOpenDir):
|
||||
defaultOpenDirPath = self.lastOpenDir
|
||||
else:
|
||||
defaultOpenDirPath = os.path.dirname(self.filePath) if self.filePath else '.'
|
||||
|
||||
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
|
||||
'%s - Open Directory' % __appname__, defaultOpenDirPath,
|
||||
QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks))
|
||||
self.importDirImages(dirpath)
|
||||
|
||||
def importDirImages(self, dirpath):
|
||||
if not self.mayContinue() or not dirpath:
|
||||
return
|
||||
|
||||
self.lastOpenDir = dirpath
|
||||
self.dirname = dirpath
|
||||
self.filePath = None
|
||||
self.fileListWidget.clear()
|
||||
@ -1109,7 +1118,7 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
if self.dirty is True:
|
||||
self.saveFile()
|
||||
else:
|
||||
self.changeSavedir()
|
||||
self.changeSavedirDialog()
|
||||
return
|
||||
|
||||
if not self.mayContinue():
|
||||
@ -1240,8 +1249,8 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
default=DEFAULT_LINE_COLOR)
|
||||
if color:
|
||||
self.lineColor = color
|
||||
# Change the color for all shape lines:
|
||||
Shape.line_color = self.lineColor
|
||||
Shape.line_color = color
|
||||
self.canvas.setDrawingColor(color)
|
||||
self.canvas.update()
|
||||
self.setDirty()
|
||||
|
||||
@ -1292,7 +1301,7 @@ class MainWindow(QMainWindow, WindowMixin):
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if self.labelHist is None:
|
||||
self.lablHist = [line]
|
||||
self.labelHist = [line]
|
||||
else:
|
||||
self.labelHist.append(line)
|
||||
|
||||
|
||||
@ -41,8 +41,9 @@ class Canvas(QWidget):
|
||||
self.current = None
|
||||
self.selectedShape = None # save the selected shape here
|
||||
self.selectedShapeCopy = None
|
||||
self.lineColor = QColor(0, 0, 255)
|
||||
self.line = Shape(line_color=self.lineColor)
|
||||
self.drawingLineColor = QColor(0, 0, 255)
|
||||
self.drawingRectColor = QColor(0, 0, 255)
|
||||
self.line = Shape(line_color=self.drawingLineColor)
|
||||
self.prevPoint = QPointF()
|
||||
self.offsets = QPointF(), QPointF()
|
||||
self.scale = 1.0
|
||||
@ -61,6 +62,10 @@ class Canvas(QWidget):
|
||||
self.setFocusPolicy(Qt.WheelFocus)
|
||||
self.verified = False
|
||||
|
||||
def setDrawingColor(self, qColor):
|
||||
self.drawingLineColor = qColor
|
||||
self.drawingRectColor = qColor
|
||||
|
||||
def enterEvent(self, ev):
|
||||
self.overrideCursor(self._cursor)
|
||||
|
||||
@ -103,7 +108,7 @@ class Canvas(QWidget):
|
||||
if self.drawing():
|
||||
self.overrideCursor(CURSOR_DRAW)
|
||||
if self.current:
|
||||
color = self.lineColor
|
||||
color = self.drawingLineColor
|
||||
if self.outOfPixmap(pos):
|
||||
# Don't allow the user to draw outside the pixmap.
|
||||
# Project the point to the pixmap's edges.
|
||||
@ -414,8 +419,7 @@ class Canvas(QWidget):
|
||||
rightBottom = self.line[1]
|
||||
rectWidth = rightBottom.x() - leftTop.x()
|
||||
rectHeight = rightBottom.y() - leftTop.y()
|
||||
color = QColor(0, 220, 0)
|
||||
p.setPen(color)
|
||||
p.setPen(self.drawingRectColor)
|
||||
brush = QBrush(Qt.BDiagPattern)
|
||||
p.setBrush(brush)
|
||||
p.drawRect(leftTop.x(), leftTop.y(), rectWidth, rectHeight)
|
||||
@ -606,9 +610,15 @@ class Canvas(QWidget):
|
||||
points = [p1+p2 for p1, p2 in zip(self.selectedShape.points, [step]*4)]
|
||||
return True in map(self.outOfPixmap, points)
|
||||
|
||||
def setLastLabel(self, text):
|
||||
def setLastLabel(self, text, line_color = None, fill_color = None):
|
||||
assert text
|
||||
self.shapes[-1].label = text
|
||||
if line_color:
|
||||
self.shapes[-1].line_color = line_color
|
||||
|
||||
if fill_color:
|
||||
self.shapes[-1].fill_color = fill_color
|
||||
|
||||
return self.shapes[-1]
|
||||
|
||||
def undoLastLine(self):
|
||||
|
||||
@ -58,6 +58,35 @@ class LabelFile(object):
|
||||
def toggleVerify(self):
|
||||
self.verified = not self.verified
|
||||
|
||||
''' ttf is disable
|
||||
def load(self, filename):
|
||||
import json
|
||||
with open(filename, 'rb') as f:
|
||||
data = json.load(f)
|
||||
imagePath = data['imagePath']
|
||||
imageData = b64decode(data['imageData'])
|
||||
lineColor = data['lineColor']
|
||||
fillColor = data['fillColor']
|
||||
shapes = ((s['label'], s['points'], s['line_color'], s['fill_color'])\
|
||||
for s in data['shapes'])
|
||||
# Only replace data after everything is loaded.
|
||||
self.shapes = shapes
|
||||
self.imagePath = imagePath
|
||||
self.imageData = imageData
|
||||
self.lineColor = lineColor
|
||||
self.fillColor = fillColor
|
||||
|
||||
def save(self, filename, shapes, imagePath, imageData, lineColor=None, fillColor=None):
|
||||
import json
|
||||
with open(filename, 'wb') as f:
|
||||
json.dump(dict(
|
||||
shapes=shapes,
|
||||
lineColor=lineColor, fillColor=fillColor,
|
||||
imagePath=imagePath,
|
||||
imageData=b64encode(imageData)),
|
||||
f, ensure_ascii=True, indent=2)
|
||||
'''
|
||||
|
||||
@staticmethod
|
||||
def isLabelFile(filename):
|
||||
fileSuffix = os.path.splitext(filename)[1].lower()
|
||||
|
||||
22
libs/lib.py
22
libs/lib.py
@ -1,5 +1,5 @@
|
||||
from math import sqrt
|
||||
|
||||
import hashlib
|
||||
try:
|
||||
from PyQt5.QtGui import *
|
||||
from PyQt5.QtCore import *
|
||||
@ -71,3 +71,23 @@ def distance(p):
|
||||
def fmtShortcut(text):
|
||||
mod, key = text.split('+', 1)
|
||||
return '<b>%s</b>+<b>%s</b>' % (mod, key)
|
||||
|
||||
|
||||
def generateColorByText(text):
|
||||
color_table = []
|
||||
color_table.append(QColor(0, 0, 50))
|
||||
color_table.append(QColor(0, 0, 255))
|
||||
color_table.append(QColor(0, 50, 0))
|
||||
color_table.append(QColor(0, 255, 0))
|
||||
color_table.append(QColor(50, 0, 0))
|
||||
color_table.append(QColor(255, 0, 0))
|
||||
color_table.append(QColor(0, 50, 50))
|
||||
color_table.append(QColor(0, 255, 255))
|
||||
color_table.append(QColor(50, 50, 0))
|
||||
color_table.append(QColor(255, 255, 0))
|
||||
color_table.append(QColor(50, 0, 50))
|
||||
color_table.append(QColor(255, 0, 255))
|
||||
color_table.append(QColor(50, 50, 50))
|
||||
color_table.append(QColor(255, 255, 255))
|
||||
colorInd = int(hashlib.sha1(text.encode('utf-8')).hexdigest(), 16) % len(color_table)
|
||||
return color_table[colorInd]
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user