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

YOLO Write

Browse Source
This commit is contained in:
Wang Yinghao 2018-03-01 21:38:22 -06:00
parent dec4885e29
commit 2783071c85
4 changed files with 188 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
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 TXT_EXT
from libs.ustr import ustr
from libs.version import __version__
@ -97,6 +98,8 @@ class MainWindow(QMainWindow, WindowMixin):
# Save as Pascal voc xml
self.defaultSaveDir = None
self.usingPascalVocFormat = True
self.usingYoloFormat = False
# For loading all image under a directory
self.mImgList = []
self.dirname = None
@ -232,6 +235,9 @@ class MainWindow(QMainWindow, WindowMixin):
save = action('&Save', self.saveFile,
'Ctrl+S', 'save', u'Save labels to file', enabled=False)
save_format = action('&Format', self.change_format,
'Ctrl+', 'format', u'Change save format', enabled=True)
saveAs = action('&Save As', self.saveFileAs,
'Ctrl+Shift+S', 'save-as', u'Save labels to a different file', enabled=False)
@ -324,7 +330,7 @@ class MainWindow(QMainWindow, WindowMixin):
self.popLabelListMenu)
# Store actions for further handling.
self.actions = struct(save=save, saveAs=saveAs, open=open, close=close, resetAll = resetAll,
self.actions = struct(save=save, save_format=save_format, saveAs=saveAs, open=open, close=close, resetAll = resetAll,
lineColor=color1, create=create, delete=delete, edit=edit, copy=copy,
createMode=createMode, editMode=editMode, advancedMode=advancedMode,
shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor,
@ -363,7 +369,7 @@ class MainWindow(QMainWindow, WindowMixin):
self.lastLabel = None
addActions(self.menus.file,
(open, opendir, changeSavedir, openAnnotation, self.menus.recentFiles, save, saveAs, close, resetAll, quit))
(open, opendir, changeSavedir, openAnnotation, self.menus.recentFiles, save, save_format, saveAs, close, resetAll, quit))
addActions(self.menus.help, (help, showInfo))
addActions(self.menus.view, (
self.autoSaving,
@ -383,11 +389,11 @@ class MainWindow(QMainWindow, WindowMixin):
self.tools = self.toolbar('Tools')
self.actions.beginner = (
open, opendir, changeSavedir, openNextImg, openPrevImg, verify, save, None, create, copy, delete, None,
open, opendir, changeSavedir, openNextImg, openPrevImg, verify, save, save_format, None, create, copy, delete, None,
zoomIn, zoom, zoomOut, fitWindow, fitWidth)
self.actions.advanced = (
open, opendir, changeSavedir, openNextImg, openPrevImg, save, None,
open, opendir, changeSavedir, openNextImg, openPrevImg, save, save_format, None,
createMode, editMode, None,
hideAll, showAll)
@ -466,6 +472,14 @@ class MainWindow(QMainWindow, WindowMixin):
## Support Functions ##
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")
def noShapes(self):
return not self.itemsToShapes
@ -733,9 +747,15 @@ class MainWindow(QMainWindow, WindowMixin):
# Can add differrent annotation formats here
try:
if self.usingPascalVocFormat is True:
annotationFilePath += XML_EXT
print ('Img: ' + self.filePath + ' -> Its xml: ' + annotationFilePath)
self.labelFile.savePascalVocFormat(annotationFilePath, shapes, self.filePath, self.imageData,
self.lineColor.getRgb(), self.fillColor.getRgb())
elif self.usingYoloFormat is True:
annotationFilePath += TXT_EXT
print ('Img: ' + self.filePath + ' -> Its txt: ' + annotationFilePath)
self.labelFile.saveYoloFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.labelHist,
self.lineColor.getRgb(), self.fillColor.getRgb())
else:
self.labelFile.save(annotationFilePath, shapes, self.filePath, self.imageData,
self.lineColor.getRgb(), self.fillColor.getRgb())
@ -1197,13 +1217,13 @@ class MainWindow(QMainWindow, WindowMixin):
if self.defaultSaveDir is not None and len(ustr(self.defaultSaveDir)):
if self.filePath:
imgFileName = os.path.basename(self.filePath)
savedFileName = os.path.splitext(imgFileName)[0] + XML_EXT
savedFileName = os.path.splitext(imgFileName)[0]
savedPath = os.path.join(ustr(self.defaultSaveDir), savedFileName)
self._saveFile(savedPath)
else:
imgFileDir = os.path.dirname(self.filePath)
imgFileName = os.path.basename(self.filePath)
savedFileName = os.path.splitext(imgFileName)[0] + XML_EXT
savedFileName = os.path.splitext(imgFileName)[0]
savedPath = os.path.join(imgFileDir, savedFileName)
self._saveFile(savedPath if self.labelFile
else self.saveFileDialog())

28
libs/labelFile.py
View File

@ -8,6 +8,7 @@ except ImportError:
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
import os.path
import sys
@ -55,6 +56,33 @@ class LabelFile(object):
writer.save(targetFile=filename)
return
def saveYoloFormat(self, filename, shapes, imagePath, imageData, classList,
lineColor=None, fillColor=None, databaseSrc=None):
imgFolderPath = os.path.dirname(imagePath)
imgFolderName = os.path.split(imgFolderPath)[-1]
imgFileName = os.path.basename(imagePath)
#imgFileNameWithoutExt = os.path.splitext(imgFileName)[0]
# Read from file path because self.imageData might be empty if saving to
# Pascal format
image = QImage()
image.load(imagePath)
imageShape = [image.height(), image.width(),
1 if image.isGrayscale() else 3]
writer = YOLOWriter(imgFolderName, imgFileName,
imageShape, localImgPath=imagePath)
writer.verified = self.verified
for shape in shapes:
points = shape['points']
label = shape['label']
# Add Chris
difficult = int(shape['difficult'])
bndbox = LabelFile.convertPoints2BndBox(points)
writer.addBndBox(bndbox[0], bndbox[1], bndbox[2], bndbox[3], label, difficult)
writer.save(targetFile=filename, classList=classList)
return
def toggleVerify(self):
self.verified = not self.verified

133
libs/yolo_io.py Normal file
View File

@ -0,0 +1,133 @@
#!/usr/bin/env python
# -*- coding: utf8 -*-
import sys
import os
from xml.etree import ElementTree
from xml.etree.ElementTree import Element, SubElement
from lxml import etree
import codecs
TXT_EXT = '.txt'
ENCODE_METHOD = 'utf-8'
class YOLOWriter:
def __init__(self, foldername, filename, imgSize, databaseSrc='Unknown', localImgPath=None):
self.foldername = foldername
self.filename = filename
self.databaseSrc = databaseSrc
self.imgSize = imgSize
self.boxlist = []
self.localImgPath = localImgPath
self.verified = False
def addBndBox(self, xmin, ymin, xmax, ymax, name, difficult):
bndbox = {'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax}
bndbox['name'] = name
bndbox['difficult'] = difficult
self.boxlist.append(bndbox)
def BndBox2YoloLine(self, box, classList=[]):
xmin = box['xmin']
xmax = box['xmax']
ymin = box['ymin']
ymax = box['ymax']
xcen = (xmin + xmax) / 2 / self.imgSize[1]
ycen = (ymin + ymax) / 2 / self.imgSize[0]
w = (xmax - xmin) / self.imgSize[1]
h = (ymax - ymin) / self.imgSize[0]
classIndex = classList.index(box['name'])
return classIndex, xcen, ycen, w, h
def save(self, classList=[], targetFile=None):
out_file = None #Update yolo .txt
out_class_file = None #Update class list .txt
if targetFile is None:
out_file = open(
self.filename + TXT_EXT, 'w', encoding=ENCODE_METHOD)
classesFile = os.path.join(os.path.dirname(os.path.abspath(self.filename)), "classes.txt")
out_class_file = open(classesFile, 'w')
else:
out_file = codecs.open(targetFile, 'w', encoding=ENCODE_METHOD)
classesFile = os.path.join(os.path.dirname(os.path.abspath(targetFile)), "classes.txt")
out_class_file = open(classesFile, 'w')
for box in self.boxlist:
classIndex, xcen, ycen, w, h = self.BndBox2YoloLine(box, classList)
print (classIndex, xcen, ycen, w, h)
out_file.write("%d %.6f %.6f %.6f %.6f\n" % (classIndex, xcen, ycen, w, h))
print (classList)
print (out_class_file)
for c in classList:
out_class_file.write(c+'\n')
out_class_file.close()
out_file.close()
class YoloReader:
def __init__(self, filepath, imgSize, classListPath=None):
# shapes type:
# [labbel, [(x1,y1), (x2,y2), (x3,y3), (x4,y4)], color, color, difficult]
self.shapes = []
self.filepath = filepath
if classListPath == None:
dir_path = os.path.dirname(os.path.realpath(self.filepath))
self.classListPath = os.path.join(dir_path, "classes.txt")
else:
self.classListPath = classListPath
classesFile = open(self.classListPath, 'r')
self.classes = classesFile.read().split('\n')
self.imgSize = imgSize
self.verified = False
try:
self.parseYoloFormat()
except:
pass
def getShapes(self):
return self.shapes
def addShape(self, label, xmin, ymin, xmax, ymax, difficult):
points = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
self.shapes.append((label, points, None, None, difficult))
def yoloLine2Shape(self, classIndex, xcen, ycen, w, h):
label = self.classes[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 = int(imgSize[1] * xmin)
xmax = int(imgSize[1] * xmax)
ymin = int(imgSize[0] * ymin)
ymax = int(imgSize[0] * ymax)
return label, xmin, ymin, xmax, ymax
def parseYoloFormat(self):
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)
# Caveat: difficult flag is discarded when saved as yolo format.
self.addShape(label, xmin, ymin, xmax, ymax, False)