| 68 | |
| 69 | |
| 70 | class GridMask(nn.Module): |
| 71 | def __init__(self, use_h, use_w, rotate = 1, offset=False, ratio = 0.5, mode=0, prob = 1.): |
| 72 | super(GridMask, self).__init__() |
| 73 | self.use_h = use_h |
| 74 | self.use_w = use_w |
| 75 | self.rotate = rotate |
| 76 | self.offset = offset |
| 77 | self.ratio = ratio |
| 78 | self.mode = mode |
| 79 | self.st_prob = prob |
| 80 | self.prob = prob |
| 81 | self.fp16_enable = False |
| 82 | def set_prob(self, epoch, max_epoch): |
| 83 | self.prob = self.st_prob * epoch / max_epoch #+ 1.#0.5 |
| 84 | @auto_fp16() |
| 85 | def forward(self, x): |
| 86 | if np.random.rand() > self.prob or not self.training: |
| 87 | return x |
| 88 | n,c,h,w = x.size() |
| 89 | x = x.view(-1,h,w) |
| 90 | hh = int(1.5*h) |
| 91 | ww = int(1.5*w) |
| 92 | d = np.random.randint(2, h) |
| 93 | self.l = min(max(int(d*self.ratio+0.5),1),d-1) |
| 94 | mask = np.ones((hh, ww), np.float32) |
| 95 | st_h = np.random.randint(d) |
| 96 | st_w = np.random.randint(d) |
| 97 | if self.use_h: |
| 98 | for i in range(hh//d): |
| 99 | s = d*i + st_h |
| 100 | t = min(s+self.l, hh) |
| 101 | mask[s:t,:] *= 0 |
| 102 | if self.use_w: |
| 103 | for i in range(ww//d): |
| 104 | s = d*i + st_w |
| 105 | t = min(s+self.l, ww) |
| 106 | mask[:,s:t] *= 0 |
| 107 | |
| 108 | r = np.random.randint(self.rotate) |
| 109 | mask = Image.fromarray(np.uint8(mask)) |
| 110 | mask = mask.rotate(r) |
| 111 | mask = np.asarray(mask) |
| 112 | mask = mask[(hh-h)//2:(hh-h)//2+h, (ww-w)//2:(ww-w)//2+w] |
| 113 | |
| 114 | mask = torch.from_numpy(mask).to(x.dtype).cuda() |
| 115 | if self.mode == 1: |
| 116 | mask = 1-mask |
| 117 | mask = mask.expand_as(x) |
| 118 | if self.offset: |
| 119 | offset = torch.from_numpy(2 * (np.random.rand(h,w) - 0.5)).to(x.dtype).cuda() |
| 120 | x = x * mask + offset * (1 - mask) |
| 121 | else: |
| 122 | x = x * mask |
| 123 | |
| 124 | return x.view(n,c,h,w) |