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

41
libs/yolo_io.py
View File

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