o  iq,ã@sÀdZddlmZddlZddlmZddlmZddl m Z m Z ddl m Z mZd0d d „Zd0d d„Zd1dd„Zd2dd„Zd3dd„Z d4d5dd„Zd6d d!„Zd7d%d&„Zd8d9d+d,„Zd8d:d.d/„ZdS);zW This script contains utility functions for developing new networks/blocks in PyTorch. é)Ú annotationsN)ÚTensor)Ú functional)Úcomplex_conj_tÚ complex_mul_t)Úfftn_centered_tÚifftn_centered_tÚxrÚreturncCsÂ|jddkrtd|jd›dƒ‚t|jƒdkr4|j\}}}}}| ddddd ¡ ¡ |d|||¡St|jƒd krY|j\}}}}}}| ddddd d¡ ¡ |d||||¡Std |j›ƒ‚) aC Swaps the complex dimension with the channel dimension so that the network treats real/imaginary parts as two separate channels. Args: x: input of shape (B,C,H,W,2) for 2D data or (B,C,H,W,D,2) for 3D data Returns: output of shape (B,C*2,H,W) for 2D data or (B,C*2,H,W,D) for 3D data éÿÿÿÿéz'last dim must be 2, but x.shape[-1] is Ú.érééééúLonly 2D (B,C,H,W,2) and 3D (B,C,H,W,D,2) data are supported but x has shape )ÚshapeÚ ValueErrorÚlenÚpermuteÚ contiguousÚview©r ÚbÚcÚhÚwÚtwoÚd©r!úo/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/reconstruction/networks/nets/utils.pyÚreshape_complex_to_channel_dims &*r#cCsÆ|jdddkrtd|jd›dƒ‚t|jƒdkr6|j\}}}}|d}| |d|||¡ ddddd¡St|jƒdkr[|j\}}}}}|d}| |d||||¡ dddddd¡Std |j›ƒ‚) a, Swaps the complex dimension with the channel dimension so that the network output has 2 as its last dimension Args: x: input of shape (B,C*2,H,W) for 2D data or (B,C*2,H,W,D) for 3D data Returns: output of shape (B,C,H,W,2) for 2D data or (B,C,H,W,D,2) for 3D data rr rz&channel dimension should be even but (z ) is odd.rrrzLonly 2D (B,C*2,H,W) and 3D (B,C*2,H,W,D) data are supported but x has shape )rrrrr)r rÚc2rrrr r!r!r"Ú#reshape_channel_complex_to_last_dim4s  $r%útuple[Tensor, int]cCsŒt|jƒdkr|j\}}}}}| ¡ ||d|||¡|fSt|jƒdkr>|j\}}}}}}| ¡ ||d||||¡|fStd|j›ƒ‚)zû Combines batch and channel dimensions. Args: x: input of shape (B,C,H,W,2) for 2D data or (B,C,H,W,D,2) for 3D data Returns: A tuple containing: (1) output of shape (B*C,1,...) (2) batch size rrrr)rrrrrrr!r!r"Úreshape_channel_to_batch_dimOs  r'Ú batch_sizeÚintc Cs„t|jƒdkr|j\}}}}}||}| |||||¡St|jƒdkr:|j\}}}}}}||}| ||||||¡Std|j›ƒ‚)zÝ Detaches batch and channel dimensions. Args: x: input of shape (B*C,1,H,W,2) for 2D data or (B*C,1,H,W,D,2) for 3D data batch_size: batch size Returns: output of shape (B,C,...) rrzPonly 2D (B*C,1,H,W,2) and 3D (B*C,1,H,W,D,2) data are supported but x has shape )rrrr) r r(ÚbcÚonerrrrr r!r!r"Ú$reshape_batch_channel_to_channel_dimhs r,útuple[Tensor, Tensor, Tensor]cCsÐt|jƒdkrk|j\}}}}| ¡ |d|d||¡}|jdd |dddd¡ |d|ddd¡ ¡ ||dd¡}|jddd |dddd¡ |d|ddd¡ ¡ ||dd¡}| ||||¡}|||||fSt|jƒdkrà|j\}}}}}| ¡ |d|d|||¡}|jdd |ddddd¡ |d|dddd¡ ¡ ||ddd¡}|jddd |ddddd¡ |d|dddd¡ ¡ ||ddd¡}| |||||¡}|||||fStd|j›ƒ‚) a¼ Performs layer mean-std normalization for complex data. Normalization is done for each batch member along each part (part refers to real and imaginary parts), separately. Args: x: input of shape (B,C,H,W) for 2D data or (B,C,H,W,D) for 3D data Returns: A tuple containing (1) normalized output of shape (B,C,H,W) for 2D data or (B,C,H,W,D) for 3D data (2) mean (3) std rr )ÚdimrF)r.ÚunbiasedrúHonly 2D (B,C,H,W) and 3D (B,C,H,W,D) data are supported but x has shape )rrrrÚmeanÚexpandÚstdr)r rrrrr1r3r r!r!r"Úcomplex_normalizes(<>"B ûr4éÚkúVtuple[Tensor, tuple[tuple[int, int], tuple[int, int], tuple[int, int], int, int, int]]cCsZt|jƒdkrJ|j\}}}}|d|dBd}|d|dBd}t||dƒ}t||dƒ} t ||| ¡}d} d} | || ||| f} || fSt|jƒdkr¥|j\}}}}} |d|dBd}|d|dBd}| d|dBd} t||dƒ}t||dƒ} t| | dƒ} t || || ¡}| || ||| f} || fStd|j›ƒ‚)ah Pad input to feed into the network (torch script compatible) Args: x: input of shape (B,C,H,W) for 2D data or (B,C,H,W,D) for 3D data k: padding factor. each padded dimension will be divisible by k. Returns: A tuple containing (1) padded input (2) pad sizes (in order to reverse padding if needed) Example: .. code-block:: python import torch # 2D data x = torch.ones([3,2,50,70]) x_pad,padding_sizes = divisible_pad_t(x, k=16) # the following line should print (3, 2, 64, 80) print(x_pad.shape) # 3D data x = torch.ones([3,2,50,70,80]) x_pad,padding_sizes = divisible_pad_t(x, k=16) # the following line should print (3, 2, 64, 80, 80) print(x_pad.shape) rrr r )r r rr0)rrÚ floor_ceilÚFÚpadr)r r6rrrrÚw_multÚh_multÚw_padÚh_padÚd_multÚd_padÚ pad_sizesr r!r!r"Údivisible_pad_t©s.!òþrBrAúGtuple[tuple[int, int], tuple[int, int], tuple[int, int], int, int, int]cCsª|\}}}}}}t|jƒdkr&|d|d||d…|d||d…fSt|jƒdkrM|d|d||d…|d||d…|d||d…fStd|j›ƒ‚)zÙ De-pad network output to match its original shape Args: x: input of shape (B,C,H,W) for 2D data or (B,C,H,W,D) for 3D data pad_sizes: padding values Returns: de-padded input r.rrrr0)rrr)r rAr>r=r@r<r;r?r!r!r"Úinverse_divisible_pad_tçs  .@rDÚnÚfloatútuple[int, int]cCst |¡t |¡fS)z« Returns floor and ceil of the input Args: n: input number Returns: A tuple containing: (1) floor(n) (2) ceil(n) )ÚmathÚfloorÚceil)rEr!r!r"r8s r8r ÚkspaceÚ sens_mapsÚ spatial_dimscCs&t||dd}t|t|ƒƒjdddS)a Reduces coil measurements to a corresponding image based on the given sens_maps. Let's say there are C coil measurements inside kspace, then this function multiplies the conjugate of each coil sensitivity map with the corresponding coil image. The result of this process will be C images. Summing those images together gives the resulting "reduced image." Args: kspace: 2D kspace (B,C,H,W,2) with the last dimension being 2 (for real/imaginary parts) and C denoting the coil dimension. 3D data will have the shape (B,C,H,W,D,2). sens_maps: sensitivity maps of the same shape as input x. spatial_dims: is 2 for 2D data and is 3 for 3D data Returns: reduction of x to (B,1,H,W,2) for 2D data or (B,1,H,W,D,2) for 3D data. T©rMÚ is_complexr)r.Úkeepdim)rrrÚsum)rKrLrMÚimgr!r!r"Úsensitivity_map_reducesrSrRcCstt||ƒ|ddS)an Expands an image to its corresponding coil images based on the given sens_maps. Let's say there are C coils. This function multiples image img with each coil sensitivity map in sens_maps and stacks the resulting C coil images along the channel dimension which is reserved for coils. Args: img: 2D image (B,1,H,W,2) with the last dimension being 2 (for real/imaginary parts). 3D data will have the shape (B,1,H,W,D,2). sens_maps: Sensitivity maps for combining coil images. The shape is (B,C,H,W,2) for 2D data or (B,C,H,W,D,2) for 3D data (C denotes the coil dimension). spatial_dims: is 2 for 2D data and is 3 for 3D data Returns: Expansion of x to (B,C,H,W,2) for 2D data and (B,C,H,W,D,2) for 3D data. The output is transferred to the frequency domain to yield coil measurements. TrN)rr)rRrLrMr!r!r"Úsensitivity_map_expand#srT)r rr r)r rr r&)r rr(r)r r)r rr r-)r5)r rr6r)r r7)r rrArCr r)rErFr rG)r )rKrrLrrMr)r r)rRrrLrrMr)r r)Ú__doc__Ú __future__rrHÚtorchrÚtorch.nnrr9Z'monai.apps.reconstruction.complex_utilsrrZ!monai.networks.blocks.fft_utils_trrr#r%r'r,r4rBrDr8rSrTr!r!r!r"Ús$        ) ÿ >