rebase & fix conflict

mmpose/datasets/transforms/__init__.py

@@ -1,12 +1,11 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .bottomup_transforms import (BottomupGetHeatmapMask, BottomupRandomAffine,
-                                  BottomupResize)
+                                  BottomupRandomChoiceResize,
+                                  BottomupRandomCrop, BottomupResize)
 from .common_transforms import (Albumentation, FilterAnnotations,
                                 GenerateTarget, GetBBoxCenterScale,
                                 PhotometricDistortion, RandomBBoxTransform,
-                                RandomFlip, RandomHalfBody, YOLOXHSVRandomAug,
-                                BottomupRandomChoiceResize,
-                                BottomupRandomCrop, BottomupResize)
+                                RandomFlip, RandomHalfBody, YOLOXHSVRandomAug)
 from .converting import KeypointConverter
 from .formatting import PackPoseInputs
 from .loading import LoadImage

mmpose/models/data_preprocessors/__init__.py

@@ -2,4 +2,7 @@
 from .batch_augmentation import BatchSyncRandomResize
 from .data_preprocessor import BatchShapeDataPreprocessor, PoseDataPreprocessor
 
-__all__ = ['PoseDataPreprocessor', 'BatchSyncRandomResize', 'BatchShapeDataPreprocessor']
+__all__ = [
+    'PoseDataPreprocessor', 'BatchSyncRandomResize',
+    'BatchShapeDataPreprocessor'
+]

mmpose/models/data_preprocessors/data_preprocessor.py

@@ -1,26 +1,17 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-<<<<<<< HEAD
-from typing import List, Optional, Sequence, Union
-
-import numpy as np
-import torch
-import torch.nn as nn
-from mmengine.model import ImgDataPreprocessor
-from mmengine.utils import is_seq_of
-=======
 import math
 from numbers import Number
-from typing import Optional, Sequence, Union
+from typing import List, Optional, Sequence, Union
 
 import numpy as np
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
-import torchvision.transforms.functional
+import torchvision
 from mmengine.model import ImgDataPreprocessor
 from mmengine.model.utils import stack_batch
 from mmengine.utils import is_seq_of
 from PIL import Image
->>>>>>> support edpose
 
 from mmpose.registry import MODELS
 
@@ -29,11 +20,90 @@
 class PoseDataPreprocessor(ImgDataPreprocessor):
     """Image pre-processor for pose estimation tasks."""
 
-<<<<<<< HEAD
     def __init__(self,
                  mean: Sequence[float] = None,
                  std: Sequence[float] = None,
-=======
+                 pad_size_divisor: int = 1,
+                 pad_value: Union[float, int] = 0,
+                 bgr_to_rgb: bool = False,
+                 rgb_to_bgr: bool = False,
+                 non_blocking: Optional[bool] = False,
+                 batch_augments: Optional[List[dict]] = None):
+        super().__init__(
+            mean=mean,
+            std=std,
+            pad_size_divisor=pad_size_divisor,
+            pad_value=pad_value,
+            bgr_to_rgb=bgr_to_rgb,
+            rgb_to_bgr=rgb_to_bgr,
+            non_blocking=non_blocking)
+        if batch_augments is not None:
+            self.batch_augments = nn.ModuleList(
+                [MODELS.build(aug) for aug in batch_augments])
+        else:
+            self.batch_augments = None
+
+    def forward(self, data: dict, training: bool = False) -> dict:
+        """Perform normalization, padding and bgr2rgb conversion based on
+        ``BaseDataPreprocessor``.
+
+        Args:
+            data (dict): Data sampled from dataloader.
+            training (bool): Whether to enable training time augmentation.
+
+        Returns:
+            dict: Data in the same format as the model input.
+        """
+        batch_pad_shape = self._get_pad_shape(data)
+        data = super().forward(data=data, training=training)
+        inputs, data_samples = data['inputs'], data['data_samples']
+        batch_input_shape = tuple(inputs[0].size()[-2:])
+        for data_sample, pad_shape in zip(data_samples, batch_pad_shape):
+            data_sample.set_metainfo({
+                'batch_input_shape': batch_input_shape,
+                'pad_shape': pad_shape
+            })
+
+        if training and self.batch_augments is not None:
+            for batch_aug in self.batch_augments:
+                inputs, data_samples = batch_aug(inputs, data_samples)
+
+        return {'inputs': inputs, 'data_samples': data_samples}
+
+    def _get_pad_shape(self, data: dict) -> List[tuple]:
+        """Get the pad_shape of each image based on data and
+        pad_size_divisor."""
+        _batch_inputs = data['inputs']
+        # Process data with `pseudo_collate`.
+        if is_seq_of(_batch_inputs, torch.Tensor):
+            batch_pad_shape = []
+            for ori_input in _batch_inputs:
+                pad_h = int(
+                    np.ceil(ori_input.shape[1] /
+                            self.pad_size_divisor)) * self.pad_size_divisor
+                pad_w = int(
+                    np.ceil(ori_input.shape[2] /
+                            self.pad_size_divisor)) * self.pad_size_divisor
+                batch_pad_shape.append((pad_h, pad_w))
+        # Process data with `default_collate`.
+        elif isinstance(_batch_inputs, torch.Tensor):
+            assert _batch_inputs.dim() == 4, (
+                'The input of `ImgDataPreprocessor` should be a NCHW tensor '
+                'or a list of tensor, but got a tensor with shape: '
+                f'{_batch_inputs.shape}')
+            pad_h = int(
+                np.ceil(_batch_inputs.shape[1] /
+                        self.pad_size_divisor)) * self.pad_size_divisor
+            pad_w = int(
+                np.ceil(_batch_inputs.shape[2] /
+                        self.pad_size_divisor)) * self.pad_size_divisor
+            batch_pad_shape = [(pad_h, pad_w)] * _batch_inputs.shape[0]
+        else:
+            raise TypeError('Output of `cast_data` should be a dict '
+                            'or a tuple with inputs and data_samples, but got'
+                            f'{type(data)}: {data}')
+        return batch_pad_shape
+
 
 @MODELS.register_module()
 class BatchShapeDataPreprocessor(ImgDataPreprocessor):
