|
|
| from typing import List, Optional, Sequence, Tuple, Union
|
|
|
| import torch
|
| from mmengine.model import BaseDataPreprocessor, stack_batch
|
|
|
| from mmaction.registry import MODELS
|
| from mmaction.utils import SampleList
|
|
|
|
|
| @MODELS.register_module()
|
| class ActionDataPreprocessor(BaseDataPreprocessor):
|
| """Data pre-processor for action recognition tasks.
|
|
|
| Args:
|
| mean (Sequence[float or int], optional): The pixel mean of channels
|
| of images or stacked optical flow. Defaults to None.
|
| std (Sequence[float or int], optional): The pixel standard deviation
|
| of channels of images or stacked optical flow. Defaults to None.
|
| to_rgb (bool): Whether to convert image from BGR to RGB.
|
| Defaults to False.
|
| to_float32 (bool): Whether to convert data to float32.
|
| Defaults to True.
|
| blending (dict, optional): Config for batch blending.
|
| Defaults to None.
|
| format_shape (str): Format shape of input data.
|
| Defaults to ``'NCHW'``.
|
| """
|
|
|
| def __init__(self,
|
| mean: Optional[Sequence[Union[float, int]]] = None,
|
| std: Optional[Sequence[Union[float, int]]] = None,
|
| to_rgb: bool = False,
|
| to_float32: bool = True,
|
| blending: Optional[dict] = None,
|
| format_shape: str = 'NCHW') -> None:
|
| super().__init__()
|
| self.to_rgb = to_rgb
|
| self.to_float32 = to_float32
|
| self.format_shape = format_shape
|
|
|
| if mean is not None:
|
| assert std is not None, 'To enable the normalization in ' \
|
| 'preprocessing, please specify both ' \
|
| '`mean` and `std`.'
|
|
|
| self._enable_normalize = True
|
| if self.format_shape == 'NCHW':
|
| normalizer_shape = (-1, 1, 1)
|
| elif self.format_shape in ['NCTHW', 'MIX2d3d']:
|
| normalizer_shape = (-1, 1, 1, 1)
|
| else:
|
| raise ValueError(f'Invalid format shape: {format_shape}')
|
|
|
| self.register_buffer(
|
| 'mean',
|
| torch.tensor(mean, dtype=torch.float32).view(normalizer_shape),
|
| False)
|
| self.register_buffer(
|
| 'std',
|
| torch.tensor(std, dtype=torch.float32).view(normalizer_shape),
|
| False)
|
| else:
|
| self._enable_normalize = False
|
|
|
| if blending is not None:
|
| self.blending = MODELS.build(blending)
|
| else:
|
| self.blending = None
|
|
|
| def forward(self,
|
| data: Union[dict, Tuple[dict]],
|
| training: bool = False) -> Union[dict, Tuple[dict]]:
|
| """Perform normalization, padding, bgr2rgb conversion and batch
|
| augmentation based on ``BaseDataPreprocessor``.
|
|
|
| Args:
|
| data (dict or Tuple[dict]): data sampled from dataloader.
|
| training (bool): Whether to enable training time augmentation.
|
|
|
| Returns:
|
| dict or Tuple[dict]: Data in the same format as the model input.
|
| """
|
| data = self.cast_data(data)
|
| if isinstance(data, dict):
|
| return self.forward_onesample(data, training=training)
|
| elif isinstance(data, (tuple, list)):
|
| outputs = []
|
| for data_sample in data:
|
| output = self.forward_onesample(data_sample, training=training)
|
| outputs.append(output)
|
| return tuple(outputs)
|
| else:
|
| raise TypeError(f'Unsupported data type: {type(data)}!')
|
|
|
| def forward_onesample(self, data, training: bool = False) -> dict:
|
| """Perform normalization, padding, bgr2rgb conversion and batch
|
| augmentation on one data sample.
|
|
|
| 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.
|
| """
|
| inputs, data_samples = data['inputs'], data['data_samples']
|
| inputs, data_samples = self.preprocess(inputs, data_samples, training)
|
| data['inputs'] = inputs
|
| data['data_samples'] = data_samples
|
| return data
|
|
|
| def preprocess(self,
|
| inputs: List[torch.Tensor],
|
| data_samples: SampleList,
|
| training: bool = False) -> Tuple:
|
|
|
| batch_inputs = stack_batch(inputs)
|
|
|
| if self.format_shape == 'MIX2d3d':
|
| if batch_inputs.ndim == 4:
|
| format_shape, view_shape = 'NCHW', (-1, 1, 1)
|
| else:
|
| format_shape, view_shape = 'NCTHW', None
|
| else:
|
| format_shape, view_shape = self.format_shape, None
|
|
|
|
|
| if self.to_rgb:
|
| if format_shape == 'NCHW':
|
| batch_inputs = batch_inputs[..., [2, 1, 0], :, :]
|
| elif format_shape == 'NCTHW':
|
| batch_inputs = batch_inputs[..., [2, 1, 0], :, :, :]
|
| else:
|
| raise ValueError(f'Invalid format shape: {format_shape}')
|
|
|
|
|
| if self._enable_normalize:
|
| if view_shape is None:
|
| batch_inputs = (batch_inputs - self.mean) / self.std
|
| else:
|
| mean = self.mean.view(view_shape)
|
| std = self.std.view(view_shape)
|
| batch_inputs = (batch_inputs - mean) / std
|
| elif self.to_float32:
|
| batch_inputs = batch_inputs.to(torch.float32)
|
|
|
|
|
| if training and self.blending is not None:
|
| batch_inputs, data_samples = self.blending(batch_inputs,
|
| data_samples)
|
|
|
| return batch_inputs, data_samples
|
|
|
