update

video_decoder.py

@@ -106,23 +106,24 @@ class VideoDecoder(PyOperator):
                 continue
             yield frame
 
-    # @staticmethod
-    # def get_video_duration(video):
-    #     print(video)
-    #     if video.duration is not None:
-    #         return float(video.duration * video.time_base)
-    #     elif video.metadata.get('DURATION') is not None:
-    #         time_str = video.metadata['DURATION']
-    #         return reduce(lambda x, y: float(x) * 60 + float(y), time_str.split(':'))
-    #     else:
-    #         return None
+    @staticmethod
+    def get_video_duration(video):
+        if video.duration is not None:
+            return float(video.duration * video.time_base)
+        elif video.metadata.get('DURATION') is not None:
+            time_str = video.metadata['DURATION']
+            return reduce(lambda x, y: float(x) * 60 + float(y), time_str.split(':'))
+        else:
+            return None
 
     def __call__(self, video_path: str) -> Generator:
         with av.open(video_path) as container:
             stream = container.streams.video[0]
-            duration = float(container.duration) / 1000000
-            image_format = 'RGB'
+            duration = VideoDecoder.get_video_duration(stream)
+            if duration is None:
+                duration = float(container.duration) / 1000000
 
+            image_format = 'RGB'
             frame_gen = VideoDecoder._decdoe(stream, container, self._start_time)
             sample_function = self.get_sample(stream, duration)
             for frame in sample_function(frame_gen):