U Ph@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 sleZdZdZddddddfdd difdd d d d f dddddddddd fdd ZdddddZZS) 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_slopeinplaceinstanceaffinegdeconvNintz Sequence[int]z str | tupleboolz float | tuplestrznn.Module | None) spatial_dimsfeaturesactnormbiasdropoutupsample pad_factorconv_netc s~t|| dkr4t|dd||||||d |_n@dd| D} | dddkrntd| ddd| |_||_dS) Nr ) r in_channels out_channelsrrrrr r!cSsg|] }|jqS)shape).0pr&r&i/home/dell461/cl/sdc2/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*r/9s&  zComplexUnet.__init__r)xreturncCsTt|}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"r0rr)r3r7meanstdZ padding_sizesr&r&r*forward]s     zComplexUnet.forward)__name__ __module__ __qualname____doc__r/r< __classcell__r&r&r5r*r s  $$r ) __future__rcollections.abcrtorch.nnnntorchr-monai.apps.reconstruction.networks.nets.utilsrrrrr Zmonai.networks.nets.basic_unetr Moduler r&r&r&r* s