code: def undistortion_factor(
      distorted_squared_radius: type_alias.TensorLike,
      distortion_coefficient: type_alias.TensorLike,
      num_iterations: int = 5,
      name: str = "quadratic_radial_distortion_undistortion_factor",
  ) -> Tuple[tf.Tensor, tf.Tensor]:
      """Calculates the inverse quadratic distortion function given squared radii.
  
      Given a vector describing a location in camera space in homogeneous
      coordinates `(x/z, y/z, 1)`, after distortion has been applied, these become
      `(x'/z, y'/z, 1)`. `distorted_squared_radius` is `(x'/z)^2 + (y'/z)^2`.
      `undistortion_factor` multiplies `x'/z` and `y'/z` to obtain the undistorted
      projective coordinates `x/z` and `y/z`.
      The undustortion factor in this function is derived from a quadratic.
      distortion function, where the distortion factor equals
      `1.0 + distortion_coefficient * squared_radius`.
  
      Note:
        In the following, A1 to An are optional batch dimensions, which must be
        broadcast compatible.
  
      Args:
        distorted_squared_radius: A tensor of shape `[A1, ..., An, H, W]` containing
          the value of  projective coordinates `(x/z)^2 + (y/z)^2`. For each pixel
          it contains the squared distance of that pixel to the center of the image
          plane. We use `distorted_squared_radius` rather than the distorted radius
          itself to avoid an unnecessary `sqrt`, which may introduce gradient
          singularities. The non-negativity of `distorted_squared_radius` is only
          enforced in debug mode.
        distortion_coefficient: A `scalar` or a tensor of shape `[A1, ..., An]`,
          which contains the distortion coefficients of each image.
        num_iterations: Number of Newton-Raphson iterations to calculate the inverse
          distortion function. Defaults to 5, which is on the high-accuracy side.
        name: A name for this op. Defaults to
          "quadratic_radial_distortion_undistortion_factor".
  
      Returns:
        undistortion: A tensor of shape `[A1, ..., An, H, W]` containing the
          correction factor that should multiply the distorted projective
          coordinates `(x'/z)` and `(y'/z)` to obtain the undistorted ones.
        overflow_mask: A `bool` tensor of shape `[A1, ..., An, H, W]`, `True` where
          `distorted_squared_radius` is beyond the range where the distortion
          function is monotonically increasing. Wherever `overflow_mask` is `True`,
          `undistortion_factor`'s value is meaningless.
  
      """
      with tf.name_scope(name):
          distorted_squared_radius = tf.convert_to_tensor(
              value=distorted_squared_radius
          )
          distortion_coefficient = tf.convert_to_tensor(
              value=distortion_coefficient
          )
  
          if distortion_coefficient.shape.ndims == 0:
              distortion_coefficient = tf.expand_dims(
                  distortion_coefficient, axis=0
              )
          shape.check_static(
              tensor=distorted_squared_radius,
              tensor_name="distorted_squared_radius",
              has_rank_greater_than=1,
          )
          shape.compare_batch_dimensions(
              tensors=(
                  distorted_squared_radius,
                  distortion_coefficient,
              ),
              tensor_names=(
                  "distorted_squared_radius",
                  "distortion_coefficient",
              ),
              last_axes=(-3, -1),
              broadcast_compatible=True,
          )
          distorted_squared_radius = asserts.assert_all_above(
              distorted_squared_radius, 0.0, open_bound=False
          )
          distortion_coefficient = tf.expand_dims(
              distortion_coefficient, axis=-1
          )
          distortion_coefficient = tf.expand_dims(
              distortion_coefficient, axis=-1
          )
          # For a distortion function of r' = (1 + ar^2)r, with a negative a, the
          # maximum r until which r'(r) is monotonically increasing is r^2 = -1/(3a).
          # At that value, r'^2 = -4 / (27a). Therefore the overflow condition for r'
          # is ar'^2 +(4/27.0) < 0. For a positive a it never holds, as it should,
          # because then r' is monotonic in r everywhere and thus never overflows.
          distortion_coefficient_times_distorted_squared_radius = (
              distortion_coefficient * distorted_squared_radius
          )
          overflow_mask = tf.less(
              4.0 / 27.0
              + distortion_coefficient_times_distorted_squared_radius,
              0.0,
          )
  
          # Newton-raphson iterations. The expression below is obtained from
          # algebrically simplifying the Newton-Raphson formula
          # (https://en.wikipedia.org/wiki/Newtons_method).
          # We initialize with the approximate formula for the undistortion function
          # given here https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4934233/.
          undistortion = (
              1.0
              - distortion_coefficient_times_distorted_squared_radius
          )
          for _ in range(num_iterations):
              two_thirds_undistortion = 2.0 * undistortion / 3.0
              undistortion = (1.0 - two_thirds_undistortion) / (
                  1.0
                  + 3.0
                  * distortion_coefficient_times_distorted_squared_radius
                  * undistortion
                  * undistortion
              ) + two_thirds_undistortion
          return undistortion, overflow_mask
  
file_path: graphics/tensorflow_graphics/rendering/camera/quadratic_radial_distortion.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Tuple
  from six.moves import range
  import tensorflow
  from tensorflow_graphics.util import asserts
  from tensorflow_graphics.util import export_api
  from tensorflow_graphics.util import shape
  from tensorflow_graphics.util import type_alias
label_desc: Use image node_id 2 for calling a global function with example usage: undistortion_factor(distorted_squared_radius, distortion_coefficient, num_iterations, name) and returns: undistortion,  overflow_mask