From af67ec1720d9e4479d1e48389e7482411ebc7645 Mon Sep 17 00:00:00 2001 From: Mika Fischer Date: Tue, 24 May 2011 21:39:18 +0200 Subject: [PATCH] Make loading of videos work --- sloth/annotations/container.py | 2 +- sloth/annotations/model.py | 56 ++++++++++++++++++---------------- 2 files changed, 31 insertions(+), 27 deletions(-) diff --git a/sloth/annotations/container.py b/sloth/annotations/container.py index 35a8fd4..f683088 100644 --- a/sloth/annotations/container.py +++ b/sloth/annotations/container.py @@ -94,7 +94,7 @@ class AnnotationContainer: for file in self.annotations_: if file['type'] == 'image': num += len(file['annotations']) - elif file['type'] == video: + elif file['type'] == 'video': for frame in file['frames']: num += len(frame['annotations']) return num diff --git a/sloth/annotations/model.py b/sloth/annotations/model.py index a1629f3..64b7c1e 100644 --- a/sloth/annotations/model.py +++ b/sloth/annotations/model.py @@ -6,11 +6,13 @@ from PyQt4.QtCore import * from functools import partial import os.path import okapy +import okapy.videoio as okv TypeRole, DataRole, ImageRole = [Qt.UserRole + i + 1 for i in range(3)] class ModelItem: - def __init__(self, parent=None): + def __init__(self, model, parent=None): + self.model_ = model self.parent_ = parent self.children_ = [] @@ -21,6 +23,9 @@ class ModelItem: # return tuple child, index of the child return [(child, index.child(row, 0)) for row, child in enumerate(self.children_)] + def model(self): + return self.model_ + def parent(self): return self.parent_ @@ -34,24 +39,24 @@ class ModelItem: return QVariant() class RootModelItem(ModelItem): - def __init__(self, files): - ModelItem.__init__(self, None) + def __init__(self, model, files): + ModelItem.__init__(self, model, None) self.files_ = files for file in files: - fmi = FileModelItem.create(file, self) + fmi = FileModelItem.create(self.model(), file, self) self.children_.append(fmi) class FileModelItem(ModelItem): - def __init__(self, file, parent): - ModelItem.__init__(self, parent) + def __init__(self, model, file, parent): + ModelItem.__init__(self, model, parent) self.file_ = file def filename(self): return self.file_['filename'] - def fullpath(self, index): - return os.path.join(index.model().basedir(), self.filename()) + def fullpath(self): + return os.path.join(self.model().basedir(), self.filename()) def data(self, index, role): if role == Qt.DisplayRole and index.column() == 0: @@ -59,15 +64,15 @@ class FileModelItem(ModelItem): return ModelItem.data(self, index, role) @staticmethod - def create(file, parent): + def create(model, file, parent): if file['type'] == 'image': - return ImageFileModelItem(file, parent) + return ImageFileModelItem(model, file, parent) elif file['type'] == 'video': - return VideoFileModelItem(file, parent) + return VideoFileModelItem(model, file, parent) class ImageFileModelItem(FileModelItem): - def __init__(self, file, parent): - FileModelItem.__init__(self, file, parent) + def __init__(self, model, file, parent): + FileModelItem.__init__(self, model, file, parent) for ann in file['annotations']: ami = AnnotationModelItem(ann, self) @@ -84,7 +89,7 @@ class ImageFileModelItem(FileModelItem): def data(self, index, role): if role == ImageRole: - return okapy.loadImage(self.fullpath(index)) + return okapy.loadImage(self.fullpath()) elif role == DataRole: return self.file_ return FileModelItem.data(self, index, role) @@ -93,19 +98,18 @@ class VideoFileModelItem(FileModelItem): _cached_vs_filename = None _cached_vs = None - def __init__(self, file, parent): - FileModelItem.__init__(self, file, parent) + def __init__(self, model, file, parent): + FileModelItem.__init__(self, model, file, parent) for frame in file['frames']: - fmi = FrameModelItem(frame, self) + fmi = FrameModelItem(self.model(), frame, self) self.children_.append(fmi) def updateCachedVideoSource(self): # have only one cached video source at a time for now # TODO: for labeling multiple synchronized videos this should # be modified, otherwise it might be awfully slow - VideoFileModelItem._cached_vs = okv.ImageSequence() - VideoFileModelItem._cached_vs.open(self.fullpath()) + VideoFileModelItem._cached_vs = okv.FFMPEGIndexedVideoSource(self.fullpath()) VideoFileModelItem._cached_vs_filename = self.fullpath() def getFrame(self, frame): @@ -116,8 +120,8 @@ class VideoFileModelItem(FileModelItem): return VideoFileModelItem._cached_vs.getImage() class FrameModelItem(ModelItem): - def __init__(self, frame, parent): - ModelItem.__init__(self, parent) + def __init__(self, model, frame, parent): + ModelItem.__init__(self, model, parent) self.frame_ = frame for ann in frame['annotations']: @@ -147,8 +151,8 @@ class FrameModelItem(ModelItem): return QVariant() class AnnotationModelItem(ModelItem): - def __init__(self, annotation, parent): - ModelItem.__init__(self, parent) + def __init__(self, model, annotation, parent): + ModelItem.__init__(self, model, parent) self.annotation_ = annotation # dummy key/value so that pyqt does not convert the dict # into a QVariantMap while communicating with the Views @@ -209,8 +213,8 @@ class AnnotationModelItem(ModelItem): return self.annotation_[key] class KeyValueModelItem(ModelItem): - def __init__(self, key, parent): - ModelItem.__init__(self, parent) + def __init__(self, model, key, parent): + ModelItem.__init__(self, model, parent) self.key_ = key def data(self, index, role): @@ -229,7 +233,7 @@ class AnnotationModel(QAbstractItemModel): def __init__(self, annotations, parent=None): QAbstractItemModel.__init__(self, parent) self.annotations_ = annotations - self.root_ = RootModelItem(self.annotations_) + self.root_ = RootModelItem(self, self.annotations_) self.dirty_ = False self.basedir_ = ""