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)

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
232
    @classmethod
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
233
    def save_function(cls, 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