diff --git a/labelImg.py b/labelImg.py index e84ac323..f2ecd617 100755 --- a/labelImg.py +++ b/labelImg.py @@ -38,7 +38,7 @@ from libs.canvas import Canvas from libs.zoomWidget import ZoomWidget from libs.labelDialog import LabelDialog from libs.colorDialog import ColorDialog -from libs.labelFile import LabelFile, LabelFileError +from libs.labelFile import LabelFile, LabelFileError, LabelFileFormat from libs.toolBar import ToolBar from libs.pascal_voc_io import PascalVocReader from libs.pascal_voc_io import XML_EXT @@ -86,8 +86,7 @@ class MainWindow(QMainWindow, WindowMixin): # Save as Pascal voc xml self.defaultSaveDir = defaultSaveDir - self.usingPascalVocFormat = settings.get(SETTING_USING_VOC_FORMAT, True) - self.usingYoloFormat = not self.usingPascalVocFormat + self.labelFileFormat = settings.get(SETTING_LABEL_FILE_FORMAT, LabelFileFormat.PASCAL_VOC) # For loading all image under a directory self.mImgList = [] @@ -229,9 +228,10 @@ class MainWindow(QMainWindow, WindowMixin): save = action(getStr('save'), self.saveFile, 'Ctrl+S', 'save', getStr('saveDetail'), enabled=False) - save_format = action('&PascalVOC' if self.usingPascalVocFormat else '&YOLO', + isUsingPascalVoc = self.labelFileFormat == LabelFileFormat.PASCAL_VOC + save_format = action('&PascalVOC' if isUsingPascalVoc else '&YOLO', self.change_format, 'Ctrl+', - 'format_voc' if self.usingPascalVocFormat else 'format_yolo', + 'format_voc' if isUsingPascalVoc else 'format_yolo', getStr('changeSaveFormat'), enabled=True) saveAs = action(getStr('saveAs'), self.saveFileAs, @@ -503,20 +503,22 @@ class MainWindow(QMainWindow, WindowMixin): if save_format == FORMAT_PASCALVOC: self.actions.save_format.setText(FORMAT_PASCALVOC) self.actions.save_format.setIcon(newIcon("format_voc")) - self.usingPascalVocFormat = True - self.usingYoloFormat = False + self.labelFileFormat = LabelFileFormat.PASCAL_VOC LabelFile.suffix = XML_EXT elif save_format == FORMAT_YOLO: self.actions.save_format.setText(FORMAT_YOLO) self.actions.save_format.setIcon(newIcon("format_yolo")) - self.usingPascalVocFormat = False - self.usingYoloFormat = True + self.labelFileFormat = LabelFileFormat.YOLO LabelFile.suffix = TXT_EXT def change_format(self): - if self.usingPascalVocFormat: self.set_format(FORMAT_YOLO) - elif self.usingYoloFormat: self.set_format(FORMAT_PASCALVOC) + if self.labelFileFormat == LabelFileFormat.PASCAL_VOC: + self.set_format(FORMAT_YOLO) + elif self.labelFileFormat == LabelFileFormat.YOLO: + self.set_format(FORMAT_PASCALVOC) + else: + raise ValueError('Unknown label file format.') def noShapes(self): return not self.itemsToShapes @@ -819,12 +821,12 @@ class MainWindow(QMainWindow, WindowMixin): shapes = [format_shape(shape) for shape in self.canvas.shapes] # Can add differrent annotation formats here try: - if self.usingPascalVocFormat is True: + if self.labelFileFormat == LabelFileFormat.PASCAL_VOC: if annotationFilePath[-4:].lower() != ".xml": annotationFilePath += XML_EXT self.labelFile.savePascalVocFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.lineColor.getRgb(), self.fillColor.getRgb()) - elif self.usingYoloFormat is True: + elif self.labelFileFormat == LabelFileFormat.YOLO: if annotationFilePath[-4:].lower() != ".txt": annotationFilePath += TXT_EXT self.labelFile.saveYoloFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.labelHist, @@ -1060,7 +1062,6 @@ class MainWindow(QMainWindow, WindowMixin): self.toggleActions(True) # Label xml file and show bound box according to its filename - # if self.usingPascalVocFormat is True: if self.defaultSaveDir is not None: basename = os.path.basename( os.path.splitext(self.filePath)[0]) @@ -1157,7 +1158,7 @@ class MainWindow(QMainWindow, WindowMixin): settings[SETTING_SINGLE_CLASS] = self.singleClassMode.isChecked() settings[SETTING_PAINT_LABEL] = self.displayLabelOption.isChecked() settings[SETTING_DRAW_SQUARE] = self.drawSquaresOption.isChecked() - settings[SETTING_USING_VOC_FORMAT] = self.usingPascalVocFormat + settings[SETTING_LABEL_FILE_FORMAT] = self.labelFileFormat settings.save() def loadRecent(self, filename): @@ -1202,7 +1203,7 @@ class MainWindow(QMainWindow, WindowMixin): path = os.path.dirname(ustr(self.filePath))\ if self.filePath else '.' - if self.usingPascalVocFormat: + if self.labelFileFormat == LabelFileFormat.PASCAL_VOC: filters = "Open Annotation XML file (%s)" % ' '.join(['*.xml']) filename = ustr(QFileDialog.getOpenFileName(self,'%s - Choose a xml file' % __appname__, path, filters)) if filename: @@ -1513,9 +1514,7 @@ def get_main_app(argv=[]): argparser = argparse.ArgumentParser() argparser.add_argument("image_dir", nargs="?") argparser.add_argument("predefined_classes_file", - default=os.path.join( - os.path.dirname(argv[0]), - "data", "predefined_classes.txt"), + default=os.path.join(os.path.dirname(__file__), "data", "predefined_classes.txt"), nargs="?") argparser.add_argument("save_dir", nargs="?") args = argparser.parse_args(argv[1:]) diff --git a/libs/constants.py b/libs/constants.py index 9f44d158..8fa4b275 100644 --- a/libs/constants.py +++ b/libs/constants.py @@ -15,5 +15,5 @@ SETTING_SINGLE_CLASS = 'singleclass' FORMAT_PASCALVOC='PascalVOC' FORMAT_YOLO='YOLO' SETTING_DRAW_SQUARE = 'draw/square' -SETTING_USING_VOC_FORMAT = 'usingVocFormat' +SETTING_LABEL_FILE_FORMAT= 'labelFileFormat' DEFAULT_ENCODING = 'utf-8' diff --git a/libs/labelFile.py b/libs/labelFile.py index 9a3c54ee..ebcca266 100644 --- a/libs/labelFile.py +++ b/libs/labelFile.py @@ -10,10 +10,16 @@ from base64 import b64encode, b64decode from libs.pascal_voc_io import PascalVocWriter from libs.yolo_io import YOLOWriter from libs.pascal_voc_io import XML_EXT +from enum import Enum import os.path import sys +class LabelFileFormat(Enum): + PASCAL_VOC= 1 + YOLO = 2 + + class LabelFileError(Exception): pass