o  i)@s|ddlmZddlZddlmZddlZddlmZddlm Z ddl m Z ddl m Z Gdd d eZGd d d eZdS) ) annotationsN)Callable)_Loss)one_hot)distance_transform_edt) LossReductioncsReZdZdZddddddejdfdfdd ZedddZ dddZ Z S) HausdorffDTLossa Compute channel-wise binary Hausdorff loss based on distance transform. It can support both multi-classes and multi-labels tasks. The data `input` (BNHW[D] where N is number of classes) is compared with ground truth `target` (BNHW[D]). Note that axis N of `input` is expected to be logits or probabilities for each class, if passing logits as input, must set `sigmoid=True` or `softmax=True`, or specifying `other_act`. And the same axis of `target` can be 1 or N (one-hot format). The original paper: Karimi, D. et. al. (2019) Reducing the Hausdorff Distance in Medical Image Segmentation with Convolutional Neural Networks, IEEE Transactions on medical imaging, 39(2), 499-513 g@FNalphafloatinclude_backgroundbool to_onehot_ysigmoidsoftmax other_actCallable | None reductionLossReduction | strbatchreturnNonec stjt|jd|durt|stdt|jdt|t|dkr+t d||_ ||_ ||_ ||_ ||_||_||_dS)a Args: include_background: if False, channel index 0 (background category) is excluded from the calculation. if the non-background segmentations are small compared to the total image size they can get overwhelmed by the signal from the background so excluding it in such cases helps convergence. to_onehot_y: whether to convert the ``target`` into the one-hot format, using the number of classes inferred from `input` (``input.shape[1]``). Defaults to False. sigmoid: if True, apply a sigmoid function to the prediction. softmax: if True, apply a softmax function to the prediction. other_act: callable function to execute other activation layers, Defaults to ``None``. for example: ``other_act = torch.tanh``. 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. batch: whether to sum the intersection and union areas over the batch dimension before the dividing. Defaults to False, a loss value is computed independently from each item in the batch before any `reduction`. Raises: TypeError: When ``other_act`` is not an ``Optional[Callable]``. ValueError: When more than 1 of [``sigmoid=True``, ``softmax=True``, ``other_act is not None``]. Incompatible values. )rNz*other_act must be None or callable but is .zXIncompatible values: more than 1 of [sigmoid=True, softmax=True, other_act is not None].)super__init__rvaluecallable TypeErrortype__name__int ValueErrorr r r rrrr) selfr r r rrrrr __class__]/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/losses/hausdorff_loss.pyr,s& zHausdorffDTLoss.__init__img torch.TensorcCs^t|}tt|D]!}||dk}|r,|s,t|}|}t|}||||<q |S)zGenerate distance transform. Args: img (np.ndarray): input mask as NCHWD or NCHW. Returns: np.ndarray: Distance field. g?)torch zeros_likerangelenanyallr)r"r'field batch_idxZfg_maskZfg_distZbg_maskZbg_distr%r%r&distance_field`s   zHausdorffDTLoss.distance_fieldinputtargetcCsn|dkr|dkrtd|jrt|}|jd}|jr0|dkr*tdnt|d}|jdur:||}|j rM|dkrGtdnt ||d}|j sn|dkrZtd n|ddddf}|ddddf}|j|jkrt d |jd |jd |j }g}t|jdD]Z}|dd|gf}|dd|gf}||} ||} ||d } | |j| |j} | | |} td t|j}|jrdg|}|| j|ddqtj|dd}|jtjjkrt|}|S|jtjjkr t |}|S|jtj!jkr.t"|jdd dgt|jd }|#|}|Std|jd)aK Args: input: the shape should be BNHW[D], where N is the number of classes. target: the shape should be BNHW[D] or B1HW[D], where N is the number of classes. Raises: ValueError: If the input is not 2D (NCHW) or 3D (NCHWD). AssertionError: When input and target (after one hot transform if set) have different shapes. ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"]. Example: >>> import torch >>> from monai.losses.hausdorff_loss import HausdorffDTLoss >>> from monai.networks.utils import one_hot >>> B, C, H, W = 7, 5, 3, 2 >>> input = torch.rand(B, C, H, W) >>> target_idx = torch.randint(low=0, high=C - 1, size=(B, H, W)).long() >>> target = one_hot(target_idx[:, None, ...], num_classes=C) >>> self = HausdorffDTLoss(reduction='none') >>> loss = self(input, target) >>> assert np.broadcast_shapes(loss.shape, input.shape) == input.shape z'Only 2D (NCHW) and 3D (NCHWD) supportedrz2single channel prediction, `softmax=True` ignored.Nz6single channel prediction, `to_onehot_y=True` ignored.) num_classesz>single channel prediction, `include_background=False` ignored.z"ground truth has different shape (z) from input ()rT)dimkeepdim)r9zUnsupported reduction: z0, available options are ["mean", "sum", "none"].)$r9r!rr)shaperwarningswarnrr rr AssertionErrordevicer+r1detachr r toaranger,tolistrappendmeancatrrMEANrSUMsumNONElistview)r"r2r3 n_pred_chr?Zall_fiZch_inputZ ch_targetZpred_dtZ target_dtZ pred_errordistanceZ running_f reduce_axisfbroadcast_shaper%r%r&forward{s`                & zHausdorffDTLoss.forward)r r r r r r rr rr rrrrrr rr)r'r(rr(r2r(r3r(rr() r __module__ __qualname____doc__rrGrr)no_gradr1rS __classcell__r%r%r#r&rs4 rcs"eZdZdZdfdd ZZS) LogHausdorffDTLossa Compute the logarithm of the Hausdorff Distance Transform Loss. This class computes the logarithm of the Hausdorff Distance Transform Loss, which is based on the distance transform. The logarithm is computed to potentially stabilize and scale the loss values, especially when the original loss values are very small. The formula for the loss is given by: log_loss = log(HausdorffDTLoss + 1) Inherits from the HausdorffDTLoss class to utilize its distance transform computation. r2r(r3rcstt||d}|S)a Compute the logarithm of the Hausdorff Distance Transform Loss. Args: input (torch.Tensor): The shape should be BNHW[D], where N is the number of classes. target (torch.Tensor): The shape should be BNHW[D] or B1HW[D], where N is the number of classes. Returns: torch.Tensor: The computed Log Hausdorff Distance Transform Loss for the given input and target. Raises: Any exceptions raised by the parent class HausdorffDTLoss. r)r)logrrS)r"r2r3Zlog_lossr#r%r&rSszLogHausdorffDTLoss.forwardrT)rrUrVrWrSrYr%r%r#r&rZs rZ) __future__rr<typingrr)torch.nn.modules.lossrmonai.networksrmonai.transforms.utilsr monai.utilsrrrZr%r%r%r&s      8