Function nerf_encoding

core/FlowFormer/common.py:14–19

x is of shape [*, 2]. The last dimension are two coordinates (x and y).

(x, L=6, NORMALIZE_FACOR=1/300)

import math

def nerf_encoding(x, L=6, NORMALIZE_FACOR=1/300):
    """
        x is of shape [*, 2]. The last dimension are two coordinates (x and y).
    """
    freq_bands = 2.** torch.linspace(0, L, L-1).to(x.device)
    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)

def sampler_gaussian(latent, mean, std, image_size, point_num=25):
    # latent [B, H*W, D]