mmpose/datasets/dataset_wrappers.py

# Copyright (c) OpenMMLab. All rights reserved.

from copy import deepcopy
from typing import Any, Callable, List, Optional, Tuple, Union, Dict

import numpy as np
from mmengine.dataset import BaseDataset
from mmengine.registry import build_from_cfg

from mmpose.registry import DATASETS
from .datasets.utils import parse_pose_metainfo


@DATASETS.register_module()
class CombinedDataset(BaseDataset):
    """A wrapper of combined dataset.

    Args:
        metainfo (dict): The meta information of combined dataset.
        datasets (list): The configs of datasets to be combined.
        pipeline (list, optional): Processing pipeline. Defaults to [].
        sample_ratio_factor (list, optional): A list of sampling ratio
            factors for each dataset. Defaults to None
    """

    def __init__(self,
                 metainfo: dict,
                 datasets: list,
                 pipeline: List[Union[dict, Callable]] = [],
                 sample_ratio_factor: Optional[List[float]] = None,
                 dataset_ratio_factor: Optional[List[float]] = None,
                 keypoints_mapping: Optional[List[Dict]] = None,
                 **kwargs):

        self.datasets = []
        self.resample = sample_ratio_factor is not None

        self.keypoints_mapping = keypoints_mapping
        self.num_joints = None
        if self.keypoints_mapping is not None:
            self.num_joints = 0
            for mapping in self.keypoints_mapping:
                self.num_joints = max(self.num_joints, max(mapping.values()) +1)


        for cfg in datasets:
            dataset = build_from_cfg(cfg, DATASETS)
            self.datasets.append(dataset)

        # For each dataset, select its random subset based on the sample_ratio_factor
        if dataset_ratio_factor is not None:
            for i, dataset in enumerate(self.datasets):
                dataset_len = len(dataset)
                random_subset = np.random.choice(
                    dataset_len,
                    int(dataset_len * dataset_ratio_factor[i]),
                    replace=False,
                )
                self.datasets[i] = dataset.get_subset(
                    random_subset.flatten().tolist(),
                )

        self._lens = [len(dataset) for dataset in self.datasets]
        if self.resample:
            assert len(sample_ratio_factor) == len(datasets), f'the length ' \
                f'of `sample_ratio_factor` {len(sample_ratio_factor)} does ' \
                f'not match the length of `datasets` {len(datasets)}'
            assert min(sample_ratio_factor) >= 0.0, 'the ratio values in ' \
                '`sample_ratio_factor` should not be negative.'
            self._lens_ori = self._lens
            self._lens = [
                round(l * sample_ratio_factor[i])
                for i, l in enumerate(self._lens_ori)
            ]

        self._len = sum(self._lens)

        super(CombinedDataset, self).__init__(pipeline=pipeline, **kwargs)
        self._metainfo = parse_pose_metainfo(metainfo)

        print("CombinedDataset initialized\n\tlen: {}\n\tlens: {}".format(self._len, self._lens))

    @property
    def metainfo(self):
        return deepcopy(self._metainfo)

    def __len__(self):
        return self._len

    def _get_subset_index(self, index: int) -> Tuple[int, int]:
        """Given a data sample's global index, return the index of the sub-
        dataset the data sample belongs to, and the local index within that
        sub-dataset.

        Args:
            index (int): The global data sample index

        Returns:
            tuple[int, int]:
            - subset_index (int): The index of the sub-dataset
            - local_index (int): The index of the data sample within
                the sub-dataset
        """
        if index >= len(self) or index < -len(self):
            raise ValueError(
                f'index({index}) is out of bounds for dataset with '
                f'length({len(self)}).')

        if index < 0:
            index = index + len(self)

        subset_index = 0
        while index >= self._lens[subset_index]:
            index -= self._lens[subset_index]
            subset_index += 1

        if self.resample:
            gap = (self._lens_ori[subset_index] -
                   1e-4) / self._lens[subset_index]
            index = round(gap * index + np.random.rand() * gap - 0.5)

        return subset_index, index

    def prepare_data(self, idx: int) -> Any:
        """Get data processed by ``self.pipeline``.The source dataset is
        depending on the index.

        Args:
            idx (int): The index of ``data_info``.

        Returns:
            Any: Depends on ``self.pipeline``.
        """

        data_info = self.get_data_info(idx)

        # the assignment of 'dataset' should not be performed within the
        # `get_data_info` function. Otherwise, it can lead to the mixed
        # data augmentation process getting stuck.
        data_info['dataset'] = self

        return self.pipeline(data_info)

    def get_data_info(self, idx: int) -> dict:
        """Get annotation by index.

        Args:
            idx (int): Global index of ``CombinedDataset``.
        Returns:
            dict: The idx-th annotation of the datasets.
        """
        subset_idx, sample_idx = self._get_subset_index(idx)
        # Get data sample processed by ``subset.pipeline``
        data_info = self.datasets[subset_idx][sample_idx]

        if 'dataset' in data_info:
            data_info.pop('dataset')

        # Add metainfo items that are required in the pipeline and the model
        metainfo_keys = [
            'upper_body_ids', 'lower_body_ids', 'flip_pairs',
            'dataset_keypoint_weights', 'flip_indices'
        ]

        for key in metainfo_keys:
            data_info[key] = deepcopy(self._metainfo[key])

        # Map keypoints based on the dataset keypoint mapping
        if self.keypoints_mapping is not None:
            mapping = self.keypoints_mapping[subset_idx]
            
            keypoints = data_info['keypoints']
            N, K, D = keypoints.shape
            keypoints_visibility = data_info.get('keypoints_visibility', np.zeros((N, K)))
            keypoints_visible = data_info.get('keypoints_visible', np.zeros((N, K)))
            
            mapped_keypoints = np.zeros((N, self.num_joints, 2))
            mapped_visibility = np.zeros((N, self.num_joints))
            mapped_visible = np.zeros((N, self.num_joints))

            map_idx = np.stack(
                [list(mapping.keys()), list(mapping.values())], axis=1)
            mapped_keypoints[:, map_idx[:, 1], :] = data_info['keypoints'][:, map_idx[:, 0], :]
            mapped_visibility[:, map_idx[:, 1]] = keypoints_visibility[:, map_idx[:, 0]]
            mapped_visible[:, map_idx[:, 1]] = keypoints_visible[:, map_idx[:, 0]]

            data_info['keypoints'] = mapped_keypoints.reshape((N, self.num_joints, 2) )
            data_info['keypoints_visibility'] = mapped_visibility.reshape((N, self.num_joints))
            data_info['keypoints_visible'] = mapped_visible.reshape((N, self.num_joints))

        # print('data_info', data_info)

        return data_info

    def full_init(self):
        """Fully initialize all sub datasets."""

        if self._fully_initialized:
            return

        for dataset in self.datasets:
            dataset.full_init()
        self._fully_initialized = True