o  i¸ã@sŠddlmZddlZddlmmZddlmZddd„Z dd d „Z dd d „Z ddd„Z dddd„Z Gdd„deƒZGdd„deƒZdS) é)Ú annotationsN)Ú_LossÚimgú torch.TensorÚreturncCs¤t|jƒdkr!t | ddd¡ }t | ddd¡ }t ||¡St|jƒdkrPt | dd d ¡ }t | d d d ¡ }t | d d d¡ }t t ||¡|¡SdS)zÕ Perform soft erosion on the input image Args: img: the shape should be BCH(WD) Adapted from: https://github.com/jocpae/clDice/blob/master/cldice_loss/pytorch/soft_skeleton.py#L6 é)éé©r r )r r)r r)rr é)rr r ©r r r )r rr)r rr )rr r)r r r)rrr N)ÚlenÚshapeÚFÚ max_pool2dÚtorchÚminÚ max_pool3d)rÚp1Úp2Úp3©rúU/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/losses/cldice.pyÚ soft_erodes  ürcCs@t|jƒdkrt |ddd¡St|jƒdkrt |ddd¡SdS)z× Perform soft dilation on the input image Args: img: the shape should be BCH(WD) Adapted from: https://github.com/jocpae/clDice/blob/master/cldice_loss/pytorch/soft_skeleton.py#L18 r)rrr r )rrrr N)r rrrr)rrrrÚ soft_dilate(s  ÿrcCst|ƒ}t|ƒ}|S)zê Wrapper function to perform soft opening on the input image Args: img: the shape should be BCH(WD) Adapted from: https://github.com/jocpae/clDice/blob/master/cldice_loss/pytorch/soft_skeleton.py#L25 )rr)rZ eroded_imageZ dilated_imagerrrÚ soft_open8s rÚiter_ÚintcCs\t|ƒ}t ||¡}t|ƒD]}t|ƒ}t|ƒ}t ||¡}|t |||¡}q|S)a> Perform soft skeletonization on the input image Adapted from: https://github.com/jocpae/clDice/blob/master/cldice_loss/pytorch/soft_skeleton.py#L29 Args: img: the shape should be BCH(WD) iter_: number of iterations for skeletonization Returns: skeletonized image )rrÚreluÚranger)rrÚimg1ÚskelÚ_ÚdeltarrrÚ soft_skelGs r$çð?Úy_trueÚy_predÚsmoothÚfloatcCstt ||dd…dd…df¡}d||t |dd…dd…df¡t |dd…dd…df¡|}d|}|S)a Function to compute soft dice loss Adapted from: https://github.com/jocpae/clDice/blob/master/cldice_loss/pytorch/cldice.py#L22 Args: y_true: the shape should be BCH(WD) y_pred: the shape should be BCH(WD) Returns: dice loss Nr .ç@r%)rÚsum)r&r'r(Ú intersectionÚcoeffÚ soft_dicerrrr._s Hr.cs.eZdZdZdd‡fd d „ Zddd„Z‡ZS)ÚSoftclDiceLossa: Compute the Soft clDice loss defined in: Shit et al. (2021) clDice -- A Novel Topology-Preserving Loss Function for Tubular Structure Segmentation. (https://arxiv.org/abs/2003.07311) Adapted from: https://github.com/jocpae/clDice/blob/master/cldice_loss/pytorch/cldice.py#L7 rr%rrr(r)rÚNonecstƒ ¡||_||_dS)z{ Args: iter_: Number of iterations for skeletonization smooth: Smoothing parameter N)ÚsuperÚ__init__Úiterr()Úselfrr(©Ú __class__rrr2~s  zSoftclDiceLoss.__init__r&rr'cCsÌt||jƒ}t||jƒ}t t ||¡dd…dd…df¡|jt |dd…dd…df¡|j}t t ||¡dd…dd…df¡|jt |dd…dd…df¡|j}dd||||}|S©Nr .r%r*)r$r3rr+Úmultiplyr()r4r&r'Ú skel_predÚ skel_trueÚtprecÚtsensÚcl_dicerrrÚforwardˆs  ( ÿ( ÿzSoftclDiceLoss.forward)rr%)rrr(r)rr0©r&rr'rrr©Ú__name__Ú __module__Ú __qualname__Ú__doc__r2r>Ú __classcell__rrr5rr/ss  r/cs.eZdZdZdd‡fd d „ Zddd„Z‡ZS)ÚSoftDiceclDiceLossa; Compute the Soft clDice loss defined in: Shit et al. (2021) clDice -- A Novel Topology-Preserving Loss Function for Tubular Structure Segmentation. (https://arxiv.org/abs/2003.07311) Adapted from: https://github.com/jocpae/clDice/blob/master/cldice_loss/pytorch/cldice.py#L38 rçà?r%rrÚalphar)r(rr0cs tƒ ¡||_||_||_dS)z© Args: iter_: Number of iterations for skeletonization smooth: Smoothing parameter alpha: Weighing factor for cldice N)r1r2r3r(rH)r4rrHr(r5rrr2 s  zSoftDiceclDiceLoss.__init__r&rr'c Csòt|||jƒ}t||jƒ}t||jƒ}t t ||¡dd…dd…df¡|jt |dd…dd…df¡|j}t t ||¡dd…dd…df¡|jt |dd…dd…df¡|j}dd||||}d|j||j|} | Sr7)r.r(r$r3rr+r8rH) r4r&r'Údicer9r:r;r<r=Z total_lossrrrr>¬s  ( ÿ( ÿzSoftDiceclDiceLoss.forward)rrGr%)rrrHr)r(r)rr0r?r@rrr5rrF•s  rF)rrrr)rrrrrr)r%)r&rr'rr(r)rr)Ú __future__rrÚtorch.nn.functionalÚnnÚ functionalrÚtorch.nn.modules.lossrrrrr$r.r/rFrrrrÚs       "