o  iAL@sddlmZddlZddlZddlmZddlmZddlm Z eddd\Z Z ed\Z Z Gd d d e Z Gd d d ejZGd ddejZd(d)ddZd*d+ddZddZGdddejZGd d!d!ejZd(d)d"d#Zd,d$d%Zd,d&d'ZdS)-) annotationsN)optional_import)StrEnumZlpipsLPIPS)name torchvisionc@s(eZdZdZdZdZdZdZdZdZ dS) PercetualNetworkTypealexvggsqueezeradimagenet_resnet50medicalnet_resnet10_23datasetsmedicalnet_resnet50_23datasetsresnet50N) __name__ __module__ __qualname__r r r r r rrrrY/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/losses/perceptual.pyrsrcsJeZdZdZejdddddddfdfdd Zd ddZd!ddZZ S)"PerceptualLossag Perceptual loss using features from pretrained deep neural networks trained. The function supports networks pretrained on: ImageNet that use the LPIPS approach from Zhang, et al. "The unreasonable effectiveness of deep features as a perceptual metric." https://arxiv.org/abs/1801.03924 ; RadImagenet from Mei, et al. "RadImageNet: An Open Radiologic Deep Learning Research Dataset for Effective Transfer Learning" https://pubs.rsna.org/doi/full/10.1148/ryai.210315 ; MedicalNet from Chen et al. "Med3D: Transfer Learning for 3D Medical Image Analysis" https://arxiv.org/abs/1904.00625 ; and ResNet50 from Torchvision: https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html . The fake 3D implementation is based on a 2.5D approach where we calculate the 2D perceptual loss on slices from all three axes and average. The full 3D approach uses a 3D network to calculate the perceptual loss. MedicalNet networks are only compatible with 3D inputs and support channel-wise loss. Args: spatial_dims: number of spatial dimensions. network_type: {``"alex"``, ``"vgg"``, ``"squeeze"``, ``"radimagenet_resnet50"``, ``"medicalnet_resnet10_23datasets"``, ``"medicalnet_resnet50_23datasets"``, ``"resnet50"``} Specifies the network architecture to use. Defaults to ``"alex"``. is_fake_3d: if True use 2.5D approach for a 3D perceptual loss. fake_3d_ratio: ratio of how many slices per axis are used in the 2.5D approach. cache_dir: path to cache directory to save the pretrained network weights. pretrained: whether to load pretrained weights. This argument only works when using networks from LIPIS or Torchvision. Defaults to ``"True"``. pretrained_path: if `pretrained` is `True`, users can specify a weights file to be loaded via using this argument. This argument only works when ``"network_type"`` is "resnet50". Defaults to `None`. pretrained_state_dict_key: if `pretrained_path` is not `None`, this argument is used to extract the expected state dict. This argument only works when ``"network_type"`` is "resnet50". Defaults to `None`. channel_wise: if True, the loss is returned per channel. Otherwise the loss is averaged over the channels. Defaults to ``False``. Tg?NF spatial_dimsint network_typestr is_fake_3dbool fake_3d_ratiofloat cache_dir str | None pretrainedpretrained_pathpretrained_state_dict_key channel_wisec st|dvr td|dks|rd|vrtd| r%d|vr%td|ttvr6tddt|rGtj |t d |d ||_ ||d kr]|d ur]t|d | d |_n"d|vrit|d d|_n|dkrwt||||d|_nt||d d|_||_||_| |_dS)N)z1Perceptual loss is implemented only in 2D and 3D.r$Z medicalnet_ziMedicalNet networks are only compatible with ``spatial_dims=3``.Argument is_fake_3d must be set to False.z>Channel-wise loss is only compatible with MedicalNet networks.zGUnrecognised criterion entered for Adversarial Loss. Must be one in: %sz, zSetting cache_dir to z@, this may change the default cache dir for all torch.hub calls.r%F)netverboser#Z radimagenet_)r&r'r)r&r r!r")r r&r')super__init__NotImplementedError ValueErrorlowerlistrjointorchhubset_dirwarningswarnrMedicalNetPerceptualSimilarityperceptual_functionRadImageNetPerceptualSimilarity$TorchvisionModelPerceptualSimilarityrrrr#) selfrrrrrr r!r"r# __class__rrr)FsN       zPerceptualLoss.__init__input torch.Tensortarget spatial_axisreturnc Csddd}gd}||d }||||ft|d }t|jd d t|jd |j|j}tj |d |d }||||ft|d } tj | d |d } t | || } | S)a Calculate perceptual loss in one of the axis used in the 2.5D approach. After the slices of one spatial axis is transformed into different instances in the batch, we compute the loss using the 2D approach. Args: input: input 5D tensor. BNHWD target: target 5D tensor. BNHWD spatial_axis: spatial axis to obtain the 2D slices. xr< fake_3d_permtupler?cSsH|d|}|d|j|d|j|d|j|d}|S)zg Transform slices from one spatial axis into different instances in the batch. )rr$r%)rpermute contiguousviewshape)r@rAslicesrrr batchify_axiss.z:PerceptualLoss._calculate_axis_loss..batchify_axisr$r%rD)r@rArN)dimindex)r@r<rArBr?r<) removerBr/randpermrHrrtodevice index_selectmeanr5) r8r;r=r>rJZpreserved_axes channel_axisZ input_slicesindicesZ target_slicesZ axis_lossrrr_calculate_axis_losss  (z#PerceptualLoss._calculate_axis_losscCs|j|jkrtd|jd|jd|jdkr:|jr:|j||dd}|j||dd}|j||dd}|||}n|||}|jrNtj| dd }|St|}|S) zx Args: input: the shape should be BNHW[D]. target: the shape should be BNHW[D]. z"ground truth has differing shape (z) from input ()r%r$)r>rLrrM) rHr+rrrWr5r#r/rTr )r8r;r=Z loss_sagittalZ loss_coronalZ loss_axiallossrrrforwards   zPerceptualLoss.forward)rrrrrrrrrrr rr!rr"rr#r)r;r<r=r<r>rr?r<r;r<r=r<r?r<) rrr__doc__rr r)rWr[ __classcell__rrr9rr$s$ <$rcs0eZdZdZ ddfd d ZdddZZS)r4a Component to perform the perceptual evaluation with the networks pretrained by Chen, et al. "Med3D: Transfer Learning for 3D Medical Image Analysis". This class uses torch Hub to download the networks from "Warvito/MedicalNet-models". Args: net: {``"medicalnet_resnet10_23datasets"``, ``"medicalnet_resnet50_23datasets"``} Specifies the network architecture to use. Defaults to ``"medicalnet_resnet10_23datasets"``. verbose: if false, mute messages from torch Hub load function. channel_wise: if True, the loss is returned per channel. Otherwise the loss is averaged over the channels. Defaults to ``False``. r Fr&rr'rr#r?NonecsRtddtj_tjjd||dd|_|||_| D]}d|_ q!dS)NcSsdS)NTr)abcrrrsz9MedicalNetPerceptualSimilarity.__init__..zwarvito/MedicalNet-modelsTmodelr' trust_repoF) r(r)r/r0_validate_not_a_forked_repoloadreevalr# parameters requires_grad)r8r&r'r#paramr9rrr)s   z'MedicalNetPerceptualSimilarity.__init__r;r<r=cCst|}t|}d}t|jdD]J}|dd|dfd}|dd|dfd}|dkrA|j|}|j|}|jd}qtj||j|gdd}tj||j|gdd}qt|} t|} | | d} |j rt | jd|jd| jd| jd| jd} t|jdD]$} | |}| d|}| dd|| |dfj dd| dd| df<qn| j dd d } t | d d } | S) a, Compute perceptual loss using MedicalNet 3D networks. The input and target tensors are inputted in the pre-trained MedicalNet that is used for feature extraction. Then, these extracted features are normalised across the channels. Finally, we compute the difference between the input and target features and calculate the mean value from the spatial dimensions to obtain the perceptual loss. Args: input: 3D input tensor with shape BCDHW. target: 3D target tensor with shape BCDHW. rrDN.rYr$r%rLTrMkeepdimrn) "medicalnet_intensity_normalisationrangerH unsqueezerer[r/catnormalize_tensorr#zerossumspatial_average_3d)r8r;r=Z feats_per_chZch_idx input_channelZtarget_channel outs_input outs_target feats_input feats_targetZ feats_diffresultsiZl_idxr_idxrrrr[s6     ( 2 z&MedicalNetPerceptualSimilarity.forward)r FF)r&rr'rr#rr?r_r\rrrr]r)r[r^rrr9rr4s  r4Tr@r<rnrr?cCs|jgd|dS)NrKrorTr@rnrrrrwrw绽|=epsrcCs&ttj|dddd}|||S)Nr$rDTrm)r/sqrtrv)r@r norm_factorrrrrts rtcCs|}|}|||S)zvBased on https://github.com/Tencent/MedicalNet/blob/18c8bb6cd564eb1b964bffef1f4c2283f1ae6e7b/datasets/brains18.py#L133)rTstd)volumerTrrrrrps rpcs.eZdZdZddfd d ZdddZZS)r6a Component to perform the perceptual evaluation with the networks pretrained on RadImagenet (pretrained by Mei, et al. "RadImageNet: An Open Radiologic Deep Learning Research Dataset for Effective Transfer Learning"). This class uses torch Hub to download the networks from "Warvito/radimagenet-models". Args: net: {``"radimagenet_resnet50"``} Specifies the network architecture to use. Defaults to ``"radimagenet_resnet50"``. verbose: if false, mute messages from torch Hub load function. r Fr&rr'rr?r_cs@ttjjd||dd|_||D]}d|_qdS)NzWarvito/radimagenet-modelsTrdF) r(r)r/r0rhrerirjrk)r8r&r'rlr9rrr)*s  z(RadImageNetPerceptualSimilarity.__init__r;r<r=cCs|jddkr|jddkr|dddd}|dddd}|ddgddf}|ddgddf}t|}t|}|j|}|j|}t|}t|}||d}t|jddddd }|S) a We expect that the input is normalised between [0, 1]. Given the preprocessing performed during the training at https://github.com/BMEII-AI/RadImageNet, we make sure that the input and target have 3 channels, reorder it from 'RGB' to 'BGR', and then remove the mean components of each input data channel. The outputs are normalised across the channels, and we obtain the mean from the spatial dimensions (similar approach to the lpips package). rDr%N)r$rDr.r$Trmro)rHrepeat subtract_meanrer[rtspatial_averagervr8r;r=ryrzr{r|r}rrrr[2s   z'RadImageNetPerceptualSimilarity.forward)r F)r&rr'rr?r_r\rrrr9rr6s r6cs6eZdZdZ    ddfdd ZdddZZS)r7a Component to perform the perceptual evaluation with TorchVision models. Currently, only ResNet50 is supported. The network structure is based on: https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html Args: net: {``"resnet50"``} Specifies the network architecture to use. Defaults to ``"resnet50"``. pretrained: whether to load pretrained weights. Defaults to `True`. pretrained_path: if `pretrained` is `True`, users can specify a weights file to be loaded via using this argument. Defaults to `None`. pretrained_state_dict_key: if `pretrained_path` is not `None`, this argument is used to extract the expected state dict. Defaults to `None`. rTNr&rr rr!rr"r?r_c stdg}||vrtd|d|d|dur*tjj|r%tjjjndd}ntjjdd}|durItj |dd}|durD||}| |d|_ tjj ||j g|_||D]}d |_q_dS) Nrz'net' z0 is not supported, please select a network from .)weightsT) weights_onlyzlayer4.2.relu_2F)r(r)r*rmodelsrResNet50_WeightsDEFAULTr/rhload_state_dict final_layerfeature_extractioncreate_feature_extractorrerirjrk) r8r&r r!r"Zsupported_networksnetwork state_dictrlr9rrr)ds,   z-TorchvisionModelPerceptualSimilarity.__init__r;r<r=cCs|jddkr|jddkr|dddd}|dddd}t|}t|}|j||j}|j||j}t|}t|}||d}t|jddddd}|S)a We expect that the input is normalised between [0, 1]. Given the preprocessing performed during the training at https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html#torchvision.models.ResNet50_Weights, we make sure that the input and target have 3 channels, and then do Z-Score normalization. The outputs are normalised across the channels, and we obtain the mean from the spatial dimensions (similar approach to the lpips package). rDr%r$Trmro) rHrtorchvision_zscore_normrer[rrtrrvrrrrr[s  z,TorchvisionModelPerceptualSimilarity.forward)rTNN) r&rr rr!rr"rr?r_r\rrrr9rr7Ts r7cCs|jddg|dS)Nr$r%rorrrrrrrrcCsgd}gd}|dddddddf|d|d|dddddddf<|dddddddf|d|d|dddddddf<|dddddddf|d|d|dddddddf<|S)N) ףp= ?v/?Cl?)gZd;O?gy&1?g?rrDr$r)r@rTrrrrrs DDDrcCsgd}|dddddddf|d8<|dddddddf|d8<|dddddddf|d8<|S)N)rrrrrDr$r)r@rTrrrrs (((r)T)r@r<rnrr?r<)r)r@r<rrr?r<)r@r<r?r<) __future__rr2r/torch.nnnn monai.utilsrmonai.utils.enumsrr_rrModulerr4rwrtrpr6r7rrrrrrrs(       M 6 O