| 32 | |
| 33 | |
| 34 | def rasterize_polygons_with_grid_sample(full_image_bit_mask, box, mask_size, threshold=0.5): |
| 35 | x0, y0, x1, y1 = box[0], box[1], box[2], box[3] |
| 36 | |
| 37 | img_h, img_w = full_image_bit_mask.shape |
| 38 | |
| 39 | mask_y = np.arange(0.0, mask_size) + 0.5 # mask y sample coords in [0.5, mask_size - 0.5] |
| 40 | mask_x = np.arange(0.0, mask_size) + 0.5 # mask x sample coords in [0.5, mask_size - 0.5] |
| 41 | mask_y = mask_y / mask_size * (y1 - y0) + y0 |
| 42 | mask_x = mask_x / mask_size * (x1 - x0) + x0 |
| 43 | |
| 44 | mask_x = (mask_x - 0.5) / (img_w - 1) * 2 + -1 |
| 45 | mask_y = (mask_y - 0.5) / (img_h - 1) * 2 + -1 |
| 46 | gy, gx = torch.meshgrid(torch.from_numpy(mask_y), torch.from_numpy(mask_x)) |
| 47 | ind = torch.stack([gx, gy], dim=-1).to(dtype=torch.float32) |
| 48 | |
| 49 | full_image_bit_mask = torch.from_numpy(full_image_bit_mask) |
| 50 | mask = F.grid_sample( |
| 51 | full_image_bit_mask[None, None, :, :].to(dtype=torch.float32), |
| 52 | ind[None, :, :, :], |
| 53 | align_corners=True, |
| 54 | ) |
| 55 | |
| 56 | return mask[0, 0] >= threshold |
| 57 | |
| 58 | |
| 59 | class TestMaskCropPaste(unittest.TestCase): |