Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2023 Alibaba Research and The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch MGP-STR model."""16 17import collections.abc18from dataclasses import dataclass19from typing import Optional, Union20 21import torch22import torch.nn.functional as F23from torch import nn24 25from ...modeling_outputs import BaseModelOutput26from ...modeling_utils import PreTrainedModel27from ...utils import ModelOutput, auto_docstring, logging28from .configuration_mgp_str import MgpstrConfig29 30 31logger = logging.get_logger(__name__)32 33 34# Copied from transformers.models.beit.modeling_beit.drop_path35def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:36    """37    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).38 39    Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks,40    however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...41    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the42    layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the43    argument.44    """45    if drop_prob == 0.0 or not training:46        return input47    keep_prob = 1 - drop_prob48    shape = (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets49    random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)50    random_tensor.floor_()  # binarize51    output = input.div(keep_prob) * random_tensor52    return output53 54 55# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->Mgpstr56class MgpstrDropPath(nn.Module):57    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""58 59    def __init__(self, drop_prob: Optional[float] = None) -> None:60        super().__init__()61        self.drop_prob = drop_prob62 63    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:64        return drop_path(hidden_states, self.drop_prob, self.training)65 66    def extra_repr(self) -> str:67        return f"p={self.drop_prob}"68 69 70@dataclass71@auto_docstring(72    custom_intro="""73    Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.74    """75)76class MgpstrModelOutput(ModelOutput):77    r"""78    logits (`tuple(torch.FloatTensor)` of shape `(batch_size, config.num_character_labels)`):79        Tuple of `torch.FloatTensor` (one for the output of character of shape `(batch_size,80        config.max_token_length, config.num_character_labels)`, + one for the output of bpe of shape `(batch_size,81        config.max_token_length, config.num_bpe_labels)`, + one for the output of wordpiece of shape `(batch_size,82        config.max_token_length, config.num_wordpiece_labels)`) .83 84        Classification scores (before SoftMax) of character, bpe and wordpiece.85    a3_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_a3_attentions=True` is passed or when `config.output_a3_attentions=True`):86        Tuple of `torch.FloatTensor` (one for the attention of character, + one for the attention of bpe`, + one87        for the attention of wordpiece) of shape `(batch_size, config.max_token_length, sequence_length)`.88 89        Attentions weights after the attention softmax, used to compute the weighted average in the self-attention90        heads.91    """92 93    logits: Optional[tuple[torch.FloatTensor]] = None94    hidden_states: Optional[tuple[torch.FloatTensor]] = None95    attentions: Optional[tuple[torch.FloatTensor]] = None96    a3_attentions: Optional[tuple[torch.FloatTensor]] = None97 98 99class MgpstrEmbeddings(nn.Module):100    """2D Image to Patch Embedding"""101 102    def __init__(self, config: MgpstrConfig):103        super().__init__()104        image_size = (105            config.image_size106            if isinstance(config.image_size, collections.abc.Iterable)107            else (config.image_size, config.image_size)108        )109        patch_size = (110            config.patch_size111            if isinstance(config.patch_size, collections.abc.Iterable)112            else (config.patch_size, config.patch_size)113        )114        self.image_size = image_size115        self.patch_size = patch_size116        self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])117        self.num_patches = self.grid_size[0] * self.grid_size[1]118        self.num_tokens = 2 if config.distilled else 1119 120        self.proj = nn.Conv2d(config.num_channels, config.hidden_size, kernel_size=patch_size, stride=patch_size)121 122        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))123 124        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + self.num_tokens, config.hidden_size))125        self.pos_drop = nn.Dropout(p=config.drop_rate)126 127    def forward(self, pixel_values):128        batch_size, channel, height, width = pixel_values.shape129        if height != self.image_size[0] or width != self.image_size[1]:130            raise ValueError(131                f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})."132            )133 134        patch_embeddings = self.proj(pixel_values)135        patch_embeddings = patch_embeddings.flatten(2).transpose(1, 2)  # BCHW -> BNC136 137        cls_tokens = self.cls_token.expand(batch_size, -1, -1)138        embedding_output = torch.cat((cls_tokens, patch_embeddings), dim=1)139        embedding_output = embedding_output + self.pos_embed140        embedding_output = self.pos_drop(embedding_output)141 142        return embedding_output143 144 145class MgpstrMlp(nn.Module):146    """MLP as used in Vision Transformer, MLP-Mixer and related networks"""147 148    def __init__(self, config: MgpstrConfig, hidden_features):149        super().__init__()150        hidden_features = hidden_features or config.hidden_size151        self.fc1 = nn.Linear(config.hidden_size, hidden_features)152        self.act = nn.GELU()153        self.fc2 = nn.Linear(hidden_features, config.hidden_size)154        self.drop = nn.Dropout(config.drop_rate)155 156    def forward(self, hidden_states):157        hidden_states = self.fc1(hidden_states)158        hidden_states = self.act(hidden_states)159        hidden_states = self.drop(hidden_states)160        hidden_states = self.fc2(hidden_states)161        hidden_states = self.drop(hidden_states)162        return hidden_states163 164 165class MgpstrAttention(nn.Module):166    def __init__(self, config: MgpstrConfig):167        super().__init__()168        self.num_heads = config.num_attention_heads169        head_dim = config.hidden_size // config.num_attention_heads170        self.scale = head_dim**-0.5171 172        self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.qkv_bias)173        self.attn_drop = nn.Dropout(config.attn_drop_rate)174        self.proj = nn.Linear(config.hidden_size, config.hidden_size)175        self.proj_drop = nn.Dropout(config.drop_rate)176 177    def forward(self, hidden_states):178        batch_size, num, channel = hidden_states.shape179        qkv = (180            self.qkv(hidden_states)181            .reshape(batch_size, num, 3, self.num_heads, channel // self.num_heads)182            .permute(2, 0, 3, 1, 4)183        )184        query, key, value = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)185 186        attention_probs = (query @ key.transpose(-2, -1)) * self.scale187        attention_probs = attention_probs.softmax(dim=-1)188        attention_probs = self.attn_drop(attention_probs)189 190        context_layer = (attention_probs @ value).transpose(1, 2).reshape(batch_size, num, channel)191        context_layer = self.proj(context_layer)192        context_layer = self.proj_drop(context_layer)193        return (context_layer, attention_probs)194 195 196class MgpstrLayer(nn.Module):197    def __init__(self, config: MgpstrConfig, drop_path=None):198        super().__init__()199        self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)200        self.attn = MgpstrAttention(config)201        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here202        self.drop_path = MgpstrDropPath(drop_path) if drop_path is not None else nn.Identity()203        self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)204        mlp_hidden_dim = int(config.hidden_size * config.mlp_ratio)205        self.mlp = MgpstrMlp(config, mlp_hidden_dim)206 207    def forward(self, hidden_states):208        self_attention_outputs = self.attn(self.norm1(hidden_states))209        attention_output = self_attention_outputs[0]210        outputs = self_attention_outputs[1]211 212        # first residual connection213        hidden_states = self.drop_path(attention_output) + hidden_states214 215        # second residual connection is done here216        layer_output = hidden_states + self.drop_path(self.mlp(self.norm2(hidden_states)))217 218        outputs = (layer_output, outputs)219        return outputs220 221 222class MgpstrEncoder(nn.Module):223    def __init__(self, config: MgpstrConfig):224        super().__init__()225        # stochastic depth decay rule226        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers, device="cpu")]227 228        self.blocks = nn.Sequential(229            *[MgpstrLayer(config=config, drop_path=dpr[i]) for i in range(config.num_hidden_layers)]230        )231 232    def forward(self, hidden_states, output_attentions=False, output_hidden_states=False, return_dict=True):233        all_hidden_states = () if output_hidden_states else None234        all_self_attentions = () if output_attentions else None235 236        for _, blk in enumerate(self.blocks):237            if output_hidden_states:238                all_hidden_states = all_hidden_states + (hidden_states,)239 240            layer_outputs = blk(hidden_states)241            hidden_states = layer_outputs[0]242 243            if output_attentions:244                all_self_attentions = all_self_attentions + (layer_outputs[1],)245 246        if output_hidden_states:247            all_hidden_states = all_hidden_states + (hidden_states,)248 249        if not return_dict:250            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)251        return BaseModelOutput(252            last_hidden_state=hidden_states,253            hidden_states=all_hidden_states,254            attentions=all_self_attentions,255        )256 257 258class MgpstrA3Module(nn.Module):259    def __init__(self, config: MgpstrConfig):260        super().__init__()261        self.token_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)262        self.tokenLearner = nn.Sequential(263            nn.Conv2d(config.hidden_size, config.hidden_size, kernel_size=(1, 1), stride=1, groups=8, bias=False),264            nn.Conv2d(config.hidden_size, config.max_token_length, kernel_size=(1, 1), stride=1, bias=False),265        )266        self.feat = nn.Conv2d(267            config.hidden_size, config.hidden_size, kernel_size=(1, 1), stride=1, groups=8, bias=False268        )269        self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)270 271    def forward(self, hidden_states):272        hidden_states = self.token_norm(hidden_states)273        hidden_states = hidden_states.transpose(1, 2).unsqueeze(-1)274        selected = self.tokenLearner(hidden_states)275        selected = selected.flatten(2)276        attentions = F.softmax(selected, dim=-1)277 278        feat = self.feat(hidden_states)279        feat = feat.flatten(2).transpose(1, 2)280        feat = torch.einsum("...si,...id->...sd", attentions, feat)281        a3_out = self.norm(feat)282 283        return (a3_out, attentions)284 285 286@auto_docstring287class MgpstrPreTrainedModel(PreTrainedModel):288    config: MgpstrConfig289    base_model_prefix = "mgp_str"290    _no_split_modules = []291 292    def _init_weights(self, module: nn.Module) -> None:293        """Initialize the weights"""294        std = self.config.initializer_range295        if isinstance(module, MgpstrEmbeddings):296            nn.init.trunc_normal_(module.pos_embed, mean=0.0, std=std)297            nn.init.trunc_normal_(module.cls_token, mean=0.0, std=std)298        elif isinstance(module, (nn.Linear, nn.Conv2d)):299            nn.init.trunc_normal_(module.weight.data, mean=0.0, std=std)300            if module.bias is not None:301                module.bias.data.zero_()302        elif isinstance(module, nn.LayerNorm):303            module.bias.data.zero_()304            module.weight.data.fill_(1.0)305 306 307@auto_docstring308class MgpstrModel(MgpstrPreTrainedModel):309    def __init__(self, config: MgpstrConfig):310        super().__init__(config)311        self.config = config312        self.embeddings = MgpstrEmbeddings(config)313        self.encoder = MgpstrEncoder(config)314 315        # Initialize weights and apply final processing316        self.post_init()317 318    def get_input_embeddings(self) -> nn.Module:319        return self.embeddings.proj320 321    @auto_docstring322    def forward(323        self,324        pixel_values: torch.FloatTensor,325        output_attentions: Optional[bool] = None,326        output_hidden_states: Optional[bool] = None,327        return_dict: Optional[bool] = None,328    ) -> Union[tuple[torch.FloatTensor], BaseModelOutput]:329        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions330        output_hidden_states = (331            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states332        )333        return_dict = return_dict if return_dict is not None else self.config.use_return_dict334 335        if pixel_values is None:336            raise ValueError("You have to specify pixel_values")337 338        embedding_output = self.embeddings(pixel_values)339 340        encoder_outputs = self.encoder(341            embedding_output,342            output_attentions=output_attentions,343            output_hidden_states=output_hidden_states,344            return_dict=return_dict,345        )346 347        if not return_dict:348            return encoder_outputs349        return BaseModelOutput(350            last_hidden_state=encoder_outputs.last_hidden_state,351            hidden_states=encoder_outputs.hidden_states,352            attentions=encoder_outputs.attentions,353        )354 355 356@auto_docstring(357    custom_intro="""358    MGP-STR Model transformer with three classification heads on top (three A^3 modules and three linear layer on top359    of the transformer encoder output) for scene text recognition (STR) .360    """361)362class MgpstrForSceneTextRecognition(MgpstrPreTrainedModel):363    config: MgpstrConfig364    main_input_name = "pixel_values"365 366    def __init__(self, config: MgpstrConfig) -> None:367        super().__init__(config)368 369        self.num_labels = config.num_labels370        self.mgp_str = MgpstrModel(config)371 372        self.char_a3_module = MgpstrA3Module(config)373        self.bpe_a3_module = MgpstrA3Module(config)374        self.wp_a3_module = MgpstrA3Module(config)375 376        self.char_head = nn.Linear(config.hidden_size, config.num_character_labels)377        self.bpe_head = nn.Linear(config.hidden_size, config.num_bpe_labels)378        self.wp_head = nn.Linear(config.hidden_size, config.num_wordpiece_labels)379 380        # Initialize weights and apply final processing381        self.post_init()382 383    @auto_docstring384    def forward(385        self,386        pixel_values: torch.FloatTensor,387        output_attentions: Optional[bool] = None,388        output_a3_attentions: Optional[bool] = None,389        output_hidden_states: Optional[bool] = None,390        return_dict: Optional[bool] = None,391    ) -> Union[tuple[torch.FloatTensor], MgpstrModelOutput]:392        r"""393        output_a3_attentions (`bool`, *optional*):394            Whether or not to return the attentions tensors of a3 modules. See `a3_attentions` under returned tensors395            for more detail.396 397        Example:398 399        ```python400        >>> from transformers import (401        ...     MgpstrProcessor,402        ...     MgpstrForSceneTextRecognition,403        ... )404        >>> import requests405        >>> from PIL import Image406 407        >>> # load image from the IIIT-5k dataset408        >>> url = "https://i.postimg.cc/ZKwLg2Gw/367-14.png"409        >>> image = Image.open(requests.get(url, stream=True).raw).convert("RGB")410 411        >>> processor = MgpstrProcessor.from_pretrained("alibaba-damo/mgp-str-base")412        >>> pixel_values = processor(images=image, return_tensors="pt").pixel_values413 414        >>> model = MgpstrForSceneTextRecognition.from_pretrained("alibaba-damo/mgp-str-base")415 416        >>> # inference417        >>> outputs = model(pixel_values)418        >>> out_strs = processor.batch_decode(outputs.logits)419        >>> out_strs["generated_text"]420        '["ticket"]'421        ```"""422        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions423        output_hidden_states = (424            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states425        )426        return_dict = return_dict if return_dict is not None else self.config.use_return_dict427 428        mgp_outputs = self.mgp_str(429            pixel_values,430            output_attentions=output_attentions,431            output_hidden_states=output_hidden_states,432            return_dict=return_dict,433        )434 435        sequence_output = mgp_outputs[0]436 437        char_a3_out, char_attention = self.char_a3_module(sequence_output)438        bpe_a3_out, bpe_attention = self.bpe_a3_module(sequence_output)439        wp_a3_out, wp_attention = self.wp_a3_module(sequence_output)440 441        char_logits = self.char_head(char_a3_out)442        bpe_logits = self.bpe_head(bpe_a3_out)443        wp_logits = self.wp_head(wp_a3_out)444 445        all_a3_attentions = (char_attention, bpe_attention, wp_attention) if output_a3_attentions else None446        all_logits = (char_logits, bpe_logits, wp_logits)447 448        if not return_dict:449            outputs = (all_logits, all_a3_attentions) + mgp_outputs[1:]450            return tuple(output for output in outputs if output is not None)451        return MgpstrModelOutput(452            logits=all_logits,453            hidden_states=mgp_outputs.hidden_states,454            attentions=mgp_outputs.attentions,455            a3_attentions=all_a3_attentions,456        )457 458 459__all__ = ["MgpstrModel", "MgpstrPreTrainedModel", "MgpstrForSceneTextRecognition"]460