| |
| |
| from abc import ABC, abstractmethod |
| import io |
| import math |
| import numpy as np |
| from typing import Optional, TypeVar, Union |
| import collections |
|
|
| try: |
| import decord |
| except ImportError: |
| _HAS_DECORD = False |
| else: |
| _HAS_DECORD = True |
|
|
| if _HAS_DECORD: |
| decord.bridge.set_bridge('native') |
|
|
| DecordDevice = TypeVar("DecordDevice") |
|
|
| |
| class VideoReaderWrapper(decord.VideoReader): |
| def __init__(self, *args, **kwargs): |
| super().__init__(*args, **kwargs) |
| self.seek(0) |
|
|
| def __getitem__(self, key): |
| frames = super().__getitem__(key) |
| self.seek(0) |
| return frames |
|
|
|
|
| class Video(ABC): |
| """ |
| Video provides an interface to access clips from a video container. |
| """ |
|
|
| @abstractmethod |
| def __init__( |
| self, |
| file: Union[str, io.IOBase], |
| video_name: Optional[str] = None, |
| decode_audio: bool = True, |
| ) -> None: |
| """ |
| Args: |
| file (BinaryIO): a file-like object (e.g. io.BytesIO or io.StringIO) that |
| contains the encoded video. |
| """ |
| pass |
|
|
| @property |
| @abstractmethod |
| def duration(self) -> float: |
| """ |
| Returns: |
| duration of the video in seconds |
| """ |
| pass |
|
|
| @abstractmethod |
| def get_clip( |
| self, start_sec: float, end_sec: float, num_samples: int |
| ): |
| """ |
| Retrieves frames from the internal video at the specified start and end times |
| in seconds (the video always starts at 0 seconds). |
| |
| Args: |
| start_sec (float): the clip start time in seconds |
| end_sec (float): the clip end time in seconds |
| Returns: |
| video_data_dictonary: A dictionary mapping strings to tensor of the clip's |
| underlying data. |
| |
| """ |
| pass |
|
|
| def close(self): |
| pass |
|
|
|
|
| class EncodedVideoDecord(Video): |
| """ |
| |
| Accessing clips from an encoded video using Decord video reading API |
| as the decoding backend. For more details, please refer to - |
| `Decord <https://github.com/dmlc/decord>` |
| """ |
|
|
| def __init__( |
| self, |
| file: Union[str, io.IOBase], |
| video_name: Optional[str] = None, |
| width: int = -1, |
| height: int = -1, |
| num_threads: int = 0, |
| fault_tol: int = -1, |
| ) -> None: |
| """ |
| Args: |
| file str: file path. |
| video_name (str): An optional name assigned to the video. |
| decode_audio (bool): If disabled, audio is not decoded. |
| 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`. |
| """ |
| self._video_name = video_name |
| if not _HAS_DECORD: |
| raise ImportError( |
| "decord is required to use EncodedVideoDecord decoder. Please " |
| "install with 'pip install decord' for CPU-only version and refer to" |
| "'https://github.com/dmlc/decord' for GPU-supported version" |
| ) |
| try: |
| self._av_reader = VideoReaderWrapper( |
| uri=file, |
| ctx=decord.cpu(0), |
| width=width, |
| height=height, |
| num_threads=num_threads, |
| fault_tol=fault_tol, |
| ) |
| except Exception as e: |
| raise RuntimeError(f"Failed to open video {video_name} with Decord. {e}") |
|
|
| self._fps = self._av_reader.get_avg_fps() |
| self._duration = float(len(self._av_reader)) / float(self._fps) |
|
|
| @property |
| def name(self) -> Optional[str]: |
| """ |
| Returns: |
| name: the name of the stored video if set. |
| """ |
| return self._video_name |
|
|
| @property |
| def duration(self) -> float: |
| """ |
| Returns: |
| duration: the video's duration/end-time in seconds. |
| """ |
| return self._duration |
|
|
| def close(self): |
| if self._av_reader is not None: |
| del self._av_reader |
| self._av_reader = None |
|
|
| def get_clip( |
| self, start_sec: float, end_sec: float, num_samples: int |
| ): |
| """ |
| Retrieves frames from the encoded video at the specified start and end times |
| in seconds (the video always starts at 0 seconds). |
| |
| Args: |
| start_sec (float): the clip start time in seconds |
| end_sec (float): the clip end time in seconds |
| Returns: |
| clip_data: |
| A dictionary mapping the entries at "video" and "audio" to a tensors. |
| |
| "video": A tensor of the clip's RGB frames with shape: |
| (channel, time, height, width). The frames are of type torch.float32 and |
| in the range [0 - 255]. |
| |
| Returns None if no video or audio found within time range. |
| |
| """ |
| if start_sec > end_sec or start_sec > self._duration: |
| raise RuntimeError( |
| f"Incorrect time window for Decord decoding for video: {self._video_name}." |
| ) |
|
|
| start_idx = math.ceil(self._fps * start_sec) |
| end_idx = math.ceil(self._fps * end_sec) |
| end_idx = min(end_idx, len(self._av_reader)) |
| |
|
|
| frame_idxs = np.linspace(start_idx, end_idx - 1, num_samples, dtype=int) |
|
|
| try: |
| outputs = self._av_reader.get_batch(frame_idxs) |
| return outputs.asnumpy(), frame_idxs - frame_idxs[0] |
| except Exception as e: |
| print(f"Failed to decode video with Decord: {self._video_name}. {e}") |
| raise e |
|
|
| try: |
| import cv2 |
| except ImportError: |
| print(f"ERR: import cv2 failed, install cv2 by 'pip install opencv-python'") |
|
|
| class EncodedVideoOpencv(): |
| def __init__( |
| self, |
| file: Union[str, io.IOBase], |
| video_name: Optional[str] = None, |
| width: int = -1, |
| height: int = -1, |
| num_threads: int = 0, |
| fault_tol: int = -1, |
| ) -> None: |
| """ |
| Args: |
| file str: file path. |
| video_name (str): An optional name assigned to the video. |
| width : Not support yet. |
| height : Not support yet. |
| num_threads : Not support yet. |
| fault_tol : Not support yet. |
| """ |
|
|
| self._video_name = video_name |
| self.cap = cv2.VideoCapture(file) |
| self._fps = self.cap.get(cv2.CAP_PROP_FPS) |
| self._vlen = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| self._duration = float(self._vlen) / float(self._fps) |
|
|
| @property |
| def name(self) -> Optional[str]: |
| """ |
| Returns: |
| name: the name of the stored video if set. |
| """ |
| return self._video_name |
|
|
| @property |
| def duration(self) -> float: |
| """ |
| Returns: |
| duration: the video's duration/end-time in seconds. |
| """ |
| return self._duration |
|
|
| def __del__(self): |
| self.close() |
|
|
| def close(self): |
| self.cap.release() |
|
|
| def get_clip( |
| self, start_sec: float, end_sec: float, num_samples: int |
| ): |
| if start_sec > end_sec or start_sec > self._duration: |
| raise RuntimeError( |
| f"Incorrect time window for Decord decoding for video: {self._video_name}." |
| ) |
| start_idx = math.ceil(self._fps * start_sec) |
| end_idx = math.ceil(self._fps * end_sec) |
| end_idx = min(end_idx, self._vlen) |
| frame_idxs = np.linspace(start_idx, end_idx - 1, num_samples, dtype=int) |
| frame_idx2freq = collections.defaultdict(int) |
| for frame_idx in frame_idxs: |
| frame_idx2freq[frame_idx] += 1 |
| try: |
| frames = [] |
| for i in range(self._vlen): |
| if i > frame_idxs[-1]: |
| break |
| ret, frame = self.cap.read() |
| if i in frame_idx2freq: |
| frames.extend([frame] * frame_idx2freq[i]) |
| frames = np.array(frames).astype(np.uint8) |
| assert len(frames) == num_samples |
| return frames, frame_idxs - frame_idxs[0] |
| except Exception as e: |
| print(f"Failed to decode video with opencv: {self._video_name}. {e}") |
| raise e |
|
|
