From fbd10c1a34ae30e0bbc01f72343bbfb8ca3c3ddf Mon Sep 17 00:00:00 2001 From: Kentaro Wada Date: Mon, 11 Jun 2018 08:07:12 +0900 Subject: [PATCH] Add image level flag annotation --- labelme/app.py | 52 ++++++++++++++++++++++-------- labelme/config/default_config.yaml | 1 + labelme/labelFile.py | 17 ++++++++-- 3 files changed, 55 insertions(+), 15 deletions(-) diff --git a/labelme/app.py b/labelme/app.py index 81459b734..3ddd0a416 100644 --- a/labelme/app.py +++ b/labelme/app.py @@ -165,6 +165,15 @@ def __init__(self, config=None, filename=None, output=None): self.labelListContainer = QtWidgets.QWidget() self.labelListContainer.setLayout(listLayout) + self.flag_dock = self.flag_widget = None + self.flag_dock = QtWidgets.QDockWidget('Flags', self) + self.flag_dock.setObjectName('Flags') + self.flag_widget = QtWidgets.QListWidget() + if config['flags']: + self.loadFlags({k: False for k in config['flags']}) + self.flag_dock.setWidget(self.flag_widget) + self.flag_widget.itemChanged.connect(self.setDirty) + self.uniqLabelList = EscapableQListWidget() self.uniqLabelList.setToolTip( "Select label to start annotating for it. " @@ -214,6 +223,7 @@ def __init__(self, config=None, filename=None, output=None): self.setCentralWidget(scrollArea) + self.addDockWidget(Qt.RightDockWidgetArea, self.flag_dock) self.addDockWidget(Qt.RightDockWidgetArea, self.labelsdock) self.addDockWidget(Qt.RightDockWidgetArea, self.dock) self.addDockWidget(Qt.RightDockWidgetArea, self.filedock) @@ -686,6 +696,14 @@ def loadLabels(self, shapes): shape.fill_color = QtGui.QColor(*fill_color) self.loadShapes(s) + def loadFlags(self, flags): + self.flag_widget.clear() + for key, flag in flags.items(): + item = QtWidgets.QListWidgetItem(key) + item.setFlags(item.flags() | Qt.ItemIsUserCheckable) + item.setCheckState(Qt.Checked if flag else Qt.Unchecked) + self.flag_widget.addItem(item) + def saveLabels(self, filename): lf = LabelFile() @@ -698,13 +716,26 @@ def format_shape(s): points=[(p.x(), p.y()) for p in s.points]) shapes = [format_shape(shape) for shape in self.labelList.shapes] + flags = {} + for i in range(self.flag_widget.count()): + item = self.flag_widget.item(i) + key = item.text() + flag = item.checkState() == Qt.Checked + flags[key] = flag try: imagePath = os.path.relpath( self.imagePath, os.path.dirname(filename)) imageData = self.imageData if self._config['store_data'] else None - lf.save(filename, shapes, imagePath, imageData, - self.lineColor.getRgb(), self.fillColor.getRgb(), - self.otherData) + lf.save( + filename=filename, + shapes=shapes, + imagePath=imagePath, + imageData=imageData, + lineColor=self.lineColor.getRgb(), + fillColor=self.fillColor.getRgb(), + otherData=self.otherData, + flags=flags, + ) self.labelFile = lf # disable allows next and previous image to proceed # self.filename = filename @@ -880,6 +911,8 @@ def loadFile(self, filename=None): self.canvas.loadPixmap(QtGui.QPixmap.fromImage(image)) if self.labelFile: self.loadLabels(self.labelFile.shapes) + if self.labelFile.flags is not None: + self.loadFlags(self.labelFile.flags) self.setClean() self.canvas.setEnabled(True) self.adjustScale(initial=True) @@ -1243,16 +1276,9 @@ def main(): output = config_from_args.pop('output') config_file = config_from_args.pop('config_file') # drop the default config - if not config_from_args['auto_save']: - config_from_args.pop('auto_save') - if config_from_args['store_data']: - config_from_args.pop('store_data') - if not config_from_args['labels']: - config_from_args.pop('labels') - if not config_from_args['sort_labels']: - config_from_args.pop('sort_labels') - if not config_from_args['validate_label']: - config_from_args.pop('validate_label') + for k, v in list(config_from_args.items()): + if v is None: + config_from_args.pop(k) config = get_config(config_from_args, config_file) app = QtWidgets.QApplication(sys.argv) diff --git a/labelme/config/default_config.yaml b/labelme/config/default_config.yaml index a62a1c623..b97a7eefe 100644 --- a/labelme/config/default_config.yaml +++ b/labelme/config/default_config.yaml @@ -1,6 +1,7 @@ auto_save: false store_data: true +flags: null labels: null sort_labels: true validate_label: null diff --git a/labelme/labelFile.py b/labelme/labelFile.py index e7b7ee891..b4b331458 100644 --- a/labelme/labelFile.py +++ b/labelme/labelFile.py @@ -24,7 +24,14 @@ def __init__(self, filename=None): self.filename = filename def load(self, filename): - keys = ['imageData', 'imagePath', 'lineColor', 'fillColor', 'shapes'] + keys = [ + 'imageData', + 'imagePath', + 'lineColor', + 'fillColor', + 'shapes', # polygonal annotations + 'flags', # image level flags + ] try: with open(filename, 'rb' if PY2 else 'r') as f: data = json.load(f) @@ -36,6 +43,7 @@ def load(self, filename): data['imagePath']) with open(imagePath, 'rb') as f: imageData = f.read() + flags = data.get('flags') imagePath = data['imagePath'] lineColor = data['lineColor'] fillColor = data['fillColor'] @@ -52,6 +60,7 @@ def load(self, filename): otherData[key] = value # Only replace data after everything is loaded. + self.flags = flags self.shapes = shapes self.imagePath = imagePath self.imageData = imageData @@ -61,12 +70,16 @@ def load(self, filename): self.otherData = otherData def save(self, filename, shapes, imagePath, imageData=None, - lineColor=None, fillColor=None, otherData=None): + lineColor=None, fillColor=None, otherData=None, + flags=None): if imageData is not None: imageData = base64.b64encode(imageData).decode('utf-8') if otherData is None: otherData = {} + if flags is None: + flags = [] data = dict( + flags=flags, shapes=shapes, lineColor=lineColor, fillColor=fillColor,