import torch


def extract_ampl_phase(fft_im):
    """
    Extracts amplitude and phase from the image
    Args:
        fft_im (Tensor): size should be bx3xhxwx2
    Returns:
        tuple: (amplitude, phase)
    """
    # fft_im: size should be bx3xhxwx2
    fft_amp = fft_im[:, :, :, :, 0]**2 + fft_im[:, :, :, :, 1]**2
    fft_amp = torch.sqrt(fft_amp)
    fft_pha = torch.atan2(fft_im[:, :, :, :, 1], fft_im[:, :, :, :, 0])
    return fft_amp, fft_pha