(strands_tbn, strands_positions, root_normals)
| 139 | #strands_normals is [B, Nr_strands, 3] |
| 140 | # @torch.compile |
| 141 | def world_to_tbn_space(strands_tbn, strands_positions, root_normals): |
| 142 | # print("strands_tbn",strands_tbn.shape) |
| 143 | nr_batch= strands_tbn.shape[0] |
| 144 | nr_strands= strands_tbn.shape[1] |
| 145 | root_pos = strands_positions[:,:,0:1,:] #[B, Nr_strands, 1, 3] |
| 146 | |
| 147 | #transform from scalp to world |
| 148 | |
| 149 | #remove the root for the positional data |
| 150 | strands_positions = strands_positions-root_pos |
| 151 | |
| 152 | #we want to map tangent to X, bitangent to Z and normal to Y, so we swap B and N |
| 153 | indices_tbn=torch.tensor([0,2,1], device="cuda").long() |
| 154 | strands_tbn=torch.index_select(strands_tbn, 3, indices_tbn) |
| 155 | #make the Tangent to be along +x |
| 156 | strands_tbn[..., 0] = -strands_tbn[..., 0] |
| 157 | |
| 158 | #rotate so that the TBN is identity |
| 159 | #TBN is basically the rotation from scalp to world, we want the inverse |
| 160 | strands_tbn_inv = strands_tbn.transpose(2,3) |
| 161 | |
| 162 | #rotate positional data [B, Nr_strands, 3, 3] x [B, Nr_strands, nr_points_per_strand, 3] |
| 163 | strands_tbn_inv = strands_tbn_inv.reshape(nr_batch, nr_strands, 1, 3, 3) |
| 164 | # print("strands_positions",strands_positions.shape) |
| 165 | strands_positions = strands_positions.reshape(nr_batch, nr_strands, -1, 3, 1) |
| 166 | strands_positions= torch.matmul(strands_tbn_inv, strands_positions) |
| 167 | strands_positions = strands_positions.reshape(nr_batch, nr_strands, -1, 3) |
| 168 | |
| 169 | #roundtrip |
| 170 | # strands_tbn = strands_tbn.reshape(nr_batch, nr_strands, 1, 3, 3) |
| 171 | # strands_positions = strands_positions.reshape(nr_batch, nr_strands, -1, 3, 1) |
| 172 | # strands_positions= torch.matmul(strands_tbn, strands_positions) |
| 173 | # strands_positions = strands_positions.reshape(nr_batch, nr_strands, -1, 3) |
| 174 | # strands_positions = strands_positions+root_pos |
| 175 | |
| 176 | #rotate normals |
| 177 | root_normals = root_normals.reshape(nr_batch, nr_strands, 3, 1) |
| 178 | strands_tbn_inv = strands_tbn_inv.reshape(nr_batch, nr_strands, 3, 3) |
| 179 | root_normals= torch.matmul(strands_tbn_inv, root_normals) |
| 180 | root_normals = root_normals.reshape(nr_batch, nr_strands, 3) |
| 181 | |
| 182 | return strands_positions, root_normals |
| 183 | |
| 184 | #incurs a copy to cpu |
| 185 | #generates uv space map where each pixel has the index of the triangle, the 3 vertex indices of the triangle and the barycentric weights |
no test coverage detected