o  i}@spddlmZddlZddlmZddlmZddlmZddl m Z ddl m Z ddl mZGd d d eZdS) ) annotationsN)CrossEntropyLoss) functional)_Loss)DiceLoss)SobelGradients)HoVerNetBranchcsNeZdZdZ      ddfd d ZdddZdddZdddZZS) HoVerNetLossa: Loss function for HoVerNet pipeline, which is combination of losses across the three branches. The NP (nucleus prediction) branch uses Dice + CrossEntropy. The HV (Horizontal and Vertical) distance from centroid branch uses MSE + MSE of the gradient. The NC (Nuclear Class prediction) branch uses Dice + CrossEntropy The result is a weighted sum of these losses. Args: lambda_hv_mse: Weight factor to apply to the HV regression MSE part of the overall loss lambda_hv_mse_grad: Weight factor to apply to the MSE of the HV gradient part of the overall loss lambda_np_ce: Weight factor to apply to the nuclei prediction CrossEntropyLoss part of the overall loss lambda_np_dice: Weight factor to apply to the nuclei prediction DiceLoss part of overall loss lambda_nc_ce: Weight factor to apply to the nuclei class prediction CrossEntropyLoss part of the overall loss lambda_nc_dice: Weight factor to apply to the nuclei class prediction DiceLoss part of the overall loss @? lambda_hv_msefloatlambda_hv_mse_grad lambda_np_celambda_np_dice lambda_nc_celambda_nc_dicereturnNonecsn||_||_||_||_||_||_ttdddddd|_ t dd|_ t ddd |_ t dd d |_dS) NTgMbP?sum)softmaxZ smooth_drZ smooth_nr reductionbatchmean)rr) kernel_size spatial_axes)r rrrrrsuper__init__rdicercersobel_vsobel_h)selfr rrrrr __class__k/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/pathology/losses/hovernet_loss.pyr-s   zHoVerNetLoss.__init__image torch.TensorcCs>||dddf}||dddf}tj||gddS)aiCompute the Sobel gradients of the horizontal vertical map (HoVerMap). More specifically, it will compute horizontal gradient of the input horizontal gradient map (channel=0) and vertical gradient of the input vertical gradient map (channel=1). Args: image: a tensor with the shape of BxCxHxW representing HoVerMap Nrrdim)r#r"torchstack)r$r)Zresult_hZresult_vr'r'r(_compute_sobelCs zHoVerNetLoss._compute_sobel predictiontargetfocuscCsb||}||}||}|ddddf}t||fd}|||}||d}|S)z[Compute the MSE loss of the gradients of the horizontal and vertical centroid distance mapsN.rg:0yE>)r/r-catr)r$r0r1r2Z pred_gradZ true_gradlossr'r'r(_mse_gradient_lossPs   zHoVerNetLoss._mse_gradient_lossdict[str, torch.Tensor]cCstjj|vr tjj|vstdtjj|vrtjj|vs tdtjj|vr0tjj|vr0tdtjj|vr@tjj|vr@td||tjj|tjj|j}|tjjjdd}| |tjj||j }||}t |tjj|tjj|j }||tjj|tjj|tjjdddf|j}||} d} tjj|vr||tjj|tjj|j} |tjjjdd}| |tjj||j} | | } | || } | S)a! Args: prediction: dictionary of predicted outputs for three branches, each of which should have the shape of BNHW. target: dictionary of ground truths for three branches, each of which should have the shape of BNHW. zrnucleus prediction (NP) and horizontal_vertical (HV) branches must be present for prediction and target parametersz_type_prediction (NC) must be present in both or neither of the prediction and target parametersrr+Nr)rNPvalueHV ValueErrorNCr rargmaxr!rFmse_lossr r5rrr)r$r0r1Z dice_loss_npZ argmax_targetZ ce_loss_npZloss_npZ loss_hv_mseZloss_hv_mse_gradZloss_hvZloss_ncZ dice_loss_ncZ ce_loss_ncr4r'r'r(forwardbsR        zHoVerNetLoss.forward)r r r r r r )r r rr rr rr rr rr rr)r)r*rr*)r0r*r1r*r2r*rr*)r0r6r1r6rr*) __name__ __module__ __qualname____doc__rr/r5r? __classcell__r'r'r%r(r s  r ) __future__rr-torch.nnrrr=Ztorch.nn.modules.lossr monai.lossesrmonai.transformsrmonai.utils.enumsrr r'r'r'r(s