diff --git a/sloth/annotations/container.py b/sloth/annotations/container.py index e832704..214934d 100644 --- a/sloth/annotations/container.py +++ b/sloth/annotations/container.py @@ -90,7 +90,7 @@ class AnnotationContainer: "AnnotationContainer.load()" ) - def save(self, filename=""): + def save(self, annotations, filename=""): """ Save the annotations. """ diff --git a/sloth/annotations/model.py b/sloth/annotations/model.py index 1980499..90505a6 100644 --- a/sloth/annotations/model.py +++ b/sloth/annotations/model.py @@ -4,6 +4,7 @@ The annotationmodel module contains the classes for the AnnotationModel. from PyQt4.QtGui import QTreeView, QSortFilterProxyModel, QAbstractItemView from PyQt4.QtCore import QModelIndex, QPersistentModelIndex, QAbstractItemModel, QVariant, Qt, pyqtSignal import os.path +import copy ItemRole, TypeRole, DataRole, ImageRole = [Qt.UserRole + i + 1 for i in range(4)] @@ -176,6 +177,9 @@ class RootModelItem(ModelItem): # TODO return 0 + def getAnnotations(self): + return [child.getAnnotations() for child in self.children()] + class FileModelItem(ModelItem): def __init__(self, fileinfo): ModelItem.__init__(self) @@ -236,6 +240,11 @@ class ImageFileModelItem(FileModelItem, ImageModelItem): return self._fileinfo return FileModelItem.data(self, role) + def getAnnotations(self): + fi = copy.deepcopy(self._fileinfo) + fi['annotations'] = [child.getAnnotations() for child in self.children()] + return fi + class VideoFileModelItem(FileModelItem): def __init__(self, fileinfo): frameinfos = fileinfo.get("frames", []) @@ -246,6 +255,11 @@ class VideoFileModelItem(FileModelItem): for frameinfo in frameinfos: self.appendChild(FrameModelItem(frameinfo)) + def getAnnotations(self): + fi = copy.deepcopy(self._fileinfo) + fi['frames'] = [child.getAnnotations() for child in self.children()] + return fi + class FrameModelItem(ImageModelItem): def __init__(self, frameinfo): if frameinfo.has_key("annotations"): @@ -264,6 +278,11 @@ class FrameModelItem(ImageModelItem): return "%d / %.3f" % (self.framenum(), self.timestamp()) return ImageModelItem.data(self, role, column) + def getAnnotations(self): + fi = copy.deepcopy(self._frameinfo) + fi['annotations'] = [child.getAnnotations() for child in self.children()] + return fi + class AnnotationModelItem(ModelItem): def __init__(self, annotation): ModelItem.__init__(self) @@ -275,6 +294,13 @@ class AnnotationModelItem(ModelItem): def type(self): return self._annotation['type'] + def getAnnotations(self): + ann = copy.deepcopy(self._annotation) + if None in ann: + del ann[None] + # TODO: Maybe it would be nice to enforce a deterministic ordering? + return ann + def annotation(self): return self._annotation diff --git a/sloth/core/labeltool.py b/sloth/core/labeltool.py index 7f594af..47fdb11 100755 --- a/sloth/core/labeltool.py +++ b/sloth/core/labeltool.py @@ -216,14 +216,12 @@ class LabelTool(QObject): # 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_ = self.container_factory_.create(fname) # Get annotations dict - ann = self._model.getAnnotations() + ann = self._model.root().getAnnotations() - self.container_.save(fname, ann) + self.container_.save(ann, 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._model.root().numFiles(), self._model.root().numAnnotations()) diff --git a/sloth/gui/labeltool.py b/sloth/gui/labeltool.py index 0a8d25b..151836d 100755 --- a/sloth/gui/labeltool.py +++ b/sloth/gui/labeltool.py @@ -195,7 +195,7 @@ class MainWindow(QMainWindow): filename = self.labeltool.getCurrentFilename() if filename is None: return self.fileSaveAs() - return self.saveAnnotations(filename) + return self.labeltool.saveAnnotations(filename) def fileSaveAs(self): fname = '.' # self.annotations.filename() or '.' @@ -205,7 +205,7 @@ class MainWindow(QMainWindow): "%s annotation files (%s)" % (APP_NAME, format_str)) if not fname.isEmpty(): - return self.saveAnnotations(str(fname)) + return self.labeltool.saveAnnotations(str(fname)) return False def addMediaFile(self):