Spaces:
Sleeping
Sleeping
File size: 4,009 Bytes
b1e87cf d38b6f8 b1e87cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | 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)
|