Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2024 University of Sydney 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 VitPose model."""16 17from dataclasses import dataclass18from typing import Optional, Union19 20import torch21from torch import nn22 23from ...modeling_outputs import BackboneOutput24from ...modeling_utils import PreTrainedModel25from ...processing_utils import Unpack26from ...utils import ModelOutput, TransformersKwargs, auto_docstring, logging27from ...utils.backbone_utils import load_backbone28from ...utils.generic import can_return_tuple29from .configuration_vitpose import VitPoseConfig30 31 32logger = logging.get_logger(__name__)33 34# General docstring35 36 37@dataclass38@auto_docstring(39    custom_intro="""40    Class for outputs of pose estimation models.41    """42)43class VitPoseEstimatorOutput(ModelOutput):44    r"""45    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):46        Loss is not supported at this moment. See https://github.com/ViTAE-Transformer/ViTPose/tree/main/mmpose/models/losses for further detail.47    heatmaps (`torch.FloatTensor` of shape `(batch_size, num_keypoints, height, width)`):48        Heatmaps as predicted by the model.49    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):50        Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +51        one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states52        (also called feature maps) of the model at the output of each stage.53    """54 55    loss: Optional[torch.FloatTensor] = None56    heatmaps: Optional[torch.FloatTensor] = None57    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None58    attentions: Optional[tuple[torch.FloatTensor, ...]] = None59 60 61@auto_docstring62class VitPosePreTrainedModel(PreTrainedModel):63    config: VitPoseConfig64    base_model_prefix = "vit"65    main_input_name = "pixel_values"66    supports_gradient_checkpointing = True67 68    def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]):69        """Initialize the weights"""70        if isinstance(module, (nn.Linear, nn.Conv2d)):71            # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid72            # `trunc_normal_cpu` not implemented in `half` issues73            module.weight.data = nn.init.trunc_normal_(74                module.weight.data.to(torch.float32), mean=0.0, std=self.config.initializer_range75            ).to(module.weight.dtype)76            if module.bias is not None:77                module.bias.data.zero_()78        elif isinstance(module, nn.LayerNorm):79            module.bias.data.zero_()80            module.weight.data.fill_(1.0)81 82 83def flip_back(output_flipped, flip_pairs, target_type="gaussian-heatmap"):84    """Flip the flipped heatmaps back to the original form.85 86    Args:87        output_flipped (`torch.tensor` of shape `(batch_size, num_keypoints, height, width)`):88            The output heatmaps obtained from the flipped images.89        flip_pairs (`torch.Tensor` of shape `(num_keypoints, 2)`):90            Pairs of keypoints which are mirrored (for example, left ear -- right ear).91        target_type (`str`, *optional*, defaults to `"gaussian-heatmap"`):92            Target type to use. Can be gaussian-heatmap or combined-target.93            gaussian-heatmap: Classification target with gaussian distribution.94            combined-target: The combination of classification target (response map) and regression target (offset map).95            Paper ref: Huang et al. The Devil is in the Details: Delving into Unbiased Data Processing for Human Pose Estimation (CVPR 2020).96 97    Returns:98        torch.Tensor: heatmaps that flipped back to the original image99    """100    if target_type not in ["gaussian-heatmap", "combined-target"]:101        raise ValueError("target_type should be gaussian-heatmap or combined-target")102 103    if output_flipped.ndim != 4:104        raise ValueError("output_flipped should be [batch_size, num_keypoints, height, width]")105    batch_size, num_keypoints, height, width = output_flipped.shape106    channels = 1107    if target_type == "combined-target":108        channels = 3109        output_flipped[:, 1::3, ...] = -output_flipped[:, 1::3, ...]110    output_flipped = output_flipped.reshape(batch_size, -1, channels, height, width)111    output_flipped_back = output_flipped.clone()112 113    # Swap left-right parts114    for left, right in flip_pairs.tolist():115        output_flipped_back[:, left, ...] = output_flipped[:, right, ...]116        output_flipped_back[:, right, ...] = output_flipped[:, left, ...]117    output_flipped_back = output_flipped_back.reshape((batch_size, num_keypoints, height, width))118    # Flip horizontally119    output_flipped_back = output_flipped_back.flip(-1)120    return output_flipped_back121 122 123class VitPoseSimpleDecoder(nn.Module):124    """125    Simple decoding head consisting of a ReLU activation, 4x upsampling and a 3x3 convolution, turning the126    feature maps into heatmaps.127    """128 129    def __init__(self, config: VitPoseConfig):130        super().__init__()131 132        self.activation = nn.ReLU()133        self.upsampling = nn.Upsample(scale_factor=config.scale_factor, mode="bilinear", align_corners=False)134        self.conv = nn.Conv2d(135            config.backbone_config.hidden_size, config.num_labels, kernel_size=3, stride=1, padding=1136        )137 138    def forward(self, hidden_state: torch.Tensor, flip_pairs: Optional[torch.Tensor] = None) -> torch.Tensor:139        # Transform input: ReLU + upsample140        hidden_state = self.activation(hidden_state)141        hidden_state = self.upsampling(hidden_state)142        heatmaps = self.conv(hidden_state)143 144        if flip_pairs is not None:145            heatmaps = flip_back(heatmaps, flip_pairs)146 147        return heatmaps148 149 150class VitPoseClassicDecoder(nn.Module):151    """152    Classic decoding head consisting of a 2 deconvolutional blocks, followed by a 1x1 convolution layer,153    turning the feature maps into heatmaps.154    """155 156    def __init__(self, config: VitPoseConfig):157        super().__init__()158 159        self.deconv1 = nn.ConvTranspose2d(160            config.backbone_config.hidden_size, 256, kernel_size=4, stride=2, padding=1, bias=False161        )162        self.batchnorm1 = nn.BatchNorm2d(256)163        self.relu1 = nn.ReLU()164 165        self.deconv2 = nn.ConvTranspose2d(256, 256, kernel_size=4, stride=2, padding=1, bias=False)166        self.batchnorm2 = nn.BatchNorm2d(256)167        self.relu2 = nn.ReLU()168 169        self.conv = nn.Conv2d(256, config.num_labels, kernel_size=1, stride=1, padding=0)170 171    def forward(self, hidden_state: torch.Tensor, flip_pairs: Optional[torch.Tensor] = None):172        hidden_state = self.deconv1(hidden_state)173        hidden_state = self.batchnorm1(hidden_state)174        hidden_state = self.relu1(hidden_state)175 176        hidden_state = self.deconv2(hidden_state)177        hidden_state = self.batchnorm2(hidden_state)178        hidden_state = self.relu2(hidden_state)179 180        heatmaps = self.conv(hidden_state)181 182        if flip_pairs is not None:183            heatmaps = flip_back(heatmaps, flip_pairs)184 185        return heatmaps186 187 188@auto_docstring(189    custom_intro="""190    The VitPose model with a pose estimation head on top.191    """192)193class VitPoseForPoseEstimation(VitPosePreTrainedModel):194    def __init__(self, config: VitPoseConfig):195        super().__init__(config)196 197        self.backbone = load_backbone(config)198 199        # add backbone attributes200        if not hasattr(self.backbone.config, "hidden_size"):201            raise ValueError("The backbone should have a hidden_size attribute")202        if not hasattr(self.backbone.config, "image_size"):203            raise ValueError("The backbone should have an image_size attribute")204        if not hasattr(self.backbone.config, "patch_size"):205            raise ValueError("The backbone should have a patch_size attribute")206 207        self.head = VitPoseSimpleDecoder(config) if config.use_simple_decoder else VitPoseClassicDecoder(config)208 209        # Initialize weights and apply final processing210        self.post_init()211 212    @can_return_tuple213    @auto_docstring214    def forward(215        self,216        pixel_values: torch.Tensor,217        dataset_index: Optional[torch.Tensor] = None,218        flip_pairs: Optional[torch.Tensor] = None,219        labels: Optional[torch.Tensor] = None,220        **kwargs: Unpack[TransformersKwargs],221    ) -> VitPoseEstimatorOutput:222        r"""223        dataset_index (`torch.Tensor` of shape `(batch_size,)`):224            Index to use in the Mixture-of-Experts (MoE) blocks of the backbone.225 226            This corresponds to the dataset index used during training, e.g. For the single dataset index 0 refers to the corresponding dataset. For the multiple datasets index 0 refers to dataset A (e.g. MPII) and index 1 refers to dataset B (e.g. CrowdPose).227        flip_pairs (`torch.tensor`, *optional*):228            Whether to mirror pairs of keypoints (for example, left ear -- right ear).229 230        Examples:231 232        ```python233        >>> from transformers import AutoImageProcessor, VitPoseForPoseEstimation234        >>> import torch235        >>> from PIL import Image236        >>> import requests237 238        >>> processor = AutoImageProcessor.from_pretrained("usyd-community/vitpose-base-simple")239        >>> model = VitPoseForPoseEstimation.from_pretrained("usyd-community/vitpose-base-simple")240 241        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"242        >>> image = Image.open(requests.get(url, stream=True).raw)243        >>> boxes = [[[412.8, 157.61, 53.05, 138.01], [384.43, 172.21, 15.12, 35.74]]]244        >>> inputs = processor(image, boxes=boxes, return_tensors="pt")245 246        >>> with torch.no_grad():247        ...     outputs = model(**inputs)248        >>> heatmaps = outputs.heatmaps249        ```"""250 251        loss = None252        if labels is not None:253            raise NotImplementedError("Training is not yet supported")254 255        outputs: BackboneOutput = self.backbone.forward_with_filtered_kwargs(256            pixel_values,257            dataset_index=dataset_index,258            **kwargs,259        )260 261        # Turn output hidden states in tensor of shape (batch_size, num_channels, height, width)262        sequence_output = outputs.feature_maps[-1]263        batch_size = sequence_output.shape[0]264        patch_height = self.config.backbone_config.image_size[0] // self.config.backbone_config.patch_size[0]265        patch_width = self.config.backbone_config.image_size[1] // self.config.backbone_config.patch_size[1]266        sequence_output = sequence_output.permute(0, 2, 1)267        sequence_output = sequence_output.reshape(batch_size, -1, patch_height, patch_width).contiguous()268 269        heatmaps = self.head(sequence_output, flip_pairs=flip_pairs)270 271        return VitPoseEstimatorOutput(272            loss=loss,273            heatmaps=heatmaps,274            hidden_states=outputs.hidden_states,275            attentions=outputs.attentions,276        )277 278 279__all__ = ["VitPosePreTrainedModel", "VitPoseForPoseEstimation"]280