def __call__()

in empose/data/noise_functions.py [0:0]

• 31 lines of code
• 4 McCabe index (conditional complexity)

    def __call__(self, batch: ABatch, **kwargs):
        if self.max_r <= 0.0:
            return batch
        if batch.marker_pos_synth is None:
            return batch

        markers = batch.marker_pos_synth
        batch_size, seq_len, n_markers = markers.shape[0], markers.shape[1], markers.shape[-1] // 3
        ms = markers.reshape(batch_size, seq_len, n_markers, 3)

        # Select `num_marker` many markers at random. For simplicity each batch entry is treated the same.
        m_ids = torch.randperm(n_markers, generator=self.rng)[:self.num_markers].to(device=markers.device)

        # Select perturbation window at random.
        window_len = int(self.ws * seq_len)
        sf = torch.randint(0, seq_len - window_len + 1, (batch_size,), generator=self.rng).to(device=markers.device)
        ef = sf + window_len

        # Choose perturbation radius at random.
        thigh_len = torch.norm(ms[0, seq_len // 2, C.T_TO_IDX_WO_ROOT[C.T_RUL]] - ms[0, 0, C.T_TO_IDX_WO_ROOT[C.T_RLL]])
        r = torch.rand(batch_size, window_len, self.num_markers).to(device=markers.device) * self.max_r * thigh_len / 2

        # Choose spherical coordinates at random.
        thetas = torch.rand((batch_size, window_len, self.num_markers), generator=self.rng).to(
            device=markers.device) * np.pi * 2
        phis = torch.rand((batch_size, window_len, self.num_markers), generator=self.rng).to(
            device=markers.device) * np.pi

        # Create displacement vector.
        xs = r * torch.cos(thetas) * torch.sin(phis)
        ys = r * torch.sin(thetas) * torch.cos(phis)
        zs = r * torch.cos(phis)

        xs = xs.to(dtype=ms.dtype, device=ms.device)
        ys = ys.to(dtype=ms.dtype, device=ms.device)
        zs = zs.to(dtype=ms.dtype, device=ms.device)

        # Apply displacement to the markers.
        # For now we're looping since this is easier then via advanced indexing.
        ms_noisy = ms.clone()
        for i in range(batch_size):
            ms_noisy[i, sf[i]:ef[i], m_ids, 0] += xs[i]
            ms_noisy[i, sf[i]:ef[i], m_ids, 1] += ys[i]
            ms_noisy[i, sf[i]:ef[i], m_ids, 2] += zs[i]

        batch.marker_pos_noisy = ms_noisy.reshape(batch_size, seq_len, -1)
        return batch