U tPh%@sbddlmZddlZddlmZddlmZdddddd ZGd d d eZGd d d eZ dS)) annotationsN)_Loss) LossReduction torch.Tensorint)xdimreturncCstdd}tdd}tdd}td}||g||g}}t||jkr\||g}||g}q8|||<|||<||||dS)a Calculate gradients on single dimension of a tensor using central finite difference. It moves the tensor along the dimension to calculate the approximate gradient dx[i] = (x[i+1] - x[i-1]) / 2. Adapted from: DeepReg (https://github.com/DeepRegNet/DeepReg) Args: x: the shape should be BCH(WD). dim: dimension to calculate gradient along. Returns: gradient_dx: the shape should be BCH(WD) Ng@)slicelenndim)rrZslice_1Z slice_2_sZ slice_2_e slice_allZ slicing_sZ slicing_erH/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/losses/deform.pyspatial_gradients     rcsBeZdZdZdejfddddfdd Zd d d d d ZZS) BendingEnergyLossaO Calculate the bending energy based on second-order differentiation of ``pred`` using central finite difference. For more information, see https://github.com/Project-MONAI/tutorials/blob/main/modules/bending_energy_diffusion_loss_notes.ipynb. Adapted from: DeepReg (https://github.com/DeepRegNet/DeepReg) FboolLossReduction | strNone normalize reductionr cstjt|jd||_dSae Args: normalize: Whether to divide out spatial sizes in order to make the computation roughly invariant to image scale (i.e. vector field sampling resolution). Defaults to False. 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. )rNsuper__init__rvaluerselfrr __class__rrr:s zBendingEnergyLoss.__init__rpredr c sjdkrtdjtjdD]0}j| ddkr(tdjddq(jdjdkrtdjdd jdfd d tdjD}|jrtjjjd ddd jdd}td}t |D]\}}|d7}|jr6|j||9}|t ||j|d}n|t ||d}t|djD]F}|jr|dt ||j|d}n|dt ||d}qXq|j t j jkrt|}n>|j t jjkrt|}n"|j t jjkrtd|j d|S)a Args: pred: the shape should be BCH(WD) Raises: ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"]. ValueError: When ``pred`` is not 3-d, 4-d or 5-d. ValueError: When any spatial dimension of ``pred`` has size less than or equal to 4. ValueError: When the number of channels of ``pred`` does not match the number of spatial dimensions. :Expecting 3-d, 4-d or 5-d pred, instead got pred of shape r r r)z;All spatial dimensions must be > 4, got spatial dimensions NGNumber of vector components, i.e. number of channels of the input DDF, /, does not match number of spatial dimensions, csg|]}t|qSrr.0rr&rr bsz-BendingEnergyLoss.forward..devicer r r rUnsupported reduction: 0, available options are ["mean", "sum", "none"].)r ValueErrorshaperangertorchtensorr4reshape enumeraterrrMEANr meanSUMsumNONE) r"r&ifirst_order_gradient spatial_dimsenergydim_1gZdim_2rr1rforwardJs< . "  zBendingEnergyLoss.forward __name__ __module__ __qualname____doc__rr@rrK __classcell__rrr#rr/s rcsBeZdZdZdejfddddfdd Zd d d d d ZZS) DiffusionLossah Calculate the diffusion based on first-order differentiation of ``pred`` using central finite difference. For the original paper, please refer to VoxelMorph: A Learning Framework for Deformable Medical Image Registration, Guha Balakrishnan, Amy Zhao, Mert R. Sabuncu, John Guttag, Adrian V. Dalca IEEE TMI: Transactions on Medical Imaging. 2019. eprint arXiv:1809.05231. For more information, see https://github.com/Project-MONAI/tutorials/blob/main/modules/bending_energy_diffusion_loss_notes.ipynb. Adapted from: VoxelMorph (https://github.com/voxelmorph/voxelmorph) Frrrrcstjt|jd||_dSrrr!r#rrrs zDiffusionLoss.__init__rr%csjdkrtdjtjdD]0}j| ddkr(tdjddq(jdjdkrtdjddjdfd d tdjD}|jrtjjjd ddd jdd }td}t |D]6\}}|d7}|jr|j||9}||d}q|j t j j krBt|}n>|j t jj kr^t|}n"|j t jj krtd|j d|S)a Args: pred: Predicted dense displacement field (DDF) with shape BCH[WD], where C is the number of spatial dimensions. Note that diffusion loss can only be calculated when the sizes of the DDF along all spatial dimensions are greater than 2. Raises: ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"]. ValueError: When ``pred`` is not 3-d, 4-d or 5-d. ValueError: When any spatial dimension of ``pred`` has size less than or equal to 2. ValueError: When the number of channels of ``pred`` does not match the number of spatial dimensions. r'r+r r z;All spatial dimensions must be > 2, got spatial dimensions Nr,r-csg|]}t|qSrr.r/r1rrr2sz)DiffusionLoss.forward..r3r5r6rr7r8)rr9r:r;rr<r=r4r>r?rrr@r rArBrCrD)r"r&rErFrG diffusionrIrJrr1rrKs2 .   zDiffusionLoss.forwardrLrrr#rrRsrR) __future__rr<torch.nn.modules.lossr monai.utilsrrrrRrrrr s   Q