|
import copy |
|
import random |
|
from dataclasses import dataclass |
|
|
|
import torch |
|
|
|
from .base import AugmentedDynamicItemDataset, DataPipe |
|
|
|
|
|
@dataclass |
|
class NoiseAugmentation(DataPipe): |
|
noise_proportion: float = 0.0 |
|
input_feat_name: str = "input_feat" |
|
output_feat_name: str = "output_feat" |
|
""" |
|
Args: |
|
noise_proportion (float): for this percentage of the time, Gaussian noise will be applied on all frames during MAM training, set to 0 for no noise |
|
input_feat_name (str): handle for the `takes` (input) |
|
output_feat_name (str): handle for the `provides` (output) |
|
""" |
|
|
|
def apply_noise_on_data(self, input_feat): |
|
with torch.no_grad(): |
|
if self.noise_proportion > 0: |
|
noised_feat = copy.deepcopy(input_feat) |
|
dice = random.random() |
|
if dice < self.noise_proportion: |
|
noise_sampler = torch.distributions.Normal(0, 0.2) |
|
noised_feat += noise_sampler.sample(noised_feat.shape) |
|
noised_feat = noised_feat.to(dtype=torch.float32) |
|
return noised_feat |
|
else: |
|
return input_feat |
|
|
|
def __call__(self, dataset: AugmentedDynamicItemDataset): |
|
dataset.add_dynamic_item( |
|
self.apply_noise_on_data, |
|
takes=self.input_feat_name, |
|
provides=self.output_feat_name, |
|
) |
|
return dataset |
|
|
