Args: img (Tensor): Tensor image of size (C, H, W). Returns: Tensor: Image with n_holes of dimension length x length cut out of it.
(self, img)
| 50 | self.p = prob |
| 51 | |
| 52 | def __call__(self, img): |
| 53 | """ |
| 54 | Args: |
| 55 | img (Tensor): Tensor image of size (C, H, W). |
| 56 | Returns: |
| 57 | Tensor: Image with n_holes of dimension length x length cut out of it. |
| 58 | """ |
| 59 | if torch.rand(1) > self.p: |
| 60 | return img |
| 61 | h = img.size(1) |
| 62 | w = img.size(2) |
| 63 | |
| 64 | mask = np.ones((h, w), np.float32) |
| 65 | |
| 66 | for n in range(self.n_holes): |
| 67 | y = np.random.randint(h) |
| 68 | x = np.random.randint(w) |
| 69 | |
| 70 | y1 = np.clip(y - self.length // 2, 0, h) |
| 71 | y2 = np.clip(y + self.length // 2, 0, h) |
| 72 | x1 = np.clip(x - self.length // 2, 0, w) |
| 73 | x2 = np.clip(x + self.length // 2, 0, w) |
| 74 | mask[y1: y2, x1: x2] = 0. |
| 75 | |
| 76 | mask = torch.from_numpy(mask) |
| 77 | mask = mask.expand_as(img) |
| 78 | img = img * mask |
| 79 | |
| 80 | return img |
| 81 | |
| 82 | def ShearX(img, v): # [-0.3, 0.3] |
| 83 | assert -0.3 <= v <= 0.3 |