o - i@srddlmZddlZddlmZmZmZmZmZddl Z ddl Z ddl m mZddl m Z ddlZddlmZmZddlmZddlmZddlmZdd lmZdd lmZmZmZed d d \ZZ ddgZ!d%d&ddZ"Gddde j#Z$Gddde j#Z%Gddde j#Z&Gddde j#Z'Gddde j#Z(Gdd d e j#Z)Gd!d"d"e j#Z*Gd#d$d$e j#Z+dS)') annotationsN)AnyCallableOptionalSequenceTuple)nn)MLPBlockUnetrBasicBlock) SegResNetDS2)convert_points_to_disc)!keep_merge_components_with_points)sample_points_from_label)optional_importunsqueeze_leftunsqueeze_righteinops rearrange)nameVISTA3D vista3d1320encoder_embed_dimint in_channelscCsBt|dd||dd}t|ddd}td|dd }t|||d }|S) a Exact VISTA3D network configuration used in https://arxiv.org/abs/2406.05285>`_. The model treats class index larger than 132 as zero-shot. Args: encoder_embed_dim: hidden dimension for encoder. in_channels: input channel number. )rrrinstancer)r blocks_downnorm out_channels init_filtersdsdepth) feature_size n_classeslast_supportedT)r'r&use_mlp) image_encoder class_head point_head)r PointMappingSAMClassMappingClassifyr)rr segresnetr,r+Zvistar0]/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/vista3d.pyr#s cseZdZdZdGfdd ZdHddZdIddZ  dJdKddZ   !dLdMd)d*ZdNd,d-Z .dOdPd4d5Z dQdRd8d9Z dSdTd=d>Z           :dUdVdEdFZ ZS)Wra VISTA3D based on: `VISTA3D: Versatile Imaging SegmenTation and Annotation model for 3D Computed Tomography `_. Args: image_encoder: image encoder backbone for feature extraction. class_head: class head used for class index based segmentation point_head: point head used for interactive segmetnation r* nn.Moduler+r,cs>t||_||_||_d|_d|_d|_d|_d|_ dS)NFii') super__init__r*r+r,image_embeddings auto_freeze point_freeze NINF_VALUEZ PINF_VALUE)selfr*r+r, __class__r0r1r4Fs  zVISTA3D.__init__pad_size list | Nonelabelstorch.Tensor | None prev_mask point_coordscCs~|dur |||fS|durtj||ddd}|dur#tj||ddd}|dur:|tj|d|d|dg|jd}|||fS) a Image has been padded by sliding window inferer. The related padding need to be performed outside of slidingwindow inferer. Args: pad_size: padding size passed from sliding window inferer. labels: image label ground truth. prev_mask: previous segmentation mask. point_coords: point click coordinates. Nconstantr)padmodevalueidevice)FrCtorchtensorrI)r9r<r>r@rAr0r0r1update_slidingwindow_paddingQs  z$VISTA3D.update_slidingwindow_padding class_vectorreturnrcCs,|dur|dur td|jdS|jdS)zAGet number of foreground classes based on class and point prompt.Nz2class_vector and point_coords cannot be both None.r) ValueErrorshape)r9rNrAr0r0r1get_foreground_class_countns   z"VISTA3D.get_foreground_class_countN9 point_label torch.Tensor label_setSequence[int] | None special_index Sequence[int]cCs|dur|S|jdt|kstdtt|D]-}|||vrFtt||D]}|||fdkr;|||fdn|||f|||f<q)q|S)a Convert point label based on its class prompt. For special classes defined in special index, the positive/negative point label will be converted from 1/0 to 3/2. The purpose is to separate those classes with ambiguous classes. Args: point_label: the point label tensor, [B, N]. label_set: the label index matching the indexes in labels. If labels are mapped to global index using RelabelID, this label_set should be global mapped index. If labels are not mapped to global index, e.g. in zero-shot evaluation, this label_set should be the original index. special_index: the special class index that needs to be converted. Nrz4point_label and label_set must have the same length.r)rQlenrPrange)r9r[r]r_ijr0r0r1convert_point_labelws 6zVISTA3D.convert_point_labelTrr patch_coordsSequence[slice] use_centerboolmapped_label_set max_ppoint max_npointc CsHt||||||j|d\}} || |} || t||jdfS)aM Sample points for patch during sliding window validation. Only used for point only validation. Args: labels: shape [1, 1, H, W, D]. patch_coords: a sequence of sliding window slice objects. label_set: local index, must match values in labels. use_center: sample points from the center. mapped_label_set: global index, it is used to identify special classes and is the global index for the sampled points. max_ppoint/max_npoint: positive points and negative points to sample. )rlrmrIrira)rrIrfrKrLto unsqueeze) r9r>rgr]rirkrlrmrA point_labelsr0r0r1sample_points_patch_vals zVISTA3D.sample_points_patch_valrpc Cs|dj|dj|djg}|dj|dj|djg}ttj||jdd}ttj||jdd}t||dkd||dkd}||}|}| ryd||<d||<|dk d} |dd| f}|dd| f}||fSdS) a Update point_coords with respect to patch coords. If point is outside of the patch, remove the coordinates and set label to -1. Args: patch_coords: a sequence of the python slice objects representing the patch coordinates during sliding window inference. This value is passed from sliding_window_inferer. point_coords: point coordinates, [B, N, 3]. point_labels: point labels, [B, N]. rFrarHrrNNN) stopstartrrKrLrI logical_andallcloneany) r9rgrArpZ patch_endsZ patch_startsZpatch_starts_tensorZpatch_ends_tensorindicesZnot_pad_indicesr0r0r1update_point_to_patchs"     zVISTA3D.update_point_to_patch?logits point_logits mapping_indexthredfloatc sbt|tjjr |n|}||g}tjdD]|t fdd| tDqt||j}t} tj|jd||k} tt|kt|d| } t|k| } t| | ||d} | |j} t| | }t| |}| rd||<t| | |j} ||d| 9<||| |7<|S) a Combine auto results with point click response. The auto results have shape [B, 1, H, W, D] which means B foreground masks from a single image patch. Out of those B foreground masks, user may add points to a subset of B1 foreground masks for editing. mapping_index represents the correspondence between B and B1. For mapping_index with point clicks, NaN values in logits will be replaced with point_logits. Meanwhile, the added/removed region in point clicks must be updated by the lcc function. Notice, if a positive point is within logits/prev_mask, the components containing the positive point will be added. Args: logits: automatic branch results, [B, 1, H, W, D]. point_logits: point branch results, [B1, 1, H, W, D]. point_coords: point coordinates, [B1, N, 3]. point_labels: point labels, [B1, N]. mapping_index: [B]. thred: the threshold to convert logits to binary. rc s2g|]}d|d|d|dfdkqS)rrr)item).0pZ_logitsrdr0r1 s$z8VISTA3D.connected_components_combine..)nan)rArprar) isinstancemonaidata MetaTensor as_tensorrcrQappendnprycpunumpyroundastyperrKrLrnrIisnan nan_to_numr8sigmoidrv logical_orrlccdtype)r9r}r~rArprrinsideZ inside_tensornan_maskZ pos_regionZdiff_posZdiff_negccZ uc_pos_region fill_maskr0rr1connected_components_combines6     z$VISTA3D.connected_components_combineradius int | NonecCs|durt|jddd}dt|jdd|||djddd}d||dk<t|tjjr4|n|}|||9<||d||7<|S) a Combine point results with auto results using gaussian. Args: logits: automatic branch results, [B, 1, H, W, D]. point_logits: point branch results, [B1, 1, H, W, D]. point_coords: point coordinates, [B1, N, 3]. point_labels: point labels, [B1, N]. mapping_index: [B]. radius: gaussian ball radius. Nrrrr)rT)keepdimsr) minrQr sumrrrrr)r9r}r~rArprrweightr0r0r1gaussian_combines zVISTA3D.gaussian_combineFr6r7cCs||jkr2t|jdr|jj||dn|jD] }| o | |_q|jD]}| |_q(||_||jkrft|jdrF|jj||dn|jD] }| oR| |_qK|jD]}| |_qZ||_dSdS)z Freeze auto-branch or point-branch. Args: auto_freeze: whether to freeze the auto branch. point_freeze: whether to freeze the point branch. set_auto_grad)r6r7N) r6hasattrr*r parameters requires_gradr+r7r,)r9r6r7paramr0r0r1r4s"       zVISTA3D.set_auto_grad input_imageslist[Sequence[slice]] | None prompt_classval_point_samplerCallable | None transposec Ks|| dd|| |\}} }|jdd}|j}|dur/|dur/|jtjddg||dS|||}|durs|durd|durd| durH|j} | ||d|\}}}|d dkr_d|d<d\}} n|durs| |d||\}}|dur|dur|dk ddk tj }|r||}||}|dur||}n|js|dur|dur|d nd\}}|dur|dur|jtj|dg||d}| r|dd}|S|jdur| d d r|dur|jd}}n|j||du|dud \}}d}tj|dur2|||\}}|dur1|j||||d }|dur(||||||| }nF||||||}n<|jtj|dg|||jd}|j||||d ||<| durn|durn|| |ddd |j|||||}| d d r|dur||_| r|dd}|S)a7 The forward function for VISTA3D. We only support single patch in training and inference. One exception is allowing sliding window batch size > 1 for automatic segmentation only case. B represents number of objects, N represents number of points for each objects. Args: input_images: [1, 1, H, W, D] point_coords: [B, N, 3] point_labels: [B, N], -1 represents padding. 0/1 means negative/positive points for regular class. 2/3 means negative/postive ponits for special supported class like tumor. class_vector: [B, 1], the global class index. prompt_class: [B, 1], the global class index. This value is associated with point_coords to identify if the points are for zero-shot or supported class. When class_vector and point_coords are both provided, prompt_class is the same as class_vector. For prompt_class[b] > 512, point_coords[b] will be considered novel class. patch_coords: a list of sequence of the python slice objects representing the patch coordinates during sliding window inference. This value is passed from sliding_window_inferer. This is an indicator for training phase or validation phase. Notice for sliding window batch size > 1 (only supported by automatic segmentation), patch_coords will inlcude coordinates of multiple patches. If point prompts are included, the batch size can only be one and all the functions using patch_coords will by default use patch_coords[0]. labels: [1, 1, H, W, D], the groundtruth label tensor, only used for point-only evaluation label_set: the label index matching the indexes in labels. If labels are mapped to global index using RelabelID, this label_set should be global mapped index. If labels are not mapped to global index, e.g. in zero-shot evaluation, this label_set should be the original index. prev_mask: [B, N, H_fullsize, W_fullsize, D_fullsize]. This is the transposed raw output from sliding_window_inferer before any postprocessing. When user click points to perform auto-results correction, this can be the auto-results. radius: single float value controling the gaussian blur when combining point and auto results. The gaussian combine is not used in VISTA3D training but might be useful for finetuning purposes. val_point_sampler: function used to sample points from labels. This is only used for point-only evaluation. transpose: bool. If true, the output will be transposed to be [1, B, H, W, D]. Required to be true if calling from sliding window inferer/point inferer. r<NrrrrHrrarsTZ keep_cacheF) with_point with_label)rNrIr)rMgetrQrIr8rKzerosrRrqrr{rrnrjryr6fill_rr5r*cuda empty_cacher+r,rrrdetach)r9rrgrArprNrr>r]r@rrrkwargs image_sizerIbsrr}outZout_auto_r~r0r0r1forwardPs2               zVISTA3D.forward)r*r2r+r2r,r2)r<r=r>r?r@r?rAr?)rNr?rAr?rOr)NrS)r[r\r]r^r_r`)TNrr)r>r\rgrhr]r`rirjrkr^rlrrmr)rgrhrAr\rpr\)r|) r}r\r~r\rAr\rpr\rr\rrN) r}r\r~r\rAr\rpr\rr\rr)FF)r6rjr7rj) NNNNNNNNNNF)rr\rgrrAr?rpr?rNr?rr?r>r?r]r^r@r?rrrrrrj)__name__ __module__ __qualname____doc__r4rMrRrfrqr{rrrr __classcell__r0r0r:r1r:s>   " !* D cs.eZdZddfd d Z ddddZZS)r- r$r%r&r max_promptr'r(c s`t|}||_ttj||ddddt|ttj||ddddt||_tjdddddd|_ t d|ddd|_ t |d|_ ttd|td|g|_td||_td||_td||_ttj||dddddt|ttj||dddd |_t|||d|_||_||_t|||_td||_td||_d S) aHInteractive point head used for VISTA3D. Adapted from segment anything: `https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/mask_decoder.py`. Args: feature_size: feature channel from encoder. max_prompt: max prompt number in each forward iteration. n_classes: number of classes the model can potentially support. This is the maximum number of class embeddings. last_supported: number of classes the model support, this value should match the trained model weights. rr)rr! kernel_sizestridepaddingr$r)depth embedding_dimmlp_dim num_heads)rrroutput_padding)rrrN)r3r4rr SequentialConv3dInstanceNorm3dGELUfeat_downsampleZmask_downsampleTwoWayTransformer transformerPositionEmbeddingRandompe_layer ModuleList Embeddingpoint_embeddingsnot_a_point_embedspecial_class_embed mask_tokensConvTranspose3doutput_upscalingMLPoutput_hypernetworks_mlpsr'r(class_embeddingszeroshot_embedsupported_embed)r9r&rr'r(Ztransformer_dimr:r0r1r4s:  zPointMappingSAM.__init__Nrr\rArprNr?c Cs||}t|jdd}d}tj|d}|j||}d||dk<||dk|jj 7<||dk|j dj 7<||dk|j dj 7<||dk|j dj |j j 7<||d k|j dj |j j 7<|j j } | d|ddd} |durtj| ||jj d|dddfdd } n(g} |D]} | |jkr| |jj q| |jj qtj| |t| fdd } g} |j}ttt| jd|D]} d \}}}tj| |t| d|| jdf}| |d|d}tj||jddd }tj||jdd d|jddd }|j\}}}}}||||\}}|ddddddf}||}|dd |||||}|!|}|j\}}}}}|| |||||}| | dd|||qt"| S) zArgs: out: feature from encoder, [1, C, H, W, C] point_coords: point coordinates, [B, N, 3] point_labels: point labels, [B, N] class_vector: class prompts, [B] rrNr|rarrrrdim)NNN)#rtuplerQrKrrrforward_with_coordsrrrrrroexpandsizecatrr(rrstackrrcrrceilrrepeat_interleaverrrviewrvstack)r9rrArprNZout_low out_shapepointspoint_embedding output_tokensZ tokens_allrrdmasksrsrcZupscaled_embeddingZhyper_inidxtokensZpos_srcbchwdhsZmask_tokens_outmaskr0r0r1r s^   $$      ,   zPointMappingSAM.forward)rr$r%)r&rrrr'rr(rr)rr\rAr\rpr\rNr?rrrr4rrr0r0r:r1r-s3r-cs.eZdZdZddfdd Zdd dZZS)r.zFClass head that performs automatic segmentation based on class vector.Tr'rr&r)rjc st||_|r tt||tdtt|||_t |||_ tt d||dddddt d||ddddd|_ dS)zArgs: n_classes: maximum number of class embedding. feature_size: class embedding size. use_mlp: use mlp to further map class embedding. rrrT) spatial_dimsrr!rr norm_name res_blockN) r3r4r)rrLinearInstanceNorm1drmlprrr image_post_mapping)r9r'r&r)r:r0r1r4Vs<     zClassMappingClassify.__init__rr\rNc Cst|j\}}}}}||}||}|jr||}|||||||} | |d|||dd} | |fS)Nrarr)rQr rr)r squeezerr) r9rrNrrrrrclass_embeddingZmasks_embeddingr0r0r1r{s   zClassMappingClassify.forward)T)r'rr&rr)rj)rr\rNr\rrrrr4rrr0r0r:r1r.Ss%r.cs.eZdZ  ddfd d ZdddZZS)rrelurrrrrr activation tuple | strattention_downsample_raterONonec szt||_||_||_||_t|_t |D]}|j t ||||||dkdqt |||d|_ t||_dS)a A transformer decoder that attends to an input image using queries whose positional embedding is supplied. Adapted from `https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/transformer.py`. Args: depth: number of layers in the transformer. embedding_dim: the channel dimension for the input embeddings. num_heads: the number of heads for multihead attention. Must divide embedding_dim. mlp_dim: the channel dimension internal to the MLP block. activation: the activation to use in the MLP block. attention_downsample_rate: the rate at which to downsample the image before projecting. r)rrrrrskip_first_layer_pedownsample_rateN)r3r4rrrrrrlayersrcrTwoWayAttentionBlock Attentionfinal_attn_token_to_image LayerNormnorm_final_attn)r9rrrrrrrdr:r0r1r4s&    zTwoWayTransformer.__init__image_embeddingr\image_per!Tuple[torch.Tensor, torch.Tensor]c Cs|dddd}|dddd}|}|}|jD] }|||||d\}}q||}||}|j|||d} || }||}||fS)a- Args: image_embedding: image to attend to. Should be shape B x embedding_dim x h x w for any h and w. image_pe: the positional encoding to add to the image. Must have the same shape as image_embedding. point_embedding: the embedding to add to the query points. Must have shape B x N_points x embedding_dim for any N_points. Returns: torch.Tensor: the processed point_embedding. torch.Tensor: the processed image_embedding. rrr)querieskeysquery_pekey_peqkv)flattenpermuterrr) r9rrrrr layerr$r%attn_outr0r0r1rs  zTwoWayTransformer.forward)rr)rrrrrrrrrrrrrOr)rr\rr\rr\rOrrr0r0r:r1rs ,rcs2eZdZ    ddfdd ZdddZZS)rrrFrrrrrrrrrjrOrcstt|||_t||_t|||d|_t||_t |||dd|_ t||_ t||_ t|||d|_ ||_dS)a A transformer block with four layers: (1) self-attention of sparse inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp block on sparse inputs, and (4) cross attention of dense inputs to sparse inputs. Adapted from `https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/transformer.py`. Args: embedding_dim: the channel dimension of the embeddings. num_heads: the number of heads in the attention layers. mlp_dim: the hidden dimension of the mlp block. activation: the activation of the mlp block. skip_first_layer_pe: skip the PE on the first layer. rvista3d) hidden_sizeract dropout_modeN)r3r4r self_attnrrnorm1cross_attn_token_to_imagenorm2r r norm3norm4cross_attn_image_to_tokenr)r9rrrrrrr:r0r1r4s       zTwoWayAttentionBlock.__init__rr\r r!r"rc Cs|jr |j|||d}n||}|j|||d}||}||}||}||}|j|||d}||}||}||}||}||}||}||}|j|||d}||}||}||fS)Nr#) rr0r1r2r3r r4r6r5) r9rr r!r"r$r*r%Zmlp_outr0r0r1rs(     zTwoWayAttentionBlock.forward)r+rrF)rrrrrrrrrrrrjrOr) rr\r r\r!r\r"r\rOrrr0r0r:r1rs&rcsBeZdZdZddfd d Zdd dZdddZdddZZS)ra An attention layer that allows for downscaling the size of the embedding after projection to queries, keys, and values. Adapted from `https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/transformer.py`. Args: embedding_dim: the channel dimension of the embeddings. num_heads: the number of heads in the attention layers. downsample_rate: the rate at which to downsample the image before projecting. rrrrrrOrcszt||_|||_||_|j|dkstdt||j|_t||j|_ t||j|_ t|j||_ dS)Nrz$num_heads must divide embedding_dim.) r3r4rZ internal_dimrrPrrq_projk_projv_projout_proj)r9rrrr:r0r1r4/s  zAttention.__init__xr\cCs,|j\}}}||||||}|ddSNrr)rQreshaper)r9r;rrnrr0r0r1_separate_heads<s  zAttention._separate_headscCs,|j\}}}}|dd}|||||Sr<)rQrr=)r9r;rn_headsZn_tokens c_per_headr0r0r1_recombine_headsBs zAttention._recombine_headsr$r%r&cCs||}||}||}|||j}|||j}|||j}|j\}}}}||dddd}|t|}t j |dd}||}| |}| |}|S)Nrrrrrar) r7r8r9r?rrQr(mathsqrtrKsoftmaxrBr:)r9r$r%r&rrAattnrr0r0r1rHs     zAttention.forward)r)rrrrrrrOr)r;r\rrrOr\r;r\rOr\)r$r\r%r\r&r\rOr\) rrrrr4r?rBrrr0r0r:r1r#s  rcsBeZdZdZddfd d ZdddZdddZdddZZS)raA Positional encoding using random spatial frequencies. Adapted from `https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/prompt_encoder.py`. Args: num_pos_feats: the number of positional encoding features. scale: the scale of the positional encoding. @N num_pos_featsrscaleOptional[float]rOrcs<t|dus |dkrd}|d|td|fdS)Nrg?#positional_encoding_gaussian_matrixr)r3r4register_bufferrKrandn)r9rIrJr:r0r1r4ks z PositionEmbeddingRandom.__init__coordstorch.torch.TensorcCsBd|d}||j}dtj|}tjt|t|gddS)z8Positionally encode points that are normalized to [0,1].rrrar)rLrpirKrsincos)r9rOr0r0r1 _pe_encodingqs  z$PositionEmbeddingRandom._pe_encodingrTuple[int, int, int]c Cs|\}}}|jj}tj|||f|tjd}|jddd}|jddd}|jddd} ||}||}| |} |tj||| gdd} | ddddS) z>Generate positional encoding for a grid of the specified size.rrrr|rrrar) rLrIrKonesfloat32cumsumrTrr() r9rrrrrIgridZx_embedZy_embedZz_embedper0r0r1r}s zPositionEmbeddingRandom.forward coords_inputrcCs|}|dddddf|d|dddddf<|dddddf|d|dddddf<|dddddf|d|dddddf<||tjS)z