Wrapper for grid_sample, uses pixel coordinates
(img, coords, mode='bilinear', mask=False)
| 139 | |
| 140 | |
| 141 | def bilinear_sampler(img, coords, mode='bilinear', mask=False): |
| 142 | """ Wrapper for grid_sample, uses pixel coordinates """ |
| 143 | H, W = img.shape[-2:] |
| 144 | xgrid, ygrid = coords.split([1,1], dim=-1) |
| 145 | xgrid = 2*xgrid/(W-1) - 1 |
| 146 | ygrid = 2*ygrid/(H-1) - 1 |
| 147 | |
| 148 | grid = torch.cat([xgrid, ygrid], dim=-1) |
| 149 | img = F.grid_sample(img, grid, align_corners=True) |
| 150 | |
| 151 | if mask: |
| 152 | mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) |
| 153 | return img, mask.float() |
| 154 | |
| 155 | return img |
| 156 | |
| 157 | |
| 158 | def coords_grid(batch, ht, wd, device='cpu'): |