diff --git a/annotationmodel.py b/annotationmodel.py index 2db5b51..aa4befe 100644 --- a/annotationmodel.py +++ b/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): diff --git a/annotations.py b/annotations.py index f453ea2..031b0a3 100644 --- a/annotations.py +++ b/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") diff --git a/annotationscene.py b/annotationscene.py index 6adff2e..3494210 100644 --- a/annotationscene.py +++ b/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): diff --git a/labeltool.py b/labeltool.py index 8b29dbf..bc34c2b 100755 --- a/labeltool.py +++ b/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() diff --git a/loaders/__init__.py b/loaders/__init__.py index e69de29..118def0 100644 --- a/loaders/__init__.py +++ b/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 + diff --git a/std_config.py b/std_config.py index 054532e..493f5ad 100644 --- a/std_config.py +++ b/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 = (