File size: 3,459 Bytes
5d21dd2 |
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 |
from . import *
from .noise_layers import *
class Random_Noise(nn.Module):
def __init__(self, layers, len_layers_R, len_layers_F):
super(Random_Noise, self).__init__()
for i in range(len(layers)):
layers[i] = eval(layers[i])
self.noise = nn.Sequential(*layers)
self.len_layers_R = len_layers_R
self.len_layers_F = len_layers_F
print(self.noise)
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
def forward(self, image_cover_mask):
image, cover_image, mask = image_cover_mask[0], image_cover_mask[1], image_cover_mask[2]
forward_image = image.clone().detach()
forward_cover_image = cover_image.clone().detach()
forward_mask = mask.clone().detach()
noised_image_C = torch.zeros_like(forward_image)
noised_image_R = torch.zeros_like(forward_image)
noised_image_F = torch.zeros_like(forward_image)
for index in range(forward_image.shape[0]):
random_noise_layer_C = np.random.choice(self.noise, 1)[0]
random_noise_layer_R = np.random.choice(self.noise[0:self.len_layers_R], 1)[0]
random_noise_layer_F = np.random.choice(self.noise[self.len_layers_R:self.len_layers_R + self.len_layers_F], 1)[0]
noised_image_C[index] = random_noise_layer_C([forward_image[index].clone().unsqueeze(0), forward_cover_image[index].clone().unsqueeze(0), forward_mask[index].clone().unsqueeze(0)])
noised_image_R[index] = random_noise_layer_R([forward_image[index].clone().unsqueeze(0), forward_cover_image[index].clone().unsqueeze(0), forward_mask[index].clone().unsqueeze(0)])
noised_image_F[index] = random_noise_layer_F([forward_image[index].clone().unsqueeze(0), forward_cover_image[index].clone().unsqueeze(0), forward_mask[index].clone().unsqueeze(0)])
'''single_image = ((noised_image_C[index].clamp(-1, 1).permute(1, 2, 0) + 1) / 2 * 255).add(0.5).clamp(0, 255).to('cpu', torch.uint8).numpy()
im = Image.fromarray(single_image)
read = np.array(im, dtype=np.uint8)
noised_image_C[index] = self.transform(read).unsqueeze(0).to(image.device)
single_image = ((noised_image_R[index].clamp(-1, 1).permute(1, 2, 0) + 1) / 2 * 255).add(0.5).clamp(0, 255).to('cpu', torch.uint8).numpy()
im = Image.fromarray(single_image)
read = np.array(im, dtype=np.uint8)
noised_image_R[index] = self.transform(read).unsqueeze(0).to(image.device)
single_image = ((noised_image_F[index].clamp(-1, 1).permute(1, 2, 0) + 1) / 2 * 255).add(0.5).clamp(0, 255).to('cpu', torch.uint8).numpy()
im = Image.fromarray(single_image)
read = np.array(im, dtype=np.uint8)
noised_image_F[index] = self.transform(read).unsqueeze(0).to(image.device)
noised_image_gap_C = noised_image_C - forward_image
noised_image_gap_R = noised_image_R - forward_image
noised_image_gap_F = noised_image_F - forward_image'''
noised_image_gap_C = noised_image_C.clamp(-1, 1) - forward_image
noised_image_gap_R = noised_image_R.clamp(-1, 1) - forward_image
noised_image_gap_F = noised_image_F.clamp(-1, 1) - forward_image
return image + noised_image_gap_C, image + noised_image_gap_R, image + noised_image_gap_F
|