o  i@snddlmZddlmZddlZddlmZddlmmZ ddl m Z ddl m Z mZGddde ZdS)) annotations)AnyN)_Loss)GaussianFilter MeanFiltercsBeZdZdZ     ddfdd ZdddZd ddZZS)!NACLLossa Neighbor-Aware Calibration Loss (NACL) is primarily developed for developing calibrated models in image segmentation. NACL computes standard cross-entropy loss with a linear penalty that enforces the logit distributions to match a soft class proportion of surrounding pixel. Murugesan, Balamurali, et al. "Trust your neighbours: Penalty-based constraints for model calibration." International Conference on Medical Image Computing and Computer-Assisted Intervention, MICCAI 2023. https://arxiv.org/abs/2303.06268 Murugesan, Balamurali, et al. "Neighbor-Aware Calibration of Segmentation Networks with Penalty-Based Constraints." https://arxiv.org/abs/2401.14487 meanl1皙??classesintdim kernel_size kernel_opsstr distance_typealphafloatsigmareturnNonecst|dvr td|dvrtd|d|dvr$td|||_||_t|_||_||_ ||_ ||dkrOt ||d |_ |j j |||j _ |d kr\t||d |_ d Sd S) am Args: classes: number of classes dim: dimension of data (supports 2d and 3d) kernel_size: size of the spatial kernel distance_type: l1/l2 distance between spatial kernel and predicted logits alpha: weightage between cross entropy and logit constraint sigma: sigma of gaussian )r gaussianz*Kernel ops must be either mean or gaussian)rzSupport 2d and 3d, got dim=.)r l2z+Distance type must be either L1 or L2, got r ) spatial_dimssizer)rrN)super__init__ ValueErrorncrnnCrossEntropyLoss cross_entropyrrksr svls_layerfilterr)selfr rrrrrr __class__X/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/losses/nacl_loss.pyr (s(  zNACLLoss.__init__mask torch.TensorcCs|jdkr tj|tj|jddddd }| |}|jdkrAtj|tj|jdddddd }| |}|S)a Converts the mask to one hot represenation and is smoothened with the selected spatial filter. Args: mask: the shape should be BH[WD]. Returns: torch.Tensor: the shape would be BNH[WD], N being number of classes. r) num_classesrr) rFone_hottotorchint64r"permute contiguousrr')r)r.Z oh_labelsrmaskr,r,r-get_constr_targetUs ,  . zNACLLoss.get_constr_targetinputstargetscCsh|||}||}|jdkr||}n|jdkr+||d}||j|}|S)a Computes standard cross-entropy loss and constraints it neighbor aware logit penalty. Args: inputs: the shape should be BNH[WD], where N is the number of classes. targets: the shape should be BH[WD]. Returns: torch.Tensor: value of the loss. Example: >>> import torch >>> from monai.losses import NACLLoss >>> B, N, H, W = 8, 3, 64, 64 >>> input = torch.rand(B, N, H, W) >>> target = torch.randint(0, N, (B, H, W)) >>> criterion = NACLLoss(classes = N, dim = 2) >>> loss = criterion(input, target) r rr)r%r;rsubabs_r pow_r)r)r<r=Zloss_ceZutargetsZ loss_conflossr,r,r-forwardks    zNACLLoss.forward)rr r r r )r rrrrrrrrrrrrrrr)r.r/rr/)r<r/r=r/rr/)__name__ __module__ __qualname____doc__r r;rB __classcell__r,r,r*r-rs -r) __future__rtypingrr6torch.nnr#torch.nn.functional functionalr3torch.nn.modules.lossrmonai.networks.layersrrrr,r,r,r-s