5
5
from .image_utils import load_labelmap
6
6
7
7
IMAGE_EXTENSIONS = {".jpg" , ".jpeg" , ".png" , ".bmp" }
8
- ANNOTATION_EXTENSIONS = {".txt" , ".json" , ".xml" }
8
+ ANNOTATION_EXTENSIONS = {".txt" , ".json" , ".xml" , ".csv" }
9
9
LABELMAPS_EXTENSIONS = {".labels" , ".yaml" , ".yml" }
10
10
11
11
@@ -126,18 +126,17 @@ def _filterIndividualAnnotations(image, annotation, format):
126
126
"iscrowd" : 0 ,
127
127
}
128
128
imgReference = imgReferences [0 ]
129
- _annotation = {
130
- "name" : "annotation.coco.json" ,
131
- "parsedType" : "coco" ,
132
- "parsed" : {
129
+ _annotation = {"name" : "annotation.coco.json" }
130
+ _annotation ["rawText" ] = json .dumps (
131
+ {
133
132
"info" : parsed ["info" ],
134
133
"licenses" : parsed ["licenses" ],
135
134
"categories" : parsed ["categories" ],
136
135
"images" : [imgReference ],
137
136
"annotations" : [a for a in parsed ["annotations" ] if a ["image_id" ] == imgReference ["id" ]]
138
137
or [fake_annotation ],
139
- },
140
- }
138
+ }
139
+ )
141
140
return _annotation
142
141
elif format == "createml" :
143
142
imgReferences = [i for i in parsed if i ["image" ] == image ["name" ]]
@@ -147,27 +146,53 @@ def _filterIndividualAnnotations(image, annotation, format):
147
146
imgReference = imgReferences [0 ]
148
147
_annotation = {
149
148
"name" : "annotation.createml.json" ,
150
- "parsedType" : "createml" ,
151
- "parsed" : [imgReference ],
149
+ "rawText" : json .dumps ([imgReference ]),
150
+ }
151
+ return _annotation
152
+ elif format == "csv" :
153
+ imgLines = [ld ["line" ] for ld in parsed ["lines" ] if ld ["file_name" ] == image ["name" ]]
154
+ if imgLines :
155
+ headers = parsed ["headers" ]
156
+ _annotation = {
157
+ "name" : "annotation.csv" ,
158
+ "rawText" : "" .join ([headers ] + imgLines ),
152
159
}
153
160
return _annotation
161
+ else :
162
+ return None
154
163
return None
155
164
156
165
157
166
def _loadAnnotations (folder , annotations ):
158
- valid_extensions = {".json" }
167
+ valid_extensions = {".json" , ".csv" }
159
168
annotations = [a for a in annotations if a ["extension" ] in valid_extensions ]
160
169
for ann in annotations :
161
170
extension = ann ["extension" ]
162
- with open (f"{ folder } { ann ['file' ]} " , "r" ) as f :
163
- parsed = json .load (f )
164
- parsedType = _guessAnnotationFileFormat (parsed , extension )
165
- if parsedType :
166
- ann ["parsed" ] = parsed
167
- ann ["parsedType" ] = parsedType
171
+ if extension == ".json" :
172
+ with open (f"{ folder } { ann ['file' ]} " , "r" ) as f :
173
+ parsed = json .load (f )
174
+ parsedType = _guessAnnotationFileFormat (parsed , extension )
175
+ if parsedType :
176
+ ann ["parsed" ] = parsed
177
+ ann ["parsedType" ] = parsedType
178
+ elif extension == ".csv" :
179
+ ann ["parsedType" ] = "csv"
180
+ ann ["parsed" ] = _parseAnnotationCSV (f"{ folder } { ann ['file' ]} " )
168
181
return annotations
169
182
170
183
184
+ def _parseAnnotationCSV (filename ):
185
+ # TODO: use a proper CSV library?
186
+ with open (filename , "r" ) as f :
187
+ lines = f .readlines ()
188
+ headers = lines [0 ]
189
+ lines = [{"file_name" : ld .split ("," )[0 ].strip (), "line" : ld } for ld in lines [1 :]]
190
+ return {
191
+ "headers" : headers ,
192
+ "lines" : lines ,
193
+ }
194
+
195
+
171
196
def _guessAnnotationFileFormat (parsed , extension ):
172
197
if extension == ".json" :
173
198
if isinstance (parsed , dict ):
0 commit comments