utils.py 6.48 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
138
139

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

    def __getitem__(self, index):
        # logger.debug("Getting frame %s from %s", index, self.path)
        if isinstance(index, int):
            idx = self.indices[index]
140
            return self.transform([self.reader[idx]])[0]
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
141
142
143
144
145
146
147
148

        if not (isinstance(index, tuple) and len(index) == self.ndim):
            raise NotImplementedError(f"Indxing like {index} is not supported yet!")

        if all(i == slice(0, 0) for i in index):
            return np.array([], dtype=self.dtype)

        if self.selection_style == "all":
149
            return self.transform(np.asarray(self.reader.load())[index])
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
150
151
152
153
154
155
156
157
158
159
160

        idx = self.indices[index[0]]
        video = []
        for i, frame in enumerate(self.reader):
            if i not in idx:
                continue
            video.append(frame)
            if i == idx[-1]:
                break

        index = (slice(len(video)),) + index[1:]
161
        return self.transform(np.asarray(video)[index])
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
162
163
164
165
166
167
168
169
170
171
172

    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
173
174
175
176
177
178
179
180
181
182
183
184
    @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
185
186
187
188
    def __len__(self):
        return len(self.data)

    def __getitem__(self, item):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
189
190
191
192
        # 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
193
194
195
196
197
        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
198
199
200
    def save(self, file):
        self.save_function(self, file)

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
201
    @classmethod
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
202
    def save_function(cls, other, file):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
203
204
205
206
207
208
        with h5py.File(file, mode="w") as f:
            f["data"] = other.data
            f["indices"] = other.indices

    @classmethod
    def load(cls, file):
209
210
211
212
213
        # 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
214
215
        self = cls(data=data, indices=indices)
        return self