-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathMITConverter.py
122 lines (96 loc) · 4.67 KB
/
MITConverter.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import os
import xml.etree.ElementTree as ET
from VOCConverter import ToVOCConverter
class MITtoVOCConverter(ToVOCConverter):
'''
MIT Street Scenes is in a similar format to Pascal VOC, but the labels are polygons not bounding boxes.
Need to convert the polygons to bounding boxes and add in the additional required info.
'''
def __init__(self,imageFolder,sourceLabelFolder,outputLabelFolder):
'''
imageFolder: folder where all the images are
sourceLabelFolder: path to folder where source data is saved
outputLabelFolder: path where ouptut files are to be saved
'''
super().__init__(imageFolder,sourceLabelFolder,outputLabelFolder)
self.imageFormat = ".JPG"
self.database = "MIT Street Scenes"
def createXMLLabelFile(self,labelFileName):
'''
Take in a MIT Street Scene label file for an image (labelFileName) and save out the Pascal VOC format label file
Read in the existing XML and add to it
'''
self.currentOutFile = os.path.join(self.outputLabelFolder,labelFileName[:labelFileName.rfind("_")] + ".xml")
self.currentImageFile = labelFileName[:labelFileName.rfind("_")] + self.imageFormat
labelXML = self.initializeXMLFile()
sourceTree = ET.parse(os.path.join(self.sourceLabelFolder,labelFileName))
sourceRoot = sourceTree.getroot()
sourceRoot = self.removeNewlines(sourceRoot)
# add the original source data to the new source data, purely to preserve data
labelXML.find("source").extend(sourceRoot.find("source"))
for objET in sourceRoot.iterfind("object"):
objectLabel = self.createXMLLabel(objET)
labelXML.append(objectLabel)
tree = ET.ElementTree(labelXML)
tree.write(self.currentOutFile)
def createXMLLabel(self,objectLabel):
'''
Create a xml style label for the line that is passed in. Return the xml data
'''
# go over every "pt" label in the object and find the min/max values
xmin = self.currentImageShape[0]
xmax = 0
ymin = self.currentImageShape[1]
ymax = 0
try:
for pt in objectLabel.find("polygon").iterfind("pt"):
x = pt.find("x").text.rstrip("\n").lstrip("\n")
x = int(x)
y = pt.find("y").text.rstrip("\n").lstrip("\n")
y = int(y)
xmin = min(xmin,x)
xmax = max(xmax,x)
ymin = min(ymin,y)
ymax = max(ymax,y)
bndbox = ET.SubElement(objectLabel,"bndbox")
ET.SubElement(bndbox,"xmin").text = str(xmin)
ET.SubElement(bndbox,"ymin").text = str(ymin)
ET.SubElement(bndbox,"xmax").text = str(xmax)
ET.SubElement(bndbox,"ymax").text = str(ymax)
except:
print("Error creating polygons for image :{}".format(self.currentImageFile))
return objectLabel
def removeNewlines(self,root):
'''
The MIT street scenes XML has leading and trailing newlines on every entry, we are going to remove them to make everything cleaner
'''
for tag in root.iter():
tag.text = tag.text.rstrip("\n").lstrip("\n")
return root
def convertDataset(self,verbose=False):
'''
Convert the entire dataset
Note that for MIT we split the label name with a "_"
'''
# find all the label files
labelFiles = os.listdir(self.sourceLabelFolder)
numLabels = len(labelFiles)
labelFileNames = [i.split("_")[0] for i in labelFiles]
labelFileNames = list(set(labelFileNames))
assert numLabels == len(labelFileNames), "Repeated label files!"
# find all image files
imageFiles = os.listdir(self.imageFolder)
imageFileNames = [i.split(".")[0] for i in imageFiles]
# verify there is a label for each image and vise-a-versa
imagesWithNoLabel = list(set(imageFileNames) - set(labelFileNames))
labelsWithNoImage = list(set(labelFileNames) - set(imageFileNames))
assert len(imagesWithNoLabel) == 0, "Images with no label file found: {}".format(imagesWithNoLabel)
assert len(labelsWithNoImage) == 0, "Labels with no image found: {}".format(labelsWithNoImage)
# call createXMLLabelFile() on each file
counter = 0
for label in labelFiles:
counter += 1
self.createXMLLabelFile(label)
if verbose and counter%100==0:
print("On image {}/{} {:.1f}% complete".format(counter,numLabels,float(counter)/float(numLabels)*100.))
print("Finished converting {} labels!".format(numLabels))