From 25d183349d2db0726dc275a31509e5b3fe3f1a71 Mon Sep 17 00:00:00 2001
From: "junjie.jiang" <junjie.jiang@zilliz.com>
Date: Thu, 21 Jul 2022 15:04:38 +0800
Subject: [PATCH] update

---
 video_decoder.py | 25 +++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/video_decoder.py b/video_decoder.py
index 08717d1..9333a41 100644
--- a/video_decoder.py
+++ b/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):