a h@s8dZddlZddlmZddlmZmZddlZddl Zddlm Z ddl m Z ddlmZdd lmZdd lmZdd lmZmZmZmZd d lmZeeZeeddGdddeZeeddGdddeZGddde j Z!Gddde j Z"Gddde j Z#Gddde j Z$d>ej%e&e'ej%dd d!Z(Gd"d#d#e j Z)Gd$d%d%eZ*Gd&d'd'e j Z+Gd(d)d)e j,Z-Gd*d+d+e j Z.Gd,d-d-e j Z/eGd.d/d/eZ0eGd0d1d1e0Z1ej%e2ej%d2d3d4Z3ej%e2e2ej%d5d6d7Z4Gd8d9d9e j Z5ed:dGd;d<dr?rEcCs|jddddf}|jd}t|d}tjsF||ksF||krtj|d||d dddd||fddd }| ddddS|d||dSdS) Nr?rrr8ZbicubicF)sizemodeZ align_corners) rNr9rrZjit is_tracingF interpolatereshaper;)r2r>r?Zpatch_pos_embedr/Zpretrain_patch_sizerrrinterpolate_pos_encodings  z)SegGptEmbeddings.interpolate_pos_encoding)r<prompt_pixel_valuesbool_masked_posembedding_typerEcCs||}||}|j\}}} } |j||| d} |d| d|| d} |d| | | }|durp|nd}||| } ||j}||j }|| }|| }|dkr|j }n|dkr|j }nt d|||}||}t j||fdd}|S)NrPrinstanceZsemanticzBEmbedding type should be either 'semantic' or 'instance', but got rdim)rMr9rHexpand unsqueezeZtype_asrVrWrIrJrKrLr:rcat)r2r<rXrYrZZinput_embeddingsZprompt_embeddingsr= patch_height patch_width_rHwZ pos_embedZtype_embeddingr@rrrrAs*     zSegGptEmbeddings.forward)NN)rrrrrr&intrTensorrWr BoolTensorstrrArBrrr4rrC~srCcseZdZdZfddZeeejejdddZejejejeje eefe eefejddd Z dejejd d d Z Z S)SegGptAttentionz=Multi-head Attention block with relative position embeddings.cst|j|j}}t|tjjr*|n||f}t|tjjrD|n||f}|d|j|d|jf}|j|j }|j |_ |d|_ t j |j|jd|j d|_t |j|j|_|j|_|jr|durtdt td|dd||_t td|dd||_dS)NrrgrbiaszBInput size must be provided if using relative positional encoding.r8)r%r&r'r(r+r,r-r.r*num_attention_headsscalerLinearZqkv_biasqkvproj use_relative_position_embeddingsr:rFrrG rel_pos_h rel_pos_w)r2r3r'r(Z input_sizeZhead_dimr4rrr&s     zSegGptAttention.__init__)q_sizek_sizerel_posrEc Cstdt||d}tj|d|jddddd|dd}|d|dd}t|dddft||d}t|dddft||d}|||dt||d}|| S) a Get relative positional embeddings according to the relative positions of query and key sizes. Args: q_size (int): size of the query. k_size (int): size of key k. rel_pos (`torch.Tensor`): relative position embeddings (L, channel). Returns: Extracted positional embeddings according to relative positions. r8rrrPZlinear)rQrRN?) remaxrTrUrVr9r;rZarangelong) r2rtrurvZ max_rel_distZrel_pos_resizedZq_coordsZk_coordsZrelative_coordsrrr get_rel_poss$$zSegGptAttention.get_rel_pos)attnqueryrrrsrtrurEcCs|\}}|\} } ||| |} ||| |} |j\} }}|| |||}td|| }td|| }|| ||| | }||dddddddddf|dddddddddf}|| ||| | }|S)a Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`. https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py Args: attn (`torch.Tensor`): attention map. query (`torch.Tensor`): query q in the attention layer with shape (batch_size, query_height * query_width, channel). rel_pos_h (`torch.Tensor`): relative position embeddings (Lh, channel) for height axis. rel_pos_w (`torch.Tensor`): relative position embeddings (Lw, channel) for width axis. q_size (tuple): spatial sequence size of query q with (query_height, query_width). k_size (tuple): spatial sequence size of key k with (key_height, key_width). Returns: attn (`torch.Tensor`): attention map with added relative positional embeddings. zbhwc,hkc->bhwkzbhwc,wkc->bhwkN)rzr9rVrZeinsum)r2r{r|rrrsrtruZ query_heightZ query_widthZ key_heightZ key_widthZrelative_position_heightZrelative_position_widthr=rcr]Zreshaped_queryZrel_hZrel_wrrradd_decomposed_rel_poss Hz&SegGptAttention.add_decomposed_rel_posFrrEc Cs:|j\}}}}|||||d|jdddddd}|d||j||dd\}} } ||j| dd} |jr| | ||j |j ||f||f} t j jj| t jdd|j} |r| ||j||d} | ||j||d} nd} | | ||j||d} | ddddd|||d} || } | | fS) NrrPr8rr)dtyper])r9rorVrlr;ZunbindrmZ transposerqr}rrrsrrrZsoftmaxfloat32torviewrp)r2routput_attentionsr=r>r?rcror|keyvalueZ attn_weightsZattn_weights_reshapedZ attn_outputrrrrA;s,  & zSegGptAttention.forward)F) rrrrr&rerrfrzrr}rArBrrr4rris "   -rics0eZdZfddZejejdddZZS) SegGptMlpcs>tt|j|j|_t|j|j|_t|j |_ dSN) r%r&rrnr*Zmlp_dimlin1lin2r hidden_actactr2r3r4rrr&cs zSegGptMlp.__init__r~cCs"||}||}||}|Sr)rrrr2rrrrrAis   zSegGptMlp.forward)rrrr&rrfrArBrrr4rrbs rF)input drop_probtrainingrEcCsd|dks |s|Sd|}|jdfd|jd}|tj||j|jd}||||}|S)aF Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the argument. rrr)rrdevice)r9ndimrZrandrrZfloor_div)rrrZ keep_probr9Z random_tensoroutputrrr drop_pathqs rcsPeZdZdZd eeddfdd ZejejdddZ e d d d Z Z S) SegGptDropPathzXDrop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).N)rrEcst||_dSr)r%r&r)r2rr4rrr&s zSegGptDropPath.__init__r~cCst||j|jSr)rrrrrrrrAszSegGptDropPath.forwardrEcCs d|jS)Nzp=)rr2rrr extra_reprszSegGptDropPath.extra_repr)N) rrrrrfloatr&rrfrArhrrBrrr4rrsrc sZeZdZeeddfdd Zd ejee e e e ejejfe ejfdddZ Z S) SegGptLayerN)r3drop_path_raterEcsdtt||_t||_|dkr.t|nt|_ tj |j |j d|_ tj |j |j d|_dS)Nreps)r%r&ri attentionrmlprrZIdentityr LayerNormr*layer_norm_epslayernorm_beforelayernorm_after)r2r3rr4rrr&s    zSegGptLayer.__init__F)r ensemble_condfeature_ensemblerrEc Cs|j|||d}|d}|dd}|r|jdd|kr|j|jdddd\}} |dkr|jdd} | d| d} | jddd| } | j|j} n| jddd| } tj|| gdd}| ||}|} | |}| |}| | |}|f|}|S) N)rrrr8r\rPT)r]Zkeepdim) rrr9splitrVmeanZ expand_asrr`rrr) r2rrrrZself_attention_outputsZattention_outputoutputspromptinputsZ num_promptsZresidualrrrrAs,    zSegGptLayer.forward)FF)rrrrrr&rrfreboolrrrArBrrr4rrs rc sHeZdZeddfdd Zd ejeeeeee e fddd Z Z S) SegGptEncoderNrDcspt|_ddtjdjjddDtfddt jD|_ tj j j d|_d|_dS) NcSsg|] }|qSr)item).0xrrr z*SegGptEncoder.__init__..rcpu)rcsg|]}t|qSr)r)rir3ZdprrrrrrF)r%r&r3rZlinspacerZnum_hidden_layersrZ ModuleListrangelayersrr*r layernormZgradient_checkpointingrr4rrr&s   "zSegGptEncoder.__init__FT)rrroutput_hidden_states return_dictrEc Cs|rdnd}|rdnd}g}t|jD]\} } |r<||f}|jj| krLdnd} | || ||} | d}| |jjkr|d|jdd||jdddd}| |jjvr||||r&|| df}q&|r||f}|stdd||||fDSt ||||dS) Nrr8rrrOcss|]}|dur|VqdSrr)rvrrr sz(SegGptEncoder.forward..)rrrr) enumeraterr3Z merge_indexr9!intermediate_hidden_state_indicesappendrrr) r2rrrrrZall_hidden_statesZall_self_attentionsrrZ layer_modulerZ layer_outputsrrrrAs:    *   zSegGptEncoder.forward)FFFT) rrrrr&rrfrrrrrArBrrr4rrs  rcs@eZdZdZdddfdd Zejejdfdd ZZS) SegGptLayerNormaALayerNorm that supports two data formats: channels_last (default) or channels_first. The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width). gư> channels_last)r data_formatc s8tj|fd|i||dvr.td|||_dS)Nr)rchannels_firstzUnsupported data format: )r%r&NotImplementedErrorr)r2normalized_shaperrkwargsr4rrr&szSegGptLayerNorm.__init__)featuresrEcsH|jdkr8|dddd}t|}|dddd}n t|}|S)z Args: features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels) rrr8rr)rr;r%rA)r2rr4rrrA s    zSegGptLayerNorm.forward) rrrrr&rrfrArBrrr4rrsrcs,eZdZfddZejdddZZS)SegGptDecoderHeadcs\ttj|j|jddd|_t|j|jdd|_t |j |_ tj|jdddd|_ dS)Nrr)r$paddingr)rrrT)r$rk) r%r&rr0decoder_hidden_sizeconvrrrrract_fctheadrr4rrr&s   zSegGptDecoderHead.__init__rcCs,||}||}||}||}|Sr)rrrrrrrrrA)s     zSegGptDecoderHead.forward)rrrr&rrrArBrrr4rrs rcs@eZdZfddZejejdddZejdddZZS) SegGptDecodercsXttj|jt|j|jd|jdd|_ t ||_ |j|_|j|_||_ dS)Nr8Trj) r%r&rrnr*lenrr(r decoder_embedr decoder_predr3rr4rrr&3s  zSegGptDecoder.__init__r~cCs`|j\}}}}|||||j|j|j}|dddddd}|j|d||j||jfd}|S) Nrrrr8rrPr9)r9rVr(rr;)r2rr=rarbrcrrr_reshape_hidden_states?sz$SegGptDecoder._reshape_hidden_statesrcCs"||}||}||}|Sr)rrrrrrrrAKs   zSegGptDecoder.forward) rrrr&rrrrArBrrr4rr2s  rc@s<eZdZUeed<dZdZdZddgZe j ddd d Z dS) SegGptPreTrainedModelr3modelr<TrCrN)modulerEcCs|jj}t|tjtjfr`tjj|jj t j d|d |jj |j_ |jdur\|jj n.t|tjtfr|jj |jj dnt|trtjj|jj t j d|d |jj |j_ tjj|jj t j d|d |jj |j_ nt|trtjj|jj t j d|d |jj |j_ t jjj|j|dt jjj|j|dt jjj|j|dt jjj|j|dt jjj|j|ddS)zInitialize the weightsr)rstdNrw)r)r3Zinitializer_ranger+rrnr0initZ trunc_normal_weightdatarrrrrkZzero_rrZfill_rirrrsrCrNZnormal_rHrIrJrKrL)r2rrrrr _init_weights[sL     z#SegGptPreTrainedModel._init_weights) rrrrrZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingZ_no_split_modulesrModulerrrrrrSs rcseZdZedfdd ZedddZeee efddd d Z e de j e j e j ee jeeeeee jeeeeeeeeefd d d ZZS) SegGptModelr3cs2t|||_t||_t||_|dSr)r%r&r3rCr@rencoder post_initrr4rrr&s    zSegGptModel.__init__rcCs|jjSr)r@rMrrrrget_input_embeddingssz SegGptModel.get_input_embeddingsN)heads_to_prunerEcCs*|D]\}}|jj|j|qdS)z Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base class PreTrainedModel N)itemsrlayerrZ prune_heads)r2rrZheadsrrr _prune_headsszSegGptModel._prune_heads r<rX prompt_masksrYrrZlabelsrrrrEc CsP|dur |n|jj}| dur | n|jj} | dur4| n|jj} |durH|nd}|jjjjj} | | }| | }t j ||fdd}|durt j ||fddnt j ||fdd}|dur|durt d|dur&|jjj} t j| dt j|jd} t j| | dt j|jd}t | |g}|d}|j||||d}|j|||| | d }|S) a prompt_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Prompt pixel values. Prompt pixel values can be obtained using [`AutoImageProcessor`]. See [`SegGptImageProcessor.__call__`] for details. prompt_masks (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Prompt mask. Prompt mask can be obtained using [`AutoImageProcessor`]. See [`SegGptImageProcessor.__call__`] for details. bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). feature_ensemble (`bool`, *optional*): Boolean indicating whether to use feature ensemble or not. If `True`, the model will use feature ensemble if we have at least two prompts. If `False`, the model will not use feature ensemble. This argument should be considered when doing few-shot inference on an input image i.e. more than one prompt for the same image. embedding_type (`str`, *optional*): Embedding type. Indicates whether the prompt is a semantic or instance embedding. Can be either instance or semantic. labels (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`, `optional`): Ground truth mask for input images. Examples: ```python >>> from transformers import SegGptImageProcessor, SegGptModel >>> from PIL import Image >>> import requests >>> image_input_url = "https://raw.githubusercontent.com/baaivision/Painter/main/SegGPT/SegGPT_inference/examples/hmbb_2.jpg" >>> image_prompt_url = "https://raw.githubusercontent.com/baaivision/Painter/main/SegGPT/SegGPT_inference/examples/hmbb_1.jpg" >>> mask_prompt_url = "https://raw.githubusercontent.com/baaivision/Painter/main/SegGPT/SegGPT_inference/examples/hmbb_1_target.png" >>> image_input = Image.open(requests.get(image_input_url, stream=True).raw) >>> image_prompt = Image.open(requests.get(image_prompt_url, stream=True).raw) >>> mask_prompt = Image.open(requests.get(mask_prompt_url, stream=True).raw).convert("L") >>> checkpoint = "BAAI/seggpt-vit-large" >>> model = SegGptModel.from_pretrained(checkpoint) >>> image_processor = SegGptImageProcessor.from_pretrained(checkpoint) >>> inputs = image_processor(images=image_input, prompt_images=image_prompt, prompt_masks=mask_prompt, return_tensors="pt") >>> outputs = model(**inputs) >>> list(outputs.last_hidden_state.shape) [1, 56, 28, 1024] ``` NFr8r\zLabels were provided, but bool_masked_pos were not. It will be set to default value. If you're training the model, make sure to provide a bool_masked_pos.rr)rZrY)rrrr)r3rruse_return_dictr@rMr1rrrrr`loggerZ warning_oncer/rGrronesr_r)r2r<rXrrYrrZrrrrZexpected_dtyper/bool_masked_pos_zerosbool_masked_pos_onesZembedding_outputZencoder_outputsrrrrAsH;     zSegGptModel.forward)NNNNNNN)rrrrr&r#rdictrelistrr rrfrrgrrhrrrrrArBrrr4rrs0  r)tensorr(rEcCsl|j\}}}}||}||}|j||||||fd}|dddddd}|j||||ddfd}|S)Nrrr8rrrr)r9rVr;)rr(r=r)r>r?rarbrrrpatchify sr)rrarbrEcCs|jd}t|jddd}|||jdkrTtd|jdd|d|d |j|||||dfd }|dd ddd d }|j|d||||fd }|S)NrrPrrOrzNumber of patches z does not match patch height (z ) and width (r7rrr8r)r9rer:rVr;)rrarbr=r(rrr unpatchifys rcs8eZdZfddZejejejejdddZZS) SegGptLosscst|j|_|j|_dSr)r%r&betar(rr4rrr&&s zSegGptLoss.__init__)rr"rrYcCstj||fdd}|dddddfdd|jdd}t||jd|j|jd|j}tj||d|jd}|| | }|S)aNComputes the L1 loss between the predicted masks and the ground truth masks. Args: prompt_masks (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Pixel values from mask prompt. pred_masks (`torch.FloatTensor` of shape `(batch_size, num_channels, 2*height, width)`): Predicted masks. labels (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Ground truth mask for input images. bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Returns: `torch.FloatTensor`: The mean L1 loss between the predicted masks and the ground truth masks. r8r\Nrrnone)Z reductionr) rr`repeatr(rr9rTZsmooth_l1_lossrsum)r2rr"rrYZ ground_truthmaskr!rrrrA+s *$zSegGptLoss.forward) rrrr&rrrgrArBrrr4rr%s  rzM SegGpt model with a decoder on top for one-shot image segmentation. csxeZdZedfdd Zedejejejeej ee ee eej ee ee ee e eefd ddZZS) SegGptForImageSegmentationrcs2t|||_t||_t||_|dSr)r%r&r3rrrdecoderrrr4rrr&Us    z#SegGptForImageSegmentation.__init__Nrc  Csx|dur |n|jj}| dur | n|jj} | dur4| n|jj} |dur|jjjj} tj | dtj |j d} tj | | dtj |j d} t | | g}|d}|j||||||||| | d }| r|jn|d}tj |dd}||}d}|dur t|j}|||||}| sd|f}| r*||df}|rL| r:dnd}|||f}|dur`|f|}|St|||j|jd S) aY prompt_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Prompt pixel values. Prompt pixel values can be obtained using [`AutoImageProcessor`]. See [`SegGptImageProcessor.__call__`] for details. prompt_masks (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): Prompt mask. Prompt mask can be obtained using [`AutoImageProcessor`]. See [`SegGptImageProcessor.__call__`] for details. bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). feature_ensemble (`bool`, *optional*): Boolean indicating whether to use feature ensemble or not. If `True`, the model will use feature ensemble if we have at least two prompts. If `False`, the model will not use feature ensemble. This argument should be considered when doing few-shot inference on an input image i.e. more than one prompt for the same image. embedding_type (`str`, *optional*): Embedding type. Indicates whether the prompt is a semantic or instance embedding. Can be either instance or semantic. labels (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`, `optional`): Ground truth mask for input images. Examples: ```python >>> from transformers import SegGptImageProcessor, SegGptForImageSegmentation >>> from PIL import Image >>> import requests >>> image_input_url = "https://raw.githubusercontent.com/baaivision/Painter/main/SegGPT/SegGPT_inference/examples/hmbb_2.jpg" >>> image_prompt_url = "https://raw.githubusercontent.com/baaivision/Painter/main/SegGPT/SegGPT_inference/examples/hmbb_1.jpg" >>> mask_prompt_url = "https://raw.githubusercontent.com/baaivision/Painter/main/SegGPT/SegGPT_inference/examples/hmbb_1_target.png" >>> image_input = Image.open(requests.get(image_input_url, stream=True).raw) >>> image_prompt = Image.open(requests.get(image_prompt_url, stream=True).raw) >>> mask_prompt = Image.open(requests.get(mask_prompt_url, stream=True).raw).convert("L") >>> checkpoint = "BAAI/seggpt-vit-large" >>> model = SegGptForImageSegmentation.from_pretrained(checkpoint) >>> image_processor = SegGptImageProcessor.from_pretrained(checkpoint) >>> inputs = image_processor(images=image_input, prompt_images=image_prompt, prompt_masks=mask_prompt, return_tensors="pt") >>> outputs = model(**inputs) >>> result = image_processor.post_process_semantic_segmentation(outputs, target_sizes=[(image_input.height, image_input.width)])[0] >>> print(list(result.shape)) [170, 297] ``` Nr8rr) r<rXrrYrrZrrrrrPr\r)r!r"rr)r3rrrrr@rMr/rrGrrrr`r_rrrr rr)r2r<rXrrYrrZrrrrr/rrrrr"r!Zloss_fnridxrrrrA_s^;        z"SegGptForImageSegmentation.forward)NNNNNNN)rrrrr&r rrfrrgrrhrrrr rArBrrr4rrOs,  r)rrr)rF)8rcollections.abcr, dataclassesrtypingrrrZtorch.utils.checkpointrZtorch.nnrrTZ activationsrZmodeling_layersr Zmodeling_utilsr utilsr r r rZconfiguration_seggptrZ get_loggerrrrr rr#rCrirrfrrrrrrrrrrrrrerrrr__all__rrrrsb        #U/=!0 *