import logging import pickle from enum import Enum from math import ceil from typing import Any, List, Optional, TypeVar, Union, overload import numpy as np import torch from mmcv.utils import print_log from detrsmpl.utils.path_utils import ( Existence, check_path_existence, check_path_suffix, ) # In T = TypeVar('T'), T can be anything. # See definition of typing.TypeVar for details. _T1 = TypeVar('_T1') _KT = TypeVar('_KT') _VT = TypeVar('_VT') _HumanData = TypeVar('_HumanData') _CPU_DEVICE = torch.device('cpu') _HumanData_SUPPORTED_KEYS = { 'image_path': { 'type': list, }, 'image_id': { 'type': list, }, 'bbox_xywh': { 'type': np.ndarray, 'shape': (-1, 5), 'dim': 0 }, 'config': { 'type': str, 'dim': None }, 'keypoints2d': { 'type': np.ndarray, 'shape': (-1, -1, 3), 'dim': 0 }, 'keypoints3d': { 'type': np.ndarray, 'shape': (-1, -1, 4), 'dim': 0 }, 'smpl': { 'type': dict, 'slice_key': 'betas', 'dim': 0 }, 'smplh': { 'type': dict, 'slice_key': 'betas', 'dim': 0 }, 'smplx': { 'type': dict, 'slice_key': 'betas', 'dim': 0 }, 'meta': { 'type': dict, }, 'keypoints2d_mask': { 'type': np.ndarray, 'shape': (-1, ), 'dim': None }, 'keypoints2d_convention': { 'type': str, 'dim': None }, 'keypoints3d_mask': { 'type': np.ndarray, 'shape': (-1, ), 'dim': None }, 'keypoints3d_convention': { 'type': str, 'dim': None }, 'vertices': { 'type': np.ndarray, 'shape': (-1, ), 'dim': None }, 'focal_length': { 'type': np.ndarray, 'shape': (-1, ), 'dim': 0 }, 'principal_point': { 'type': np.ndarray, 'shape': (-1, ), 'dim': 0 }, 'misc': { 'type': dict, }, } class _KeyCheck(Enum): PASS = 0 WARN = 1 ERROR = 2 class HumanData(dict): logger = None SUPPORTED_KEYS = _HumanData_SUPPORTED_KEYS WARNED_KEYS = [] def __new__(cls: _HumanData, *args: Any, **kwargs: Any) -> _HumanData: """New an instance of HumanData. Args: cls (HumanData): HumanData class. Returns: HumanData: An instance of HumanData. """ ret_human_data = super().__new__(cls, args, kwargs) setattr(ret_human_data, '__data_len__', -1) setattr(ret_human_data, '__key_strict__', False) setattr(ret_human_data, '__keypoints_compressed__', False) return ret_human_data @classmethod def set_logger(cls, logger: Union[logging.Logger, str, None] = None): """Set logger of HumanData class. Args: logger (logging.Logger | str | None, optional): The way to print summary. See `mmcv.utils.print_log()` for details. Defaults to None. """ cls.logger = logger @classmethod def fromfile(cls, npz_path: str) -> _HumanData: """Construct a HumanData instance from an npz file. Args: npz_path (str): Path to a dumped npz file. Returns: HumanData: A HumanData instance load from file. """ ret_human_data = cls() ret_human_data.load(npz_path) return ret_human_data @classmethod def new(cls, source_dict: dict = None, key_strict: bool = False) -> _HumanData: """Construct a HumanData instance from a dict. Args: source_dict (dict, optional): A dict with items in HumanData fashion. Defaults to None. key_strict (bool, optional): Whether to raise error when setting unsupported keys. Defaults to False. Returns: HumanData: A HumanData instance. """ if source_dict is None: ret_human_data = cls() else: ret_human_data = cls(source_dict) ret_human_data.set_key_strict(key_strict) return ret_human_data def get_key_strict(self) -> bool: """Get value of attribute key_strict. Returns: bool: Whether to raise error when setting unsupported keys. """ return self.__key_strict__ def set_key_strict(self, value: bool): """Set value of attribute key_strict. Args: value (bool, optional): Whether to raise error when setting unsupported keys. Defaults to True. """ former__key_strict__ = self.__key_strict__ self.__key_strict__ = value if former__key_strict__ is False and \ value is True: self.pop_unsupported_items() def check_keypoints_compressed(self) -> bool: """Check whether the keypoints are compressed. Returns: bool: Whether the keypoints are compressed. """ return self.__keypoints_compressed__ def load(self, npz_path: str): """Load data from npz_path and update them to self. Args: npz_path (str): Path to a dumped npz file. """ supported_keys = self.__class__.SUPPORTED_KEYS with np.load(npz_path, allow_pickle=True) as npz_file: tmp_data_dict = dict(npz_file) for key, value in list(tmp_data_dict.items()): if isinstance(value, np.ndarray) and\ len(value.shape) == 0: # value is not an ndarray before dump value = value.item() elif key in supported_keys and\ type(value) != supported_keys[key]['type']: value = supported_keys[key]['type'](value) if value is None: tmp_data_dict.pop(key) elif key == '__key_strict__' or \ key == '__data_len__' or\ key == '__keypoints_compressed__': self.__setattr__(key, value) # pop the attributes to keep dict clean tmp_data_dict.pop(key) elif key == 'bbox_xywh' and value.shape[1] == 4: value = np.hstack([value, np.ones([value.shape[0], 1])]) tmp_data_dict[key] = value else: tmp_data_dict[key] = value self.update(tmp_data_dict) self.__set_default_values__() def dump(self, npz_path: str, overwrite: bool = True): """Dump keys and items to an npz file. Args: npz_path (str): Path to a dumped npz file. overwrite (bool, optional): Whether to overwrite if there is already a file. Defaults to True. Raises: ValueError: npz_path does not end with '.npz'. FileExistsError: When overwrite is False and file exists. """ if not check_path_suffix(npz_path, ['.npz']): raise ValueError('Not an npz file.') if not overwrite: if check_path_existence(npz_path, 'file') == Existence.FileExist: raise FileExistsError dict_to_dump = { '__key_strict__': self.__key_strict__, '__data_len__': self.__data_len__, '__keypoints_compressed__': self.__keypoints_compressed__, } dict_to_dump.update(self) np.savez_compressed(npz_path, **dict_to_dump) def get_sliced_cache(self, slice_size=10) -> List: """Slice the whole HumanData into pieces for HumanDataCacheWriter. Args: slice_size (int, optional): The length of each unit in HumanData cache. Defaults to 10. Returns: List: Two dicts for HumanDataCacheWriter. Init HumanDataCacheWriter by HumanDataCacheWriter(**Returns[0]) and set data by human_data_cache_writer.update_sliced_dict(Returns[1]). """ keypoints_info = {} non_sliced_data = {} sliced_data = {} slice_num = ceil(self.__data_len__ / slice_size) for slice_index in range(slice_num): sliced_data[str(slice_index)] = {} dim_dict = self.__get_slice_dim__() for key, dim in dim_dict.items(): # no dim to slice if dim is None: if key.startswith('keypoints') and\ (key.endswith('_mask') or key.endswith('_convention')): keypoints_info[key] = self[key] else: non_sliced_data[key] = self[key] elif isinstance(dim, dict): value_dict = self.get_raw_value(key) non_sliced_sub_dict = {} for sub_key in value_dict.keys(): sub_value = value_dict[sub_key] if dim[sub_key] is None: non_sliced_sub_dict[sub_key] = sub_value else: sub_dim = dim[sub_key] for slice_index in range(slice_num): slice_start = slice_index * slice_size slice_end = min((slice_index + 1) * slice_size, self.__data_len__) slice_range = slice(slice_start, slice_end) sliced_sub_value = \ HumanData.__get_sliced_result__( sub_value, sub_dim, slice_range ) if key not in sliced_data[str(slice_index)]: sliced_data[str(slice_index)][key] = {} sliced_data[str(slice_index)][key][sub_key] = \ sliced_sub_value if len(non_sliced_sub_dict) > 0: non_sliced_data[key] = non_sliced_sub_dict else: value = self.get_raw_value(key) # slice as ndarray if isinstance(value, np.ndarray): slice_list = [ slice(None), ] * len(value.shape) for slice_index in range(slice_num): slice_start = slice_index * slice_size slice_end = min((slice_index + 1) * slice_size, self.__data_len__) slice_list[dim] = slice(slice_start, slice_end) sliced_value = value[tuple(slice_list)] sliced_data[str(slice_index)][key] = sliced_value # slice as list/tuple else: for slice_index in range(slice_num): slice_start = slice_index * slice_size slice_end = min((slice_index + 1) * slice_size, self.__data_len__) sliced_value = value[slice(slice_start, slice_end)] sliced_data[str(slice_index)][key] = sliced_value writer_args_dict = { 'slice_size': slice_size, 'keypoints_info': keypoints_info, 'data_len': self.data_len, 'non_sliced_data': non_sliced_data, 'key_strict': self.get_key_strict() } return writer_args_dict, sliced_data def to(self, device: Optional[Union[torch.device, str]] = _CPU_DEVICE, dtype: Optional[torch.dtype] = None, non_blocking: Optional[bool] = False, copy: Optional[bool] = False, memory_format: Optional[torch.memory_format] = None) -> dict: """Convert values in numpy.ndarray type to torch.Tensor, and move Tensors to the target device. All keys will exist in the returned dict. Args: device (Union[torch.device, str], optional): A specified device. Defaults to CPU_DEVICE. dtype (torch.dtype, optional): The data type of the expected torch.Tensor. If dtype is None, it is decided according to numpy.ndarry. Defaults to None. non_blocking (bool, optional): When non_blocking, tries to convert asynchronously with respect to the host if possible, e.g., converting a CPU Tensor with pinned memory to a CUDA Tensor. Defaults to False. copy (bool, optional): When copy is set, a new Tensor is created even when the Tensor already matches the desired conversion. No matter what value copy is, Tensor constructed from numpy will not share the same memory with the source numpy.ndarray. Defaults to False. memory_format (torch.memory_format, optional): The desired memory format of returned Tensor. Not supported by pytorch-cpu. Defaults to None. Returns: dict: A dict with all numpy.ndarray values converted into torch.Tensor and all Tensors moved to the target device. """ ret_dict = {} for key in self.keys(): raw_value = self.get_raw_value(key) tensor_value = None if isinstance(raw_value, np.ndarray): tensor_value = torch.from_numpy(raw_value).clone() elif isinstance(raw_value, torch.Tensor): tensor_value = raw_value if tensor_value is None: ret_dict[key] = raw_value else: if memory_format is None: ret_dict[key] = \ tensor_value.to(device, dtype, non_blocking, copy) else: ret_dict[key] = \ tensor_value.to(device, dtype, non_blocking, copy, memory_format=memory_format) return ret_dict def __getitem__(self, key: _KT) -> _VT: """Get value defined by HumanData. This function will be called by self[key]. In keypoints_compressed mode, if the key contains 'keypoints', an array with zero-padding at absent keypoint will be returned. Call self.get_raw_value(k) to get value without padding. Args: key (_KT): Key in HumanData. Returns: _VT: Value to the key. """ value = super().__getitem__(key) if self.__keypoints_compressed__: mask_key = f'{key}_mask' if key in self and \ isinstance(value, np.ndarray) and \ 'keypoints' in key and \ mask_key in self: mask_array = np.asarray(super().__getitem__(mask_key)) value = \ self.__class__.__add_zero_pad__(value, mask_array) return value def get_raw_value(self, key: _KT) -> _VT: """Get raw value from the dict. It acts the same as dict.__getitem__(k). Args: key (_KT): Key in dict. Returns: _VT: Value to the key. """ value = super().__getitem__(key) return value def get_value_in_shape(self, key: _KT, shape: Union[list, tuple], padding_constant: int = 0) -> np.ndarray: """Get value in a specific shape. For each dim, if the required shape is smaller than current shape, ndarray will be sliced. Otherwise, it will be padded with padding_constant at the end. Args: key (_KT): Key in dict. The value of this key must be an instance of numpy.ndarray. shape (Union[list, tuple]): Shape of the returned array. Its length must be equal to value.ndim. Set -1 for a dimension if you do not want to edit it. padding_constant (int, optional): The value to set the padded values for each axis. Defaults to 0. Raises: ValueError: A value in shape is neither positive integer nor -1. Returns: np.ndarray: An array in required shape. """ value = self.get_raw_value(key) assert isinstance(value, np.ndarray) assert value.ndim == len(shape) pad_width_list = [] slice_list = [] for dim_index in range(len(shape)): if shape[dim_index] == -1: # no pad or slice pad_width_list.append((0, 0)) slice_list.append(slice(None)) elif shape[dim_index] > 0: # valid shape value wid = shape[dim_index] - value.shape[dim_index] if wid > 0: pad_width_list.append((0, wid)) else: pad_width_list.append((0, 0)) slice_list.append(slice(0, shape[dim_index])) else: # invalid raise ValueError pad_value = np.pad(value, pad_width=pad_width_list, mode='constant', constant_values=padding_constant) return pad_value[tuple(slice_list)] @overload def get_slice(self, stop: int): """Slice [0, stop, 1] of all sliceable values.""" ... @overload def get_slice(self, start: int, stop: int): """Slice [start, stop, 1] of all sliceable values.""" ... @overload def get_slice(self, start: int, stop: int, step: int): """Slice [start, stop, step] of all sliceable values.""" ... def get_slice(self, arg_0: int, arg_1: Union[int, Any] = None, step: int = 1) -> _HumanData: """Slice all sliceable values along major_dim dimension. Args: arg_0 (int): When arg_1 is None, arg_0 is stop and start=0. When arg_1 is not None, arg_0 is start. arg_1 (Union[int, Any], optional): None or where to stop. Defaults to None. step (int, optional): Length of step. Defaults to 1. Returns: HumanData: A new HumanData instance with sliced values. """ ret_human_data = \ HumanData.new(key_strict=self.get_key_strict()) if arg_1 is None: start = 0 stop = arg_0 else: start = arg_0 stop = arg_1 slice_index = slice(start, stop, step) dim_dict = self.__get_slice_dim__() for key, dim in dim_dict.items(): # keys not expected be sliced if dim is None: ret_human_data[key] = self[key] elif isinstance(dim, dict): value_dict = self.get_raw_value(key) sliced_dict = {} for sub_key in value_dict.keys(): sub_value = value_dict[sub_key] if dim[sub_key] is None: sliced_dict[sub_key] = sub_value else: sub_dim = dim[sub_key] sliced_sub_value = \ HumanData.__get_sliced_result__( sub_value, sub_dim, slice_index) sliced_dict[sub_key] = sliced_sub_value ret_human_data[key] = sliced_dict else: value = self[key] sliced_value = \ HumanData.__get_sliced_result__( value, dim, slice_index) ret_human_data[key] = sliced_value # check keypoints compressed if self.check_keypoints_compressed(): ret_human_data.compress_keypoints_by_mask() return ret_human_data def __get_slice_dim__(self) -> dict: """For each key in this HumanData, get the dimension for slicing. 0 for default, if no other value specified. Returns: dict: Keys are self.keys(). Values indicate where to slice. None for not expected to be sliced or failed. """ supported_keys = self.__class__.SUPPORTED_KEYS ret_dict = {} for key in self.keys(): # keys not expected be sliced if key in supported_keys and \ 'dim' in supported_keys[key] and \ supported_keys[key]['dim'] is None: ret_dict[key] = None else: value = self[key] if isinstance(value, dict) and len(value) > 0: ret_dict[key] = {} for sub_key in value.keys(): try: sub_value_len = len(value[sub_key]) if sub_value_len != self.__data_len__: ret_dict[key][sub_key] = None elif 'dim' in value: ret_dict[key][sub_key] = value['dim'] else: ret_dict[key][sub_key] = 0 except TypeError: ret_dict[key][sub_key] = None continue # instance cannot be sliced without len method try: value_len = len(value) except TypeError: ret_dict[key] = None continue # slice on dim 0 by default slice_dim = 0 if key in supported_keys and \ 'dim' in supported_keys[key]: slice_dim = \ supported_keys[key]['dim'] data_len = value_len if slice_dim == 0 \ else value.shape[slice_dim] # dim not for slice if data_len != self.__data_len__: ret_dict[key] = None continue else: ret_dict[key] = slice_dim return ret_dict def __setitem__(self, key: _KT, val: _VT) -> None: """Set self[key] to value. Only be called when using human_data[key] = val. Methods like update won't call __setitem__. In keypoints_compressed mode, if the key contains 'keypoints', and f'{key}_mask' is in self.keys(), invalid zeros will be removed before setting value. Args: key (_KT): Key in HumanData. Better be an element in HumanData.SUPPORTED_KEYS. If not, an Error will be raised in key_strict mode. val (_VT): Value to the key. Raises: KeyError: self.get_key_strict() is True and key cannot be found in HumanData.SUPPORTED_KEYS. ValueError: Value is supported but doesn't match definition. ValueError: self.check_keypoints_compressed() is True and mask of a keypoint item is missing. """ self.__check_key__(key) self.__check_value__(key, val) # if it can be compressed by mask if self.__keypoints_compressed__: class_logger = self.__class__.logger if 'keypoints' in key and \ '_mask' in key: msg = 'Mask cannot be modified ' +\ 'in keypoints_compressed mode.' print_log(msg=msg, logger=class_logger, level=logging.WARN) return elif isinstance(val, np.ndarray) and \ 'keypoints' in key and \ '_mask' not in key: mask_key = f'{key}_mask' if mask_key in self: mask_array = np.asarray(super().__getitem__(mask_key)) val = \ self.__class__.__remove_zero_pad__(val, mask_array) else: msg = f'Mask for {key} has not been set.' +\ f' Please set {mask_key} before compression.' print_log(msg=msg, logger=class_logger, level=logging.ERROR) raise ValueError dict.__setitem__(self, key, val) def set_raw_value(self, key: _KT, val: _VT) -> None: """Set the raw value of self[key] to val after key check. It acts the same as dict.__setitem__(self, key, val) if the key satisfied constraints. Args: key (_KT): Key in dict. val (_VT): Value to the key. Raises: KeyError: self.get_key_strict() is True and key cannot be found in HumanData.SUPPORTED_KEYS. ValueError: Value is supported but doesn't match definition. """ self.__check_key__(key) self.__check_value__(key, val) dict.__setitem__(self, key, val) def pop_unsupported_items(self) -> None: """Find every item with a key not in HumanData.SUPPORTED_KEYS, and pop it to save memory.""" for key in list(self.keys()): if key not in self.__class__.SUPPORTED_KEYS: self.pop(key) def __check_key__(self, key: Any) -> _KeyCheck: """Check whether the key matches definition in HumanData.SUPPORTED_KEYS. Args: key (Any): Key in HumanData. Returns: _KeyCheck: PASS, WARN or ERROR. Raises: KeyError: self.get_key_strict() is True and key cannot be found in HumanData.SUPPORTED_KEYS. """ ret_key_check = _KeyCheck.PASS if self.get_key_strict(): if key not in self.__class__.SUPPORTED_KEYS: ret_key_check = _KeyCheck.ERROR else: if key not in self.__class__.SUPPORTED_KEYS and \ key not in self.__class__.WARNED_KEYS: # log warning message at the first time ret_key_check = _KeyCheck.WARN self.__class__.WARNED_KEYS.append(key) if ret_key_check == _KeyCheck.ERROR: raise KeyError(self.__class__.__get_key_error_msg__(key)) elif ret_key_check == _KeyCheck.WARN: class_logger = self.__class__.logger if class_logger == 'silent': pass else: print_log(msg=self.__class__.__get_key_warn_msg__(key), logger=class_logger, level=logging.WARN) return ret_key_check def __check_value__(self, key: Any, val: Any) -> bool: """Check whether the value matches definition in HumanData.SUPPORTED_KEYS. Args: key (Any): Key in HumanData. val (Any): Value to the key. Returns: bool: True for matched, ortherwise False. Raises: ValueError: Value is supported but doesn't match definition. """ ret_bool = self.__check_value_type__(key, val) and\ self.__check_value_shape__(key, val) and\ self.__check_value_len__(key, val) if not ret_bool: raise ValueError(self.__class__.__get_value_error_msg__()) return ret_bool def __check_value_type__(self, key: Any, val: Any) -> bool: """Check whether the type of val matches definition in HumanData.SUPPORTED_KEYS. Args: key (Any): Key in HumanData. val (Any): Value to the key. Returns: bool: If type doesn't match, return False. Else return True. """ ret_bool = True supported_keys = self.__class__.SUPPORTED_KEYS # check definition if key in supported_keys: # check type if type(val) != supported_keys[key]['type']: ret_bool = False if not ret_bool: expected_type = supported_keys[key]['type'] err_msg = 'Type check Failed:\n' err_msg += f'key={str(key)}\n' err_msg += f'type(val)={type(val)}\n' err_msg += f'expected type={expected_type}\n' print_log(msg=err_msg, logger=self.__class__.logger, level=logging.ERROR) return ret_bool def __check_value_shape__(self, key: Any, val: Any) -> bool: """Check whether the shape of val matches definition in HumanData.SUPPORTED_KEYS. Args: key (Any): Key in HumanData. val (Any): Value to the key. Returns: bool: If expected shape is defined and doesn't match, return False. Else return True. """ ret_bool = True supported_keys = self.__class__.SUPPORTED_KEYS # check definition if key in supported_keys: # check shape if 'shape' in supported_keys[key]: val_shape = val.shape for shape_ind in range(len(supported_keys[key]['shape'])): # length not match if shape_ind >= len(val_shape): ret_bool = False break expect_val = supported_keys[key]['shape'][shape_ind] # value not match if expect_val > 0 and \ expect_val != val_shape[shape_ind]: ret_bool = False break if not ret_bool: expected_shape = str(supported_keys[key]['shape']) expected_shape = expected_shape.replace('-1', 'Any') err_msg = 'Shape check Failed:\n' err_msg += f'key={str(key)}\n' err_msg += f'val.shape={val_shape}\n' err_msg += f'expected shape={expected_shape}\n' print_log(msg=err_msg, logger=self.__class__.logger, level=logging.ERROR) return ret_bool @property def data_len(self) -> int: """Get the temporal length of this HumanData instance. Returns: int: Number of frames related to this instance. """ return self.__data_len__ @data_len.setter def data_len(self, value: int): """Set the temporal length of this HumanData instance. Args: value (int): Number of frames related to this instance. """ self.__data_len__ = value def __check_value_len__(self, key: Any, val: Any) -> bool: """Check whether the temporal length of val matches other values. Args: key (Any): Key in HumanData. val (Any): Value to the key. Returns: bool: If temporal dim is defined and temporal length doesn't match, return False. Else return True. """ ret_bool = True supported_keys = self.__class__.SUPPORTED_KEYS # check definition if key in supported_keys: # check temporal length if 'dim' in supported_keys[key] and \ supported_keys[key]['dim'] is not None: val_slice_dim = supported_keys[key]['dim'] if supported_keys[key]['type'] == dict: slice_key = supported_keys[key]['slice_key'] val_data_len = val[slice_key].shape[val_slice_dim] else: val_data_len = val.shape[val_slice_dim] if self.data_len < 0: # no data_len yet, assign a new one self.data_len = val_data_len else: # check if val_data_len matches recorded data_len if self.data_len != val_data_len: ret_bool = False if not ret_bool: err_msg = 'Temporal check Failed:\n' err_msg += f'key={str(key)}\n' err_msg += f'val\'s data_len={val_data_len}\n' err_msg += f'expected data_len={self.data_len}\n' print_log(msg=err_msg, logger=self.__class__.logger, level=logging.ERROR) return ret_bool def generate_mask_from_confidence(self, keys=None) -> None: """Generate mask from keypoints' confidence. Keypoints that have zero confidence in all occurrences will have a zero mask. Note that the last value of the keypoint is assumed to be confidence. Args: keys: None, str, or list of str. None: all keys with `keypoint` in it will have mask generated from their confidence. str: key of the keypoint, the mask has name f'{key}_name' list of str: a list of keys of the keypoints. Generate mask for multiple keypoints. Defaults to None. Returns: None Raises: KeyError: A key is not not found """ if keys is None: keys = [] for key in self.keys(): val = self.get_raw_value(key) if isinstance(val, np.ndarray) and \ 'keypoints' in key and \ '_mask' not in key: keys.append(key) elif isinstance(keys, str): keys = [keys] elif isinstance(keys, list): for key in keys: assert isinstance(key, str) else: raise TypeError(f'`Keys` must be None, str, or list of str, ' f'got {type(keys)}.') update_dict = {} for kpt_key in keys: kpt_array = self.get_raw_value(kpt_key) num_joints = kpt_array.shape[-2] # if all conf of a joint are zero, this joint is masked joint_conf = kpt_array[..., -1].reshape(-1, num_joints) mask_array = (joint_conf > 0).astype(np.uint8).max(axis=0) assert len(mask_array) == num_joints # generate mask update_dict[f'{kpt_key}_mask'] = mask_array self.update(update_dict) def compress_keypoints_by_mask(self) -> None: """If a key contains 'keypoints', and f'{key}_mask' is in self.keys(), invalid zeros will be removed and f'{key}_mask' will be locked. Raises: KeyError: A key contains 'keypoints' has been found but its corresponding mask is missing. """ assert self.__keypoints_compressed__ is False key_pairs = [] for key in self.keys(): mask_key = f'{key}_mask' val = self.get_raw_value(key) if isinstance(val, np.ndarray) and \ 'keypoints' in key and \ '_mask' not in key and 'has' not in key: if mask_key in self: key_pairs.append([key, mask_key]) else: msg = f'Mask for {key} has not been set.' +\ f'Please set {mask_key} before compression.' raise KeyError(msg) compressed_dict = {} for kpt_key, mask_key in key_pairs: kpt_array = self.get_raw_value(kpt_key) mask_array = np.asarray(self.get_raw_value(mask_key)) compressed_kpt = \ self.__class__.__remove_zero_pad__(kpt_array, mask_array) compressed_dict[kpt_key] = compressed_kpt # set value after all pairs are compressed self.update(compressed_dict) self.__keypoints_compressed__ = True def decompress_keypoints(self) -> None: """If a key contains 'keypoints', and f'{key}_mask' is in self.keys(), invalid zeros will be inserted to the right places and f'{key}_mask' will be unlocked. Raises: KeyError: A key contains 'keypoints' has been found but its corresponding mask is missing. """ assert self.__keypoints_compressed__ is True key_pairs = [] for key in self.keys(): mask_key = f'{key}_mask' val = self.get_raw_value(key) if isinstance(val, np.ndarray) and \ 'keypoints' in key and \ '_mask' not in key: if mask_key in self: key_pairs.append([key, mask_key]) else: class_logger = self.__class__.logger msg = f'Mask for {key} has not been found.' +\ f'Please remove {key} before decompression.' print_log(msg=msg, logger=class_logger, level=logging.ERROR) raise KeyError decompressed_dict = {} for kpt_key, mask_key in key_pairs: mask_array = np.asarray(self.get_raw_value(mask_key)) compressed_kpt = self.get_raw_value(kpt_key) kpt_array = \ self.__class__.__add_zero_pad__(compressed_kpt, mask_array) decompressed_dict[kpt_key] = kpt_array # set value after all pairs are decompressed self.update(decompressed_dict) self.__keypoints_compressed__ = False def dump_by_pickle(self, pkl_path: str, overwrite: bool = True) -> None: """Dump keys and items to a pickle file. It's a secondary dump method, when a HumanData instance is too large to be dumped by self.dump() Args: pkl_path (str): Path to a dumped pickle file. overwrite (bool, optional): Whether to overwrite if there is already a file. Defaults to True. Raises: ValueError: npz_path does not end with '.pkl'. FileExistsError: When overwrite is False and file exists. """ if not check_path_suffix(pkl_path, ['.pkl']): raise ValueError('Not an pkl file.') if not overwrite: if check_path_existence(pkl_path, 'file') == Existence.FileExist: raise FileExistsError dict_to_dump = { '__key_strict__': self.__key_strict__, '__data_len__': self.__data_len__, '__keypoints_compressed__': self.__keypoints_compressed__, } dict_to_dump.update(self) with open(pkl_path, 'wb') as f_writeb: pickle.dump(dict_to_dump, f_writeb, protocol=pickle.HIGHEST_PROTOCOL) def load_by_pickle(self, pkl_path: str) -> None: """Load data from pkl_path and update them to self. When a HumanData Instance was dumped by self.dump_by_pickle(), use this to load. Args: npz_path (str): Path to a dumped npz file. """ with open(pkl_path, 'rb') as f_readb: tmp_data_dict = pickle.load(f_readb) for key, value in list(tmp_data_dict.items()): if value is None: tmp_data_dict.pop(key) elif key == '__key_strict__' or \ key == '__data_len__' or\ key == '__keypoints_compressed__': self.__setattr__(key, value) # pop the attributes to keep dict clean tmp_data_dict.pop(key) elif key == 'bbox_xywh' and value.shape[1] == 4: value = np.hstack([value, np.ones([value.shape[0], 1])]) tmp_data_dict[key] = value else: tmp_data_dict[key] = value self.update(tmp_data_dict) self.__set_default_values__() def __set_default_values__(self) -> None: """For older versions of HumanData, call this method to apply missing values (also attributes).""" supported_keys = self.__class__.SUPPORTED_KEYS if self.__data_len__ == -1: for key in supported_keys: if key in self and \ 'dim' in supported_keys[key] and\ supported_keys[key]['dim'] is not None: if 'slice_key' in supported_keys[key] and\ supported_keys[key]['type'] == dict: sub_key = supported_keys[key]['slice_key'] slice_dim = supported_keys[key]['dim'] self.__data_len__ = \ self[key][sub_key].shape[slice_dim] else: slice_dim = supported_keys[key]['dim'] self.__data_len__ = self[key].shape[slice_dim] break for key in list(self.keys()): convention_key = f'{key}_convention' if key.startswith('keypoints') and \ not key.endswith('_mask') and \ not key.endswith('_convention') and \ convention_key not in self: self[convention_key] = 'human_data' @classmethod def concatenate(cls, human_data_0: _HumanData, human_data_1: _HumanData) -> _HumanData: """Concatenate two human_data. All keys will be kept it the returned human_data. If either value from human_data_0 or human_data_1 matches data_len from its HumanData, the two values will be concatenated as a single value. If not, postfix will be added to the key to specify source of the value. Args: human_data_0 (_HumanData) human_data_1 (_HumanData) Returns: _HumanData: A new human_data instance with all concatenated data. """ ret_human_data = cls.new(key_strict=False) set_0 = set(human_data_0.keys()) set_1 = set(human_data_1.keys()) common_keys = set_0.intersection(set_1) dim_dict_0 = human_data_0.__get_slice_dim__() dim_dict_1 = human_data_1.__get_slice_dim__() for key in common_keys: value_0 = human_data_0[key] value_1 = human_data_1[key] # align type value_0 = list(value_0) if isinstance(value_0, tuple)\ else value_0 value_1 = list(value_1) if isinstance(value_1, tuple)\ else value_1 assert type(value_0) == type(value_1) # align convention if key.startswith('keypoints') and\ key.endswith('_convention'): assert value_0 == value_1 ret_human_data[key] = value_0 continue # mask_0 and mask_1 elif key.startswith('keypoints') and\ key.endswith('_mask'): new_mask = value_0 * value_1 ret_human_data[key] = new_mask continue # go through the sub dict if isinstance(value_0, dict): sub_dict = {} for sub_key, sub_value_0 in value_0.items(): # only found in value_0 if sub_key not in value_1: sub_dict[sub_key] = sub_value_0 # found in both values else: sub_value_1 = value_1[sub_key] concat_sub_dict = cls.__concat_value__( key=sub_key, value_0=sub_value_0, dim_0=dim_dict_0[key][sub_key], value_1=sub_value_1, dim_1=dim_dict_1[key][sub_key]) sub_dict.update(concat_sub_dict) for sub_key, sub_value_1 in value_1.items(): if sub_key not in value_0: sub_dict[sub_key] = sub_value_1 ret_human_data[key] = sub_dict # try concat else: concat_dict = cls.__concat_value__(key=key, value_0=value_0, dim_0=dim_dict_0[key], value_1=value_1, dim_1=dim_dict_1[key]) ret_human_data.update(concat_dict) # check exclusive keys for key, value in human_data_0.items(): if key not in common_keys: # value not for concat and slice if dim_dict_0[key] is None: ret_human_data[key] = value # value aligned with data_len of HumanData_0 else: ret_human_data[f'{key}_0'] = value for key, value in human_data_1.items(): if key not in common_keys: # same as above if dim_dict_1[key] is None: ret_human_data[key] = value else: ret_human_data[f'{key}_1'] = value return ret_human_data @classmethod def __concat_value__(cls, key: Any, value_0: Any, value_1: Any, dim_0: Union[None, int], dim_1: Union[None, int]) -> dict: """Concat two values from two different HumanData. Args: key (Any): The common key of the two values. value_0 (Any): Value from 0. value_1 (Any): Value from 1. dim_0 (Union[None, int]): The dim for concat and slice. None for N/A. dim_1 (Union[None, int]): The dim for concat and slice. None for N/A. Returns: dict: Dict for concatenated result. """ ret_dict = {} if dim_0 is None or dim_1 is None: ret_dict[f'{key}_0'] = value_0 ret_dict[f'{key}_1'] = value_1 elif isinstance(value_0, list): ret_dict[key] = value_0 + value_1 # elif isinstance(value_0, np.ndarray): else: ret_dict[key] = np.concatenate((value_0, value_1), axis=dim_0) return ret_dict @classmethod def __add_zero_pad__(cls, compressed_array: np.ndarray, mask_array: np.ndarray) -> np.ndarray: """Pad zeros to a compressed keypoints array. Args: compressed_array (np.ndarray): A compressed keypoints array. mask_array (np.ndarray): The mask records compression relationship. Returns: np.ndarray: A keypoints array in full-size. """ assert mask_array.sum() == compressed_array.shape[1] data_len, _, dim = compressed_array.shape mask_len = mask_array.shape[0] ret_value = np.zeros(shape=[data_len, mask_len, dim], dtype=compressed_array.dtype) valid_mask_index = np.where(mask_array == 1)[0] ret_value[:, valid_mask_index, :] = compressed_array return ret_value @classmethod def __remove_zero_pad__(cls, zero_pad_array: np.ndarray, mask_array: np.ndarray) -> np.ndarray: """Remove zero-padding from a full-size keypoints array. Args: zero_pad_array (np.ndarray): A keypoints array in full-size. mask_array (np.ndarray): The mask records compression relationship. Returns: np.ndarray: A compressed keypoints array. """ assert mask_array.shape[0] == zero_pad_array.shape[1] valid_mask_index = np.where(mask_array == 1)[0] ret_value = np.take(zero_pad_array, valid_mask_index, axis=1) return ret_value @classmethod def __get_key_warn_msg__(cls, key: Any) -> str: """Get the warning message when a key fails the check. Args: key (Any): The key with wrong. Returns: str: The warning message. """ class_name = cls.__name__ warn_message = \ f'{key} is absent in' +\ f' {class_name}.SUPPORTED_KEYS.\n' suggestion_message = \ 'Ignore this if you know exactly' +\ ' what you are doing.\n' +\ 'Otherwise, Call self.set_key_strict(True)' +\ ' to avoid wrong keys.\n' return warn_message + suggestion_message @classmethod def __get_key_error_msg__(cls, key: Any) -> str: """Get the error message when a key fails the check. Args: key (Any): The key with wrong. Returns: str: The error message. """ class_name = cls.__name__ absent_message = \ f'{key} is absent in' +\ f' {class_name}.SUPPORTED_KEYS.\n' suggestion_message = \ 'Call self.set_key_strict(False)' +\ ' to allow unsupported keys.\n' return absent_message + suggestion_message @classmethod def __get_value_error_msg__(cls) -> str: """Get the error message when a value fails the check. Returns: str: The error message. """ error_message = \ 'An supported value doesn\'t ' +\ 'match definition.\n' suggestion_message = \ 'See error log for details.\n' return error_message + suggestion_message @classmethod def __get_sliced_result__( cls, input_data: Union[np.ndarray, list, tuple], slice_dim: int, slice_range: slice) -> Union[np.ndarray, list, tuple]: """Slice input_data along slice_dim with slice_range. Args: input_data (Union[np.ndarray, list, tuple]): Data to be sliced. slice_dim (int): Dimension to be sliced. slice_range (slice): An instance of class slice. Returns: Union[np.ndarray, list, tuple]: A slice of input_data. """ if isinstance(input_data, np.ndarray): slice_list = [ slice(None), ] * len(input_data.shape) slice_list[slice_dim] = slice_range sliced_data = input_data[tuple(slice_list)] else: sliced_data = \ input_data[slice_range] return sliced_data