o  i%@sZddlmZddlZddlmZddlmZdd d ZGd d d eZGdddeZ dS)) annotationsN)_Loss) LossReductionx torch.TensordimintreturncCstdd}tdd}tdd}td}||g||g}}t||jkr4||g}||g}t||jks#|||<|||<||||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_erU/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/losses/deform.pyspatial_gradients     rc4eZdZdZdejfdfd d Zdd dZZS)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) F normalizebool reductionLossReduction | strr Nonectjt|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 : zBendingEnergyLoss.__init__predrc sjdvr tdjtjdD]}j| ddkr,tdjddqjdjdkrGtdjdd jdfd d tdjD}|jrntjjjd ddd jdd}td}t |D]V\}}|d7}|jr|j||9}|t ||j|d}n |t ||d}t|djD]!}|jr|dt ||j|d}q|dt ||d}qqw|j t j jkrt|}|S|j t jjkrt|}|S|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, cg|]}t|qSrr.0rr'rr bz-BendingEnergyLoss.forward..devicer r r rUnsupported reduction: 0, available options are ["mean", "sum", "none"].)r ValueErrorshaperangertorchtensorr7reshape enumeraterrrMEANr!meanSUMsumNONE) r#r'ifirst_order_gradient spatial_dimsenergydim_1gZdim_2rr3rforwardJsH . "  zBendingEnergyLoss.forwardrrrrr rr'rr r __name__ __module__ __qualname____doc__rrCr rN __classcell__rrr$rr/s rcr) 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) Frrrrr rcrrrr"r$rrr r&zDiffusionLoss.__init__r'rcsjdvr tdjtjdD]}j| ddkr,tdjddqjdjdkrGtdjddjdfd d tdjD}|jrntjjjd ddd jdd }td}t |D]\}}|d7}|jr|j||9}||d}qw|j t j j krt|}|S|j t jj krt|}|S|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.cr/rr0r1r3rrr4r5z)DiffusionLoss.forward..r6r8r9rr:r;)rr<r=r>rr?r@r7rArBrrrCr!rDrErFrG)r#r'rHrIrJZ diffusionrLrMrr3rrNs< .   zDiffusionLoss.forwardrOrPrQrrr$rrWsrW)rrrrr r) __future__rr?torch.nn.modules.lossr monai.utilsrrrrWrrrrs    Q