SMPL(X) body mesh vertex segmentation.
| 3 | |
| 4 | |
| 5 | class body_segmentation(object): |
| 6 | """SMPL(X) body mesh vertex segmentation.""" |
| 7 | def __init__(self, model_type='smpl') -> None: |
| 8 | if model_type == 'smpl': |
| 9 | self.DICT = SMPL_SEGMENTATION_DICT |
| 10 | self.super_set = SMPL_SUPER_SET |
| 11 | self.NUM_VERTS = 6890 |
| 12 | elif model_type == 'smplx': |
| 13 | self.DICT = SMPLX_SEGMENTATION_DICT |
| 14 | self.super_set = SMPLX_SUPER_SET |
| 15 | self.NUM_VERTS = 10475 |
| 16 | else: |
| 17 | raise ValueError(f'Wrong model_type: {model_type}.' |
| 18 | f' Should be in {["smpl", "smplx"]}') |
| 19 | self.model_type = model_type |
| 20 | self.len = len(list(self.DICT)) |
| 21 | |
| 22 | def items(self, ): |
| 23 | return zip(self.keys(), [self.__getitem__(key) for key in self.keys()]) |
| 24 | |
| 25 | def keys(self, ): |
| 26 | return self.DICT.keys() |
| 27 | |
| 28 | def values(self, ): |
| 29 | return [self.__getitem__(key) for key in self.keys()] |
| 30 | |
| 31 | def __len__(self, ): |
| 32 | return self.len |
| 33 | |
| 34 | def __getitem__(self, key): |
| 35 | if key in self.DICT.keys(): |
| 36 | part_segmentation = [] |
| 37 | raw_segmentation = self.DICT[key] |
| 38 | for continuous in raw_segmentation: |
| 39 | if len(continuous) == 2: |
| 40 | part_segmentation.extend( |
| 41 | list(range(continuous[0], continuous[1] + 1))) |
| 42 | elif len(continuous) == 1: |
| 43 | part_segmentation.extend(continuous) |
| 44 | return part_segmentation |
| 45 | elif key in self.super_set.keys(): |
| 46 | super_part_segmentation = [] |
| 47 | for body_part_key in self.super_set[key]: |
| 48 | super_part_segmentation += self.__getitem__(body_part_key) |
| 49 | return super_part_segmentation |
| 50 | elif key.lower() == 'all': |
| 51 | return list(range(self.NUM_VERTS)) |
| 52 | else: |
| 53 | raise KeyError(f'{key} not in {self.model_type} conventions.') |
| 54 | |
| 55 | |
| 56 | def _preprocess_segmentation_dict(segmentation_dict): |
no outgoing calls
no test coverage detected