Wrapper for grid_sample, uses pixel coordinates
(img, coords, mask=False)
| 72 | return img |
| 73 | |
| 74 | def indexing(img, coords, mask=False): |
| 75 | """ Wrapper for grid_sample, uses pixel coordinates """ |
| 76 | """ |
| 77 | TODO: directly indexing features instead of sampling |
| 78 | """ |
| 79 | H, W = img.shape[-2:] |
| 80 | xgrid, ygrid = coords.split([1,1], dim=-1) |
| 81 | xgrid = 2*xgrid/(W-1) - 1 |
| 82 | ygrid = 2*ygrid/(H-1) - 1 |
| 83 | |
| 84 | grid = torch.cat([xgrid, ygrid], dim=-1) |
| 85 | img = F.grid_sample(img, grid, align_corners=True, mode='nearest') |
| 86 | |
| 87 | if mask: |
| 88 | mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) |
| 89 | return img, mask.float() |
| 90 | |
| 91 | return img |
| 92 | |
| 93 | def coords_grid(batch, ht, wd): |
| 94 | coords = torch.meshgrid(torch.arange(ht), torch.arange(wd)) |