diff --git a/annotationmodel.py b/annotationmodel.py
index 8193ae0..88a3f8d 100644
--- a/annotationmodel.py
+++ b/annotationmodel.py
@@ -364,14 +364,13 @@ class AnnotationModel(QAbstractItemModel):
         return True
 
     def removeAnnotation(self, annidx):
-        # TODO fix to work with ImageFileModelItem
         annidx = QModelIndex(annidx)  # explicitly convert from QPersistentModelIndex
         item = self.itemFromIndex(annidx)
         assert isinstance(item, AnnotationModelItem)
 
         parent = item.parent_
         parentidx = annidx.parent()
-        assert isinstance(parent, FrameModelItem)
+        assert isinstance(parent, FrameModelItem) or isinstance(parent, ImageFileModelItem)
 
         pos = parent.rowOfChild(item)
         self.beginRemoveRows(parentidx, pos, pos)
diff --git a/annotationscene.py b/annotationscene.py
index 246ef43..677ca9e 100644
--- a/annotationscene.py
+++ b/annotationscene.py
@@ -195,6 +195,7 @@ class AnnotationScene(QGraphicsScene):
 
         num_items = self.model_.rowCount(self.root_)
         self.insertItems(0, num_items)
+        self.update()
 
     def insertItems(self, first, last):
         assert self.model_ is not None
@@ -303,8 +304,8 @@ class AnnotationScene(QGraphicsScene):
             # selection mode
             if event.key() == Qt.Key_Delete:
                 for item in self.selectedItems():
-                    index = self.model_.mapFromSource(QModelIndex(item.index()))
-                    self.model_.removeRow(index.row(), QModelIndex(index.parent()))
+                    index = item.index()
+                    index.model().removeAnnotation(index)
                 event.accept()
                 return