o , ia2@sddlmZddlZddlmZddlmZmZddlm Z ddl m Z ddl m Z ddlmZddlmZmZGd d d ejZGd d d e ZGd ddejZdS)) annotationsN)CABlock FeedForward) Convolution) DownSample)UpSample)Norm)DownsampleMode UpsampleModecs2eZdZdZ  ddfd d ZdddZZS)MDTATransformerBlockaBasic transformer unit combining MDTA and GDFN with skip connections. Unlike standard transformers that use LayerNorm, this block uses Instance Norm for better adaptation to image restoration tasks. Args: spatial_dims: Number of spatial dimensions (2D or 3D) dim: Number of input channels num_heads: Number of attention heads ffn_expansion_factor: Expansion factor for feed-forward network bias: Whether to use bias in attention layers layer_norm_use_bias: Whether to use bias in layer normalization. Defaults to False. flash_attention: Whether to use flash attention optimization. Defaults to False. F spatial_dimsintdim num_headsffn_expansion_factorfloatbiasboollayer_norm_use_biasflash_attentioncs`tttj|f||d|_t||||||_ttj|f||d|_t|||||_ dS)N)affine) super__init__rINSTANCEnorm1rattnnorm2rffn)selfr rrrrrr __class___/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/restormer.pyr's zMDTATransformerBlock.__init__x torch.TensorreturncCs,||||}||||}|SN)rrrrrr#r!r!r"forward7szMDTATransformerBlock.forward)FF)r r rr rr rrrrrrrrr#r$r%r$__name__ __module__ __qualname____doc__rr( __classcell__r!r!rr"r s r cs2eZdZdZddfd d Zdfdd ZZS)OverlapPatchEmbedaInitial feature extraction using overlapped convolutions. Unlike standard patch embeddings that use non-overlapping patches, this approach maintains spatial continuity through 3x3 convolutions. Args: spatial_dims: Number of spatial dimensions (2D or 3D) in_channels: Number of input channels embed_dim: Dimension of embedded features. Defaults to 48. bias: Whether to use bias in convolution layer. Defaults to False. 0Fr r in_channels embed_dimrrc s tj|||ddd|dddS)Nr1Tr r3 out_channels kernel_sizestridespaddingr conv_only)rr)rr r3r4rrr!r"rIs zOverlapPatchEmbed.__init__r#r$r%cst|}|Sr&)rr(r'rr!r"r(Us zOverlapPatchEmbed.forward)r1r2F)r r r3r r4r rrr)r*r!r!rr"r0=s  r0csFeZdZdZ           d"d#fdd Zd$d d!ZZS)% Restormera{Restormer: Efficient Transformer for High-Resolution Image Restoration. Implements a U-Net style architecture with transformer blocks, combining: - Multi-scale feature processing through progressive down/upsampling - Efficient attention via MDTA blocks - Local feature mixing through GDFN - Skip connections for preserving spatial details Architecture: - Encoder: Progressive feature downsampling with increasing channels - Latent: Deep feature processing at lowest resolution - Decoder: Progressive upsampling with skip connections - Refinement: Final feature enhancement r1r2r5r5r5r5HzG@FTr r r3r7r num_blockstuple[int, ...]headsnum_refinement_blocksrrrrrdual_pixel_taskrr%Nonec  st t|dksJdt|tksJdtdd|Ds)Jdt |||_t|_t|_ t|_ t|_ t|_ t|d |_ |_d d} t D]7|d| }|jtj fdd t|D|j t|j|tjdd q_|d tj fd d t| D|_tt D]T|d|dd}|j t|j|tjdd d dkr|j t|j|dddn||j tj fdd t|Dqtj fdd t|D|_| |_|jr:t|j||dddd|_t|j|d|ddddd|_dS)Nr5z'Number of blocks must be greater than 1z(Number of blocks and heads must be equalcss|]}|dkVqdS)rNr!).0nr!r!r" sz%Restormer.__init__..z'Number of blocks must be greater than 0r=c &g|]}tdqSr rrrrrrr rG_)r current_dimrrrCrrHr r!r"  z&Restormer.__init__..)r r3r7mode scale_factorrc s&g|]}tdqSrKrMrN)rrrrC latent_dimr num_stepsr r!r"rQrRF)r r3r7rSrTrapply_pad_poolrT)r r3r7r8rr;c rJrKrMrN)r decoder_dimrrrCrrHr r!r"rQrRc s&g|]}tddqS)rrLrMrN)rrXrrrCrr r!r"rQrRr1r6)rrlenallr0 patch_embednn ModuleListencoder_levels downsamplesdecoder_levels upsamplesreduce_channelsrVr rangeappend Sequentialrr PIXELUNSHUFFLElatentreversedrr PIXELSHUFFLEr refinementrE skip_convoutput)rr r3r7rrArCrDrrrrErZspatial_multiplierZnext_dimr) rrPrXrrrCrUrrHrVr r"rjs                   zRestormer.__init__r#r$cs&tfddtdjdDsJdg}ttjjD]\}\}}|||q& tt j D]-}j |t ||d gd|t j dkrlj|j |qFjr|dSS)aForward pass of Restormer. Processes input through encoder-decoder architecture with skip connections. Args: inp_img: Input image tensor of shape (B, C, H, W, [D]) Returns: Restored image tensor of shape (B, C, H, W, [D]) c3s&|]}j| djkVqdS)r=N)shaperV)rGir'r!r"rI/s z$Restormer.forward..r5z=All spatial dimensions should be larger than 2^number_of_stepr)rZrcr r[ enumeratezipr^r_rdrgrYr`ratorchconcatrbrjrErkrl)rr#Zskip_connections_idxencoder downsampleidxr!r'r"r(&s2         zRestormer.forward) r=r1r1r2r>r>r?r@FTFF)r r r3r r7r rr rArBrCrBrDr rrrrrrrErrrr%rFr)r*r!r!rr"r<Zs"=r<) __future__rrqtorch.nnr\Zmonai.networks.blocks.cablockrr"monai.networks.blocks.convolutionsrZ monai.networks.blocks.downsamplermonai.networks.blocks.upsamplermonai.networks.layers.factoriesrmonai.utils.enumsr r Moduler r0r<r!r!r!r"s      %