apoorva/greenhouse
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

use enum for label file formats

Browse Source
This commit is contained in:
tzutalin 2020-08-02 00:14:18 -07:00
parent 196bd46828
commit f387c5932b
3 changed files with 25 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
labelImg.py
View File

@ -38,7 +38,7 @@ from libs.canvas import Canvas
from libs.zoomWidget import ZoomWidget from libs.zoomWidget import ZoomWidget
from libs.labelDialog import LabelDialog from libs.labelDialog import LabelDialog
from libs.colorDialog import ColorDialog 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.toolBar import ToolBar
from libs.pascal_voc_io import PascalVocReader from libs.pascal_voc_io import PascalVocReader
from libs.pascal_voc_io import XML_EXT from libs.pascal_voc_io import XML_EXT
@ -86,8 +86,7 @@ class MainWindow(QMainWindow, WindowMixin):
# Save as Pascal voc xml # Save as Pascal voc xml
self.defaultSaveDir = defaultSaveDir self.defaultSaveDir = defaultSaveDir
self.usingPascalVocFormat = settings.get(SETTING_USING_VOC_FORMAT, True) self.labelFileFormat = settings.get(SETTING_LABEL_FILE_FORMAT, LabelFileFormat.PASCAL_VOC)
self.usingYoloFormat = not self.usingPascalVocFormat
# For loading all image under a directory # For loading all image under a directory
self.mImgList = [] self.mImgList = []
@ -229,9 +228,10 @@ class MainWindow(QMainWindow, WindowMixin):
save = action(getStr('save'), self.saveFile, save = action(getStr('save'), self.saveFile,
'Ctrl+S', 'save', getStr('saveDetail'), enabled=False) '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+', self.change_format, 'Ctrl+',
'format_voc' if self.usingPascalVocFormat else 'format_yolo', 'format_voc' if isUsingPascalVoc else 'format_yolo',
getStr('changeSaveFormat'), enabled=True) getStr('changeSaveFormat'), enabled=True)
saveAs = action(getStr('saveAs'), self.saveFileAs, saveAs = action(getStr('saveAs'), self.saveFileAs,
@ -503,20 +503,22 @@ class MainWindow(QMainWindow, WindowMixin):
if save_format == FORMAT_PASCALVOC: if save_format == FORMAT_PASCALVOC:
self.actions.save_format.setText(FORMAT_PASCALVOC) self.actions.save_format.setText(FORMAT_PASCALVOC)
self.actions.save_format.setIcon(newIcon("format_voc")) self.actions.save_format.setIcon(newIcon("format_voc"))
self.usingPascalVocFormat = True self.labelFileFormat = LabelFileFormat.PASCAL_VOC
self.usingYoloFormat = False
LabelFile.suffix = XML_EXT LabelFile.suffix = XML_EXT
elif save_format == FORMAT_YOLO: elif save_format == FORMAT_YOLO:
self.actions.save_format.setText(FORMAT_YOLO) self.actions.save_format.setText(FORMAT_YOLO)
self.actions.save_format.setIcon(newIcon("format_yolo")) self.actions.save_format.setIcon(newIcon("format_yolo"))
self.usingPascalVocFormat = False self.labelFileFormat = LabelFileFormat.YOLO
self.usingYoloFormat = True
LabelFile.suffix = TXT_EXT LabelFile.suffix = TXT_EXT
def change_format(self): def change_format(self):
if self.usingPascalVocFormat: self.set_format(FORMAT_YOLO) if self.labelFileFormat == LabelFileFormat.PASCAL_VOC:
elif self.usingYoloFormat: self.set_format(FORMAT_PASCALVOC) 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): def noShapes(self):
return not self.itemsToShapes return not self.itemsToShapes
@ -819,12 +821,12 @@ class MainWindow(QMainWindow, WindowMixin):
shapes = [format_shape(shape) for shape in self.canvas.shapes] shapes = [format_shape(shape) for shape in self.canvas.shapes]
# Can add differrent annotation formats here # Can add differrent annotation formats here
try: try:
if self.usingPascalVocFormat is True: if self.labelFileFormat == LabelFileFormat.PASCAL_VOC:
if annotationFilePath[-4:].lower() != ".xml": if annotationFilePath[-4:].lower() != ".xml":
annotationFilePath += XML_EXT annotationFilePath += XML_EXT
self.labelFile.savePascalVocFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.labelFile.savePascalVocFormat(annotationFilePath, shapes, self.filePath, self.imageData,
self.lineColor.getRgb(), self.fillColor.getRgb()) self.lineColor.getRgb(), self.fillColor.getRgb())
elif self.usingYoloFormat is True: elif self.labelFileFormat == LabelFileFormat.YOLO:
if annotationFilePath[-4:].lower() != ".txt": if annotationFilePath[-4:].lower() != ".txt":
annotationFilePath += TXT_EXT annotationFilePath += TXT_EXT
self.labelFile.saveYoloFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.labelHist, self.labelFile.saveYoloFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.labelHist,
@ -1060,7 +1062,6 @@ class MainWindow(QMainWindow, WindowMixin):
self.toggleActions(True) self.toggleActions(True)
# Label xml file and show bound box according to its filename # Label xml file and show bound box according to its filename
# if self.usingPascalVocFormat is True:
if self.defaultSaveDir is not None: if self.defaultSaveDir is not None:
basename = os.path.basename( basename = os.path.basename(
os.path.splitext(self.filePath)[0]) os.path.splitext(self.filePath)[0])
@ -1157,7 +1158,7 @@ class MainWindow(QMainWindow, WindowMixin):
settings[SETTING_SINGLE_CLASS] = self.singleClassMode.isChecked() settings[SETTING_SINGLE_CLASS] = self.singleClassMode.isChecked()
settings[SETTING_PAINT_LABEL] = self.displayLabelOption.isChecked() settings[SETTING_PAINT_LABEL] = self.displayLabelOption.isChecked()
settings[SETTING_DRAW_SQUARE] = self.drawSquaresOption.isChecked() settings[SETTING_DRAW_SQUARE] = self.drawSquaresOption.isChecked()
settings[SETTING_USING_VOC_FORMAT] = self.usingPascalVocFormat settings[SETTING_LABEL_FILE_FORMAT] = self.labelFileFormat
settings.save() settings.save()
def loadRecent(self, filename): def loadRecent(self, filename):
@ -1202,7 +1203,7 @@ class MainWindow(QMainWindow, WindowMixin):
path = os.path.dirname(ustr(self.filePath))\ path = os.path.dirname(ustr(self.filePath))\
if self.filePath else '.' if self.filePath else '.'
if self.usingPascalVocFormat: if self.labelFileFormat == LabelFileFormat.PASCAL_VOC:
filters = "Open Annotation XML file (%s)" % ' '.join(['*.xml']) filters = "Open Annotation XML file (%s)" % ' '.join(['*.xml'])
filename = ustr(QFileDialog.getOpenFileName(self,'%s - Choose a xml file' % __appname__, path, filters)) filename = ustr(QFileDialog.getOpenFileName(self,'%s - Choose a xml file' % __appname__, path, filters))
if filename: if filename:
@ -1513,9 +1514,7 @@ def get_main_app(argv=[]):
argparser = argparse.ArgumentParser() argparser = argparse.ArgumentParser()
argparser.add_argument("image_dir", nargs="?") argparser.add_argument("image_dir", nargs="?")
argparser.add_argument("predefined_classes_file", argparser.add_argument("predefined_classes_file",
default=os.path.join( default=os.path.join(os.path.dirname(__file__), "data", "predefined_classes.txt"),
os.path.dirname(argv[0]),
"data", "predefined_classes.txt"),
nargs="?") nargs="?")
argparser.add_argument("save_dir", nargs="?") argparser.add_argument("save_dir", nargs="?")
args = argparser.parse_args(argv[1:]) args = argparser.parse_args(argv[1:])

