diff --git a/sloth/annotations/container.py b/sloth/annotations/container.py index 325b9e5..71450d1 100644 --- a/sloth/annotations/container.py +++ b/sloth/annotations/container.py @@ -23,6 +23,7 @@ except: pass try: import okapy + import okapy.videoio as okv _use_pil = False except: try: @@ -86,6 +87,7 @@ class AnnotationContainer: def clear(self): self._annotations = [] self._filename = None + self._video_cache = {} def load(self, filename): """ @@ -165,12 +167,22 @@ class AnnotationContainer: ``frame_number``. In the default implementation this will try to load the video from a path relative to the label files directory. """ - fullpath = self._fullpath(filename) + fullpath = str(self._fullpath(filename)) if not os.path.exists(fullpath): LOG.warn("Video file %s does not exists." % fullpath) return None - #TODO load video + # get video source from cache or load from file + if fullpath in self._video_cache: + vidsrc = self._video_cache[fullpath] + else: + vidsrc = okv.createVideoSourceFromString(fullpath) + vidsrc = okv.toRandomAccessVideoSource(vidsrc) + self._video_cache[fullpath] = vidsrc + + # get requested frame + vidsrc.getFrame(frame_number) + return vidsrc.getImage() class PickleContainer(AnnotationContainer): diff --git a/sloth/core/labeltool.py b/sloth/core/labeltool.py index e09bc06..a457660 100755 --- a/sloth/core/labeltool.py +++ b/sloth/core/labeltool.py @@ -321,8 +321,11 @@ class LabelTool(QObject): self.currentImageChanged.emit() def getImage(self, item): - # TODO: Also handle video frames - return self._container.loadImage(item['filename']) + if item['class'] == 'frame': + video = item.parent() + return self._container.loadFrame(video['filename'], item['num']) + else: + return self._container.loadImage(item['filename']) def getAnnotationFilePatterns(self): return self._container_factory.patterns() @@ -345,13 +348,17 @@ class LabelTool(QObject): # 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) + LOG.info("Importing frames from %s. This may take a while..." % fname) + video = okv.createVideoSourceFromString(fname) + video = okv.toRandomAccessVideoSource(video) i = 0 while video.getNextFrame(): + LOG.debug("Adding frame %d" % i) ts = video.getTimestamp() frame = { 'annotations': [], 'num': i, 'timestamp': ts, + 'class': 'frame' } fileitem['frames'].append(frame) i += 1