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

YOLO READ

Browse Source
This commit is contained in:
Wang Yinghao 2018-03-01 22:34:57 -06:00
parent 2783071c85
commit 201a54d0f9
2 changed files with 73 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
labelImg.py
View File

@ -38,6 +38,7 @@ from libs.labelFile import LabelFile, LabelFileError
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
from libs.yolo_io import YoloReader
from libs.yolo_io import TXT_EXT from libs.yolo_io import TXT_EXT
from libs.ustr import ustr from libs.ustr import ustr
from libs.version import __version__ from libs.version import __version__
@ -471,14 +472,20 @@ class MainWindow(QMainWindow, WindowMixin):
self.openDirDialog(dirpath=self.filePath) self.openDirDialog(dirpath=self.filePath)
## Support Functions ## ## 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): def change_format(self):
self.usingPascalVocFormat = not self.usingPascalVocFormat if self.usingPascalVocFormat: self.set_format("YOLO")
self.usingYoloFormat = not self.usingYoloFormat if self.usingYoloFormat: self.set_format("PascalVOC")
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")
def noShapes(self): def noShapes(self):
return not self.itemsToShapes return not self.itemsToShapes
@ -968,16 +975,27 @@ 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.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]) + XML_EXT os.path.splitext(self.filePath)[0])
xmlPath = os.path.join(self.defaultSaveDir, basename) 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) self.loadPascalXMLByFilename(xmlPath)
else: elif os.path.isfile(txtPath):
xmlPath = os.path.splitext(filePath)[0] + XML_EXT self.loadYOLOTXTByFilename(txtPath)
if os.path.isfile(xmlPath): else:
self.loadPascalXMLByFilename(xmlPath) 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) self.setWindowTitle(__appname__ + ' ' + filePath)
@ -1340,11 +1358,27 @@ class MainWindow(QMainWindow, WindowMixin):
if os.path.isfile(xmlPath) is False: if os.path.isfile(xmlPath) is False:
return return
self.set_format("PascalVOC")
tVocParseReader = PascalVocReader(xmlPath) tVocParseReader = PascalVocReader(xmlPath)
shapes = tVocParseReader.getShapes() shapes = tVocParseReader.getShapes()
self.loadLabels(shapes) self.loadLabels(shapes)
self.canvas.verified = tVocParseReader.verified 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): def inverted(color):
return QColor(*[255 - v for v in color.getRgb()]) return QColor(*[255 - v for v in color.getRgb()])

41
libs/yolo_io.py
View File

@ -77,28 +77,35 @@ class YOLOWriter:
class YoloReader: class YoloReader:
def __init__(self, filepath, imgSize, classListPath=None): def __init__(self, filepath, image, classListPath=None):
# shapes type: # shapes type:
# [labbel, [(x1,y1), (x2,y2), (x3,y3), (x4,y4)], color, color, difficult] # [labbel, [(x1,y1), (x2,y2), (x3,y3), (x4,y4)], color, color, difficult]
self.shapes = [] self.shapes = []
self.filepath = filepath self.filepath = filepath
if classListPath == None: if classListPath is None:
dir_path = os.path.dirname(os.path.realpath(self.filepath)) dir_path = os.path.dirname(os.path.realpath(self.filepath))
self.classListPath = os.path.join(dir_path, "classes.txt") self.classListPath = os.path.join(dir_path, "classes.txt")
else: else:
self.classListPath = classListPath self.classListPath = classListPath
print (filepath, self.classListPath)
classesFile = open(self.classListPath, 'r') 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.imgSize = imgSize
self.verified = False self.verified = False
try: # try:
self.parseYoloFormat() self.parseYoloFormat()
except: # except:
pass # pass
def getShapes(self): def getShapes(self):
return self.shapes return self.shapes
@ -109,17 +116,17 @@ class YoloReader:
self.shapes.append((label, points, None, None, difficult)) self.shapes.append((label, points, None, None, difficult))
def yoloLine2Shape(self, classIndex, xcen, ycen, w, h): 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) xmin = max(float(xcen) - float(w) / 2, 0)
xmax = max(float(xcen) + float(w) / 2, 1) xmax = min(float(xcen) + float(w) / 2, 1)
ymin = min(float(ycen) - float(h) / 2, 0) ymin = max(float(ycen) - float(h) / 2, 0)
ymax = max(float(ycen) + float(h) / 2, 1) ymax = min(float(ycen) + float(h) / 2, 1)
xmin = int(imgSize[1] * xmin) xmin = int(self.imgSize[1] * xmin)
xmax = int(imgSize[1] * xmax) xmax = int(self.imgSize[1] * xmax)
ymin = int(imgSize[0] * ymin) ymin = int(self.imgSize[0] * ymin)
ymax = int(imgSize[0] * ymax) ymax = int(self.imgSize[0] * ymax)
return label, xmin, ymin, xmax, ymax return label, xmin, ymin, xmax, ymax
@ -127,7 +134,7 @@ class YoloReader:
bndBoxFile = open(self.filepath, 'r') bndBoxFile = open(self.filepath, 'r')
for bndBox in bndBoxFile: for bndBox in bndBoxFile:
classIndex, xcen, ycen, w, h = bndBox.split(' ') 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. # Caveat: difficult flag is discarded when saved as yolo format.
self.addShape(label, xmin, ymin, xmax, ymax, False) self.addShape(label, xmin, ymin, xmax, ymax, False)