implement loaders

annotationmodel.py

@@ -233,7 +233,7 @@ class AnnotationModel(QAbstractItemModel):
         return self.basedir_
 
     def setBasedir(self, dir):
-        print "setBasedir", dir
+        print "setBasedir: \"" + dir + "\"" 
         self.basedir_ = dir
 
     def itemFromIndex(self, index):

annotations.py

@@ -1,9 +1,9 @@
 import os
-import glob
 
 class AnnotationContainer:
-    def __init__(self):
+    def __init__(self, loaders):
         self.clear()
+        self.loaders_ = dict(loaders)
 
     def filename(self):
         return self.filename_
@@ -13,69 +13,17 @@ class AnnotationContainer:
         self.filename_    = None
 
     def load(self, filename):
-        self.loadRectFormat(filename)
+        ext = os.path.splitext(filename)[1]
+        if ext.startswith('.'):
+            ext = ext[1:]
 
-    def loadRectFormat(self, filename):
-        self.annotations_ = []
+        loader = self.loaders_[ext]()
+        self.annotations_ = loader.load(filename)
         self.filename_ = filename
-        filename_helper = {}
 
-        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:
-                    rects = []
-
-                for i in range(0, 4*num_rects, 4):
-                    rects.append({
-                        'type': 'rect',
-                        'x': s[i+2],
-                        'y': s[i+3],
-                        '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 = {
-                            'type':     'image',
-                            'filename': image_fname,
-                            'annotations': rects
-                            }
-                    self.annotations_.append(annotation)
-                    filename_helper[image_fname] = annotation
-        except Exception as e:
-            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)
+        pass
+        #self.saveRectFormat(filename)
 
     def saveRectFormat(self, filename):
         f = open(filename, "w")

annotationscene.py

@@ -90,7 +90,6 @@ class AnnotationScene(QGraphicsScene):
     def insertItems(self, first, last):
         assert self.model_ is not None
         assert self.root_.isValid()
-        print "insertItems"
 
         # create a graphics item for each model index
         for row in range(first, last+1):

labeltool.py

@@ -9,7 +9,6 @@ from buttonarea import *
 from annotationmodel import *
 from annotationscene import *
 from frameviewer import *
-from loaders import *
 from optparse import OptionParser
 import annotations
 from conf import config
@@ -29,9 +28,7 @@ class MainWindow(QMainWindow):
             # load config
             config.update(options.config)
 
-        self.loaders_ = []
-
-        self.anno_container = annotations.AnnotationContainer()
+        self.anno_container = annotations.AnnotationContainer(config.LOADERS)
         self.current_index_ = None
 
         self.setupGui()
@@ -125,7 +122,7 @@ class MainWindow(QMainWindow):
             self.anno_container.load(fname)
             msg = "Successfully loaded %s (%d files, %d annotations)" % \
                     (fname, self.anno_container.numFiles(), self.anno_container.numAnnotations())
-        except Exception as e:
+        except Exception, e:
             msg = "Error: Loading failed (%s)" % str(e)
         self.updateStatus(msg)
         self.updateViews()

loaders/__init__.py

@@ -0,0 +1,80 @@
+import os
+
+class Loader:
+    def __init__(self):
+        pass
+
+    def load(self, filename):
+        pass
+
+    def save(self, filename, annotations):
+        pass
+
+class RectIdLoader(Loader):
+    def load(self, filename):
+        annotations = []
+
+        print "loading ", filename
+        filename_helper = {}
+
+        f = open(filename)
+        basedir = ""
+        try:
+            for line in f:
+                s = line.split()
+                print s
+                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 not in filename_helper:
+                    annotation = {
+                            'type':     'image',
+                            'filename': image_fname,
+                            'annotations': []
+                            }
+                    annotations.append(annotation)
+                    filename_helper[image_fname] = annotation
+
+                rects = filename_helper[image_fname]['annotations']
+
+                for i in range(0, 4*num_rects, 4):
+                    rects.append({
+                        'type': 'rect',
+                        'x': s[i+2],
+                        'y': s[i+3],
+                        'width':  s[i+4],
+                        'height': s[i+5]
+                        })
+                    if len(s) > i+6:
+                        rects[-1]['id'] = s[i+6]
+        except Exception, e:
+            print e
+
+        print annotations
+        return annotations
+
+class FeretLoader(Loader):
+    def load(self, filename):
+        basedir = os.path.dirname(filename)
+        f = open(filename)
+
+        annotations = []
+        for line in f:
+            s = line.split()
+            fileitem = {
+                'filename': os.path.join(basedir, s[0]+".bmp"),
+                'type': 'image',
+            }
+            fileitem['annotations'] = [
+                {'type': 'point', 'class': 'left_eye',  'x': int(s[1]), 'y': int(s[2])},
+                {'type': 'point', 'class': 'right_eye', 'x': int(s[3]), 'y': int(s[4])},
+                {'type': 'point', 'class': 'mouth',     'x': int(s[5]), 'y': int(s[6])}
+            ]
+            annotations.append(fileitem)
+
+        return annotations
+

std_config.py

@@ -1,5 +1,6 @@
 from items import AnnotationGraphicsRectItem as RectItem 
-from items import RectItemInserter
+from items import RectItemInserter, AnnotationGraphicsPointItem
+from loaders import FeretLoader, RectIdLoader
 from bboxitem import *
 
 RATIOS = ["0.5", "1", "2"]
@@ -21,6 +22,7 @@ HOTKEYS = (
 
 ITEMS = (
     ("rect",     RectItem),
+    ("point",    AnnotationGraphicsPointItem),
     ("bodybbox", BodyBoundingboxItem),
 )
 
@@ -30,6 +32,7 @@ INSERTERS = (
 )
 
 LOADERS = (
+    ('txt', RectIdLoader),
 )
 
 PLUGINS = (