2
libs/constants.py
View File

@ -15,5 +15,5 @@ SETTING_SINGLE_CLASS = 'singleclass'
FORMAT_PASCALVOC='PascalVOC' FORMAT_PASCALVOC='PascalVOC'
FORMAT_YOLO='YOLO' FORMAT_YOLO='YOLO'
SETTING_DRAW_SQUARE = 'draw/square' SETTING_DRAW_SQUARE = 'draw/square'
SETTING_USING_VOC_FORMAT = 'usingVocFormat' SETTING_LABEL_FILE_FORMAT= 'labelFileFormat'
DEFAULT_ENCODING = 'utf-8' DEFAULT_ENCODING = 'utf-8'

6
libs/labelFile.py
View File

@ -10,10 +10,16 @@ from base64 import b64encode, b64decode
from libs.pascal_voc_io import PascalVocWriter from libs.pascal_voc_io import PascalVocWriter
from libs.yolo_io import YOLOWriter from libs.yolo_io import YOLOWriter
from libs.pascal_voc_io import XML_EXT from libs.pascal_voc_io import XML_EXT
from enum import Enum
import os.path import os.path
import sys import sys
class LabelFileFormat(Enum):
PASCAL_VOC= 1
YOLO = 2
class LabelFileError(Exception): class LabelFileError(Exception):
pass pass