text: """
  Copyright (c) Facebook, Inc. and its affiliates.
  
  This source code is licensed under the MIT license found in the
  LICENSE file in the root directory of this source tree.
  """
  
  from typing import Dict, Optional, Sequence, Tuple, Union
  
  import fastmri
  import numpy as np
  import torch
  
  from .subsample import MaskFunc
  
  
  def to_tensor(data: np.ndarray) -> torch.Tensor:
      """
      Convert numpy array to PyTorch tensor.
  
      For complex arrays, the real and imaginary parts are stacked along the last
      dimension.
  
      Args:
          data: Input numpy array.
  
      Returns:
          PyTorch version of data.
      """
      if np.iscomplexobj(data):
          data = np.stack((data.real, data.imag), axis=-1)
  
      return torch.from_numpy(data)
  
  
  def tensor_to_complex_np(data: torch.Tensor) -> np.ndarray:
      """
      Converts a complex torch tensor to numpy array.
  
      Args:
          data: Input data to be converted to numpy.
  
      Returns:
          Complex numpy version of data.
      """
      data = data.numpy()
  
      return data[..., 0] + 1j * data[..., 1]
  
  
  def apply_mask(
      data: torch.Tensor,
      mask_func: MaskFunc,
      seed: Optional[Union[int, Tuple[int, ...]]] = None,
      padding: Optional[Sequence[int]] = None,
      rev=False,
      hamming=False,
      cuda=False
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Subsample given k-space by multiplying with a mask.
  
      Args:
          data: The input k-space data. This should have at least 3 dimensions,
              where dimensions -3 and -2 are the spatial dimensions, and the
              final dimension has size 2 (for complex values).
          mask_func: A function that takes a shape (tuple of ints) and a random
              number seed and returns a mask.
          seed: Seed for the random number generator.
          padding: Padding value to apply for mask.
  
      Returns:
          tuple containing:
              masked data: Subsampled k-space data
              mask: The generated mask
      """
      shape = np.array(data.shape)
      shape[:-3] = 1
      mask = mask_func(shape, seed)
      if padding is not None:
          mask[:, :, : padding[0]] = 0
          mask[:, :, padding[1] :] = 0  # padding value inclusive on right of zeros
  
      # masked_data = data * mask + 0.0  # the + 0.0 removes the sign of the zeros
      if cuda:
          mask = mask.cuda()
          if not rev:
              if not hamming:
                  return torch.where(mask == 0, torch.Tensor([0]).cuda(), data), mask
              else:
                  return mask*data, mask
          else:
              if not hamming:
                  return torch.where(mask != 0, torch.Tensor([0]).cuda(), data), mask
              else:
                  return (1.-mask)*data, mask
      else:
          if not rev:
              if not hamming:
                  return torch.where(mask == 0, torch.Tensor([0]), data), mask
              else:
                  return mask*data, mask
          else:
              if not hamming:
                  return torch.where(mask != 0, torch.Tensor([0]), data), mask
              else:
                  return (1.-mask)*data, mask
  
  
      return masked_data, mask
  
  
  def mask_center(x: torch.Tensor, mask_from: int, mask_to: int) -> torch.Tensor:
      """
      Initializes a mask with the center filled in.
  
      Args:
          mask_from: Part of center to start filling.
          mask_to: Part of center to end filling.
  
      Returns:
          A mask with the center filled.
      """
      mask = torch.zeros_like(x)
      mask[:, :, :, mask_from:mask_to] = x[:, :, :, mask_from:mask_to]
  
      return mask
  
  
  def center_crop(data: torch.Tensor, shape: Tuple[int, int]) -> torch.Tensor:
      """
      Apply a center crop to the input real image or batch of real images.
  
      Args:
          data: The input tensor to be center cropped. It should
              have at least 2 dimensions and the cropping is applied along the
              last two dimensions.
          shape: The output shape. The shape should be smaller
              than the corresponding dimensions of data.
  
      Returns:
          The center cropped image.
      """
      if not (0 < shape[0] <= data.shape[-2] and 0 < shape[1] <= data.shape[-1]):
          raise ValueError("Invalid shapes.")
  
      w_from = (data.shape[-2] - shape[0]) // 2
      h_from = (data.shape[-1] - shape[1]) // 2
      w_to = w_from + shape[0]
      h_to = h_from + shape[1]
  
      return data[..., w_from:w_to, h_from:h_to]
  
  
  def complex_center_crop(data: torch.Tensor, shape: Tuple[int, int]) -> torch.Tensor:
      """
      Apply a center crop to the input image or batch of complex images.
  
      Args:
          data: The complex input tensor to be center cropped. It should have at
              least 3 dimensions and the cropping is applied along dimensions -3
              and -2 and the last dimensions should have a size of 2.
          shape: The output shape. The shape should be smaller than the
              corresponding dimensions of data.
  
      Returns:
          The center cropped image
      """
      if not (0 < shape[0] <= data.shape[-3] and 0 < shape[1] <= data.shape[-2]):
          raise ValueError("Invalid shapes.")
  
      w_from = (data.shape[-3] - shape[0]) // 2
      h_from = (data.shape[-2] - shape[1]) // 2
      w_to = w_from + shape[0]
      h_to = h_from + shape[1]
  
      return data[..., w_from:w_to, h_from:h_to, :]
  
  
  def center_crop_to_smallest(
      x: torch.Tensor, y: torch.Tensor
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Apply a center crop on the larger image to the size of the smaller.
  
      The minimum is taken over dim=-1 and dim=-2. If x is smaller than y at
      dim=-1 and y is smaller than x at dim=-2, then the returned dimension will
      be a mixture of the two.
  
      Args:
          x: The first image.
          y: The second image.
  
      Returns:
          tuple of tensors x and y, each cropped to the minimim size.
      """
      smallest_width = min(x.shape[-1], y.shape[-1])
      smallest_height = min(x.shape[-2], y.shape[-2])
      x = center_crop(x, (smallest_height, smallest_width))
      y = center_crop(y, (smallest_height, smallest_width))
  
      return x, y
  
  
  def normalize(
      data: torch.Tensor,
      mean: Union[float, torch.Tensor],
      stddev: Union[float, torch.Tensor],
      eps: Union[float, torch.Tensor] = 0.0,
  ) -> torch.Tensor:
      """
      Normalize the given tensor.
  
      Applies the formula (data - mean) / (stddev + eps).
  
      Args:
          data: Input data to be normalized.
          mean: Mean value.
          stddev: Standard deviation.
          eps: Added to stddev to prevent dividing by zero.
  
      Returns:
          Normalized tensor.
      """
      return (data - mean) / (stddev + eps)
  
  
  def normalize_instance(
      data: torch.Tensor, eps: Union[float, torch.Tensor] = 0.0
  ) -> Tuple[torch.Tensor, Union[torch.Tensor], Union[torch.Tensor]]:
      """
      Normalize the given tensor  with instance norm/
  
      Applies the formula (data - mean) / (stddev + eps), where mean and stddev
      are computed from the data itself.
  
      Args:
          data: Input data to be normalized
          eps: Added to stddev to prevent dividing by zero.
  
      Returns:
          torch.Tensor: Normalized tensor
      """
      mean = data.mean()
      std = data.std()
  
      return normalize(data, mean, std, eps), mean, std
  
  
  class UnetDataTransform:
      """
      Data Transformer for training U-Net models.
      """
  
      def __init__(
          self,
          which_challenge: str,
          mask_func: Optional[MaskFunc] = None,
          use_seed: bool = True,
      ):
          """
          Args:
              which_challenge: Challenge from ("singlecoil", "multicoil").
              mask_func: Optional; A function that can create a mask of
                  appropriate shape.
              use_seed: If true, this class computes a pseudo random number
                  generator seed from the filename. This ensures that the same
                  mask is used for all the slices of a given volume every time.
          """
          if which_challenge not in ("singlecoil", "multicoil", "Img"):
              raise ValueError("Challenge should either be 'singlecoil' or 'multicoil'")
  
          self.mask_func = mask_func
          self.which_challenge = which_challenge
          self.use_seed = use_seed
  
      def __call__(
          self, 
          target, 
          fname, 
          slice_num=0, 
          attrs=None,
          seed=None
          ):
          # Preprocess the data here
          # target shape: [H, W, 1] or [H, W, 3]
          img = target
          if target.shape[2] != 2:
              img = np.concatenate((target, np.zeros_like(target)), axis=2)
          assert img.shape[-1] == 2
          img = to_tensor(img)
          kspace = fastmri.fft2c(img) 
  
          center_kspace, _ = apply_mask(kspace, self.mask_func, hamming=True, seed=seed)
          img_LF = fastmri.complex_abs(fastmri.ifft2c(center_kspace))
          img_LF = img_LF.unsqueeze(0)
          image, mean, std = normalize_instance(img_LF, eps=1e-11)
          image = image.clamp(-6, 6)
          # img_LF tensor should have shape [H, W, ?]
          target = to_tensor(np.transpose(target, (2, 0, 1)))  # target shape [1, H, W]
          target = normalize(target, mean, std, eps=1e-11)
          target = target.clamp(-6, 6)
          target = target.squeeze(0)
          # check for max value
          max_value = 0.0
          # print('traget shape', target.shape)
          # print('image shape', image.shape)
          return image, target, mean, std, fname, slice_num, max_value
  '''
      def __call__(
          self,
          kspace: np.ndarray,
          mask: np.ndarray,
          target: np.ndarray,
          attrs: Dict,
          fname: str,
          slice_num: int,
      ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, str, int, float]:
          """
          Args:
              kspace: Input k-space of shape (num_coils, rows, cols) for
                  multi-coil data or (rows, cols) for single coil data.
              mask: Mask from the test dataset.
              target: Target image.
              attrs: Acquisition related information stored in the HDF5 object.
              fname: File name.
              slice_num: Serial number of the slice.
  
          Returns:
              tuple containing:
                  image: Zero-filled input image.
                  target: Target image converted to a torch.Tensor.
                  mean: Mean value used for normalization.
                  std: Standard deviation value used for normalization.
                  fname: File name.
                  slice_num: Serial number of the slice.
          """
          kspace = to_tensor(kspace)
  
          # check for max value
          max_value = attrs["max"] if "max" in attrs.keys() else 0.0
  
          # apply mask
          if self.mask_func:
              seed = None if not self.use_seed else tuple(map(ord, fname))
              masked_kspace, mask = apply_mask(kspace, self.mask_func, seed)
          else:
              masked_kspace = kspace
  
          # inverse Fourier transform to get zero filled solution
          image = fastmri.ifft2c(masked_kspace)
  
          # crop input to correct size
          if target is not None:
              crop_size = (target.shape[-2], target.shape[-1])
          else:
              crop_size = (attrs["recon_size"][0], attrs["recon_size"][1])
  
          # check for FLAIR 203
          if image.shape[-2] < crop_size[1]:
              crop_size = (image.shape[-2], image.shape[-2])
  
          image = complex_center_crop(image, crop_size)
  
          # absolute value
          image = fastmri.complex_abs(image)
  
          # apply Root-Sum-of-Squares if multicoil data
          if self.which_challenge == "multicoil":
              image = fastmri.rss(image)
  
          # normalize input
          image, mean, std = normalize_instance(image, eps=1e-11)
          image = image.clamp(-6, 6)
  
          # normalize target
          if target is not None:
              target = to_tensor(target)
              target = center_crop(target, crop_size)
              target = normalize(target, mean, std, eps=1e-11)
              target = target.clamp(-6, 6)
          else:
              target = torch.Tensor([0])
  
          return image, target, mean, std, fname, slice_num, max_value
  '''
  
  class VarNetDataTransform:
      """
      Data Transformer for training VarNet models.
      """
  
      def __init__(self, mask_func: Optional[MaskFunc] = None, use_seed: bool = True):
          """
          Args:
              mask_func: Optional; A function that can create a mask of
                  appropriate shape. Defaults to None.
              use_seed: If True, this class computes a pseudo random number
                  generator seed from the filename. This ensures that the same
                  mask is used for all the slices of a given volume every time.
          """
          self.mask_func = mask_func
          self.use_seed = use_seed
  
      def __call__(
          self,
          kspace: np.ndarray,
          mask: np.ndarray,
          target: np.ndarray,
          attrs: Dict,
          fname: str,
          slice_num: int,
      ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, str, int, float, torch.Tensor]:
          """
          Args:
              kspace: Input k-space of shape (num_coils, rows, cols) for
                  multi-coil data.
              mask: Mask from the test dataset.
              target: Target image.
              attrs: Acquisition related information stored in the HDF5 object.
              fname: File name.
              slice_num: Serial number of the slice.
  
          Returns:
              tuple containing:
                  masked_kspace: k-space after applying sampling mask.
                  mask: The applied sampling mask
                  target: The target image (if applicable).
                  fname: File name.
                  slice_num: The slice index.
                  max_value: Maximum image value.
                  crop_size: The size to crop the final image.
          """
          if target is not None:
              target = to_tensor(target)
              max_value = attrs["max"]
          else:
              target = torch.tensor(0)
              max_value = 0.0
  
          kspace = to_tensor(kspace)
          seed = None if not self.use_seed else tuple(map(ord, fname))
          acq_start = attrs["padding_left"]
          acq_end = attrs["padding_right"]
  
          crop_size = torch.tensor([attrs["recon_size"][0], attrs["recon_size"][1]])
  
          if self.mask_func:
              masked_kspace, mask = apply_mask(
                  kspace, self.mask_func, seed, (acq_start, acq_end)
              )
          else:
              masked_kspace = kspace
              shape = np.array(kspace.shape)
              num_cols = shape[-2]
              shape[:-3] = 1
              mask_shape = [1] * len(shape)
              mask_shape[-2] = num_cols
              mask = torch.from_numpy(mask.reshape(*mask_shape).astype(np.float32))
              mask = mask.reshape(*mask_shape)
              mask[:, :, :acq_start] = 0
              mask[:, :, acq_end:] = 0
  
          return (
              masked_kspace,
              mask.byte(),
              target,
              fname,
              slice_num,
              max_value,
              crop_size,
          )
  