U tPhd<@sddlmZddlZddlmZddlmZddlm Z m Z ddl m Z ddl mZdd d d d Zdd d d dZdd d ddZeeedZGdddeZGdddeZdS)) annotationsN) functional)_Loss) gaussian_1dseparable_filtering) LossReduction)look_up_optionint torch.Tensor) kernel_sizereturncCs t|S)N)torchones)r rU/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/losses/image_dissimilarity.pymake_rectangular_kernelsrcCsf|dd}|ddkr |d8}tjdd|ftjd|}||d|d}tj|||ddS)Nrdtype)padding)r rfloatdivFconv1dreshape)r fsizefrrrrmake_triangular_kernels   rcCs6t|d}t||ddddd|}|d|S)Ng@rsampledF)sigma truncatedapprox normalizeg @)r tensorr)r r!kernelrrrmake_gaussian_kernel$s r') rectangular triangulargaussianc s\eZdZdZdddejddfddddd d d d fd d ZddZddddddZZ S)#LocalNormalizedCrossCorrelationLossa Local squared zero-normalized cross-correlation. The loss is based on a moving kernel/window over the y_true/y_pred, within the window the square of zncc is calculated. The kernel can be a rectangular / triangular / gaussian window. The final loss is the averaged loss over all windows. Adapted from: https://github.com/voxelmorph/voxelmorph/blob/legacy/src/losses.py DeepReg (https://github.com/DeepRegNet/DeepReg) r(gh㈵>r strLossReduction | strrNone) spatial_dimsr kernel_type reduction smooth_nr smooth_drr cstjt|jd||_|jdkr6td|jd||_|jddkrZtd|jt|t}||j|_ d|j _ | |_ t ||_t ||_d S) a1 Args: spatial_dims: number of spatial dimensions, {``1``, ``2``, ``3``}. Defaults to 3. kernel_size: kernel spatial size, must be odd. kernel_type: {``"rectangular"``, ``"triangular"``, ``"gaussian"``}. Defaults to ``"rectangular"``. reduction: {``"none"``, ``"mean"``, ``"sum"``} Specifies the reduction to apply to the output. Defaults to ``"mean"``. - ``"none"``: no reduction will be applied. - ``"mean"``: the sum of the output will be divided by the number of elements in the output. - ``"sum"``: the output will be summed. smooth_nr: a small constant added to the numerator to avoid nan. smooth_dr: a small constant added to the denominator to avoid nan. r3>rrr,zUnsupported ndim: z/-d, only 1-d, 2-d, and 3-d inputs are supportedrrzkernel_size must be odd, got FN)super__init__rvaluendim ValueErrorr r kernel_dictr&Z require_gradsget_kernel_vol kernel_volrr4r5)selfr1r r2r3r4r5_kernel __class__rrr8@s     z,LocalNormalizedCrossCorrelationLoss.__init__cCs>|j}t|jdD]}t|d|jd}qt|S)Nrrr)r&ranger:r matmul unsqueezesum)r?vol_rrrr=jsz2LocalNormalizedCrossCorrelationLoss.get_kernel_volr predtargetr cCs|jd|jkr(td|jd|j|j|jkrNtd|jd|jd||||||}}}|j||j|}}|g|j}t||d} t||d} t||d} t||d} t||d} | |}| |}| || }t| || tj |j | j | j d}t| || tj |j | j | j d}|||j ||}|jtjjkrht|S|jtjjkr|S|jtjjkrt|Std |jd d S) z Args: pred: the shape should be BNH[WD]. target: the shape should be BNH[WD]. Raises: ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"]. rzexpecting pred with z' spatial dimensions, got pred of shape "ground truth has differing shape ( ) from pred ())kernels)rdeviceUnsupported reduction: 0, available options are ["mean", "sum", "none"].N)r:r;shaper&tor>rr max as_tensorr5rrPr4r3rSUMr9rFnegNONEMEANmean)r?rJrKt2p2tpr&r>rOZt_sump_sumZt2_sumZp2_sumZtp_sumZt_avgZp_avgcrossZt_varZp_varZnccrrrforwardps>          z+LocalNormalizedCrossCorrelationLoss.forward) __name__ __module__ __qualname____doc__rrZr8r=ra __classcell__rrrArr+3s *r+c seZdZdZdddejddfdddd ddd d fd d ZddddddZddddddZdddddZ ddddddZ Z S)GlobalMutualInformationLossz Differentiable global mutual information loss via Parzen windowing method. Reference: https://dspace.mit.edu/handle/1721.1/123142, Section 3.1, equation 3.1-3.5, Algorithm 1 r*?gHz>r.r rr/r0)r2num_bins sigma_ratior3r4r5r c stjt|jd|dkr$tdtdd|}t|dd|dd|}t|d d g|_ ||_ ||_ |j d krdd |d |_ |d |_ t ||_t ||_dS) a Args: kernel_type: {``"gaussian"``, ``"b-spline"``} ``"gaussian"``: adapted from DeepReg Reference: https://dspace.mit.edu/handle/1721.1/123142, Section 3.1, equation 3.1-3.5, Algorithm 1. ``"b-spline"``: based on the method of Mattes et al [1,2] and adapted from ITK References: [1] "Nonrigid multimodality image registration" D. Mattes, D. R. Haynor, H. Vesselle, T. Lewellen and W. Eubank Medical Imaging 2001: Image Processing, 2001, pp. 1609-1620. [2] "PET-CT Image Registration in the Chest Using Free-form Deformations" D. Mattes, D. R. Haynor, H. Vesselle, T. Lewellen and W. Eubank IEEE Transactions in Medical Imaging. Vol.22, No.1, January 2003. pp.120-128. num_bins: number of bins for intensity sigma_ratio: a hyper param for gaussian function reduction: {``"none"``, ``"mean"``, ``"sum"``} Specifies the reduction to apply to the output. Defaults to ``"mean"``. - ``"none"``: no reduction will be applied. - ``"mean"``: the sum of the output will be divided by the number of elements in the output. - ``"sum"``: the output will be summed. smooth_nr: a small constant added to the numerator to avoid nan. smooth_dr: a small constant added to the denominator to avoid nan. r6rz!num_bins must > 0, got {num_bins}r-g?rNrr*b-spliner)NN.)r7r8rr9r;r linspacer[rr2rjpreterm bin_centersrr4r5) r?r2rjrkr3r4r5ror!rArrr8s#"   z$GlobalMutualInformationLoss.__init__r z=tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]rIcCsh|jdkr(||\}}||\}}n4|jdkrX|j|dd\}}|j|dd\}}nt||||fS)Nr*rlr,)orderr)r2parzen_windowing_gaussianparzen_windowing_b_spliner;)r?rJrKZ pred_weightZpred_probabilityZ target_weightZtarget_probabilityrrrparzen_windowings  z,GlobalMutualInformationLoss.parzen_windowingz!tuple[torch.Tensor, torch.Tensor])imgrpr c Cspt|t|}}d}|||jd|}t|||}t|||}t|||j|d}||jddd}tj|j|j dddd} t | |} tj | tj d} |dkr| | dk| dkd} nl|dkr4| d d | dd| d| dkd } | d| d| dk| dkd } nt d |d | tj| dd d} tj| dd d} | | fS)z Parzen windowing with b-spline kernel (adapted from ITK) Args: img: the shape should be B[NDHW]. order: int. rrrr)rPrrir,zDo not support b-spline z-order parzen windowingTdimkeepdim)r rUminrjrclamprrSarangerPabs zeros_likerr;rFr[) r?rtrp_max_minrbin_sizenorm_minZ window_termbinsZsample_bin_matrixweight probabilityrrrrrs& ,&z5GlobalMutualInformationLoss.parzen_windowing_b_spline)rtr cCsvt|dd}||jddd}t|j| ||j|d}|tj|ddd}tj |ddd}||fS)z Parzen windowing with gaussian kernel (adapted from DeepReg implementation) Note: the input is expected to range between 0 and 1 Args: img: the shape should be B[NDHW]. rrrrTrwrz) r r|rrSexprnrTrorFr[)r?rtrrrrrrq!s z5GlobalMutualInformationLoss.parzen_windowing_gaussianc Cs|j|jkr&td|jd|jd|||\}}}}t|ddd|||jd}t|ddd||}tj|t ||j ||j |j dd} |j t jjkrt| S|j t jjkr| S|j t jjkrt| Std |j d d S) z Args: pred: the shape should be B[NDHW]. target: the shape should be same as the pred shape. Raises: ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"]. rLrMrNrrr)rr)rxrQrRN)rSr;rsr bmmpermuterTrrFlogr4r5r3rrWr9rXrYrZr[) r?rJrKwapawbpbZpabZpapbmirrrra1s  ("z#GlobalMutualInformationLoss.forward) rbrcrdrerrZr8rsrrrqrarfrrrArrgs  13rg) __future__rr torch.nnrrtorch.nn.modules.lossrmonai.networks.layersrr monai.utilsrmonai.utils.modulerrrr'r<r+rgrrrr s       s