Turn 3D meshes into sequential tokens: [ , , ], ...
(self, data_dict: dict)
| 45 | |
| 46 | |
| 47 | def tokenize(self, data_dict: dict) -> dict: |
| 48 | ''' |
| 49 | Turn 3D meshes into sequential tokens: <bos> [<x>, <y>, <z>], ... <eos> |
| 50 | ''' |
| 51 | |
| 52 | ### 3D mesh face parsing |
| 53 | vertices = data_dict['vertices'] # batch x nv x 3 |
| 54 | faces = data_dict['faces'] # batch x nf x 3 |
| 55 | face_mask = reduce(faces != self.pad_id, 'b nf c -> b nf', 'all') # batch x nf |
| 56 | |
| 57 | batch, num_vertices, num_coors = vertices.shape |
| 58 | _, num_faces, _ = faces.shape |
| 59 | |
| 60 | # fill padding tokens with 0, to prevent gather idx error |
| 61 | face_without_pad = faces.masked_fill(~rearrange(face_mask, 'b nf -> b nf 1'), 0) |
| 62 | |
| 63 | # collect vertice coordinates per-face: b x nf x nv x c |
| 64 | faces_vertices = repeat(face_without_pad, 'b nf nv -> b nf nv c', c = num_coors) |
| 65 | vertices = repeat(vertices, 'b nv c -> b nf nv c', nf = num_faces) |
| 66 | face_coords = vertices.gather(-2, faces_vertices.long()) |
| 67 | |
| 68 | # continuous to discrete face coords: b x nf x nv x c |
| 69 | discrete_face_coords = discretize( |
| 70 | face_coords, |
| 71 | continuous_range=self.coor_continuous_range, |
| 72 | num_discrete=self.num_discrete_coors |
| 73 | ) |
| 74 | |
| 75 | # pad invalid faces with <pad_id>: batch x nf x nv x c |
| 76 | discrete_padded_coords = discrete_face_coords.masked_fill( |
| 77 | ~rearrange(face_mask, 'b nf -> b nf 1 1'), |
| 78 | self.pad_id |
| 79 | ) |
| 80 | |
| 81 | |
| 82 | ### mesh to sequence convertion: batch x ntokens |
| 83 | input_ids = discrete_padded_coords.reshape(batch, -1) |
| 84 | attention_mask = (input_ids != self.pad_id).float() |
| 85 | # reserve two spots: |
| 86 | # input_ids: <bos> ... <eos> <pad> ... => <pad> ... <pad> <pad> ... |
| 87 | # attn_mask: 1 ... 1 0 ... => 1 ... 1 0 ... |
| 88 | place_holder = torch.ones_like(input_ids[:, [0]]) # batch x 1 |
| 89 | input_ids = torch.cat((place_holder * self.pad_id, input_ids, place_holder * self.pad_id), dim=1) |
| 90 | attention_mask = torch.cat((place_holder, place_holder, attention_mask), dim=1) |
| 91 | |
| 92 | ### meshXL inputs |
| 93 | data_dict['input_ids'] = input_ids.long() # batch x (nf * 3 * 3 + 2) |
| 94 | data_dict['attention_mask'] = attention_mask.float() # batch x (nf * 3 * 3 + 2) |
| 95 | |
| 96 | # discard <bos> and <eos> tokens |
| 97 | data_dict['codes'] = discrete_padded_coords.long() # batch x (nf * 3 * 3) |
| 98 | data_dict['discrete_face_coords'] = discrete_face_coords |
| 99 | |
| 100 | return data_dict |
| 101 | |
| 102 | |
| 103 | def detokenize(self, input_ids: Tensor) -> dict: |
no test coverage detected