def visualize_depth(depth,
                    acc=None,
                    near=None,
                    far=None,
                    ignore_frac=0,
                    curve_fn=lambda x: -jnp.log(x + jnp.finfo(jnp.float32).eps),
                    modulus=0,
                    colormap=None):
  """Visualize a depth map.

  Args:
    depth: A depth map.
    acc: An accumulation map, in [0, 1].
    near: The depth of the near plane, if None then just use the min().
    far: The depth of the far plane, if None then just use the max().
    ignore_frac: What fraction of the depth map to ignore when automatically
      generating `near` and `far`. Depends on `acc` as well as `depth'.
    curve_fn: A curve function that gets applied to `depth`, `near`, and `far`
      before the rest of visualization. Good choices: x, 1/(x+eps), log(x+eps).
        Note that the default choice will flip the sign of depths, so that the
        default colormap (turbo) renders "near" as red and "far" as blue.
    modulus: If > 0, mod the normalized depth by `modulus`. Use (0, 1].
    colormap: A colormap function. If None (default), will be set to
      matplotlib's turbo if modulus==0, sinebow otherwise.

  Returns:
    An RGB visualization of `depth`.
  """
  if acc is None:
    acc = jnp.ones_like(depth)
  acc = jnp.where(jnp.isnan(depth), jnp.zeros_like(acc), acc)

  # Sort `depth` and `acc` according to `depth`, then identify the depth values
  # that span the middle of `acc`, ignoring `ignore_frac` fraction of `acc`.
  sortidx = jnp.argsort(depth.reshape([-1]))
  depth_sorted = depth.reshape([-1])[sortidx]
  acc_sorted = acc.reshape([-1])[sortidx]
  cum_acc_sorted = jnp.cumsum(acc_sorted)
  mask = ((cum_acc_sorted >= cum_acc_sorted[-1] * ignore_frac) &
          (cum_acc_sorted <= cum_acc_sorted[-1] * (1 - ignore_frac)))
  depth_keep = depth_sorted[mask]

  # If `near` or `far` are None, use the highest and lowest non-NaN values in
  # `depth_keep` as automatic near/far planes.
  eps = jnp.finfo(jnp.float32).eps
  near = near or depth_keep[0] - eps
  far = far or depth_keep[-1] + eps

  # Curve all values.
  depth, near, far = [curve_fn(x) for x in [depth, near, far]]

  # Wrap the values around if requested.
  if modulus > 0:
    value = jnp.mod(depth, modulus) / modulus
    colormap = colormap or sinebow
  else:
    # Scale to [0, 1].
    value = jnp.nan_to_num(
        jnp.clip((depth - jnp.minimum(near, far)) / jnp.abs(far - near), 0, 1))
    colormap = colormap or cm.get_cmap('turbo')

  vis = colormap(value)[:, :, :3]

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

  return vis
Visualize a depth map.

  Args:
    depth: A depth map.
    acc: An accumulation map, in [0, 1].
    near: The depth of the near plane, if None then just use the min().
    far: The depth of the far plane, if None then just use the max().
    ignore_frac: What fraction of the depth map to ignore when automatically
      generating `near` and `far`. Depends on `acc` as well as `depth'.
    curve_fn: A curve function that gets applied to `depth`, `near`, and `far`
      before the rest of visualization. Good choices: x, 1/(x+eps), log(x+eps).
        Note that the default choice will flip the sign of depths, so that the
        default colormap (turbo) renders "near" as red and "far" as blue.
    modulus: If > 0, mod the normalized depth by `modulus`. Use (0, 1].
    colormap: A colormap function. If None (default), will be set to
      matplotlib's turbo if modulus==0, sinebow otherwise.

  Returns:
    An RGB visualization of `depth`.
  
internal/vis.py