Spaces:
Sleeping
Sleeping
| # We modified the original AVReader class of decord to solve the problem of memory leak. | |
| # For more details, refer to: https://github.com/dmlc/decord/issues/208 | |
| import numpy as np | |
| from decord.video_reader import VideoReader | |
| from decord.audio_reader import AudioReader | |
| from decord.ndarray import cpu | |
| from decord import ndarray as _nd | |
| from decord.bridge import bridge_out | |
| class AVReader(object): | |
| """Individual audio video reader with convenient indexing function. | |
| Parameters | |
| ---------- | |
| uri: str | |
| Path of file. | |
| ctx: decord.Context | |
| The context to decode the file, can be decord.cpu() or decord.gpu(). | |
| sample_rate: int, default is -1 | |
| Desired output sample rate of the audio, unchanged if `-1` is specified. | |
| mono: bool, default is True | |
| Desired output channel layout of the audio. `True` is mono layout. `False` is unchanged. | |
| width : int, default is -1 | |
| Desired output width of the video, unchanged if `-1` is specified. | |
| height : int, default is -1 | |
| Desired output height of the video, unchanged if `-1` is specified. | |
| num_threads : int, default is 0 | |
| Number of decoding thread, auto if `0` is specified. | |
| fault_tol : int, default is -1 | |
| The threshold of corupted and recovered frames. This is to prevent silent fault | |
| tolerance when for example 50% frames of a video cannot be decoded and duplicate | |
| frames are returned. You may find the fault tolerant feature sweet in many cases, | |
| but not for training models. Say `N = # recovered frames` | |
| If `fault_tol` < 0, nothing will happen. | |
| If 0 < `fault_tol` < 1.0, if N > `fault_tol * len(video)`, raise `DECORDLimitReachedError`. | |
| If 1 < `fault_tol`, if N > `fault_tol`, raise `DECORDLimitReachedError`. | |
| """ | |
| def __init__( | |
| self, uri, ctx=cpu(0), sample_rate=44100, mono=True, width=-1, height=-1, num_threads=0, fault_tol=-1 | |
| ): | |
| self.__audio_reader = AudioReader(uri, ctx, sample_rate, mono) | |
| self.__audio_reader.add_padding() | |
| if hasattr(uri, "read"): | |
| uri.seek(0) | |
| self.__video_reader = VideoReader(uri, ctx, width, height, num_threads, fault_tol) | |
| self.__video_reader.seek(0) | |
| def __len__(self): | |
| """Get length of the video. Note that sometimes FFMPEG reports inaccurate number of frames, | |
| we always follow what FFMPEG reports. | |
| Returns | |
| ------- | |
| int | |
| The number of frames in the video file. | |
| """ | |
| return len(self.__video_reader) | |
| def __getitem__(self, idx): | |
| """Get audio samples and video frame at `idx`. | |
| Parameters | |
| ---------- | |
| idx : int or slice | |
| The frame index, can be negative which means it will index backwards, | |
| or slice of frame indices. | |
| Returns | |
| ------- | |
| (ndarray/list of ndarray, ndarray) | |
| First element is samples of shape CxS or a list of length N containing samples of shape CxS, | |
| where N is the number of frames, C is the number of channels, | |
| S is the number of samples of the corresponding frame. | |
| Second element is Frame of shape HxWx3 or batch of image frames with shape NxHxWx3, | |
| where N is the length of the slice. | |
| """ | |
| assert self.__video_reader is not None and self.__audio_reader is not None | |
| if isinstance(idx, slice): | |
| return self.get_batch(range(*idx.indices(len(self.__video_reader)))) | |
| if idx < 0: | |
| idx += len(self.__video_reader) | |
| if idx >= len(self.__video_reader) or idx < 0: | |
| raise IndexError("Index: {} out of bound: {}".format(idx, len(self.__video_reader))) | |
| audio_start_idx, audio_end_idx = self.__video_reader.get_frame_timestamp(idx) | |
| audio_start_idx = self.__audio_reader._time_to_sample(audio_start_idx) | |
| audio_end_idx = self.__audio_reader._time_to_sample(audio_end_idx) | |
| results = (self.__audio_reader[audio_start_idx:audio_end_idx], self.__video_reader[idx]) | |
| self.__video_reader.seek(0) | |
| return results | |
| def get_batch(self, indices): | |
| """Get entire batch of audio samples and video frames. | |
| Parameters | |
| ---------- | |
| indices : list of integers | |
| A list of frame indices. If negative indices detected, the indices will be indexed from backward | |
| Returns | |
| ------- | |
| (list of ndarray, ndarray) | |
| First element is a list of length N containing samples of shape CxS, | |
| where N is the number of frames, C is the number of channels, | |
| S is the number of samples of the corresponding frame. | |
| Second element is Frame of shape HxWx3 or batch of image frames with shape NxHxWx3, | |
| where N is the length of the slice. | |
| """ | |
| assert self.__video_reader is not None and self.__audio_reader is not None | |
| indices = self._validate_indices(indices) | |
| audio_arr = [] | |
| prev_video_idx = None | |
| prev_audio_end_idx = None | |
| for idx in list(indices): | |
| frame_start_time, frame_end_time = self.__video_reader.get_frame_timestamp(idx) | |
| # timestamp and sample conversion could have some error that could cause non-continuous audio | |
| # we detect if retrieving continuous frame and make the audio continuous | |
| if prev_video_idx and idx == prev_video_idx + 1: | |
| audio_start_idx = prev_audio_end_idx | |
| else: | |
| audio_start_idx = self.__audio_reader._time_to_sample(frame_start_time) | |
| audio_end_idx = self.__audio_reader._time_to_sample(frame_end_time) | |
| audio_arr.append(self.__audio_reader[audio_start_idx:audio_end_idx]) | |
| prev_video_idx = idx | |
| prev_audio_end_idx = audio_end_idx | |
| results = (audio_arr, self.__video_reader.get_batch(indices)) | |
| self.__video_reader.seek(0) | |
| return results | |
| def _get_slice(self, sl): | |
| audio_arr = np.empty(shape=(self.__audio_reader.shape()[0], 0), dtype="float32") | |
| for idx in list(sl): | |
| audio_start_idx, audio_end_idx = self.__video_reader.get_frame_timestamp(idx) | |
| audio_start_idx = self.__audio_reader._time_to_sample(audio_start_idx) | |
| audio_end_idx = self.__audio_reader._time_to_sample(audio_end_idx) | |
| audio_arr = np.concatenate( | |
| (audio_arr, self.__audio_reader[audio_start_idx:audio_end_idx].asnumpy()), axis=1 | |
| ) | |
| results = (bridge_out(_nd.array(audio_arr)), self.__video_reader.get_batch(sl)) | |
| self.__video_reader.seek(0) | |
| return results | |
| def _validate_indices(self, indices): | |
| """Validate int64 integers and convert negative integers to positive by backward search""" | |
| assert self.__video_reader is not None and self.__audio_reader is not None | |
| indices = np.array(indices, dtype=np.int64) | |
| # process negative indices | |
| indices[indices < 0] += len(self.__video_reader) | |
| if not (indices >= 0).all(): | |
| raise IndexError("Invalid negative indices: {}".format(indices[indices < 0] + len(self.__video_reader))) | |
| if not (indices < len(self.__video_reader)).all(): | |
| raise IndexError("Out of bound indices: {}".format(indices[indices >= len(self.__video_reader)])) | |
| return indices | |