ffmpeg
copied
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Readme
Files and versions
51 lines
1.6 KiB
51 lines
1.6 KiB
2 years ago
|
import math
|
||
|
|
import logging
|
||
|
|
import av
|
||
|
|
|
||
|
|
from towhee.types.video_frame import VideoFrame
|
||
|
|
|
||
|
|
|
||
|
|
logger = logging.getLogger()
|
||
|
|
|
||
|
|
|
||
|
|
class PyAVDecode:
|
||
|
|
def __init__(self, video_path, start_time=None, time_step=None) -> None:
|
||
|
|
self._container = av.open(video_path)
|
||
|
|
self._stream = self._container.streams.video[0]
|
||
|
|
self._start_time = start_time if start_time is not None else 0
|
||
|
|
self._time_step = time_step
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
self._container.close()
|
||
|
|
|
||
|
|
def time_step_decode(self):
|
||
|
|
ts = self._start_time
|
||
|
|
is_end = False
|
||
|
|
while not is_end:
|
||
|
|
is_end = True
|
||
|
|
offset = int(math.floor(ts / self._stream.time_base))
|
||
|
|
self._container.seek(offset, stream=self._stream)
|
||
|
|
for f in self._container.decode(self._stream):
|
||
|
|
if f.time < ts:
|
||
|
|
continue
|
||
|
|
yield self.av_frame_to_video_frame(f)
|
||
|
|
is_end = False
|
||
|
|
break
|
||
|
|
ts += self._time_step
|
||
|
|
|
||
|
|
def av_frame_to_video_frame(self, frame):
|
||
|
|
timestamp = int(round(frame.time * 1000))
|
||
|
|
ndarray = frame.to_ndarray(format='rgb24')
|
||
|
|
img = VideoFrame(ndarray, 'RGB', timestamp, frame.key_frame)
|
||
|
|
return img
|
||
|
|
|
||
|
|
def decode(self):
|
||
|
|
if self._start_time > 0:
|
||
|
|
offset = int(math.floor(self._start_time / self._stream.time_base))
|
||
|
|
self._container.seek(offset, any_frame=False, backward=True, stream=self._stream)
|
||
|
|
|
||
|
|
for frame in self._container.decode(self._stream):
|
||
|
|
if frame.time < self._start_time:
|
||
|
|
continue
|
||
|
|
yield self.av_frame_to_video_frame(frame)
|