diff --git a/sloth/annotations/model.py b/sloth/annotations/model.py index ead48d3..5d43202 100644 --- a/sloth/annotations/model.py +++ b/sloth/annotations/model.py @@ -47,6 +47,15 @@ class RootModelItem(ModelItem): fmi = FileModelItem.create(self.model(), file, self) self.children_.append(fmi) + def addFile(self, file): + fmi = FileModelItem.create(self.model(), file, self) + next = len(self.children_) + index = self.model().index(0, 0, QModelIndex()) + self.model().beginInsertRows(index, next, next) + self.children_.append(fmi) + self.model().endInsertRows() + self.model().emit(SIGNAL("dataChanged(QModelIndex,QModelIndex)"), index, index) + class FileModelItem(ModelItem): def __init__(self, model, file, parent): ModelItem.__init__(self, model, parent) diff --git a/sloth/gui/annotationscene.py b/sloth/gui/annotationscene.py index db8b479..182921c 100644 --- a/sloth/gui/annotationscene.py +++ b/sloth/gui/annotationscene.py @@ -90,8 +90,10 @@ class AnnotationScene(QGraphicsScene): self.update() def insertItems(self, first, last): + if not self.root_.isValid(): + return + assert self.model_ is not None - assert self.root_.isValid() # create a graphics item for each model index for row in range(first, last+1): @@ -235,6 +237,9 @@ class AnnotationScene(QGraphicsScene): # this is the implemenation of the scene as a view of the model #______________________________________________________________________________________________________ def dataChanged(self, indexFrom, indexTo): + if not self.root_.isValid(): + return + if self.root_ != indexFrom.parent() or self.root_ != indexTo.parent(): return diff --git a/sloth/gui/labeltool.py b/sloth/gui/labeltool.py index 5d1dc15..559bb91 100755 --- a/sloth/gui/labeltool.py +++ b/sloth/gui/labeltool.py @@ -1,6 +1,7 @@ #!/usr/bin/python import sys, os import functools, importlib +import fnmatch from optparse import OptionParser from PyQt4.QtGui import * from PyQt4.QtCore import * @@ -15,6 +16,8 @@ from sloth.gui.controlbuttons import * from sloth.conf import config from sloth import APP_NAME, ORGANIZATION_NAME, ORGANIZATION_DOMAIN, VERSION +import okapy.videoio as okv + GUIDIR=os.path.join(os.path.dirname(__file__)) class MainWindow(QMainWindow): @@ -120,7 +123,7 @@ class MainWindow(QMainWindow): self.ui.action_About.triggered.connect(self.about) ## Navigation - #self.ui.action_Add_Image.triggered.connect(self.addImage) + self.ui.action_Add_Image.triggered.connect(self.addMediaFile) self.ui.actionNext. triggered.connect(self.gotoNext) self.ui.actionPrevious. triggered.connect(self.gotoPrevious) self.ui.actionZoom_In. triggered.connect(functools.partial(self.view.setScaleRelative, 1.2)) @@ -301,6 +304,60 @@ class MainWindow(QMainWindow): if index != self.treeview.currentIndex(): self.treeview.setCurrentIndex(self.current_index_) + def addImageFile(self, fname): + fileitem = { + 'filename': fname, + 'type': 'image', + 'annotations': [ ], + } + self.model_.root_.addFile(fileitem) + + def addVideoFile(self, fname): + fileitem = { + 'filename': fname, + 'type': 'video', + 'frames': [ ], + } + + # FIXME: OKAPI should provide a method to get all timestamps at once + # FIXME: Some dialog should be displayed, telling the user that the + # video is being loaded/indexed and that this might take a while + video = okv.FFMPEGIndexedVideoSource(fname) + i = 0 + while video.getNextFrame(): + ts = video.getTimestamp() + frame = { 'annotations': [], + 'num': i, + 'timestamp': ts, + } + fileitem['frames'].append(frame) + i += 1 + + self.model_.root_.addFile(fileitem) + + def addMediaFile(self): + path = '.' + if (self.container_.filename() is not None) and \ + (len(self.container_.filename()) > 0): + path = QFileInfo(self.container_.filename()).path() + + image_types = [ '*.jpg', '*.bmp', '*.png', '*.pgm', '*.ppm', '*.ppm', '*.tif', '*.gif' ] + video_types = [ '*.mp4', '*.mpg', '*.mpeg', '*.avi', '*.mov', '*.vob' ] + format_str = ' '.join(image_types + video_types) + fname = QFileDialog.getOpenFileName(self, "%s - Add Media File" % APP_NAME, path, "Media files (%s)" % (format_str, )) + + if fname.isEmpty(): + return + + fname = str(fname) + + for pattern in image_types: + if fnmatch.fnmatch(fname, pattern): + return self.addImageFile(fname) + + return self.addVideoFile(fname) + + ### ### global event handling ###______________________________________________________________________________