import torch


def from_4D_image(image: torch.Tensor, ndims: int) -> torch.Tensor:
  """Convert `image` to `ndims`-D tensors

  Args:
    image (torch.Tensor): 4D image tensors in shape (b, c, h, w).
    ndims (int): the original rank of the image.

  Returns:
    torch.Tensor: `ndims`-D tensors
  """
  _ndims = len(image.shape)
  assert _ndims == 4, f"`image` must be a 4D tensor, while {_ndims}-D are given."
  if ndims == 2:
    return image[0, 0, :, :] # -b, -c
  elif ndims == 3:
    return image[0, :, :, :] # -b
  else:
    return image