class_docstr: Pyramidal optical flow estimation.
  
  This estimates the optical flow at `levels` different pyramidal scales. It
  should be trained in conjunction with a pyramidal optical flow supervision
  signal.
  
  Args:
      in_channels: Number of input channels. Typically uses two 3-channel
          images, plus an initial flow with 2 more channels.
      filter_counts: Number of filters for each stage of pyramid. Defaults
          to ``[32, 64, 32, 16, 2]``.
      kernel_size: Kernel size of filters.
      levels: Number of pyramid levels.
class_docstr_tok: ["Pyramidal","optical","flow","estimation",".","This","estimates","the","optical","flow","at","`","levels","`","different","pyramidal","scales",".","It","should","be","trained","in","conjunction","with","a","pyramidal","optical","flow","supervision","signal",".","Args",":","in_channels",":","Number","of","input","channels",".","Typically","uses","two","3-channel","images",",","plus","an","initial","flow","with","2","more","channels",".","filter_counts",":","Number","of","filters","for","each","stage","of","pyramid",".","Defaults","to","``","[","32",",","64",",","32",",","16",",","2","]","``",".","kernel_size",":","Kernel","size","of","filters",".","levels",":","Number","of","pyramid","levels","."]
docstr: Calculate optical flow and return image pairs for each pyramid level.
  
  During training, we optimize the optical flow by trying to match the
  flow at each pyramid scale to its target. This function computes the
  images output by the flow at each level of the pyramid and returns them
  all in addition to the final flow.
  
  Args:
      image1: The based image - a flow map will be computed to transform
          ``image1`` into ``image2``.
      image2: The target image.
  
  Returns:
      A 2-tuple containing:
          The pyramidal flow estimate.
          Base/target image pairs from each level of the pyramid that can
              be input into a distortion function.
code: def calculate_flow_with_image_pairs(
      self, image1: Tensor, image2: Tensor
  ) -> Tuple[Tensor, List[Tuple[Tensor, Tensor]]]:
      """
      Calculate optical flow and return image pairs for each pyramid level.
  
      During training, we optimize the optical flow by trying to match the
      flow at each pyramid scale to its target. This function computes the
      images output by the flow at each level of the pyramid and returns them
      all in addition to the final flow.
  
      Args:
          image1: The based image - a flow map will be computed to transform
              ``image1`` into ``image2``.
          image2: The target image.
  
      Returns:
          A 2-tuple containing:
              The pyramidal flow estimate.
              Base/target image pairs from each level of the pyramid that can
                  be input into a distortion function.
      """
      if not image1.shape == image2.shape:
          raise ValueError("Image shapes must match.")
  
      images_1, images_2 = self._decompose_images_to_pyramids(
          image1, image2
      )
      flow = torch.zeros_like(images_1[-1])[
          :, :2
      ]  # just the first two channels
  
      # estimate flows at all levels and return warped images
      optical_flow_pairs = []
      for model_level in self.model_levels:
          image1 = images_1.pop()
          image2 = images_2.pop()
          flow = F.interpolate(flow, image1.shape[2:], mode="bilinear")
          flow = flow + model_level(
              torch.cat(
                  (
                      ncF.dense_image_warp(
                          image1, flow.permute(0, 2, 3, 1)
                      ),
                      image2,
                      flow,
                  ),
                  dim=1,
              )
          )
          optical_flow_pairs.append(
              (
                  ncF.dense_image_warp(
                      image1, flow.permute(0, 2, 3, 1)
                  ),
                  image2,
              )
          )
  
      # we have to return all of the warped images at different levels for
      # training
      return flow, optical_flow_pairs
  
imports: from typing import Any, List, NamedTuple, Optional, Sequence, Tuple
  import torch
  import torch.nn.functional
  from compressai.entropy_models import EntropyBottleneck
  from torch import Tensor, nn
  import neuralcompression.functional
  from neuralcompression.layers import SimplifiedGDN, SimplifiedInverseGDN
label_desc: Use image node_id 3 for calling the DVCPyramidFlowEstimator obj's underlying member method code with example usage: obj.calculate_flow_with_image_pairs(image1, image2) and returns: flow,  optical_flow_pairs