U kPh L@sddlmZddlZddlZddlmZddlmZddlm Z eddd\Z Z ed\Z Z Gd d d e Z Gd d d ejZGd ddejZd'ddddddZd(ddddddZddZGdddejZGdddejZd)ddddd d!Zddd"d#d$Zddd"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 rrrrL/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/losses/perceptual.pyrsrc speZdZdZejdddddddfdddd d dd d dd fd d ZdddddddZddddddZZ 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?NFintstrboolfloat str | None) spatial_dims network_type is_fake_3d fake_3d_ratio cache_dir pretrainedpretrained_pathpretrained_state_dict_key channel_wisec st|dkrtd|dks&|r6d|kr6td| rJd|krJtd|ttkrltddt|rtj |t d |d ||_ ||d kr|d krt|d | d |_nDd|krt|d d|_n,|dkrt||||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)FsR      zPerceptualLoss.__init__ torch.Tensor)inputtarget spatial_axisreturnc Csdddddd}dddg}||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. r;tuple)x fake_3d_permr?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)rArBslicesrrr batchify_axiss.z:PerceptualLoss._calculate_axis_loss..batchify_axisr$r%rD)rArBrN)dimindex) remover@r/randpermrHrrtodevice index_selectmeanr5) r8r<r=r>rJZpreserved_axes channel_axisZ input_slicesindicesZ target_slicesZ axis_lossrrr_calculate_axis_losss   (z#PerceptualLoss._calculate_axis_lossr<r=r?cCs|j|jkr&td|jd|jd|jdkrt|jrt|j||dd}|j||dd}|j||dd}|||}n |||}|jrtj| dd }n t|}|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>rKrrL) rHr+rrrVr5r#r/rSr )r8r<r=Z loss_sagittalZ loss_coronalZ loss_axiallossrrrforwards   zPerceptualLoss.forward) rrr__doc__rr r)rVr[ __classcell__rrr9rr$s$$<$rcs@eZdZdZddddddfdd Zd d d d d d ZZS)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 FrrNone)r&r'r#r?csPtddtj_tjjd||d|_|||_| D] }d|_ q@dS)NcSsdS)NTr)abcrrrz9MedicalNetPerceptualSimilarity.__init__..zwarvito/MedicalNet-modelsmodelr'F) r(r)r/r0_validate_not_a_forked_repoloadreevalr# parameters requires_grad)r8r&r'r#paramr9rrr)s   z'MedicalNetPerceptualSimilarity.__init__r;rWcCst|}t|}d}t|jdD]}|dd|dfd}|dd|dfd}|dkr|j|}|j|}|jd}q"tj||j|gdd}tj||j|gdd}q"t|} t|} | | d} |j rht | jd|jd| jd| jd| jd} t|jdD]J} | |}| 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%rKTrLkeepdimrm) "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[s<      6 z&MedicalNetPerceptualSimilarity.forward)r FFrrrr\r)r[r]rrr9rr4s  r4Tr;r)rArmr?cCs|jdddg|dS)Nr$r%rKrnrSrArmrrrrvsrv绽|=r)rAepsr?cCs&ttj|dddd}|||S)Nr$rDTrl)r/sqrtru)rAr norm_factorrrrrssrscCs|}|}|||S)zvBased on https://github.com/Tencent/MedicalNet/blob/18c8bb6cd564eb1b964bffef1f4c2283f1ae6e7b/datasets/brains18.py#L133)rSstd)volumerSrrrrrosrocs>eZdZdZdddddfdd Zd d d d d d ZZS)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 Frrr^)r&r'r?cs>ttjjd||d|_||D] }d|_q.dS)NzWarvito/radimagenet-modelsrdF) r(r)r/r0rgrerhrirj)r8r&r'rkr9rrr)*s   z(RadImageNetPerceptualSimilarity.__init__r;rWcCs|jddkr<|jddkr<|dddd}|dddd}|dddddgdf}|dddddgdf}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%Nr$r.Trlrn)rHrepeat subtract_meanrer[rsspatial_averagerur8r<r=rxryrzr{r|rrrr[2s   z'RadImageNetPerceptualSimilarity.forward)r Frrrr9rr6s r6csBeZdZdZddddddd fd d Zd d d d ddZZS)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`. rTNrrrr^)r&r r!r"r?c stdg}||kr.td|d|d|dkrTtjj|rJtjjjndd}n:tjjdd}|dkrt |}|dk r||}| |d|_ tjj ||j g|_||D] }d|_qdS) Nrz'net' z0 is not supported, please select a network from .)weightsTzlayer4.2.relu_2F)r(r)r*rmodelsrResNet50_WeightsDEFAULTr/rgload_state_dict final_layerfeature_extractioncreate_feature_extractorrerhrirj) r8r&r r!r"Zsupported_networksnetwork state_dictrkr9rrr)ds*    z-TorchvisionModelPerceptualSimilarity.__init__r;rWcCs|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$Trlrn) rHrtorchvision_zscore_normrer[rrsrrurrrrr[s  z,TorchvisionModelPerceptualSimilarity.forward)rTNNrrrr9rr7Ts r7cCs|jddg|dS)Nr$r%rnrrrrrrsr)rAr?cCsdddg}dddg}|dddddddf|d|d|dddddddf<|dddddddf|d|d|dddddddf<|ddd ddddf|d |d |ddd ddddf<|S) N ףp= ?v/?Cl?gZd;O?gy&1?g?rrDr$r)rArSrrrrrs   DDDrcCsdddg}|dddddddf|d8<|dddddddf|d8<|dddddddf|d8<|S)NrrrrrDr$r)rArSrrrrs  (((r)T)r)T) __future__rr2r/torch.nnnn monai.utilsrmonai.utils.enumsrr_rrModulerr4rvrsror6r7rrrrrrr s&      M6O