(batch, scalp_mesh_data)
| 136 | |
| 137 | |
| 138 | def horizontally_flip(batch, scalp_mesh_data): |
| 139 | |
| 140 | # tbn=batch['full_strands']["tbn"].cuda() |
| 141 | positions=batch['full_strands']["positions"][0].numpy() |
| 142 | # print("positions", positions.shape) |
| 143 | |
| 144 | #flip |
| 145 | positions[:,:,0]*=-1 |
| 146 | |
| 147 | root_position = positions[:,0,:] #nr_strands x 3 |
| 148 | |
| 149 | #it's more error prone to flip the tbn so we just recompute it |
| 150 | |
| 151 | mesh_verts=scalp_mesh_data["verts"] |
| 152 | mesh_faces=scalp_mesh_data["faces"] |
| 153 | mesh_v_tangents=scalp_mesh_data["v_tangents"] |
| 154 | mesh_v_bitangents=scalp_mesh_data["v_bitangents"] |
| 155 | mesh_v_normals=scalp_mesh_data["v_normals"] |
| 156 | |
| 157 | mesh_faces=mesh_faces.astype(np.int32) |
| 158 | closest_points, barys, vertex_idxs, face_idxs=closest_point_barycentrics(root_position, mesh_verts, mesh_faces) |
| 159 | |
| 160 | # root_tangent, root_bitangent, root_normal = interpolate_tbn(self.root_position, barys, vertex_idxs, mesh_v_tangents, mesh_v_bitangents, mesh_v_normals) |
| 161 | root_tangent, root_bitangent, root_normal = interpolate_tbn(barys, vertex_idxs, mesh_v_tangents, mesh_v_bitangents, mesh_v_normals) |
| 162 | #replace the normals because it's smoother |
| 163 | root_normal=root_normal |
| 164 | tbn = np.stack((root_tangent,root_bitangent,root_normal),axis=2) |
| 165 | |
| 166 | #put it back on cuda |
| 167 | tbn=torch.from_numpy(tbn).unsqueeze(0).cuda() |
| 168 | positions=torch.from_numpy(positions).unsqueeze(0).cuda() |
| 169 | root_normal=root_normal.unsqueeze(0).cuda() |
| 170 | |
| 171 | |
| 172 | batch['full_strands']["tbn"]=tbn |
| 173 | batch['full_strands']["positions"]=positions |
| 174 | batch['full_strands']["root_normal"]=root_normal |
| 175 | batch["full_strands"]["root_uv"][:,:,0] = 1 - batch["full_strands"]["root_uv"][:,:,0] |
| 176 | |
| 177 | |
| 178 | return batch |
| 179 | |
| 180 | |
| 181 | def main(): |
no test coverage detected