From c84e02befed69f5b296f096c969ece558640cf4b Mon Sep 17 00:00:00 2001 From: Martin Baeuml Date: Fri, 13 May 2011 11:07:03 +0200 Subject: [PATCH] - move annotation model and containers to package annotations - implement AnnotationContainerFactory --- annotations.py | 59 ------------- annotations/__init__.py | 0 annotations/container.py | 97 ++++++++++++++++++++++ annotationmodel.py => annotations/model.py | 0 annotationscene.py | 2 +- conf/default_config.py | 10 +-- core/__init__.py | 0 core/exceptions.py | 7 ++ labeltool.py | 54 ++++++------ 9 files changed, 140 insertions(+), 89 deletions(-) delete mode 100644 annotations.py create mode 100644 annotations/__init__.py create mode 100644 annotations/container.py rename annotationmodel.py => annotations/model.py (100%) create mode 100644 core/__init__.py create mode 100644 core/exceptions.py diff --git a/annotations.py b/annotations.py deleted file mode 100644 index 031b0a3..0000000 --- a/annotations.py +++ /dev/null @@ -1,59 +0,0 @@ -import os - -class AnnotationContainer: - def __init__(self, loaders): - self.clear() - self.loaders_ = dict(loaders) - - def filename(self): - return self.filename_ - - def clear(self): - self.annotations_ = [] - self.filename_ = None - - def load(self, filename): - ext = os.path.splitext(filename)[1] - if ext.startswith('.'): - ext = ext[1:] - - loader = self.loaders_[ext]() - self.annotations_ = loader.load(filename) - self.filename_ = filename - - def save(self, filename): - pass - #self.saveRectFormat(filename) - - def saveRectFormat(self, filename): - f = open(filename, "w") - - for file in self.annotations_: - if file['type'] == 'image': - 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 %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 - - def asDict(self): - return self.annotations_ - - def numFiles(self): - return len(self.annotations_) - - def numAnnotations(self): - if self.annotations_ is None: - return 0 - num = 0 - for file in self.annotations_: - if file['type'] == 'image': - num += len(file['annotations']) - elif file['type'] == video: - for frame in file['frames']: - num += len(frame['annotations']) - return num - diff --git a/annotations/__init__.py b/annotations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/annotations/container.py b/annotations/container.py new file mode 100644 index 0000000..f5b69c1 --- /dev/null +++ b/annotations/container.py @@ -0,0 +1,97 @@ +import os +import fnmatch +from core.exceptions import ImproperlyConfigured + +class AnnotationContainerFactory: + def __init__(self, containers): + """ + Initialize the factory with the mappings between file pattern and + the container. + + Parameters + ========== + containers: tuple of tuples (str, str/class) + The mapping between file pattern and container class responsible + for loading/saving. + """ + self.containers_ = containers + + def create(self, filename, *args, **kwargs): + """ + Create a container for the filename. + + Parameters + ========== + filename: str + Filename for which a matching container should be created. + *args, **kwargs: + Arguments passed to constructor of the container. + """ + for pattern, container in self.containers_: + if fnmatch.fnmatch(filename, pattern): + return container(*args, **kwargs) + raise ImproperlyConfigured("No container registered for filename %s" % filename) + +class AnnotationContainer: + def __init__(self): + self.clear() + + def filename(self): + return self.filename_ + + def clear(self): + self.annotations_ = [] + self.filename_ = None + + def load(self, filename): + """ + Load the annotations. Must be implemented in the subclass. + """ + pass + + def save(self, filename): + """ + Save the annotations. Must be implemented in the subclass. + """ + pass + + def annotations(self): + """ + Returns the current annotations as python dictionary. + """ + return self.annotations_ + + def loadImage(self, filename): + """ + Load the image referenced to by the filename. In the default + implementation this will try to load the image from a path + relative to the label files directory. + """ + pass + + def loadVideo(self, filename): + """ + Load the video referenced to by the filename. In the default + implementation this will try to load the video from a path + relative to the label files directory. + """ + pass + + def setAnnotations(self, annotations): + self.annotations_ = annotations + + def numFiles(self): + return len(self.annotations_) + + def numAnnotations(self): + if self.annotations_ is None: + return 0 + num = 0 + for file in self.annotations_: + if file['type'] == 'image': + num += len(file['annotations']) + elif file['type'] == video: + for frame in file['frames']: + num += len(frame['annotations']) + return num + diff --git a/annotationmodel.py b/annotations/model.py similarity index 100% rename from annotationmodel.py rename to annotations/model.py diff --git a/annotationscene.py b/annotationscene.py index 0d0fe18..af669e3 100644 --- a/annotationscene.py +++ b/annotationscene.py @@ -2,7 +2,7 @@ from PyQt4.QtGui import * from PyQt4.QtCore import * from items import * -from annotationmodel import TypeRole, ImageRole +from annotations.model import TypeRole, ImageRole import math import okapy from okapy.guiqt.utilities import toQImage diff --git a/conf/default_config.py b/conf/default_config.py index 332d144..ca185b7 100644 --- a/conf/default_config.py +++ b/conf/default_config.py @@ -5,16 +5,16 @@ HOTKEYS = ( ) ITEMS = { - "rect": "items.RectItem", - "point": "items.PointItem", + 'rect': 'items.RectItem', + 'point': 'items.PointItem', } INSERTERS = { - "rect": "items.RectItemInserter", - "point": "items.PointItemInserter", + 'rect': 'items.RectItemInserter', + 'point': 'items.PointItemInserter', } -LOADERS = ( +CONTAINERS = ( ) PLUGINS = ( diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/core/exceptions.py b/core/exceptions.py new file mode 100644 index 0000000..5a32282 --- /dev/null +++ b/core/exceptions.py @@ -0,0 +1,7 @@ +""" +Label tool exception classes. +""" + +class ImproperlyConfigured(Exception): + """There is an error in the configuration.""" + pass diff --git a/labeltool.py b/labeltool.py index 368027f..56698e1 100755 --- a/labeltool.py +++ b/labeltool.py @@ -9,11 +9,11 @@ from PyQt4.QtCore import * import PyQt4.uic as uic import qrc_icons from buttonarea import * -from annotationmodel import * +from annotations.model import * +from annotations.container import AnnotationContainerFactory, AnnotationContainer from annotationscene import * from frameviewer import * from optparse import OptionParser -import annotations from conf import config APP_NAME = """labeltool""" @@ -31,7 +31,8 @@ class MainWindow(QMainWindow): # load config config.update(options.config) - self.anno_container = annotations.AnnotationContainer(config.LOADERS) + self.container_factory_ = AnnotationContainerFactory(config.CONTAINERS) + self.container_ = AnnotationContainer() self.current_index_ = None self.setupGui() @@ -109,8 +110,8 @@ class MainWindow(QMainWindow): settings.setValue("MainWindow/Size", QVariant(self.size())) settings.setValue("MainWindow/Position", QVariant(self.pos())) settings.setValue("MainWindow/State", QVariant(self.saveState())) - if self.anno_container.filename() is not None: - filename = QVariant(QString(self.anno_container.filename())) + if self.container_.filename() is not None: + filename = QVariant(QString(self.container_.filename())) else: filename = QVariant() settings.setValue("LastFile", filename) @@ -122,9 +123,10 @@ class MainWindow(QMainWindow): def loadAnnotations(self, fname): fname = str(fname) # convert from QString try: - self.anno_container.load(fname) + self.container_ = self.container_factory_.create(fname) + self.container_.load(fname) msg = "Successfully loaded %s (%d files, %d annotations)" % \ - (fname, self.anno_container.numFiles(), self.anno_container.numAnnotations()) + (fname, self.container_.numFiles(), self.container_.numAnnotations()) except Exception, e: msg = "Error: Loading failed (%s)" % str(e) self.updateStatus(msg) @@ -133,10 +135,17 @@ class MainWindow(QMainWindow): def saveAnnotations(self, fname): success = False try: - self.anno_container.save(fname) + # create new container if the filename is different + if fname != self.container_.filename(): + # TODO: skip if it is the same class + newcontainer = self.container_factory_.create(fname) + newcontainer.setAnnotations(self.container_.annotations()) + self.container_ = newcontainer + + self.container_.save(fname) #self.model_.writeback() # write back changes that are cached in the model itself, e.g. mask updates msg = "Successfully saved %s (%d files, %d annotations)" % \ - (fname, self.anno_container.numFiles(), self.anno_container.numAnnotations()) + (fname, self.container_.numFiles(), self.container_.numAnnotations()) success = True self.model_.setDirty(False) except Exception as e: @@ -170,7 +179,7 @@ class MainWindow(QMainWindow): def fileNew(self): if not self.okToContinue(): return - self.anno_container.clear() + self.container_.clear() self.updateStatus() self.updateViews() @@ -178,11 +187,11 @@ class MainWindow(QMainWindow): if not self.okToContinue(): return path = '.' - if (self.anno_container.filename() is not None) and \ - (len(self.anno_container.filename()) > 0): - path = QFileInfo(self.anno_container.filename()).path() + if (self.container_.filename() is not None) and \ + (len(self.container_.filename()) > 0): + path = QFileInfo(self.container_.filename()).path() - #format_str = ' '.join(['*.'+fmt for fmt in self.anno_container.formats()]) + # TODO: compile a list from all the patterns in self.container_factory_ format_str = ' '.join(['*.txt']) fname = QFileDialog.getOpenFileName(self, "%s - Load Annotations" % APP_NAME, path, @@ -191,21 +200,18 @@ class MainWindow(QMainWindow): self.loadAnnotations(fname) def fileSave(self): - if self.anno_container.filename() is None: + if self.container_.filename() is None: return self.fileSaveAs() - return self.saveAnnotations(self.anno_container.filename()) + return self.saveAnnotations(self.container_.filename()) def fileSaveAs(self): fname = '.' # self.annotations.filename() or '.' - #format_str = ' '.join(['*.'+fmt for fmt in self.anno_container.formats()]) format_str = ' '.join(['*.txt']) fname = QFileDialog.getSaveFileName(self, "%s - Save Annotations" % APP_NAME, fname, "%s annotation files (%s)" % (APP_NAME, format_str)) if not fname.isEmpty(): - #if not fname.contains("."): - #fname += ".yaml" return self.saveAnnoations(str(fname)) return False @@ -223,17 +229,17 @@ class MainWindow(QMainWindow): def updateStatus(self, message=''): self.statusBar().showMessage(message, 5000) - if self.anno_container.filename() is not None: + if self.container_.filename() is not None: self.setWindowTitle("%s - %s[*]" % \ - (APP_NAME, QFileInfo(self.anno_container.filename()).fileName())) + (APP_NAME, QFileInfo(self.container_.filename()).fileName())) else: self.setWindowTitle("%s - Unnamed[*]" % APP_NAME) self.updateModified() def updateViews(self): - self.model_ = AnnotationModel(self.anno_container.asDict()) - if self.anno_container.filename() is not None: - self.model_.setBasedir(os.path.dirname(self.anno_container.filename())) + self.model_ = AnnotationModel(self.container_.annotations()) + if self.container_.filename() is not None: + self.model_.setBasedir(os.path.dirname(self.container_.filename())) else: self.model_.setBasedir("") self.model_.dirtyChanged.connect(self.updateModified)