U uPh @sddlmZddlmZmZddlZddlmZddlm Z ddddddZ dd ddd d d dddddZ Gddde Z dS)) annotations)CallableOptionalN)_Loss torch.Tensor)xyreturncCsj|s"t|dddd}|sDt|dddd}t||}tjj|t |ddS)a First compute the difference in the complex domain, then get the absolute value and take the mse Args: x, y - B, 2, H, W real valued tensors representing complex numbers or B,1,H,W complex valued tensors Returns: l2_loss - scalar rmean reduction) is_complextorchview_as_complexpermute contiguousabsnn functionalmse_loss zeros_like)rrdiffrK/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/losses/sure_loss.pycomplex_diff_abs_losss rFrOptional[torch.Tensor]Optional[float]Optional[bool])operatorr y_pseudo_gty_refeps perturb_noise complex_inputr c Cs|dkrt|}|dkr0tt|d}|dkr@||}|||}||}td|t|dddd||} |rt||} nt j j ||d d } | | |j d|j d|j d} | S) aU Args: operator (function): The operator function that takes in an input tensor x and returns an output tensor y. We will use this to compute the divergence. More specifically, we will perturb the input x by a small amount and compute the divergence between the perturbed output and the reference output x (torch.Tensor): The input tensor of shape (B, C, H, W) to the operator. For complex input, the shape is (B, 2, H, W) aka C=2 real. For real input, the shape is (B, 1, H, W) real. y_pseudo_gt (torch.Tensor): The pseudo ground truth tensor of shape (B, C, H, W) used to compute the L2 loss. For complex input, the shape is (B, 2, H, W) aka C=2 real. For real input, the shape is (B, 1, H, W) real. y_ref (torch.Tensor, optional): The reference output tensor of shape (B, C, H, W) used to compute the divergence. Defaults to None. For complex input, the shape is (B, 2, H, W) aka C=2 real. For real input, the shape is (B, 1, H, W) real. eps (float, optional): The perturbation scalar. Set to -1 to set it automatically estimated based on y_pseudo_gtk perturb_noise (torch.Tensor, optional): The noise vector of shape (B, C, H, W). Defaults to None. For complex input, the shape is (B, 2, H, W) aka C=2 real. For real input, the shape is (B, 1, H, W) real. complex_input(bool, optional): Whether the input is complex or not. Defaults to False. Returns: sure_loss (torch.Tensor): The SURE loss scalar. Nrig?rr r r r r) r randn_likefloatrmaxsummatmulrrrrrshape) r"rr#r$r%r&r'Z x_perturbedZ y_perturbedZ divergenceZl2_loss sure_lossrrrsure_loss_function)s-  * &r/csFeZdZdZdddddfdd Zdd d d dd d d ddZZS)SURELossa Calculate the Stein's Unbiased Risk Estimator (SURE) loss for a given operator. This is a differentiable loss function that can be used to train/guide an operator (e.g. neural network), where the pseudo ground truth is available but the reference ground truth is not. For example, in the MRI reconstruction, the pseudo ground truth is the zero-filled reconstruction and the reference ground truth is the fully sampled reconstruction. Often, the reference ground truth is not available due to the lack of fully sampled data. The original SURE loss is proposed in [1]. The SURE loss used for guiding the diffusion model based MRI reconstruction is proposed in [2]. Reference [1] Stein, C.M.: Estimation of the mean of a multivariate normal distribution. Annals of Statistics [2] B. Ozturkler et al. SMRD: SURE-based Robust MRI Reconstruction with Diffusion Models. (https://arxiv.org/pdf/2310.01799.pdf) Nrr None)r&r%r cst||_||_dS)aX Args: perturb_noise (torch.Tensor, optional): The noise vector of shape (B, C, H, W). Defaults to None. For complex input, the shape is (B, 2, H, W) aka C=2 real. For real input, the shape is (B, 1, H, W) real. eps (float, optional): The perturbation scalar. Defaults to None. N)super__init__r&r%)selfr&r% __class__rrr3s zSURELoss.__init__Frrr!)r"rr#r$r'r cCs|dkr td|d|dkr@td|d|dk rh|dkrhtd|d|j|jkrtd|jd|jd|dk r|j|jkrtd |jd |jdt|||||j|j|}|S) a: Args: operator (function): The operator function that takes in an input tensor x and returns an output tensor y. We will use this to compute the divergence. More specifically, we will perturb the input x by a small amount and compute the divergence between the perturbed output and the reference output x (torch.Tensor): The input tensor of shape (B, C, H, W) to the operator. C=1 or 2: For complex input, the shape is (B, 2, H, W) aka C=2 real. For real input, the shape is (B, 1, H, W) real. y_pseudo_gt (torch.Tensor): The pseudo ground truth tensor of shape (B, C, H, W) used to compute the L2 loss. C=1 or 2: For complex input, the shape is (B, 2, H, W) aka C=2 real. For real input, the shape is (B, 1, H, W) real. y_ref (torch.Tensor, optional): The reference output tensor of the same shape as y_pseudo_gt Returns: sure_loss (torch.Tensor): The SURE loss scalar. z!Input tensor x should be 4D, got .z/Input tensor y_pseudo_gt should be 4D, but got Nz)Input tensor y_ref should be 4D, but got zKInput tensor x and y_pseudo_gt should have the same shape, but got x shape z, y_pseudo_gt shape zYInput tensor y_pseudo_gt and y_ref should have the same shape, but got y_pseudo_gt shape z, y_ref shape )dim ValueErrorr-r/r%r&)r4r"rr#r$r'lossrrrforwards !   zSURELoss.forward)NN)NF)__name__ __module__ __qualname____doc__r3r< __classcell__rrr5rr0os r0)NrNF) __future__rtypingrrrtorch.nnrtorch.nn.modules.lossrrr/r0rrrr s   F