Add WaNet trigger to DatasetFolder images. Args: identity_grid (orch.Tensor): the poisoned pattern shape. noise_grid (orch.Tensor): the noise pattern. noise (bool): turn on noise mode, default is False. s (int or float): The strength of the noise grid. Default is
| 45 | |
| 46 | |
| 47 | class AddDatasetFolderTrigger(AddTrigger): |
| 48 | """Add WaNet trigger to DatasetFolder images. |
| 49 | |
| 50 | Args: |
| 51 | identity_grid (orch.Tensor): the poisoned pattern shape. |
| 52 | noise_grid (orch.Tensor): the noise pattern. |
| 53 | noise (bool): turn on noise mode, default is False. |
| 54 | s (int or float): The strength of the noise grid. Default is 0.5. |
| 55 | grid_rescale (int or float): Scale :attr:`grid` to avoid pixel values going out of [-1, 1]. |
| 56 | Default is 1. |
| 57 | noise_rescale (int or float): Scale the random noise from a uniform distribution on the |
| 58 | interval [0, 1). Default is 2. |
| 59 | """ |
| 60 | |
| 61 | def __init__(self, identity_grid, noise_grid, noise=False, s=0.5, grid_rescale=1, noise_rescale=2): |
| 62 | super(AddDatasetFolderTrigger, self).__init__() |
| 63 | |
| 64 | self.identity_grid = deepcopy(identity_grid) |
| 65 | self.noise_grid = deepcopy(noise_grid) |
| 66 | self.h = self.identity_grid.shape[2] |
| 67 | self.noise = noise |
| 68 | self.s = s |
| 69 | self.grid_rescale = grid_rescale |
| 70 | grid = self.identity_grid + self.s * self.noise_grid / self.h |
| 71 | self.grid = torch.clamp(grid * self.grid_rescale, -1, 1) |
| 72 | self.noise_rescale = noise_rescale |
| 73 | |
| 74 | |
| 75 | |
| 76 | |
| 77 | def __call__(self, img): |
| 78 | """Get the poisoned image. |
| 79 | |
| 80 | Args: |
| 81 | img (PIL.Image.Image | numpy.ndarray | torch.Tensor): If img is numpy.ndarray or torch.Tensor, the shape should be (H, W, C) or (H, W). |
| 82 | Returns: |
| 83 | torch.Tensor: The poisoned image. |
| 84 | """ |
| 85 | if type(img) == PIL.Image.Image: |
| 86 | img = F.pil_to_tensor(img) |
| 87 | img = F.convert_image_dtype(img, torch.float) |
| 88 | img = self.add_trigger(img, noise=self.noise) |
| 89 | # 1 x H x W |
| 90 | if img.size(0) == 1: |
| 91 | img = img.squeeze().numpy() |
| 92 | img = Image.fromarray(np.clip(img*255,0,255).round().astype(np.uint8), mode='L') |
| 93 | # 3 x H x W |
| 94 | elif img.size(0) == 3: |
| 95 | img = img.numpy().transpose(1, 2, 0) |
| 96 | img = Image.fromarray(np.clip(img*255,0,255).round().astype(np.uint8)) |
| 97 | else: |
| 98 | raise ValueError("Unsupportable image shape.") |
| 99 | return img |
| 100 | elif type(img) == np.ndarray: |
| 101 | # H x W |
| 102 | if len(img.shape) == 2: |
| 103 | img = torch.from_numpy(img) |
| 104 | img = F.convert_image_dtype(img, torch.float) |