From 28b698196bff283383f6abd62cd00cdf9ff9e4d9 Mon Sep 17 00:00:00 2001
From: Martin Baeuml <baeuml@kit.edu>
Date: Mon, 31 Jan 2011 12:34:58 +0100
Subject: [PATCH] add reading/writing of ids

---
 annotations.py | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/annotations.py b/annotations.py
index 0358541..f453ea2 100644
--- a/annotations.py
+++ b/annotations.py
@@ -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