(feat, center, window_size)
| 261 | return weighted_latent, sampled_weights |
| 262 | |
| 263 | def sampler(feat, center, window_size): |
| 264 | # feat [B, C, H, W] |
| 265 | # center [B, 2, H, W] |
| 266 | center = center.permute(0, 2, 3, 1) # [B, H, W, 2] |
| 267 | B, H, W, C = center.shape |
| 268 | |
| 269 | radius = window_size // 2 |
| 270 | dx = torch.linspace(-radius, radius, 2*radius+1) |
| 271 | dy = torch.linspace(-radius, radius, 2*radius+1) |
| 272 | delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(center.device) # [B*H*W, window_size, point_num**0.5, 2] |
| 273 | |
| 274 | center = center.reshape(B*H*W, 1, 1, 2) |
| 275 | coords = center + delta |
| 276 | |
| 277 | coords = rearrange(coords, '(b h w) r1 r2 c -> b (h w) (r1 r2) c', b=B, h=H, w=W) |
| 278 | sampled_latents = bilinear_sampler(feat, coords) # [B*H*W, dim, window_size, window_size] |
| 279 | # sampled_latents = sampled_latents.permute(0, 2, 3, 1) |
| 280 | |
| 281 | return sampled_latents |
| 282 | |
| 283 | def retrieve_tokens(feat, center, window_size, sampler): |
| 284 | # feat [B, C, H, W] |
nothing calls this directly
no test coverage detected