@@ -78,17 +148,12 @@ class BatchShapeDataPreprocessor(ImgDataPreprocessor):
     def __init__(self,
                  mean: Sequence[Number] = None,
                  std: Sequence[Number] = None,
->>>>>>> support edpose
                  pad_size_divisor: int = 1,
                  pad_value: Union[float, int] = 0,
                  bgr_to_rgb: bool = False,
                  rgb_to_bgr: bool = False,
                  non_blocking: Optional[bool] = False,
-<<<<<<< HEAD
-                 batch_augments: Optional[List[dict]] = None):
-=======
                  normalize_bakend: str = 'cv2'):
->>>>>>> support edpose
         super().__init__(
             mean=mean,
             std=std,
@@ -97,21 +162,10 @@ def __init__(self,
             bgr_to_rgb=bgr_to_rgb,
             rgb_to_bgr=rgb_to_bgr,
             non_blocking=non_blocking)
-<<<<<<< HEAD
-        if batch_augments is not None:
-            self.batch_augments = nn.ModuleList(
-                [MODELS.build(aug) for aug in batch_augments])
-        else:
-            self.batch_augments = None
-
-    def forward(self, data: dict, training: bool = False) -> dict:
-        """Perform normalization, padding and bgr2rgb conversion based on
-=======
         self.normalize_bakend = normalize_bakend
 
     def forward(self, data: dict, training: bool = False) -> dict:
         """Perform normalization、padding and bgr2rgb conversion based on
->>>>>>> support edpose
         ``BaseDataPreprocessor``.
 
         Args:
@@ -121,39 +175,6 @@ def forward(self, data: dict, training: bool = False) -> dict:
         Returns:
             dict: Data in the same format as the model input.
         """
-<<<<<<< HEAD
-        batch_pad_shape = self._get_pad_shape(data)
-        data = super().forward(data=data, training=training)
-        inputs, data_samples = data['inputs'], data['data_samples']
-        batch_input_shape = tuple(inputs[0].size()[-2:])
-        for data_sample, pad_shape in zip(data_samples, batch_pad_shape):
-            data_sample.set_metainfo({
-                'batch_input_shape': batch_input_shape,
-                'pad_shape': pad_shape
-            })
-
-        if training and self.batch_augments is not None:
-            for batch_aug in self.batch_augments:
-                inputs, data_samples = batch_aug(inputs, data_samples)
-
-        return {'inputs': inputs, 'data_samples': data_samples}
-
-    def _get_pad_shape(self, data: dict) -> List[tuple]:
-        """Get the pad_shape of each image based on data and
-        pad_size_divisor."""
-        _batch_inputs = data['inputs']
-        # Process data with `pseudo_collate`.
-        if is_seq_of(_batch_inputs, torch.Tensor):
-            batch_pad_shape = []
-            for ori_input in _batch_inputs:
-                pad_h = int(
-                    np.ceil(ori_input.shape[1] /
-                            self.pad_size_divisor)) * self.pad_size_divisor
-                pad_w = int(
-                    np.ceil(ori_input.shape[2] /
-                            self.pad_size_divisor)) * self.pad_size_divisor
-                batch_pad_shape.append((pad_h, pad_w))
-=======
         if self.normalize_bakend == 'cv2':
             data = super().forward(data=data, training=training)
         else:
@@ -215,27 +236,12 @@ def normalize_pillow(self,
             # Pad and stack Tensor.
             batch_inputs = stack_batch(batch_inputs, self.pad_size_divisor,
                                        self.pad_value)
->>>>>>> support edpose
         # Process data with `default_collate`.
         elif isinstance(_batch_inputs, torch.Tensor):
             assert _batch_inputs.dim() == 4, (
                 'The input of `ImgDataPreprocessor` should be a NCHW tensor '
                 'or a list of tensor, but got a tensor with shape: '
                 f'{_batch_inputs.shape}')
-<<<<<<< HEAD
-            pad_h = int(
-                np.ceil(_batch_inputs.shape[1] /
-                        self.pad_size_divisor)) * self.pad_size_divisor
-            pad_w = int(
-                np.ceil(_batch_inputs.shape[2] /
-                        self.pad_size_divisor)) * self.pad_size_divisor
-            batch_pad_shape = [(pad_h, pad_w)] * _batch_inputs.shape[0]
-        else:
-            raise TypeError('Output of `cast_data` should be a dict '
-                            'or a tuple with inputs and data_samples, but got'
-                            f'{type(data)}: {data}')
-        return batch_pad_shape
-=======
             if self._channel_conversion:
                 _batch_inputs = _batch_inputs[:, [2, 1, 0], ...]
             # Convert to float after channel conversion to ensure
@@ -266,8 +272,7 @@ def normalize_pillow(self,
         else:
             raise TypeError('Output of `cast_data` should be a dict of '
                             'list/tuple with inputs and data_samples, '
-                            f'but got {type(data)}： {data}')
+                            f'but got {type(data)}: {data}')
         data['inputs'] = batch_inputs
         data.setdefault('data_samples', None)
         return data
->>>>>>> support edpose