latentsync/utils/av_reader.py

# 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