6
6
from xml .etree .ElementTree import Element , SubElement
7
7
from lxml import etree
8
8
import codecs
9
+ from libs .constants import DEFAULT_ENCODING
9
10
10
11
TXT_EXT = '.txt'
11
- ENCODE_METHOD = 'utf-8'
12
+ ENCODE_METHOD = DEFAULT_ENCODING
12
13
13
14
class YOLOWriter :
14
15
@@ -39,7 +40,12 @@ def BndBox2YoloLine(self, box, classList=[]):
39
40
w = float ((xmax - xmin )) / self .imgSize [1 ]
40
41
h = float ((ymax - ymin )) / self .imgSize [0 ]
41
42
42
- classIndex = classList .index (box ['name' ])
43
+ # PR387
44
+ boxName = box ['name' ]
45
+ if boxName not in classList :
46
+ classList .append (boxName )
47
+
48
+ classIndex = classList .index (boxName )
43
49
44
50
return classIndex , xcen , ycen , w , h
45
51
@@ -62,11 +68,11 @@ def save(self, classList=[], targetFile=None):
62
68
63
69
for box in self .boxlist :
64
70
classIndex , xcen , ycen , w , h = self .BndBox2YoloLine (box , classList )
65
- print (classIndex , xcen , ycen , w , h )
71
+ # print (classIndex, xcen, ycen, w, h)
66
72
out_file .write ("%d %.6f %.6f %.6f %.6f\n " % (classIndex , xcen , ycen , w , h ))
67
73
68
- print (classList )
69
- print (out_class_file )
74
+ # print (classList)
75
+ # print (out_class_file)
70
76
for c in classList :
71
77
out_class_file .write (c + '\n ' )
72
78
@@ -89,12 +95,12 @@ def __init__(self, filepath, image, classListPath=None):
89
95
else :
90
96
self .classListPath = classListPath
91
97
92
- print (filepath , self .classListPath )
98
+ # print (filepath, self.classListPath)
93
99
94
100
classesFile = open (self .classListPath , 'r' )
95
101
self .classes = classesFile .read ().strip ('\n ' ).split ('\n ' )
96
102
97
- print (self .classes )
103
+ # print (self.classes)
98
104
99
105
imgSize = [image .height (), image .width (),
100
106
1 if image .isGrayscale () else 3 ]
0 commit comments