Sample the feature map at the uv values. Args: uv: (b, *p, 2) values between [0, 1] feature_map: (b, h, w, dim), boundary corresponded to u=0, u=1, v=0, v=1. (0,0) at top left, u to right, v to down pad_one_outsize: If true, val
(
uv: torch.Tensor, # (b, *p, 2)
feature_map: torch.Tensor, # (b, h, w, dim)
mode: str = 'bilinear',
padding_mode: str = 'zeros',
uv_normalized: bool = True,
)
| 1311 | |
| 1312 | |
| 1313 | def uv_sampling( |
| 1314 | uv: torch.Tensor, # (b, *p, 2) |
| 1315 | feature_map: torch.Tensor, # (b, h, w, dim) |
| 1316 | mode: str = 'bilinear', |
| 1317 | padding_mode: str = 'zeros', |
| 1318 | uv_normalized: bool = True, |
| 1319 | ): |
| 1320 | """ |
| 1321 | Sample the feature map at the uv values. |
| 1322 | |
| 1323 | Args: |
| 1324 | uv: |
| 1325 | (b, *p, 2) values between [0, 1] |
| 1326 | feature_map: |
| 1327 | (b, h, w, dim), boundary corresponded to u=0, u=1, v=0, v=1. (0,0) at top left, u to right, v to down |
| 1328 | pad_one_outsize: |
| 1329 | If true, values outside feature_map will be set as 1, else 0 |
| 1330 | mode: |
| 1331 | mode used by grid_sample |
| 1332 | padding_mode: |
| 1333 | padding mode used by grid_sample. "zeros", "border", "reflection" |
| 1334 | uv_normalized: |
| 1335 | whether uv is normalized to [0, 1]. if None, uv is in the range of [0, w] [0, h] |
| 1336 | |
| 1337 | Returns: |
| 1338 | resampled_feature: |
| 1339 | (b, *p, dim) |
| 1340 | """ |
| 1341 | |
| 1342 | if not uv_normalized: |
| 1343 | b, h, w, dim = feature_map.shape |
| 1344 | uv = uv.clone() |
| 1345 | uv[..., 0] = uv[..., 0] / w |
| 1346 | uv[..., 1] = uv[..., 1] / h |
| 1347 | |
| 1348 | # [0, 1] -> [-1, 1] used by grid_sampling |
| 1349 | uv = 2 * uv - 1 # (b, *p, 2) |
| 1350 | |
| 1351 | b, *p_shape, _2 = uv.shape |
| 1352 | assert _2 == 2 |
| 1353 | uv = uv.reshape(b, 1, -1, 2) # (b, 1, p, 2) |
| 1354 | |
| 1355 | # (b, h, w, dim) -> (b, dim, h, w) |
| 1356 | feature_map = feature_map.permute(0, 3, 1, 2) # (b, dim, h, w) |
| 1357 | |
| 1358 | resampled_feature = torch.nn.functional.grid_sample( |
| 1359 | input=feature_map, # (b, dim, h, w) |
| 1360 | grid=uv, # (b, 1, p, 2) |
| 1361 | mode=mode, |
| 1362 | padding_mode=padding_mode, |
| 1363 | align_corners=False, |
| 1364 | ) |
| 1365 | # (b, dim, 1, p) -> (b, 1, p, dim) |
| 1366 | resampled_feature = resampled_feature.permute(0, 2, 3, 1) # (b, 1, p, dim) |
| 1367 | resampled_feature = resampled_feature.reshape(b, *p_shape, resampled_feature.size(-1)) # (b, *p, dim) |
| 1368 | |
| 1369 | return resampled_feature # (b, *p, dim) |
| 1370 |
no test coverage detected