ezhang94 / gimbal Goto Github PK

def opencv_triangulate(Ps, ys, camera_pairs=[]):
    C = len(Ps)
    batch_shape = ys.shape[1:-1]
    if not camera_pairs:
        camera_pairs = [(i,j) for i in range(C) for j in range(i+1, C)]
    Xs = jnp.empty((len(camera_pairs), *batch_shape, 3))
    for i, (c0, c1) in enumerate(camera_pairs):
        Xs = Xs.at[i].set(opencv_triangulate_dlt(Ps[(c0,c1),:,:], ys[(c0,c1),...]))
    return jnp.median(Xs, axis=0)


def opencv_triangulate_dlt(Ps, ys):
    batch_shape = ys.shape[1:-1]
    Ps = np.array(Ps)
    ys = np.array(ys).reshape(2,-1,2)
    X_hom = cv2.triangulatePoints(Ps[0],Ps[1],ys[0].T,ys[1].T)
    X = (X_hom[:3] / X_hom[3]).T.reshape(*batch_shape,3)
    return jnp.array(X)


print('Observations shape:', ys.shape)

t = time.time()
X1 = opencv_triangulate(Ps, ys, camera_pairs=[])
print('Runtime with opencv:', time.time()-t)

t = time.time()
X2 = gimbal.util.triangulate(Ps, ys, camera_pairs=[]).block_until_ready()
print('Runtime with jax:', time.time()-t)

print('Mean squared difference of output:', jnp.nanmean((X1-X2)**2))

Observations shape: (6, 5000, 23, 2)
Runtime with opencv: 6.574772834777832
Runtime with jax: 68.68624544143677
Mean squared difference of output: 9.889472634621758e-23