Function construct

robust_loss_pytorch/wavelet.py:294–339

num_levels a wavelet decomposition of an image. Args: im: A numpy or TF tensor of single or double precision floats of size (batch_size, width, height) num_levels: The number of levels (or scales) of the wavelet decomposition to apply. A value of 0 returns a "wavelet decomposi

(im, num_levels, wavelet_type)

def construct(im, num_levels, wavelet_type):
  """num_levels a wavelet decomposition of an image.

  Args:
    im: A numpy or TF tensor of single or double precision floats of size
      (batch_size, width, height)
    num_levels: The number of levels (or scales) of the wavelet decomposition to
      apply. A value of 0 returns a "wavelet decomposition" that is just the
      image.
    wavelet_type: The kind of wavelet to use, see generate_filters().

  Returns:
    A wavelet decomposition of `im` that has `num_levels` levels (not including
    the coarsest residual level) and is of type `wavelet_type`. This
    decomposition is represented as a tuple of 3-tuples, with the final element
    being a tensor:
      ((band00, band01, band02), (band10, band11, band12), ..., resid)
    Where band** and resid are TF tensors. Each element of these nested tuples
    is of shape [batch_size, width * 2^-(level+1), height * 2^-(level+1)],
    though the spatial dimensions may be off by 1 if width and height are not
    factors of 2. The residual image is of the same (rough) size as the last set
    of bands. The floating point precision of these tensors matches that of
    `im`.
  """
  if len(im.shape) != 3:
    raise ValueError(
        'Expected `im` to have a rank of 3, but is of size {}'.format(im.shape))
  im = torch.as_tensor(im)
  if num_levels == 0:
    return (im,)
  max_num_levels = get_max_num_levels(im.shape)
  assert max_num_levels >= num_levels
  filters = generate_filters(wavelet_type)
  pyr = []
  for _ in range(num_levels):
    # import pdb
    # pdb.set_trace()
    hi = _downsample(im, filters.analysis_hi, 0, 1)
    lo = _downsample(im, filters.analysis_lo, 0, 0)
    pyr.append((_downsample(hi, filters.analysis_hi, 1,
                            1), _downsample(lo, filters.analysis_hi, 1, 1),
                _downsample(hi, filters.analysis_lo, 1, 0)))
    im = _downsample(lo, filters.analysis_lo, 1, 0)
  pyr.append(im)
  pyr = tuple(pyr)
  return pyr


def collapse(pyr, wavelet_type):