implement loading of videos in the AnnotationContainer

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

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