From 201a54d0f9328c8f9e835af1ec646e427881a0e9 Mon Sep 17 00:00:00 2001 From: Wang Yinghao Date: Thu, 1 Mar 2018 22:34:57 -0600 Subject: [PATCH] YOLO READ --- labelImg.py | 64 +++++++++++++++++++++++++++++++++++++------------ libs/yolo_io.py | 41 ++++++++++++++++++------------- 2 files changed, 73 insertions(+), 32 deletions(-) diff --git a/labelImg.py b/labelImg.py index 973c7686..1993482e 100755 --- a/labelImg.py +++ b/labelImg.py @@ -38,6 +38,7 @@ from libs.labelFile import LabelFile, LabelFileError from libs.toolBar import ToolBar from libs.pascal_voc_io import PascalVocReader from libs.pascal_voc_io import XML_EXT +from libs.yolo_io import YoloReader from libs.yolo_io import TXT_EXT from libs.ustr import ustr from libs.version import __version__ @@ -471,14 +472,20 @@ class MainWindow(QMainWindow, WindowMixin): self.openDirDialog(dirpath=self.filePath) ## Support Functions ## + def set_format(self, save_format): + if save_format == 'PascalVOC': + self.actions.save_format.setText("PascalVOC") + self.usingPascalVocFormat = True + self.usingYoloFormat = False + + elif save_format == 'YOLO': + if self.usingYoloFormat: self.actions.save_format.setText("YOLO") + self.usingPascalVocFormat = False + self.usingYoloFormat = True def change_format(self): - self.usingPascalVocFormat = not self.usingPascalVocFormat - self.usingYoloFormat = not self.usingYoloFormat - print ("changing_format") - print(self.actions.save_format) - if self.usingPascalVocFormat: self.actions.save_format.setText("PascalVOC") - if self.usingYoloFormat: self.actions.save_format.setText("YOLO") + if self.usingPascalVocFormat: self.set_format("YOLO") + if self.usingYoloFormat: self.set_format("PascalVOC") def noShapes(self): return not self.itemsToShapes @@ -968,16 +975,27 @@ 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]) + XML_EXT - xmlPath = os.path.join(self.defaultSaveDir, basename) + # if self.usingPascalVocFormat is True: + if self.defaultSaveDir is not None: + basename = os.path.basename( + os.path.splitext(self.filePath)[0]) + xmlPath = os.path.join(self.defaultSaveDir, basename + XML_EXT) + txtPath = os.path.join(self.defaultSaveDir, basename + TXT_EXT) + + """Annotation file priority: + PascalXML > YOLO + """ + if os.path.isfile(xmlPath): self.loadPascalXMLByFilename(xmlPath) - else: - xmlPath = os.path.splitext(filePath)[0] + XML_EXT - if os.path.isfile(xmlPath): - self.loadPascalXMLByFilename(xmlPath) + elif os.path.isfile(txtPath): + self.loadYOLOTXTByFilename(txtPath) + else: + xmlPath = os.path.splitext(filePath)[0] + XML_EXT + txtPath = os.path.splitext(filePath)[0] + TXT_EXT + if os.path.isfile(xmlPath): + self.loadPascalXMLByFilename(xmlPath) + elif os.path.isfile(txtPath): + self.loadYOLOTXTByFilename(txtPath) self.setWindowTitle(__appname__ + ' ' + filePath) @@ -1340,11 +1358,27 @@ class MainWindow(QMainWindow, WindowMixin): if os.path.isfile(xmlPath) is False: return + self.set_format("PascalVOC") + tVocParseReader = PascalVocReader(xmlPath) shapes = tVocParseReader.getShapes() self.loadLabels(shapes) self.canvas.verified = tVocParseReader.verified + def loadYOLOTXTByFilename(self, txtPath): + if self.filePath is None: + return + if os.path.isfile(txtPath) is False: + return + + self.set_format("YOLO") + tYoloParseReader = YoloReader(txtPath, self.image) + shapes = tYoloParseReader.getShapes() + print (shapes) + self.loadLabels(shapes) + self.canvas.verified = tYoloParseReader.verified + + def inverted(color): return QColor(*[255 - v for v in color.getRgb()]) diff --git a/libs/yolo_io.py b/libs/yolo_io.py index 36e75927..f585af08 100644 --- a/libs/yolo_io.py +++ b/libs/yolo_io.py @@ -77,28 +77,35 @@ class YOLOWriter: class YoloReader: - def __init__(self, filepath, imgSize, classListPath=None): + def __init__(self, filepath, image, classListPath=None): # shapes type: # [labbel, [(x1,y1), (x2,y2), (x3,y3), (x4,y4)], color, color, difficult] self.shapes = [] self.filepath = filepath - if classListPath == None: + if classListPath is None: dir_path = os.path.dirname(os.path.realpath(self.filepath)) self.classListPath = os.path.join(dir_path, "classes.txt") else: self.classListPath = classListPath + print (filepath, self.classListPath) + classesFile = open(self.classListPath, 'r') - self.classes = classesFile.read().split('\n') + self.classes = classesFile.read().strip('\n').split('\n') + + print (self.classes) + + imgSize = [image.height(), image.width(), + 1 if image.isGrayscale() else 3] self.imgSize = imgSize self.verified = False - try: - self.parseYoloFormat() - except: - pass + # try: + self.parseYoloFormat() + # except: + # pass def getShapes(self): return self.shapes @@ -109,17 +116,17 @@ class YoloReader: self.shapes.append((label, points, None, None, difficult)) def yoloLine2Shape(self, classIndex, xcen, ycen, w, h): - label = self.classes[classIndex] + label = self.classes[int(classIndex)] - xmin = min(float(xcen) - float(w) / 2, 0) - xmax = max(float(xcen) + float(w) / 2, 1) - ymin = min(float(ycen) - float(h) / 2, 0) - ymax = max(float(ycen) + float(h) / 2, 1) + xmin = max(float(xcen) - float(w) / 2, 0) + xmax = min(float(xcen) + float(w) / 2, 1) + ymin = max(float(ycen) - float(h) / 2, 0) + ymax = min(float(ycen) + float(h) / 2, 1) - xmin = int(imgSize[1] * xmin) - xmax = int(imgSize[1] * xmax) - ymin = int(imgSize[0] * ymin) - ymax = int(imgSize[0] * ymax) + xmin = int(self.imgSize[1] * xmin) + xmax = int(self.imgSize[1] * xmax) + ymin = int(self.imgSize[0] * ymin) + ymax = int(self.imgSize[0] * ymax) return label, xmin, ymin, xmax, ymax @@ -127,7 +134,7 @@ class YoloReader: bndBoxFile = open(self.filepath, 'r') for bndBox in bndBoxFile: classIndex, xcen, ycen, w, h = bndBox.split(' ') - label, xmin, ymin, xmax, ymax = yoloLine2Shape(classIndex, xcen, ycen, w, h) + label, xmin, ymin, xmax, ymax = self.yoloLine2Shape(classIndex, xcen, ycen, w, h) # Caveat: difficult flag is discarded when saved as yolo format. self.addShape(label, xmin, ymin, xmax, ymax, False)