U Phq,@sdZddlmZddlZddlmZddlmZddl m Z m Z ddl m Z mZddd d d Zddd d d Zddd ddZddddddZddd ddZd-ddddddZdddddd Zd!d"d#d$d%Zd.ddddd'd(d)Zd/ddddd*d+d,ZdS)0zW This script contains utility functions for developing new networks/blocks in PyTorch. ) annotationsN)Tensor) functional)complex_conj_t complex_mul_t)fftn_centered_tifftn_centered_tr)xreturncCs|jddkr$td|jddt|jdkrh|j\}}}}}|ddddd |d|||St|jd kr|j\}}}}}}|ddddd d|d||||Std |jd S) 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 .rLonly 2D (B,C,H,W,2) and 3D (B,C,H,W,D,2) data are supported but x has shape N)shape ValueErrorlenpermute contiguousviewr bchwtwodr!b/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/reconstruction/networks/nets/utils.pyreshape_complex_to_channel_dims &*r#cCs|jdddkr(td|jddt|jdkrl|j\}}}}|d}||d|||dddddSt|jdkr|j\}}}}}|d}||d||||ddddddStd |jd S) 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 N)rrrrr)r rc2rrrr r!r!r"#reshape_channel_complex_to_last_dim4s  $r%ztuple[Tensor, int]cCst|jdkr<|j\}}}}}|||d||||fSt|jdkr||j\}}}}}}|||d|||||fStd|jdS)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 rrrrN)rrrrrrr!r!r"reshape_channel_to_batch_dimOs  r&int)r batch_sizer c Cst|jdkr8|j\}}}}}||}||||||St|jdkrt|j\}}}}}}||}|||||||Std|jdS)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 N)rrrr) r r(bconerrrrr r!r!r"$reshape_batch_channel_to_channel_dimhs r+ztuple[Tensor, Tensor, Tensor]cCst|jdkr|j\}}}}||d|d||}|jdd|dddd|d|ddd||dd}|jddd|dddd|d|ddd||dd}|||||}|||||fSt|jdkr|j\}}}}}||d|d|||}|jdd|ddddd|d|dddd||ddd}|jddd|ddddd|d|dddd||ddd}||||||}|||||fStd|jd S) 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,unbiasedrHonly 2D (B,C,H,W) and 3D (B,C,H,W,D) data are supported but x has shape N)rrrrmeanexpandstdr)r rrrrr/r1r r!r!r"complex_normalizesH<>"Br2zVtuple[Tensor, tuple[tuple[int, int], tuple[int, int], tuple[int, int], int, int, int]])r kr cCsXt|jdkr|j\}}}}|d|dBd}|d|dBd}t||d}t||d} t||| }d} d} | || ||| f} nt|jdkr@|j\}}}}} |d|dBd}|d|dBd}| d|dBd} t||d}t||d} t| | d} t|| || }| || ||| f} ntd|j|| fS)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 rr.)rr floor_ceilFpadr)r r4rrrrw_multh_multw_padh_padd_multd_pad pad_sizesr r!r!r"divisible_pad_ts,!r?zGtuple[tuple[int, int], tuple[int, int], tuple[int, int], int, int, int])r r>r cCs|\}}}}}}t|jdkrL|d|d||d|d||dfSt|jdkr|d|d||d|d||d|d||dfStd|jdS)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.rrrr.N)rrr)r r>r;r:r=r9r8r<r!r!r"inverse_divisible_pad_ts .@r@floatztuple[int, int])nr 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) )mathfloorceil)rBr!r!r"r5s r5r )kspace sens_maps spatial_dimsr cCs&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. TrH is_complexr)r,keepdim)rrrsum)rFrGrHimgr!r!r"sensitivity_map_reducesrN)rMrGrHr cCstt|||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. TrI)rr)rMrGrHr!r!r"sensitivity_map_expand#srO)r3)r )r )__doc__ __future__rrCtorchrtorch.nnrr6Z'monai.apps.reconstruction.complex_utilsrrZ!monai.networks.blocks.fft_utils_trrr#r%r&r+r2r?r@r5rNrOr!r!r!r" s"   )>