mirror of
https://github.com/wassname/sloth.git
synced 2026-07-11 22:51:51 +08:00
add reading/writing of ids
This commit is contained in:
+29
-2
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user