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.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:])

2
libs/constants.py
View File

@ -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'

6
libs/labelFile.py
View File

@ -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