o  i@snddlmZddlmZddlmZddlmZddlm Z m Z m Z m Z m Z ddlmZGdddejZdS) ) annotations)SequenceN)Tensor)complex_normalizedivisible_pad_tinverse_divisible_pad_t#reshape_channel_complex_to_last_dimreshape_complex_to_channel_dim) BasicUNetc sReZdZdZddddddfdd difdd d d d f d%fdd Zd&d#d$ZZS)' ComplexUneta This variant of U-Net handles complex-value input/output. It can be used as a model to learn sensitivity maps in multi-coil MRI data. It is built based on :py:class:`monai.networks.nets.BasicUNet` by default but the user can input their convolutional model as well. ComplexUnet also applies default normalization to the input which makes it more stable to train. The data being a (complex) 2-channel tensor is a requirement for using this model. Modified and adopted from: https://github.com/facebookresearch/fastMRI Args: spatial_dims: number of spatial dimensions. features: six integers as numbers of features. denotes number of channels in each layer. act: activation type and arguments. Defaults to LeakyReLU. norm: feature normalization type and arguments. Defaults to instance norm. bias: whether to have a bias term in convolution blocks. Defaults to True. dropout: dropout ratio. Defaults to 0.0. upsample: upsampling mode, available options are ``"deconv"``, ``"pixelshuffle"``, ``"nontrainable"``. pad_factor: an integer denoting the number which each padded dimension will be divisible to. For example, 16 means each dimension will be divisible by 16 after padding conv_net: the learning model used inside the ComplexUnet. The default is :py:class:`monai.networks.nets.basic_unet`. The only requirement on the model is to have 2 as input and output number of channels. ) r @r LeakyReLUg?T)negative_slopeinplaceinstanceaffinegdeconvN spatial_dimsintfeatures Sequence[int]act str | tuplenormbiasbooldropout float | tupleupsamplestr pad_factorconv_netnn.Module | Nonec s~t|| durt|dd||||||d |_n dd| D} | dddkr7td| ddd| |_||_dS) Nr ) r in_channels out_channelsrrrrr!r#cSsg|]}|jqS)shape).0pr*r*v/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/reconstruction/networks/nets/complex_unet.py Vsz(ComplexUnet.__init__..rz!in_channels should be 2 but it's .)super__init__r unet parameters ValueErrorr%) selfrrrrrr!r#r%r&params __class__r*r.r39s&   zComplexUnet.__init__xrreturncCsTt|}t|\}}}t||jd\}}||}t||}|||}t|}|S)z 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,H,W,2) for 2D data or (B,C,H,W,D,2) for 3D data )k)r rrr%r4rr)r7r;meanstdZ padding_sizesr*r*r.forward]s    zComplexUnet.forward)rrrrrrrrrr r!r"r#r$r%rr&r')r;rr<r)__name__ __module__ __qualname____doc__r3r@ __classcell__r*r*r9r.r s  $r ) __future__rcollections.abcrtorch.nnnntorchr-monai.apps.reconstruction.networks.nets.utilsrrrrr Zmonai.networks.nets.basic_unetr Moduler r*r*r*r.s