Turn sequential tokens: [ , , ], ... into 3D meshes
(self, input_ids: Tensor)
| 101 | |
| 102 | |
| 103 | def detokenize(self, input_ids: Tensor) -> dict: |
| 104 | ''' |
| 105 | Turn sequential tokens: <bos> [<x>, <y>, <z>], ... <eos> into 3D meshes |
| 106 | ''' |
| 107 | # input_ids: b (n q) or b n q, without <bos> or <eos> |
| 108 | input_ids = input_ids.reshape(input_ids.shape[0], -1) |
| 109 | # batch x nface |
| 110 | face_mask = reduce( |
| 111 | input_ids != self.pad_id, 'b (nf c) -> b nf', 'all', c = 9 |
| 112 | ) |
| 113 | |
| 114 | # batch x (nface x 9) -> batch x nface x 3 x 3 |
| 115 | pred_face_coords = input_ids.reshape(input_ids.shape[0], -1, 9) |
| 116 | pred_face_coords = rearrange( |
| 117 | pred_face_coords, '... (v c) -> ... v c', v = 3 |
| 118 | ) |
| 119 | |
| 120 | # back to continuous space |
| 121 | continuous_coors = undiscretize( |
| 122 | pred_face_coords, |
| 123 | num_discrete = self.num_discrete_coors, |
| 124 | continuous_range = self.coor_continuous_range |
| 125 | ) |
| 126 | # mask padding coordinates out with nan |
| 127 | continuous_coors = continuous_coors.masked_fill( |
| 128 | ~rearrange(face_mask, 'b nf -> b nf 1 1'), |
| 129 | float('nan') |
| 130 | ) |
| 131 | output_dict = {} |
| 132 | output_dict['recon_faces'] = continuous_coors |
| 133 | |
| 134 | return output_dict |
| 135 | |
| 136 | |
| 137 | def forward(self, data_dict: dict) -> dict: |
no test coverage detected