add reading/writing of ids

This commit is contained in:
Martin Baeuml
2011-01-31 12:34:58 +01:00
parent 7d8532ea76
commit 28b698196b
+29 -2
View File
@@ -1,3 +1,6 @@
import os
import glob
class AnnotationContainer:
def __init__(self):
self.clear()
@@ -19,10 +22,16 @@ class AnnotationContainer:
try:
f = open(filename)
basedir = ""
for line in f:
s = line.split()
image_fname = s[0]
if basedir == "":
basedir = os.path.dirname(image_fname)
else:
assert basedir == os.path.dirname(image_fname)
num_rects = int(s[1])
assert num_rects == 0 or num_rects == 1
if image_fname in filename_helper:
rects = filename_helper[image_fname]['annotations']
else:
@@ -36,6 +45,8 @@ class AnnotationContainer:
'width': s[i+4],
'height': s[i+5]
})
if len(s) > i+6:
rects[-1]['id'] = s[i+6]
if image_fname not in filename_helper:
annotation = {
@@ -46,8 +57,22 @@ class AnnotationContainer:
self.annotations_.append(annotation)
filename_helper[image_fname] = annotation
except Exception as e:
self.annotations = []
self.annotations_ = []
raise e
basedir1 = basedir
basedir = os.path.join(os.path.dirname(filename), basedir)
images = glob.glob(os.path.join(basedir, "*.png"))
for fname in images:
fname = os.path.join(basedir1, os.path.basename(fname))
if fname not in filename_helper:
print fname, "not in filename_helper"
annotation = {
'type': 'image',
'filename': fname,
'annotations': []
}
self.annotations_.append(annotation)
def save(self, filename):
self.saveRectFormat(filename)
@@ -60,7 +85,9 @@ class AnnotationContainer:
rect_anns = [ann for ann in file['annotations'] if ann['type'] == 'rect']
#f.write("%s %d", file['filename'], len(rect_anns))
for ann in rect_anns:
f.write('%s 1 %s %s %s %s\n' % (file['filename'], str(ann['x']), str(ann['y']), str(ann['width']), str(ann['height'])))
f.write('%s 1 %s %s %s %s %s\n' % (file['filename'], str(ann['x']), str(ann['y']), str(ann['width']), str(ann['height']), str(ann.get('id', ""))))
if len(rect_anns) == 0:
f.write('%s 0\n' % (file['filename']))
self.filename_ = filename