Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn.functional as F | |
| class HandDrawnAugmentation: | |
| def __init__(self, max_drift_scale=0.05, max_wobble_scale=0.01, p=0.8): | |
| """ | |
| :param max_drift_scale: Максимальное плавное искажение | |
| :param max_wobble_scale: Максимальные мелкие неровности | |
| :param p: Вероятность применения аугментации (иногда полезно видеть идеальные данные) | |
| """ | |
| self.max_drift_scale = max_drift_scale | |
| self.max_wobble_scale = max_wobble_scale | |
| self.p = p | |
| def __call__(self, points: torch.Tensor) -> torch.Tensor: | |
| if torch.rand(1).item() > self.p: | |
| return points | |
| y = points[0:1, :].clone() | |
| mask = points[1:2, :].clone() | |
| seq_len = y.size(1) | |
| device = y.device | |
| min_strength = 0.5 | |
| drift_strength = ( | |
| torch.empty(1, device=device).uniform_(min_strength, 1.0).item() | |
| ) | |
| wobble_strength = ( | |
| torch.empty(1, device=device).uniform_(min_strength, 1.0).item() | |
| ) | |
| current_drift = self.max_drift_scale * drift_strength | |
| current_wobble = self.max_wobble_scale * wobble_strength | |
| drift_freq = torch.randint(2, 6, (1,)).item() | |
| wobble_freq = torch.randint(10, 30, (1,)).item() | |
| low_freq_noise = torch.randn(1, 1, 1, drift_freq, device=device) # type: ignore | |
| drift = F.interpolate( | |
| low_freq_noise, size=(1, seq_len), mode="bicubic", align_corners=True | |
| ).view(1, seq_len) | |
| mid_freq_noise = torch.randn(1, 1, 1, wobble_freq, device=device) # type: ignore | |
| wobble = F.interpolate( | |
| mid_freq_noise, size=(1, seq_len), mode="bicubic", align_corners=True | |
| ).view(1, seq_len) | |
| total_distortion = (drift * current_drift) + (wobble * current_wobble) | |
| y = y + total_distortion * mask | |
| x_freq = torch.randint(5, 15, (1,)).item() | |
| x_noise = torch.randn(1, 1, 1, x_freq, device=device) # type: ignore | |
| x_drift = F.interpolate( | |
| x_noise, size=(1, seq_len), mode="bicubic", align_corners=True | |
| ).view(seq_len) | |
| envelope = F.interpolate( | |
| torch.rand(1, 1, 1, 3, device=device), | |
| size=(1, seq_len), | |
| mode="bicubic", | |
| align_corners=True, | |
| ).view(seq_len) | |
| x_drift = x_drift * envelope | |
| x_shift_amplitude = torch.rand(1, device=device).item() * 3.0 + 1.0 | |
| new_indices = ( | |
| torch.arange(seq_len, device=device).float() + x_drift * x_shift_amplitude | |
| ) | |
| new_indices = torch.clamp(new_indices, 0, seq_len - 1) | |
| idx_floor = new_indices.long() | |
| idx_ceil = torch.clamp(idx_floor + 1, 0, seq_len - 1) | |
| weight = new_indices - idx_floor.float() | |
| mask_floor = mask[:, idx_floor] | |
| mask_ceil = mask[:, idx_ceil] | |
| valid_transition = mask_floor * mask_ceil | |
| idx_nearest = torch.round(new_indices).long() | |
| y_nearest = y[:, idx_nearest] | |
| y_interp = y[:, idx_floor] * (1.0 - weight) + y[:, idx_ceil] * weight | |
| y = torch.where(valid_transition == 1.0, y_interp, y_nearest) | |
| mask = mask[:, idx_nearest] | |
| if torch.rand(1).item() > 0.5: | |
| shear_factor = (torch.rand(1, device=device) - 0.5) * 0.1 | |
| tiling = torch.linspace(-1, 1, seq_len, device=device) * shear_factor | |
| y = y + tiling * mask | |
| y_valid_mask = mask.bool().squeeze(0) | |
| if y_valid_mask.any(): | |
| y_valid_points = y[0, y_valid_mask] | |
| min_val = torch.min(y_valid_points) | |
| max_val = torch.max(y_valid_points) | |
| if max_val > min_val: | |
| y[0, y_valid_mask] = (y[0, y_valid_mask] - min_val) / ( | |
| max_val - min_val | |
| ) | |
| else: | |
| y[0, y_valid_mask] = 0.5 | |
| return torch.cat((y, mask), dim=0) | |