utils.py 8.25 KB
Newer Older
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
1 2
import logging

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
3 4
from bob.bio.base import selected_indices
from bob.io.video import reader
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
5 6 7 8 9 10
import h5py
import numpy as np

logger = logging.getLogger(__name__)


11 12 13
def select_frames(
    count, max_number_of_frames=None, selection_style=None, step_size=None
):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
    """Returns indices of the frames to be selected given the parameters.

    Different selection styles are supported:

    * first : The first frames are selected
    * spread : Frames are selected to be taken from the whole video with equal spaces in
      between.
    * step : Frames are selected every ``step_size`` indices, starting at
      ``step_size/2`` **Think twice if you want to have that when giving FrameContainer
      data!**
    * all : All frames are selected unconditionally.

    Parameters
    ----------
    count : int
        Total number of frames that are available
    max_number_of_frames : int
        The maximum number of frames to be selected. Ignored when selection_style is
        "all".
    selection_style : str
        One of (``first``, ``spread``, ``step``, ``all``). See above.
    step_size : int
        Only useful when ``selection_style`` is ``step``.

    Returns
    -------
    range
        A range of frames to be selected.

    Raises
    ------
    ValueError
        If ``selection_style`` is not one of the supported ones.
    """
    # default values
    if max_number_of_frames is None:
        max_number_of_frames = 20
    if selection_style is None:
        selection_style = "spread"
    if step_size is None:
        step_size = 10

    if selection_style == "first":
        # get the first frames (limited by all frames)
        indices = range(0, min(count, max_number_of_frames))
    elif selection_style == "spread":
        # get frames lineraly spread over all frames
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
61
        indices = selected_indices(count, max_number_of_frames)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
