dlc2action.transformer.kinematic

  1#
  2# Copyright 2020-2022 by A. Mathis Group and contributors. All rights reserved.
  3#
  4# This project and all its files are licensed under GNU AGPLv3 or later version. A copy is included in dlc2action/LICENSE.AGPL.
  5#
  6"""
  7Kinematic transformer
  8"""
  9
 10from typing import Dict, Tuple, List, Set, Iterable
 11import torch
 12import numpy as np
 13from copy import copy, deepcopy
 14from random import getrandbits
 15from collections import defaultdict
 16from dlc2action.transformer.base_transformer import Transformer
 17from dlc2action.utils import rotation_matrix_2d, rotation_matrix_3d
 18
 19
 20class KinematicTransformer(Transformer):
 21    """
 22    A transformer that augments the output of the Kinematic feature extractor
 23
 24    The available augmentations are `'rotate'`, `'mirror'`, `'shift'`, `'add_noise'` and `'zoom'`
 25    """
 26
 27    def __init__(
 28        self,
 29        model_name: str,
 30        augmentations: List = None,
 31        use_default_augmentations: bool = False,
 32        rotation_limits: List = None,
 33        mirror_dim: Set = None,
 34        noise_std: float = 0.05,
 35        zoom_limits: List = None,
 36        masking_probability: float = 0.05,
 37        dim: int = 2,
 38        graph_features: bool = False,
 39        bodyparts_order: List = None,
 40        canvas_shape: List = None,
 41        move_around_image_center: bool = True,
 42        **kwargs,
 43    ) -> None:
 44        """
 45        Parameters
 46        ----------
 47        augmentations : list, optional
 48            list of augmentation names to use ("rotate", "mirror", "shift", "add_noise", "zoom")
 49        use_default_augmentations : bool, default False
 50            if `True` and augmentations are not passed, default augmentations will be applied; otherwise no augmentations
 51        rotation_limits : list, default [-pi/2, pi/2]
 52            list of float rotation angle limits (`[low, high]``, or `[[low_x, high_x], [low_y, high_y], [low_z, high_z]]`
 53            for 3D data)
 54        mirror_dim : set, default {0}
 55            set of integer indices of dimensions that can be mirrored
 56        noise_std : float, default 0.05
 57            standard deviation of noise
 58        zoom_limits : list, default [0.5, 1.5]
 59            list of float zoom limits ([low, high])
 60        masking_probability : float, default 0.1
 61            the probability of masking a joint
 62        dim : int, default 2
 63            the dimensionality of the input data
 64        **kwargs : dict
 65            other parameters for the base transformer class
 66        """
 67
 68        if augmentations is None:
 69            augmentations = []
 70        if canvas_shape is None:
 71            canvas_shape = [1, 1]
 72        self.dim = int(dim)
 73
 74        self.offset = [0 for _ in range(self.dim)]
 75        self.scale = canvas_shape[1] / canvas_shape[0]
 76        self.image_center = move_around_image_center
 77        # if canvas_shape is None:
 78        #     self.offset = [0.5 for _ in range(self.dim)]
 79        # else:
 80        #     self.offset = [0.5 * canvas_shape[i] / canvas_shape[0] for i in range(self.dim)]
 81
 82        self.blank = 0  # the value that nan values are set to (shouldn't be changed in augmentations)
 83        super().__init__(
 84            model_name,
 85            augmentations,
 86            use_default_augmentations,
 87            graph_features=graph_features,
 88            bodyparts_order=bodyparts_order,
 89            **kwargs,
 90        )
 91        if rotation_limits is None:
 92            rotation_limits = [-np.pi / 2, np.pi / 2]
 93        if mirror_dim is None:
 94            mirror_dim = [0]
 95        if zoom_limits is None:
 96            zoom_limits = [0.5, 1.5]
 97        self.rotation_limits = rotation_limits
 98        self.n_bodyparts = None
 99        self.mirror_dim = mirror_dim
