def depth_to_normals(depth):
  """Assuming `depth` is orthographic, linearize it to a set of normals."""
  f_blur = jnp.array([1, 2, 1]) / 4
  f_edge = jnp.array([-1, 0, 1]) / 2
  dy = convolve2d(depth, f_blur[None, :] * f_edge[:, None])
  dx = convolve2d(depth, f_blur[:, None] * f_edge[None, :])
  inv_denom = 1 / jnp.sqrt(1 + dx**2 + dy**2)
  normals = jnp.stack([dx * inv_denom, dy * inv_denom, inv_denom], -1)
  return normals
Assuming `depth` is orthographic, linearize it to a set of normals.
internal/vis.py