| 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 | |