(latent, mean, std, image_size, point_num=25)
| 19 | return torch.cat([x*NORMALIZE_FACOR, torch.sin(3.14*x[..., -2:-1]*freq_bands*NORMALIZE_FACOR), torch.cos(3.14*x[..., -2:-1]*freq_bands*NORMALIZE_FACOR), torch.sin(3.14*x[..., -1:]*freq_bands*NORMALIZE_FACOR), torch.cos(3.14*x[..., -1:]*freq_bands*NORMALIZE_FACOR)], dim=-1) |
| 20 | |
| 21 | def sampler_gaussian(latent, mean, std, image_size, point_num=25): |
| 22 | # latent [B, H*W, D] |
| 23 | # mean [B, 2, H, W] |
| 24 | # std [B, 1, H, W] |
| 25 | H, W = image_size |
| 26 | B, HW, D = latent.shape |
| 27 | STD_MAX = 20 |
| 28 | latent = rearrange(latent, 'b (h w) c -> b c h w', h=H, w=W) # latent = latent.view(B, H, W, D).permute(0, 3, 1, 2) |
| 29 | mean = mean.permute(0, 2, 3, 1) # [B, H, W, 2] |
| 30 | |
| 31 | dx = torch.linspace(-1, 1, int(point_num**0.5)) |
| 32 | dy = torch.linspace(-1, 1, int(point_num**0.5)) |
| 33 | delta = torch.stack(torch.meshgrid(dy, dx), axis=-1).to(mean.device) # [B*H*W, point_num**0.5, point_num**0.5, 2] |
| 34 | delta_3sigma = F.sigmoid(std.permute(0, 2, 3, 1).reshape(B*HW, 1, 1, 1)) * STD_MAX * delta * 3 # [B*H*W, point_num**0.5, point_num**0.5, 2] |
| 35 | |
| 36 | centroid = mean.reshape(B*H*W, 1, 1, 2) |
| 37 | coords = centroid + delta_3sigma |
| 38 | |
| 39 | coords = rearrange(coords, '(b h w) r1 r2 c -> b (h w) (r1 r2) c', b=B, h=H, w=W) |
| 40 | sampled_latents = bilinear_sampler(latent, coords) # [B*H*W, dim, point_num**0.5, point_num**0.5] |
| 41 | sampled_latents = sampled_latents.permute(0, 2, 3, 1) |
| 42 | sampled_weights = -(torch.sum(delta.pow(2), dim=-1)) |
| 43 | |
| 44 | return sampled_latents, sampled_weights |
| 45 | |
| 46 | def sampler_gaussian_zy(latent, mean, std, image_size, point_num=25, return_deltaXY=False, beta=1): |
| 47 | # latent [B, H*W, D] |
nothing calls this directly
no test coverage detected