62 63 64 65 66 67 68 69 70 71 72
    elif selection_style == "step":
        indices = range(step_size // 2, count, step_size)[:max_number_of_frames]
    elif selection_style == "all":
        indices = range(0, count)
    else:
        raise ValueError(f"Invalid selection style: {selection_style}")

    return indices


class VideoAsArray:
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
73
    """A memory efficient class to load only select video frames.
74 75
    It also supports efficient conversion to dask arrays.
    """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
76

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
77 78 79 80 81 82
    def __init__(
        self,
        path,
        selection_style=None,
        max_number_of_frames=None,
        step_size=None,
83
        transform=None,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
84 85
        **kwargs,
    ):
86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
        """init

        Parameters
        ----------
        path : str
            Path to the video file
        selection_style : str, optional
            See :any:`select_frames`, by default None
        max_number_of_frames : int, optional
            See :any:`select_frames`, by default None
        step_size : int, optional
            See :any:`select_frames`, by default None
        transform : callable, optional
            A function that transforms the loaded video. This function should
            not change the video shape or its dtype. For example, you may flip
            the frames horizontally using this function, by default None
        """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
103 104
        super().__init__(**kwargs)
        self.path = path
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
105
        self.reader = reader(self.path)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
106 107 108 109 110 111 112 113 114 115 116
        self.dtype, shape = self.reader.video_type[:2]
        self.ndim = len(shape)
        self.selection_style = selection_style
        indices = select_frames(
            count=self.reader.number_of_frames,
            max_number_of_frames=max_number_of_frames,
            selection_style=selection_style,
            step_size=step_size,
        )
        self.indices = indices
        self.shape = (len(indices),) + shape[1:]
117
        if transform is None:
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
118

119 120
            def transform(x):
                return x
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
121

122
        self.transform = transform
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
123 124 125 126 127 128 129 130

    def __getstate__(self):
        d = self.__dict__.copy()
        d.pop("reader")
        return d

    def __setstate__(self, state):
        self.__dict__.update(state)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
131
        self.reader = reader(self.path)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
132 133 134 135 136 137

    def __len__(self):
        return self.shape[0]

    def __getitem__(self, index):
        # logger.debug("Getting frame %s from %s", index, self.path)
138 139 140 141 142 143 144 145 146 147

        # In this method, someone is requesting indices thinking this video has
        # the shape of self.shape but self.shape is determined through
        # select_frames parameters. What we want to do here is to translate
        # ``index`` to real indices of the video file given that we want to load
        # only the selected frames. List of the selected frames are stored in
        # self.indices

        # If only one frame is requested, first translate the index to the real
        # frame number in the video file and load that
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
148 149
        if isinstance(index, int):
            idx = self.indices[index]
150
            return self.transform([self.reader[idx]])[0]
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
151

152 153 154 155 156
        if not (
            isinstance(index, tuple)
            and len(index) == self.ndim
            and all(isinstance(idx, slice) for idx in index)
        ):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
157 158
            raise NotImplementedError(f"Indxing like {index} is not supported yet!")

159
        # dask.array.from_array sometimes requests empty arrays
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
160 161 162
        if all(i == slice(0, 0) for i in index):
            return np.array([], dtype=self.dtype)

163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192
        def _frames_generator():
            # read the frames one by one and yield them
            real_frame_numbers = self.indices[index[0]]
            for i, frame in enumerate(self.reader):
                if i not in real_frame_numbers:
                    continue
                # make sure arrays are loaded in C order because we reshape them
                # by C order later. Also, index into the frames here
                frame = np.ascontiguousarray(frame)[index[1:]]
                # return a tuple of flat array to match what is expected by
                # field_dtype
                yield (frame.ravel(),)
                if i == real_frame_numbers[-1]:
                    break

        iterable = _frames_generator()
        # compute the final shape given self.shape and index
        # see https://stackoverflow.com/a/36188683/1286165
        shape = [len(range(*idx.indices(dim))) for idx, dim in zip(index, self.shape)]
        # field_dtype contains information about dtype and shape of each frame
        # numpy black magic: https://stackoverflow.com/a/12473478/1286165 allows
        # us to yield frame by frame in _frames_generator which greatly speeds
        # up loading the video
        field_dtype = [("", (self.dtype, (np.prod(shape[1:]),)))]
        total_number_of_frames = shape[0]
        video = np.fromiter(iterable, field_dtype, total_number_of_frames)
        # view the array as self.dtype to remove the field_dtype
        video = np.reshape(video.view(self.dtype), shape, order="C")

        return self.transform(video)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
193 194 195 196 197 198 199 200 201 202 203

    def __repr__(self):
        return f"{self.reader!r} {self.dtype!r} {self.ndim!r} {self.shape!r} {self.indices!r}"


class VideoLikeContainer:
    def __init__(self, data, indices, **kwargs):
        super().__init__(**kwargs)
        self.data = data
        self.indices = indices

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
204 205 206 207 208 209 210 211 212 213 214 215
    @property
    def dtype(self):
        return self.data.dtype

    @property
    def shape(self):
        return self.data.shape

    @property
    def ndim(self):
        return self.data.ndim

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
216 217 218 219
    def __len__(self):
        return len(self.data)

    def __getitem__(self, item):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
220 221 222 223
        # we need to throw IndexErrors here because h5py throws ValueErrors
        # instead and this breaks loops on this class
        if isinstance(item, int) and item >= len(self):
            raise IndexError(f"Index ({item}) out of range (0-{len(self)-1})")
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
224 225 226 227 228
        return self.data[item]

    def __array__(self, dtype=None, *args, **kwargs):
        return np.asarray(self.data, dtype, *args, **kwargs)

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
229 230 231
    def save(self, file):
        self.save_function(self, file)

232 233
    @staticmethod
    def save_function(other, file):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
234 235 236 237 238 239
        with h5py.File(file, mode="w") as f:
            f["data"] = other.data
            f["indices"] = other.indices

    @classmethod
    def load(cls, file):
240 241 242 243 244
        # weak closing of the hdf5 file so we don't load all the data into
        # memory https://docs.h5py.org/en/stable/high/file.html#closing-files
        f = h5py.File(file, mode="r")
        data = f["data"]
        indices = f["indices"]
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
245 246
        self = cls(data=data, indices=indices)
        return self