| 5 | from mmcv.runner import force_fp32, auto_fp16 |
| 6 | |
| 7 | class Grid(object): |
| 8 | def __init__(self, use_h, use_w, rotate = 1, offset=False, ratio = 0.5, mode=0, prob = 1.): |
| 9 | self.use_h = use_h |
| 10 | self.use_w = use_w |
| 11 | self.rotate = rotate |
| 12 | self.offset = offset |
| 13 | self.ratio = ratio |
| 14 | self.mode=mode |
| 15 | self.st_prob = prob |
| 16 | self.prob = prob |
| 17 | |
| 18 | def set_prob(self, epoch, max_epoch): |
| 19 | self.prob = self.st_prob * epoch / max_epoch |
| 20 | |
| 21 | def __call__(self, img, label): |
| 22 | if np.random.rand() > self.prob: |
| 23 | return img, label |
| 24 | h = img.size(1) |
| 25 | w = img.size(2) |
| 26 | self.d1 = 2 |
| 27 | self.d2 = min(h, w) |
| 28 | hh = int(1.5*h) |
| 29 | ww = int(1.5*w) |
| 30 | d = np.random.randint(self.d1, self.d2) |
| 31 | if self.ratio == 1: |
| 32 | self.l = np.random.randint(1, d) |
| 33 | else: |
| 34 | self.l = min(max(int(d*self.ratio+0.5),1),d-1) |
| 35 | mask = np.ones((hh, ww), np.float32) |
| 36 | st_h = np.random.randint(d) |
| 37 | st_w = np.random.randint(d) |
| 38 | if self.use_h: |
| 39 | for i in range(hh//d): |
| 40 | s = d*i + st_h |
| 41 | t = min(s+self.l, hh) |
| 42 | mask[s:t,:] *= 0 |
| 43 | if self.use_w: |
| 44 | for i in range(ww//d): |
| 45 | s = d*i + st_w |
| 46 | t = min(s+self.l, ww) |
| 47 | mask[:,s:t] *= 0 |
| 48 | |
| 49 | r = np.random.randint(self.rotate) |
| 50 | mask = Image.fromarray(np.uint8(mask)) |
| 51 | mask = mask.rotate(r) |
| 52 | mask = np.asarray(mask) |
| 53 | mask = mask[(hh-h)//2:(hh-h)//2+h, (ww-w)//2:(ww-w)//2+w] |
| 54 | |
| 55 | mask = torch.from_numpy(mask).float() |
| 56 | if self.mode == 1: |
| 57 | mask = 1-mask |
| 58 | |
| 59 | mask = mask.expand_as(img) |
| 60 | if self.offset: |
| 61 | offset = torch.from_numpy(2 * (np.random.rand(h,w) - 0.5)).float() |
| 62 | offset = (1 - mask) * offset |
| 63 | img = img * mask + offset |
| 64 | else: |
nothing calls this directly
no outgoing calls
no test coverage detected