How to compute the scene flow

[m43]

Thanks for the nice work! Is the code for computing the scene flow already in the codebase? If not, could you please share a short snippet of how you've computed the scene flow? Thanks!

[wei-mao-2019]

Already replied via email. Here is a snapshot of the code in case others are also interested.

import torch

def compute_scene_flow(points, normals, t, model):
    """_summary_
    Args:

        points_torch: [N, 3] surface points (uniformly sampled from marching cube results) at time step t
        normals_torch: [N, 3] surface normal (from the marching cube results) at time step t
        t: [1,1] a scalar range from -1 to 1 indicating the time
        model: the trained model
        In our experiments/visualization we use N = 2000
    Returns:
        scene_flow: [N, 3] scene flow computed from sdf flow
    """
    sdf_flow = model.implicit_network.forward_sdf(points, t)
    sdf_flow = sdf_flow[:,:1]

    dist = torch.cdist(points,points,p=2)
    dist[dist>0.4] = 10.0
    knn = torch.topk(dist,dim=1,k=250,largest=False)

    points_knn = points[knn[1]]
    normals_knn = normals[knn[1]]
    sdf_flow_knn = sdf_flow[knn[1]]

    A = torch.cross(points_knn, normals_knn, dim=-1)
    A = torch.cat([A, normals_knn], dim=-1)  # [bs,N,6]
    AtA = torch.matmul(A.transpose(1, 2), A)
    Atb = torch.matmul(A.transpose(1, 2), sdf_flow_knn)
    velocities = -torch.matmul(torch.pinverse(AtA), Atb)

    scene_flow = torch.cross(velocities[...,:3,0],points,dim=1) + velocities[...,3:,0]
    return scene_flow

[m43]

Thanks again!