o & i@s^ddlmZddlZddlmZddlmZddlmZdgZ Gdddej j Z dd Z dS) ) annotationsN)softmax) PHLFilter) meshgrid_ijCRFcs>eZdZdZ        ddfdd ZdddZZS)ra Conditional Random Field: Combines message passing with a class compatibility convolution into an iterative process designed to successively minimise the energy of the class labeling. In this implementation, the message passing step is a weighted combination of a gaussian filter and a bilateral filter. The bilateral term is included to respect existing structure within the reference tensor. See: https://arxiv.org/abs/1502.03240 ?@?@N iterationsintbilateral_weightfloatgaussian_weightbilateral_spatial_sigmabilateral_color_sigmagaussian_spatial_sigma update_factorcompatibility_matrixtorch.Tensor | Nonec s>t||_||_||_||_||_||_||_||_ dS)a Args: iterations: the number of iterations. bilateral_weight: the weighting of the bilateral term in the message passing step. gaussian_weight: the weighting of the gaussian term in the message passing step. bilateral_spatial_sigma: standard deviation in spatial coordinates for the bilateral term. bilateral_color_sigma: standard deviation in color space for the bilateral term. gaussian_spatial_sigma: standard deviation in spatial coordinates for the gaussian term. update_factor: determines the magnitude of each update. compatibility_matrix: a matrix describing class compatibility, should be NxN where N is the number of classes. N) super__init__r rrrrrrr) selfr rrrrrrr __class__[/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/crf.pyr&s  z CRF.__init__ input_tensor torch.Tensorreference_tensorc Cst|}tj||j||jgdd}||j}t|dd}t|jD]E}t ||}t ||} |j ||j | } |j dur]| jddddd} t| |j } | ddd| j} t||j| dd}q#|S)z Args: input_tensor: tensor containing initial class logits. reference_tensor: the reference tensor used to guide the message passing. Returns: output (torch.Tensor): output tensor. dimN) start_dimr)_create_coordinate_tensortorchcatrrrrranger rapplyrrrflattenpermutematmulreshapeshaper) rrr Zspatial_featuresZbilateral_featuresZgaussian_features output_tensor_Zbilateral_outputZgaussian_outputZcombined_outputflatrrrforwardGs       z CRF.forward)rrrr r r r N)r r rrrrrrrrrrrrrr)rrr r)__name__ __module__ __qualname____doc__rr3 __classcell__rrrrrs!csVfddtdD}t|}t|jjjd}tjd|gddS)Ncsg|] }t|qSr)r'arangesize).0itensorrr tsz-_create_coordinate_tensor..r$)devicedtyperr") r)r#rr'stacktor@rAr:)r>axesgridscoordsrr=rr&ssr&) __future__rr'torch.nn.functionalrZmonai.networks.layers.filteringrmonai.networks.utilsr__all__nnModulerr&rrrrs     \