File size: 982 Bytes
0a9cf85 3dc5131 |
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 |
import os
import torch
import random
import numpy as np
def set_seed(seed):
random.seed(seed)
np.random.seed(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
def generate_centered_gaussian_noise(shape=(1, 1, 28, 28), sigma=5.0, mu=0):
B, C, H, W = shape
assert C == 1, "only image gray"
yy, xx = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='ij')
yy = yy.to(torch.float32)
xx = xx.to(torch.float32)
center_y, center_x = H / 2, W / 2
gauss = torch.exp(-((yy - center_y)**2 + (xx - center_x)**2) / (2 * sigma**2))
gauss = gauss / gauss.max() # Normalization to [0, 1]
gauss = gauss.unsqueeze(0).unsqueeze(0).expand(B, C, H, W)
noise = mu + torch.randn(B, C, H, W) # Noise with mean mu
localized_noise = noise * gauss + mu * (1 - gauss)
return localized_noise
|