logo
Browse Source

update

Signed-off-by: video-decode <video-decode@towhee.io>
main
video-decode 2 years ago
parent
commit
342306e394
  1. 39
      video_decoder.py

39
video_decoder.py

@ -33,17 +33,15 @@ class VideoDecoder(PyOperator):
self._sample_type = sample_type self._sample_type = sample_type
self._args = args if args is not None else {} self._args = args if args is not None else {}
def get_sample(self, stream):
def get_sample(self, stream, duration):
if self._sample_type is None: if self._sample_type is None:
return self._no_sample return self._no_sample
elif self._sample_type.lower() == SAMPLE_TYPE.UNIFORM_TEMPORAL_SUBSAMPLE: elif self._sample_type.lower() == SAMPLE_TYPE.UNIFORM_TEMPORAL_SUBSAMPLE:
duration = VideoDecoder.get_video_duration(stream)
end_time = self._end_time if self._end_time is not None and self._end_time <= duration else duration end_time = self._end_time if self._end_time is not None and self._end_time <= duration else duration
start_time = self._start_time if self._start_time is not None else 0 start_time = self._start_time if self._start_time is not None else 0
nums = int(stream.rate * (end_time - start_time)) nums = int(stream.rate * (end_time - start_time))
return partial(self._uniform_temporal_subsample, total_frames=nums) return partial(self._uniform_temporal_subsample, total_frames=nums)
elif self._sample_type.lower() == SAMPLE_TYPE.TIME_STEP_SAMPLE: elif self._sample_type.lower() == SAMPLE_TYPE.TIME_STEP_SAMPLE:
duration = VideoDecoder.get_video_duration(stream)
start_time = self._start_time if self._start_time is not None else 0 start_time = self._start_time if self._start_time is not None else 0
end_time = self._end_time if self._end_time is not None and self._end_time <= duration else duration end_time = self._end_time if self._end_time is not None and self._end_time <= duration else duration
return partial(self._time_step_sample, start_time=start_time, end_time=end_time) return partial(self._time_step_sample, start_time=start_time, end_time=end_time)
@ -108,22 +106,25 @@ class VideoDecoder(PyOperator):
continue continue
yield frame yield frame
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
# @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
def __call__(self, video_path: str) -> Generator: def __call__(self, video_path: str) -> Generator:
with av.open(video_path) as container: with av.open(video_path) as container:
stream = container.streams.video[0] stream = container.streams.video[0]
duration = float(container.duration) / 1000000
image_format = 'RGB' image_format = 'RGB'
frame_gen = VideoDecoder._decdoe(stream, container, self._start_time) frame_gen = VideoDecoder._decdoe(stream, container, self._start_time)
sample_function = self.get_sample(stream)
sample_function = self.get_sample(stream, duration)
for frame in sample_function(frame_gen): for frame in sample_function(frame_gen):
timestamp = int(frame.time * 1000) timestamp = int(frame.time * 1000)
ndarray = frame.to_ndarray(format='rgb24') ndarray = frame.to_ndarray(format='rgb24')
@ -136,7 +137,8 @@ class VideoDecoder(PyOperator):
# video_path = "/home/junjie.jiangjjj/workspace/video/[The Rock] [1996] [Trailer] [#2]-16-l-rO5B64.mkv" # video_path = "/home/junjie.jiangjjj/workspace/video/[The Rock] [1996] [Trailer] [#2]-16-l-rO5B64.mkv"
# video_path1 = "/home/junjie.jiangjjj/workspace/video/'Eagle Eye' Trailer (2008)-_wkqo_Rd3_Q.mp4" # video_path1 = "/home/junjie.jiangjjj/workspace/video/'Eagle Eye' Trailer (2008)-_wkqo_Rd3_Q.mp4"
# video_path2 = "/home/junjie.jiangjjj/workspace/video/2001 - A Space Odyssey - Trailer [1968] HD-Z2UWOeBcsJI.webm" # video_path2 = "/home/junjie.jiangjjj/workspace/video/2001 - A Space Odyssey - Trailer [1968] HD-Z2UWOeBcsJI.webm"
# video_path3 = "/home/zhangchen/zhangchen_workspace/dataset/MSRVTT/msrvtt_data/MSRVTT_Videos/video9991.mp4"
# # video_path3 = "/home/zhangchen/zhangchen_workspace/dataset/MSRVTT/msrvtt_data/MSRVTT_Videos/video9991.mp4"
# video_path3 = "/home/junjie.jiangjjj/e2adc784b83446ae775f698b9d17c9fd392b2f75.flv"
# def d(video_path): # def d(video_path):
# d = VideoDecoder(10, 17, 'time_step_sample', {'time_step': 1}) # d = VideoDecoder(10, 17, 'time_step_sample', {'time_step': 1})
@ -144,7 +146,12 @@ class VideoDecoder(PyOperator):
# for f in fs: # for f in fs:
# print(f.mode, f.key_frame, f.timestamp) # print(f.mode, f.key_frame, f.timestamp)
# d(video_path1)
# print('#' * 100)
# d(video_path2)
# d(video_path)
# # print('#' * 100)
# # with av.open(video_path) as container:
# # print(container.duration)
# # stream = container.streams.video[0]
# # print(stream.time_base)

Loading…
Cancel
Save