| 11 | |
| 12 | |
| 13 | class HumanDataCacheReader(): |
| 14 | def __init__(self, npz_path: str): |
| 15 | self.npz_path = npz_path |
| 16 | npz_file = np.load(npz_path, allow_pickle=True) |
| 17 | self.slice_size = npz_file['slice_size'].item() |
| 18 | self.data_len = npz_file['data_len'].item() |
| 19 | self.keypoints_info = npz_file['keypoints_info'].item() |
| 20 | self.non_sliced_data = None |
| 21 | self.npz_file = None |
| 22 | |
| 23 | def __del__(self): |
| 24 | if self.npz_file is not None: |
| 25 | self.npz_file.close() |
| 26 | |
| 27 | def get_item(self, index, required_keys: List[str] = []): |
| 28 | if self.npz_file is None: |
| 29 | self.npz_file = np.load(self.npz_path, allow_pickle=True) |
| 30 | cache_key = str(int(index / self.slice_size)) |
| 31 | base_data = self.npz_file[cache_key].item() |
| 32 | base_data.update(self.keypoints_info) |
| 33 | for key in required_keys: |
| 34 | non_sliced_value = self.get_non_sliced_data(key) |
| 35 | if isinstance(non_sliced_value, dict) and\ |
| 36 | key in base_data and\ |
| 37 | isinstance(base_data[key], dict): |
| 38 | base_data[key].update(non_sliced_value) |
| 39 | else: |
| 40 | base_data[key] = non_sliced_value |
| 41 | ret_human_data = HumanData.new(source_dict=base_data) |
| 42 | # data in cache is compressed |
| 43 | ret_human_data.__keypoints_compressed__ = True |
| 44 | # set missing values and attributes by default method |
| 45 | ret_human_data.__set_default_values__() |
| 46 | return ret_human_data |
| 47 | |
| 48 | def get_non_sliced_data(self, key: str): |
| 49 | if self.non_sliced_data is None: |
| 50 | if self.npz_file is None: |
| 51 | npz_file = np.load(self.npz_path, allow_pickle=True) |
| 52 | self.non_sliced_data = npz_file['non_sliced_data'].item() |
| 53 | else: |
| 54 | self.non_sliced_data = self.npz_file['non_sliced_data'].item() |
| 55 | return self.non_sliced_data[key] |
| 56 | |
| 57 | |
| 58 | class HumanDataCacheWriter(): |