Fix saving (closes: #2)

sloth/annotations/container.py

@@ -90,7 +90,7 @@ class AnnotationContainer:
             "AnnotationContainer.load()"
         )
 
-    def save(self, filename=""):
+    def save(self, annotations, filename=""):
         """
         Save the annotations.
         """

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
 

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())

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):