def visualize_normals(depth, acc, scaling=None):
  """Visualize fake normals of `depth` (optionally scaled to be isotropic)."""
  if scaling is None:
    mask = ~jnp.isnan(depth)
    x, y = jnp.meshgrid(
        jnp.arange(depth.shape[1]), jnp.arange(depth.shape[0]), indexing='xy')
    xy_var = (jnp.var(x[mask]) + jnp.var(y[mask])) / 2
    z_var = jnp.var(depth[mask])
    scaling = jnp.sqrt(xy_var / z_var)

  scaled_depth = scaling * depth
  normals = depth_to_normals(scaled_depth)
  vis = jnp.isnan(normals) + jnp.nan_to_num((normals + 1) / 2, 0)

  # Set non-accumulated pixels to white.
  if acc is not None:
    vis = vis * acc[:, :, None] + (1 - acc)[:, :, None]

  return vis
Visualize fake normals of `depth` (optionally scaled to be isotropic).
internal/vis.py