100        self.noise_std = noise_std
101        self.zoom_limits = zoom_limits
102        self.masking_probability = masking_probability
103
104    def _apply_augmentations(
105        self,
106        main_input: Dict,
107        ssl_inputs: List = None,
108        ssl_targets: List = None,
109        augment: bool = False,
110        ssl: bool = False,
111    ) -> Tuple:
112
113        if ssl_targets is None:
114            ssl_targets = [None]
115        if ssl_inputs is None:
116            ssl_inputs = [None]
117
118        keys_main = self._get_keys(
119            (
120                "coords",
121                "speed_joints",
122                "speed_direction",
123                "speed_bones",
124                "acc_bones",
125                "bones",
126                "coord_diff",
127            ),
128            main_input,
129        )
130        if len(keys_main) > 0:
131            key = keys_main[0]
132            final_shape = main_input[key].shape
133            # self._get_bodyparts(final_shape)
134            batch = final_shape[0]
135            s = main_input[key].reshape((batch, -1, self.dim, final_shape[-1]))
136            x_shape = s.shape
137            # print(f'{x_shape=}, {self.dim=}')
138            self.n_bodyparts = x_shape[1]
139            if self.dim == 3:
140                if len(self.rotation_limits) == 2:
141                    self.rotation_limits = [[0, 0], [0, 0], self.rotation_limits]
142            dicts = [main_input] + ssl_inputs + ssl_targets
143            for x in dicts:
144                if x is not None:
145                    keys = self._get_keys(
146                        (
147                            "coords",
148                            "speed_joints",
149                            "speed_direction",
150                            "speed_bones",
151                            "acc_bones",
152                            "bones",
153                            "coord_diff",
154                        ),
155                        x,
156                    )
157                    for key in keys:
158                        x[key] = x[key].reshape(x_shape)
159        if len(self._get_keys(("center"), main_input)) > 0:
160            dicts = [main_input] + ssl_inputs + ssl_targets
161            for x in dicts:
162                if x is not None:
163                    key_bases = ["center"]
164                    keys = self._get_keys(
165                        key_bases,
166                        x,
167                    )
168                    for key in keys:
169                        x[key] = x[key].reshape(
170                            (x[key].shape[0], 1, -1, x[key].shape[-1])
171                        )
172        main_input, ssl_inputs, ssl_targets = super()._apply_augmentations(
173            main_input, ssl_inputs, ssl_targets, augment
174        )
175        if len(keys_main) > 0:
176            dicts = [main_input] + ssl_inputs + ssl_targets
177            for x in dicts:
178                if x is not None:
179                    keys = self._get_keys(
180                        (
181                            "coords",
182                            "speed_joints",
183                            "speed_direction",
184                            "speed_bones",
185                            "acc_bones",
186                            "bones",
187                            "coord_diff",
188                        ),
189                        x,
190                    )
191                    for key in keys:
192                        x[key] = x[key].reshape(final_shape)
193        if len(self._get_keys(("center"), main_input)) > 0:
194            dicts = [main_input] + ssl_inputs + ssl_targets
195            for x in dicts:
196                if x is not None:
197                    key_bases = ["center"]
198                    keys = self._get_keys(
199                        key_bases,
200                        x,
201                    )
202                    for key in keys:
203                        x[key] = x[key].reshape((x[key].shape[0], -1, x[key].shape[-1]))
204        return main_input, ssl_inputs, ssl_targets
205
206    def _rotate(
207        self, main_input: Dict, ssl_inputs: List, ssl_targets: List
208    ) -> Tuple[Dict, List, List]:
209        """
210        Rotate the "coords" and "speed_joints" features of the input to a random degree
211        """
212
213        keys = self._get_keys(
214            (
215                "coords",
216                "bones",
217                "coord_diff",
218                "center",
219                "speed_joints",
220                "speed_direction",
221                "center",
222            ),
223            main_input,
224        )
225        if len(keys) == 0:
226            return main_input, ssl_inputs, ssl_targets
227        batch = main_input[keys[0]].shape[0]
228        if self.dim == 2:
229            angles = torch.FloatTensor(batch).uniform_(*self.rotation_limits)
230            R = rotation_matrix_2d(angles).to(main_input[keys[0]].device)
231        else:
232            angles_x = torch.FloatTensor(batch).uniform_(*self.rotation_limits[0])
233            angles_y = torch.FloatTensor(batch).uniform_(*self.rotation_limits[1])
234            angles_z = torch.FloatTensor(batch).uniform_(*self.rotation_limits[2])
235            R = rotation_matrix_3d(angles_x, angles_y, angles_z).to(
236                main_input[keys[0]].device
237            )
238        dicts = [main_input] + ssl_inputs + ssl_targets
239        for x in dicts:
240            if x is not None:
241                keys = self._get_keys(
242                    (
243                        "coords",
244                        "speed_joints",
245                        "speed_direction",
246                        "speed_bones",
247                        "acc_bones",
248                        "bones",
249                        "coord_diff",
250                        "center",
251                    ),
252                    x,
253                )
254                for key in keys:
255                    if key in x:
256                        mask = x[key] == self.blank
257                        if (
258                            any([key.startswith(x) for x in ["coords", "center"]])
259                            and not self.image_center
260                        ):
261                            offset = x[key].mean(1).unsqueeze(1)
262                        else:
263                            offset = 0
264                        x[key] = (
265                            torch.einsum(
266                                "abjd,aij->abid",
267                                x[key] - offset,
268                                R,
269                            )
270                            + offset
271                        )
272                        x[key][mask] = self.blank
273        return main_input, ssl_inputs, ssl_targets
274
275    def _mirror(
276        self, main_input: Dict, ssl_inputs: List, ssl_targets: List
277    ) -> Tuple[Dict, List, List]:
278        """
279        Mirror the "coords" and "speed_joints" features of the input randomly
280        """
281
282        mirror = []
283        for i in range(3):
284            if i in self.mirror_dim and bool(getrandbits(1)):
285                mirror.append(i)
286        if len(mirror) == 0:
287            return main_input, ssl_inputs, ssl_targets
288        dicts = [main_input] + ssl_inputs + ssl_targets
289        for x in dicts:
290            if x is not None:
291                keys = self._get_keys(
292                    (
293                        "coords",
294                        "speed_joints",
295                        "speed_direction",
296                        "speed_bones",
297                        "acc_bones",
298                        "bones",
299                        "center",
300                        "coord_diff",
301                    ),
302                    x,
303                )
304                for key in keys:
305                    if key in x:
306                        mask = x[key] == self.blank
307                        y = deepcopy(x[key])
308                        if not self.image_center and any(
309                            [key.startswith(x) for x in ["coords", "center"]]
310                        ):
311                            mean = y.mean(1).unsqueeze(1)
312                            for dim in mirror:
313                                y[:, :, dim, :] = (
314                                    2 * mean[:, :, dim, :] - y[:, :, dim, :]
315                                )
316                        else:
317                            for dim in mirror:
318                                y[:, :, dim, :] *= -1
319                        x[key] = y
320                        x[key][mask] = self.blank
321        return main_input, ssl_inputs, ssl_targets
322
323    def _shift(
324        self, main_input: Dict, ssl_inputs: List, ssl_targets: List
325    ) -> Tuple[Dict, List, List]:
326        """
327        Shift the "coords" features of the input randomly
328        """
329
330        keys = self._get_keys(("coords", "center"), main_input)
331        if len(keys) == 0:
332            return main_input, ssl_inputs, ssl_targets
333        batch = main_input[keys[0]].shape[0]
334        dim = main_input[keys[0]].shape[2]
335        device = main_input[keys[0]].device
336        coords = torch.cat(
337            [main_input[key].transpose(-1, -2).reshape(batch, -1, dim) for key in keys],
338            dim=1,
339        )
340        minimask = coords[:, :, 0] != self.blank
341        min_x = -0.5 - torch.min(coords[:, :, 0], dim=-1)[0]
342        min_y = -0.5 * self.scale - torch.min(coords[:, :, 1], dim=-1)[0]
343        max_x = 0.5 - torch.max(coords[:, :, 0][minimask], dim=-1)[0]
344        max_y = 0.5 * self.scale - torch.max(coords[:, :, 1][minimask], dim=-1)[0]
345        del coords
346        shift_x = min_x + torch.FloatTensor(batch).uniform_().to(device) * (
347            max_x - min_x
348        )
349        shift_y = min_y + torch.FloatTensor(batch).uniform_().to(device) * (
350            max_y - min_y
351        )
352        shift_x = shift_x.unsqueeze(-1).unsqueeze(-1)
353        shift_y = shift_y.unsqueeze(-1).unsqueeze(-1)
354        dicts = [main_input] + ssl_inputs + ssl_targets
355        for x in dicts:
356            if x is not None:
357                keys = self._get_keys(("coords", "center"), x)
358                for key in keys:
359                    y = deepcopy(x[key])
360                    mask = y == self.blank
361                    y[:, :, 0, :] += shift_x
362                    y[:, :, 1, :] += shift_y
363                    x[key] = y
364                    x[key][mask] = self.blank
365        return main_input, ssl_inputs, ssl_targets
366
367    def _add_noise(
368        self, main_input: Dict, ssl_inputs: List, ssl_targets: List
369    ) -> Tuple[Dict, List, List]:
370        """
371        Add normal noise to all features of the input
372        """
373
374        dicts = [main_input] + ssl_inputs + ssl_targets
375        for x in dicts:
376            if x is not None:
377                keys = self._get_keys(
378                    (
379                        "coords",
380                        "speed_joints",
381                        "speed_direction",
382                        "intra_distance",
383                        "acc_joints",
384                        "angle_speeds",
385                        "inter_distance",
386                        "speed_bones",
387                        "acc_bones",
388                        "bones",
389                        "coord_diff",
390                        "center",
391                    ),
392                    x,
393                )
394                for key in keys:
395                    mask = x[key] == self.blank
396                    x[key] = x[key] + torch.FloatTensor(x[key].shape).normal_(
397                        std=self.noise_std
398                    ).to(x[key].device)
399                    x[key][mask] = self.blank
400        return main_input, ssl_inputs, ssl_targets
401
402    def _zoom(
403        self, main_input: Dict, ssl_inputs: List, ssl_targets: List
404    ) -> Tuple[Dict, List, List]:
405        """
406        Add random zoom to all features of the input
407        """
408
409        key = list(main_input.keys())[0]
410        batch = main_input[key].shape[0]
411        device = main_input[key].device
412        zoom = torch.FloatTensor(batch).uniform_(*self.zoom_limits).to(device)
413        dicts = [main_input] + ssl_inputs + ssl_targets
414        for x in dicts:
415            if x is not None:
416                keys = self._get_keys(
417                    (
418                        "speed_joints",
419                        "intra_distance",
420                        "acc_joints",
421                        "inter_distance",
422                        "speed_bones",
423                        "acc_bones",
424                        "bones",
425                        "coord_diff",
426                        "center",
427                        "speed_value",
428                    ),
429                    x,
430                )
431                for key in keys:
432                    mask = x[key] == self.blank
433                    shape = np.array(x[key].shape)
434                    shape[1:] = 1
435                    y = deepcopy(x[key])
436                    y *= zoom.reshape(list(shape))
437                    x[key] = y
438                    x[key][mask] = self.blank
439                keys = self._get_keys("coords", x)
440                for key in keys:
441                    mask = x[key] == self.blank
442                    shape = np.array(x[key].shape)
443                    shape[1:] = 1
444                    zoom = zoom.reshape(list(shape))
445                    x[key][mask] = 10
446                    min_x = x[key][:, :, 0, :].min(1)[0].unsqueeze(1)
447                    min_y = x[key][:, :, 1, :].min(1)[0].unsqueeze(1)
448                    center = torch.stack([min_x, min_y], dim=2)
449                    coords = (x[key] - center) * zoom + center
450                    x[key] = coords
451                    x[key][mask] = self.blank
452
453        return main_input, ssl_inputs, ssl_targets
454
455    def _mask_joints(
456        self, main_input: Dict, ssl_inputs: List, ssl_targets: List
457    ) -> Tuple[Dict, List, List]:
458        """
459        Mask joints randomly
460        """
461
462        key = list(main_input.keys())[0]
463        batch, *_, frames = main_input[key].shape
464        masked_joints = (
465            torch.FloatTensor(batch, self.n_bodyparts).uniform_()
466            < self.masking_probability
467        )
468        dicts = [main_input] + ssl_inputs + ssl_targets
469        for x in [y for y in dicts if y is not None]:
470            keys = self._get_keys(("intra_distance", "inter_distance"), x)
471            for key in keys:
472                mask = (
473                    masked_joints.repeat(self.n_bodyparts, frames, 1, 1)
474                    .transpose(0, 2)
475                    .transpose(1, 3)
476                )
477                indices = torch.triu_indices(self.n_bodyparts, self.n_bodyparts, 1)
478
479                X = torch.zeros((batch, self.n_bodyparts, self.n_bodyparts, frames)).to(
480                    x[key].device
481                )
482                X[:, indices[0], indices[1], :] = x[key]
483                X[mask] = self.blank
484                X[mask.transpose(1, 2)] = self.blank
485                x[key] = X[:, indices[0], indices[1], :].reshape(batch, -1, frames)
486            keys = self._get_keys(
487                (
488                    "speed_joints",
489                    "speed_direction",
490                    "coords",
491                    "acc_joints",
492                    "speed_bones",
493                    "acc_bones",
494                    "bones",
495                    "coord_diff",
496                ),
497                x,
498            )
499            for key in keys:
500                mask = (
501                    masked_joints.repeat(self.dim, frames, 1, 1)
502                    .transpose(0, 2)
503                    .transpose(1, 3)
504                )
505                x[key][mask] = (
506                    x[key].mean(1).unsqueeze(1).repeat(1, x[key].shape[1], 1, 1)[mask]
507                )
508            keys = self._get_keys("angle_speeds", x)
509            for key in keys:
510                mask = (
511                    masked_joints.repeat(frames, 1, 1).transpose(0, 1).transpose(1, 2)
512                )
513                x[key][mask] = (
514                    x[key].mean(1).unsqueeze(1).repeat(1, x[key].shape[1], 1)[mask]
515                )
516
517        return main_input, ssl_inputs, ssl_targets
518
519    def _switch(
520        self, main_input: Dict, ssl_inputs: List, ssl_targets: List
521    ) -> Tuple[Dict, List, List]:
522        if bool(getrandbits(1)):
523            return main_input, ssl_inputs, ssl_targets
524        individuals = set()
525        ind_dict = defaultdict(lambda: set())
526        for key in main_input:
527            if len(key.split("---")) != 2:
528                continue
529            ind = key.split("---")[1]
530            if "+" in ind:
531                continue
532            individuals.add(ind)
533            ind_dict[ind].add(key.split("---")[0])
534        individuals = list(individuals)
535        if len(individuals) < 2:
536            return main_input, ssl_inputs, ssl_targets
537        for x, y in zip(individuals[::2], individuals[1::2]):
538            for key in ind_dict[x]:
539                if key in ind_dict[y]:
540                    main_input[f"{key}---{x}"], main_input[f"{key}---{y}"] = (
541                        main_input[f"{key}---{y}"],
542                        main_input[f"{key}---{x}"],
543                    )
544                for d in ssl_inputs + ssl_targets:
545                    if d is None:
546                        continue
547                    if f"{key}---{x}" in d and f"{key}---{y}" in d:
548                        d[f"{key}---{x}"], d[f"{key}---{y}"] = (
549                            d[f"{key}---{y}"],
550                            d[f"{key}---{x}"],
551                        )
552        return main_input, ssl_inputs, ssl_targets
553
554    def _get_keys(self, key_bases: Iterable, x: Dict) -> List:
555        """
556        Get the keys of x that start with one of the strings from key_bases
557        """
558
559        keys = []
560        if isinstance(key_bases, str):
561            key_bases = [key_bases]
562        for key in x:
563            if any([x == key.split("---")[0] for x in key_bases]):
564                keys.append(key)
565        return keys
566
567    def _augmentations_dict(self) -> Dict:
568        """
569        Get the mapping from augmentation names to functions
570        """
571
572        return {
573            "mirror": self._mirror,
574            "shift": self._shift,
575            "add_noise": self._add_noise,
576            "zoom": self._zoom,
577            "rotate": self._rotate,
578            "mask": self._mask_joints,
579            "switch": self._switch,
580        }
581
582    def _default_augmentations(self) -> List:
583        """
584        Get the list of default augmentation names
585        """
586
587        return ["mirror", "shift", "add_noise"]
588
589    def _get_bodyparts(self, shape: Tuple) -> None:
590        """
591        Set the number of bodyparts from the data if it is not known
592        """
593
594        if self.n_bodyparts is None:
595            N, B, F = shape
596            self.n_bodyparts = B // 2