Coverage for trimesh/registration.py: 93%

1""" 

2registration.py 

3--------------- 

4 

5Functions for registering (aligning) point clouds with meshes. 

6""" 

7 

8import numpy as np 

9 

10from . import bounds, transformations, util 

11from .geometry import weighted_vertex_normals 

12from .points import PointCloud, plane_fit 

13from .transformations import transform_points 

14from .triangles import angles, cross, normals 

15from .typed import ArrayLike, Integer, Seed 

16 

17try: 

18 import scipy.sparse as sparse 

19 from scipy.spatial import cKDTree 

20except BaseException as E: 

21 # wrapping just ImportError fails in some cases 

22 # will raise the error when someone tries to use KDtree 

23 from . import exceptions 

24 

25 cKDTree = exceptions.ExceptionWrapper(E) 

26 sparse = exceptions.ExceptionWrapper(E) 

27 

28 

29# permutations of cube rotations 

30# the principal inertia transform has arbitrary sign 

31# along the 3 major axis so try all combinations of 

32# 180 degree rotations with a quick first ICP pass 

33_cube_diagonals = np.array( 

34 [ 

35 [1, 1, 1, 1], 

36 [1, 1, -1, 1], 

37 [1, -1, 1, 1], 

38 [-1, 1, 1, 1], 

39 [-1, -1, 1, 1], 

40 [-1, 1, -1, 1], 

41 [1, -1, -1, 1], 

42 [-1, -1, -1, 1], 

43 ], 

44 dtype=np.float64, 

45) 

46 

47 

48def mesh_other( 

49 mesh, 

50 other, 

51 samples: Integer = 500, 

52 scale: bool = False, 

53 icp_first: Integer = 10, 

54 icp_final: Integer = 50, 

55 reflection: bool = True, 

56 seed: Seed = None, 

57 **kwargs, 

58): 

59 """ 

60 Align a mesh with another mesh or a PointCloud using 

61 the principal axes of inertia as a starting point which 

62 is refined by iterative closest point. 

63 

64 Parameters 

65 ------------ 

66 mesh : trimesh.Trimesh object 

67 Mesh to align with other 

68 other : trimesh.Trimesh or (n, 3) float 

69 Mesh or points in space 

70 samples : int 

71 Number of samples from mesh surface to align 

72 scale : bool 

73 Allow scaling in transform 

74 icp_first : int 

75 How many ICP iterations for the 9 possible 

76 combinations of sign flippage 

77 icp_final : int 

78 How many ICP iterations for the closest 

79 candidate from the wider search 

80 seed : None or int 

81 Seed the surface sampling this uses to pick key points: 

82 pass an integer for deterministic results. 

83 kwargs : dict 

84 Passed through to `icp`, which passes through to `procrustes` 

85 

86 Returns 

87 ----------- 

88 mesh_to_other : (4, 4) float 

89 Transform to align mesh to the other object 

90 cost : float 

91 Average squared distance per point 

92 """ 

93 

94 def key_points(m, count): 

95 """ 

96 Return a combination of mesh vertices and surface samples 

97 with vertices chosen by likelihood to be important 

98 to registration. 

99 """ 

100 if len(m.vertices) < (count / 2): 

101 return np.vstack((m.vertices, m.sample(count - len(m.vertices), seed=seed))) 

102 else: 

103 return m.sample(count, seed=seed) 

104 

105 if not util.is_instance_named(mesh, "Trimesh"): 

106 raise ValueError("mesh must be Trimesh object!") 

107 

108 inverse = True 

109 search = mesh 

110 # if both are meshes use the smaller one for searching 

111 if util.is_instance_named(other, "Trimesh"): 

112 if len(mesh.vertices) > len(other.vertices): 

113 # do the expensive tree construction on the 

114 # smaller mesh and query the others points 

115 search = other 

116 inverse = False 

117 points = key_points(m=mesh, count=samples) 

118 points_mesh = mesh 

119 else: 

120 points_mesh = other 

121 points = key_points(m=other, count=samples) 

122 

123 if points_mesh.is_volume: 

124 points_PIT = points_mesh.principal_inertia_transform 

125 else: 

126 points_PIT = points_mesh.bounding_box_oriented.principal_inertia_transform 

127 

128 elif util.is_shape(other, (-1, 3)): 

129 # case where other is just points 

130 points = other 

131 points_PIT = bounds.oriented_bounds(points)[0] 

132 else: 

133 raise ValueError("other must be mesh or (n, 3) points!") 

134 

135 # get the transform that aligns the search mesh principal 

136 # axes of inertia with the XYZ axis at the origin 

137 if search.is_volume: 

138 search_PIT = search.principal_inertia_transform 

139 else: 

140 search_PIT = search.bounding_box_oriented.principal_inertia_transform 

141 

142 # transform that moves the principal axes of inertia 

143 # of the search mesh to be aligned with the best- guess 

144 # principal axes of the points 

145 search_to_points = np.dot(np.linalg.inv(points_PIT), search_PIT) 

146 if not reflection: 

147 # drop the negative-determinant seeds — icp can only refine 

148 # proper rotations so a reflected seed stays reflected, #2482 

149 diagonals = _cube_diagonals[np.prod(_cube_diagonals, axis=1) > 0.0] 

150 else: 

151 diagonals = _cube_diagonals 

152 

153 # expand the diagonals into (n, 4, 4) transforms 

154 cubes = np.eye(4) * diagonals[:, None, :] 

155 

156 # loop through permutations and run iterative closest point 

157 costs = np.ones(len(cubes)) * np.inf 

158 transforms = [None] * len(cubes) 

159 centroid = search.centroid 

160 

161 for i, flip in enumerate(cubes): 

162 # transform from points to search mesh 

163 # flipped around the centroid of search 

164 a_to_b = np.dot( 

165 transformations.transform_around(flip, centroid), 

166 np.linalg.inv(search_to_points), 

167 ) 

168 

169 # run first pass ICP 

170 matrix, _junk, cost = icp( 

171 a=points, 

172 b=search, 

173 initial=a_to_b, 

174 max_iterations=int(icp_first), 

175 scale=scale, 

176 reflection=reflection, 

177 **kwargs, 

178 ) 

179 

180 # save transform and costs from ICP 

181 transforms[i] = matrix 

182 costs[i] = cost 

183 

184 # run a final ICP refinement step 

185 matrix, _junk, cost = icp( 

186 a=points, 

187 b=search, 

188 initial=transforms[np.argmin(costs)], 

189 max_iterations=int(icp_final), 

190 scale=scale, 

191 reflection=reflection, 

192 **kwargs, 

193 ) 

194 

195 # convert to per- point distance average 

196 cost /= len(points) 

197 

198 # we picked the smaller mesh to construct the tree 

199 # on so we may have calculated a transform backwards 

200 # to save computation, so just invert matrix here 

201 if inverse: 

202 mesh_to_other = np.linalg.inv(matrix) 

203 else: 

204 mesh_to_other = matrix 

205 

206 return mesh_to_other, cost 

207 

208 

209def procrustes( 

210 a: ArrayLike, 

211 b: ArrayLike, 

212 weights: ArrayLike | None = None, 

213 reflection: bool = True, 

214 translation: bool = True, 

215 scale: bool = True, 

216 return_cost: bool = True, 

217): 

218 """ 

219 Perform Procrustes' analysis to quickly align two corresponding 

220 point clouds subject to constraints. This is much cheaper than 

221 any other registration method but only applies if the two inputs 

222 correspond in order. 

223 

224 Finds the transformation T mapping a to b which minimizes the 

225 square sum distances between Ta and b, also called the cost. 

226 

227 Optionally filter the points in a and b via a binary weights array. 

228 Non-uniform weights are also supported, but won't yield the optimal rotation. 

229 

230 Parameters 

231 ---------- 

232 a : (n,3) float 

233 List of points in space 

234 b : (n,3) float 

235 List of points in space 

236 weights : (n,) float 

237 List of floats representing how much weight is assigned to each point. 

238 Binary entries can be used to filter the arrays; normalization is not required. 

239 Translation and scaling are adjusted according to the weighting. 

240 Note, however, that this method does not yield the optimal rotation for 

241 non-uniform weighting, 

242 as this would require an iterative, nonlinear optimization approach. 

243 reflection : bool 

244 If the transformation is allowed reflections 

245 translation : bool 

246 If the transformation is allowed translation and rotation. 

247 scale : bool 

248 If the transformation is allowed scaling 

249 return_cost : bool 

250 Whether to return the cost and transformed a as well 

251 

252 Returns 

253 ---------- 

254 matrix : (4,4) float 

255 The transformation matrix sending a to b 

256 transformed : (n,3) float 

257 The image of a under the transformation 

258 cost : float 

259 The cost of the transformation 

260 """ 

261 

262 a_original = np.asanyarray(a, dtype=np.float64) 

263 b_original = np.asanyarray(b, dtype=np.float64) 

264 if not util.is_shape(a_original, (-1, 3)) or not util.is_shape(b_original, (-1, 3)): 

265 raise ValueError("points must be (n,3)!") 

266 if len(a_original) != len(b_original): 

267 raise ValueError("a and b must contain same number of points!") 

268 # weights are set to uniform if not provided. 

269 if weights is None: 

270 weights = np.ones(len(a_original)) 

271 w = np.maximum(np.asanyarray(weights, dtype=np.float64), 0) 

272 if len(w) != len(a): 

273 raise ValueError("weights must have same length as a and b!") 

274 w_norm = (w / w.sum()).reshape((-1, 1)) 

275 

276 # All zero entries are removed from further computations. 

277 # If weights is a binary array, the optimal solution can still be found by 

278 # simply removing the zero entries. 

279 nonzero_weights = w_norm[:, 0] > 0.0 

280 a_nonzero = a_original[nonzero_weights] 

281 b_nonzero = b_original[nonzero_weights] 

282 w_norm = w_norm[nonzero_weights] 

283 

284 # Remove translation component 

285 if translation: 

286 # centers are (weighted) averages of the individual points. 

287 acenter = (a_nonzero * w_norm).sum(axis=0) 

288 bcenter = (b_nonzero * w_norm).sum(axis=0) 

289 else: 

290 acenter = np.zeros(a_nonzero.shape[1]) 

291 bcenter = np.zeros(b_nonzero.shape[1]) 

292 

293 # Remove scale component 

294 if scale: 

295 # scale is the square root of the (weighted) average of the 

296 # squared difference between each point and the center. 

297 ascale = np.sqrt((((a_nonzero - acenter) ** 2) * w_norm).sum()) 

298 bscale = np.sqrt((((b_nonzero - bcenter) ** 2) * w_norm).sum()) 

299 else: 

300 ascale = 1 

301 bscale = 1 

302 

303 # Use SVD to find optimal orthogonal matrix R 

304 # constrained to det(R) = 1 if necessary. 

305 

306 target = np.dot(((b_nonzero - bcenter) / bscale).T, ((a_nonzero - acenter) / ascale)) 

307 

308 u, _s, vh = np.linalg.svd(target) 

309 

310 if reflection: 

311 R = np.dot(u, vh) 

312 else: 

313 # no reflection allowed, so determinant must be 1.0 

314 R = np.dot(np.dot(u, np.diag([1, 1, np.linalg.det(np.dot(u, vh))])), vh) 

315 

316 # Compute our 4D transformation matrix encoding 

317 # a -> (R @ (a - acenter)/ascale) * bscale + bcenter 

318 # = (bscale/ascale)R @ a + (bcenter - (bscale/ascale)R @ acenter) 

319 translation = bcenter - (bscale / ascale) * np.dot(R, acenter) 

320 matrix = np.hstack((bscale / ascale * R, translation.reshape(-1, 1))) 

321 matrix = np.vstack((matrix, np.array([0.0] * (a.shape[1]) + [1.0]).reshape(1, -1))) 

322 

323 if return_cost: 

324 # Transform the original input array, including zero-weighted points 

325 transformed = transform_points(a_original, matrix) 

326 # The cost is the (weighted) sum of the euclidean distances between 

327 # the transformed source points and the target points. 

328 cost = (((b_nonzero - transformed[nonzero_weights]) ** 2) * w_norm).sum() 

329 return matrix, transformed, cost 

330 

331 return matrix 

332 

333 

334def icp(a, b, initial=None, threshold=1e-5, max_iterations=20, **kwargs): 

335 """ 

336 Apply the iterative closest point algorithm to align a point cloud with 

337 another point cloud or mesh. Will only produce reasonable results if the 

338 initial transformation is roughly correct. Initial transformation can be 

339 found by applying Procrustes' analysis to a suitable set of landmark 

340 points (often picked manually). 

341 

342 Parameters 

343 ---------- 

344 a : (n,3) float 

345 List of points in space. 

346 b : (m,3) float or Trimesh 

347 List of points in space or mesh. 

348 initial : (4,4) float 

349 Initial transformation. 

350 threshold : float 

351 Stop when change in cost is less than threshold 

352 max_iterations : int 

353 Maximum number of iterations 

354 kwargs : dict 

355 Args to pass to procrustes 

356 

357 Returns 

358 ---------- 

359 matrix : (4,4) float 

360 The transformation matrix sending a to b 

361 transformed : (n,3) float 

362 The image of a under the transformation 

363 cost : float 

364 The cost of the transformation 

365 """ 

366 

367 a = np.asanyarray(a, dtype=np.float64) 

368 if not util.is_shape(a, (-1, 3)): 

369 raise ValueError("points must be (n,3)!") 

370 

371 if initial is None: 

372 initial = np.eye(4) 

373 

374 is_mesh = util.is_instance_named(b, "Trimesh") 

375 if not is_mesh: 

376 b = np.asanyarray(b, dtype=np.float64) 

377 if not util.is_shape(b, (-1, 3)): 

378 raise ValueError("points must be (n,3)!") 

379 btree = cKDTree(b) 

380 

381 # transform a under initial_transformation 

382 a = transform_points(a, initial) 

383 total_matrix = initial 

384 

385 # start with infinite cost 

386 old_cost = np.inf 

387 

388 # avoid looping forever by capping iterations 

389 for _ in range(max_iterations): 

390 # Closest point in b to each point in a 

391 if is_mesh: 

392 closest, _distance, _faces = b.nearest.on_surface(a) 

393 else: 

394 _distances, ix = btree.query(a, 1) 

395 closest = b[ix] 

396 

397 # align a with closest points 

398 matrix, transformed, cost = procrustes(a=a, b=closest, **kwargs) 

399 

400 # update a with our new transformed points 

401 a = transformed 

402 total_matrix = np.dot(matrix, total_matrix) 

403 

404 if old_cost - cost < threshold: 

405 break 

406 else: 

407 old_cost = cost 

408 

409 return total_matrix, transformed, cost 

410 

411 

412def _normalize_by_source(source_mesh, target_geometry, target_positions): 

413 # Utility function to put the source mesh in [-1, 1]^3 and transform 

414 # target geometry accordingly 

415 if not util.is_instance_named(target_geometry, "Trimesh") and not isinstance( 

416 target_geometry, PointCloud 

417 ): 

418 vertices = np.asanyarray(target_geometry) 

419 target_geometry = PointCloud(vertices) 

420 centroid, scale = source_mesh.centroid, source_mesh.scale 

421 source_mesh.vertices = (source_mesh.vertices - centroid[None, :]) / scale 

422 # Dont forget to also transform the target positions 

423 target_geometry.vertices = (target_geometry.vertices - centroid[None, :]) / scale 

424 if target_positions is not None: 

425 target_positions = (target_positions - centroid[None, :]) / scale 

426 return target_geometry, target_positions, centroid, scale 

427 

428 

429def _denormalize_by_source( 

430 source_mesh, target_geometry, target_positions, result, centroid, scale 

431): 

432 # Utility function to transform source mesh from 

433 # [-1, 1]^3 to its original transform 

434 # and transform target geometry accordingly 

435 source_mesh.vertices = scale * source_mesh.vertices + centroid[None, :] 

436 target_geometry.vertices = scale * target_geometry.vertices + centroid[None, :] 

437 if target_positions is not None: 

438 target_positions = scale * target_positions + centroid[None, :] 

439 if isinstance(result, list): 

440 result = [scale * x + centroid[None, :] for x in result] 

441 else: 

442 result = scale * result + centroid[None, :] 

443 return result 

444 

445 

446def nricp_amberg( 

447 source_mesh, 

448 target_geometry, 

449 source_landmarks=None, 

450 target_positions=None, 

451 steps=None, 

452 eps=0.0001, 

453 gamma=1, 

454 distance_threshold=0.1, 

455 return_records=False, 

456 use_faces=True, 

457 use_vertex_normals=True, 

458 neighbors_count=8, 

459): 

460 """ 

461 Non Rigid Iterative Closest Points 

462 

463 Implementation of "Amberg et al. 2007: Optimal Step 

464 Nonrigid ICP Algorithms for Surface Registration." 

465 Allows to register non-rigidly a mesh on another or 

466 on a point cloud. The core algorithm is explained 

467 at the end of page 3 of the paper. 

468 

469 Comparison between nricp_amberg and nricp_sumner: 

470 * nricp_amberg fits to the target mesh in less steps 

471 * nricp_amberg can generate sharp edges 

472 * only vertices and their neighbors are considered 

473 * nricp_sumner tend to preserve more the original shape 

474 * nricp_sumner parameters are easier to tune 

475 * nricp_sumner solves for triangle positions whereas 

476 nricp_amberg solves for vertex transforms 

477 * nricp_sumner is less optimized when wn > 0 

478 

479 Parameters 

480 ---------- 

481 source_mesh : Trimesh 

482 Source mesh containing both vertices and faces. 

483 target_geometry : Trimesh or PointCloud or (n, 3) float 

484 Target geometry. It can contain no faces or be a PointCloud. 

485 source_landmarks : (n,) int or ((n,) int, (n, 3) float) 

486 n landmarks on the the source mesh. 

487 Represented as vertex indices (n,) int. 

488 It can also be represented as a tuple of triangle 

489 indices and barycentric coordinates ((n,) int, (n, 3) float,). 

490 target_positions : (n, 3) float 

491 Target positions assigned to source landmarks 

492 steps : Core parameters of the algorithm 

493 Iterable of iterables (ws, wl, wn, max_iter,). 

494 ws is smoothness term, wl weights landmark importance, wn normal importance 

495 and max_iter is the maximum number of iterations per step. 

496 eps : float 

497 If the error decrease if inferior to this value, the current step ends. 

498 gamma : float 

499 Weight the translation part against the rotational/skew part. 

500 Recommended value : 1. 

501 distance_threshold : float 

502 Distance threshold to account for a vertex match or not. 

503 return_records : bool 

504 If True, also returns all the intermediate results. It can help debugging 

505 and tune the parameters to match a specific case. 

506 use_faces : bool 

507 If True and if target geometry has faces, use proximity.closest_point to find 

508 matching points. Else use scipy's cKDTree object. 

509 use_vertex_normals : bool 

510 If True and if target geometry has faces, interpolate the normals of the target 

511 geometry matching points. 

512 Else use face normals or estimated normals if target geometry has no faces. 

513 neighbors_count : int 

514 number of neighbors used for normal estimation. Only used if target geometry has 

515 no faces or if use_faces is False. 

516 

517 Returns 

518 ---------- 

519 result : (n, 3) float or list[(n, 3) float] 

520 The vertices positions of source_mesh such that it is registered non-rigidly 

521 onto the target geometry. 

522 If return_records is True, it returns the list of the vertex positions at each 

523 iteration. 

524 """ 

525 

526 def _solve_system(M_kron_G, D, vertices_weight, nearest, ws, nE, nV, Dl, Ul, wl): 

527 # Solve for Eq. 12 

528 U = nearest * vertices_weight[:, None] 

529 use_landmarks = Dl is not None and Ul is not None 

530 A_stack = [ws * M_kron_G, D.multiply(vertices_weight[:, None])] 

531 B_shape = (4 * nE + nV, 3) 

532 if use_landmarks: 

533 A_stack.append(wl * Dl) 

534 B_shape = (4 * nE + nV + Ul.shape[0], 3) 

535 A = sparse.csr_matrix(sparse.vstack(A_stack)) 

536 B = sparse.lil_matrix(B_shape, dtype=np.float32) 

537 B[4 * nE : (4 * nE + nV), :] = U 

538 if use_landmarks: 

539 B[4 * nE + nV : (4 * nE + nV + Ul.shape[0]), :] = Ul * wl 

540 X = sparse.linalg.spsolve(A.T * A, A.T * B).toarray() 

541 return X 

542 

543 def _node_arc_incidence(mesh, do_weight): 

544 # Computes node-arc incidence matrix of mesh (Eq.10) 

545 nV = mesh.edges.max() + 1 

546 nE = len(mesh.edges) 

547 rows = np.repeat(np.arange(nE), 2) 

548 cols = mesh.edges.flatten() 

549 data = np.ones(2 * nE, np.float32) 

550 data[1::2] = -1 

551 if do_weight: 

552 edge_lengths = np.linalg.norm( 

553 mesh.vertices[mesh.edges[:, 0]] - mesh.vertices[mesh.edges[:, 1]], axis=-1 

554 ) 

555 data *= np.repeat(1 / edge_lengths, 2) 

556 return sparse.coo_matrix((data, (rows, cols)), shape=(nE, nV)) 

557 

558 def _create_D(vertex_3d_data): 

559 # Create Data matrix (Eq. 8) 

560 nV = len(vertex_3d_data) 

561 rows = np.repeat(np.arange(nV), 4) 

562 cols = np.arange(4 * nV) 

563 data = np.concatenate((vertex_3d_data, np.ones((nV, 1))), axis=-1).flatten() 

564 return sparse.csr_matrix((data, (rows, cols)), shape=(nV, 4 * nV)) 

565 

566 def _create_X(nV): 

567 # Create Unknowns Matrix (Eq. 1) 

568 X_ = np.concatenate((np.eye(3), np.array([[0, 0, 0]])), axis=0) 

569 return np.tile(X_, (nV, 1)) 

570 

571 def _create_Dl_Ul(D, source_mesh, source_landmarks, target_positions): 

572 # Create landmark terms (Eq. 11) 

573 Dl, Ul = None, None 

574 

575 if source_landmarks is None or target_positions is None: 

576 # If no landmarks are provided, return None for both 

577 return Dl, Ul 

578 

579 if isinstance(source_landmarks, tuple): 

580 source_tids, source_barys = source_landmarks 

581 source_tri_vids = source_mesh.faces[source_tids] 

582 # u * x1, v * x2 and w * x3 combined 

583 Dl = D[source_tri_vids.flatten(), :] 

584 Dl.data *= source_barys.flatten().repeat(np.diff(Dl.indptr)) 

585 x0 = source_mesh.vertices[source_tri_vids[:, 0]] 

586 x1 = source_mesh.vertices[source_tri_vids[:, 1]] 

587 x2 = source_mesh.vertices[source_tri_vids[:, 2]] 

588 Ul0 = ( 

589 target_positions 

590 - x1 * source_barys[:, 1, None] 

591 - x2 * source_barys[:, 2, None] 

592 ) 

593 Ul1 = ( 

594 target_positions 

595 - x0 * source_barys[:, 0, None] 

596 - x2 * source_barys[:, 2, None] 

597 ) 

598 Ul2 = ( 

599 target_positions 

600 - x0 * source_barys[:, 0, None] 

601 - x1 * source_barys[:, 1, None] 

602 ) 

603 Ul = np.zeros((Ul0.shape[0] * 3, 3)) 

604 Ul[0::3] = Ul0 # y - v * x2 + w * x3 

605 Ul[1::3] = Ul1 # y - u * x1 + w * x3 

606 Ul[2::3] = Ul2 # y - u * x1 + v * x2 

607 else: 

608 Dl = D[source_landmarks, :] 

609 Ul = target_positions 

610 return Dl, Ul 

611 

612 target_geometry, target_positions, centroid, scale = _normalize_by_source( 

613 source_mesh, target_geometry, target_positions 

614 ) 

615 

616 # Number of edges and vertices in source mesh 

617 nE = len(source_mesh.edges) 

618 nV = len(source_mesh.vertices) 

619 

620 # Initialize transformed vertices 

621 transformed_vertices = source_mesh.vertices.copy() 

622 # Node-arc incidence (M in Eq. 10) 

623 M = _node_arc_incidence(source_mesh, True) 

624 # G (Eq. 10) 

625 G = np.diag([1, 1, 1, gamma]) 

626 # M kronecker G (Eq. 10) 

627 M_kron_G = sparse.kron(M, G) 

628 # D (Eq. 8) 

629 D = _create_D(source_mesh.vertices) 

630 # D but for normal computation from the transformations X 

631 DN = _create_D(source_mesh.vertex_normals) 

632 # Unknowns 4x3 transformations X (Eq. 1) 

633 X = _create_X(nV) 

634 # Landmark related terms (Eq. 11) 

635 Dl, Ul = _create_Dl_Ul(D, source_mesh, source_landmarks, target_positions) 

636 

637 # Parameters of the algorithm (Eq. 6) 

638 # order : Alpha, Beta, normal weighting, and max iteration for step 

639 if steps is None: 

640 steps = [ 

641 [0.01, 10, 0.5, 10], 

642 [0.02, 5, 0.5, 10], 

643 [0.03, 2.5, 0.5, 10], 

644 [0.01, 0, 0.0, 10], 

645 ] 

646 if return_records: 

647 records = [transformed_vertices] 

648 

649 # Main loop 

650 for ws, wl, wn, max_iter in steps: 

651 # If normals are estimated from points and if there are less 

652 # than 3 points per query, avoid normal estimation 

653 if not use_faces and neighbors_count < 3: 

654 wn = 0 

655 

656 last_error = np.finfo(np.float32).max 

657 error = np.finfo(np.float16).max 

658 cpt_iter = 0 

659 

660 # Current step iterations loop 

661 while last_error - error > eps and (max_iter is None or cpt_iter < max_iter): 

662 qres = _from_mesh( 

663 target_geometry, 

664 transformed_vertices, 

665 from_vertices_only=not use_faces, 

666 return_normals=wn > 0, 

667 return_interpolated_normals=wn > 0 and use_vertex_normals, 

668 neighbors_count=neighbors_count, 

669 ) 

670 

671 # Data weighting 

672 vertices_weight = np.ones(nV) 

673 vertices_weight[qres["distances"] > distance_threshold] = 0 

674 

675 if wn > 0 and "normals" in qres: 

676 target_normals = qres["normals"] 

677 if use_vertex_normals and "interpolated_normals" in qres: 

678 target_normals = qres["interpolated_normals"] 

679 # Normal weighting = multiplying weights by cosines^wn 

680 source_normals = DN * X 

681 dot = util.diagonal_dot(source_normals, target_normals) 

682 # Normal orientation is only known for meshes as target 

683 dot = np.clip(dot, 0, 1) if use_faces else np.abs(dot) 

684 vertices_weight = vertices_weight * dot**wn 

685 

686 # Actual system solve 

687 X = _solve_system( 

688 M_kron_G, D, vertices_weight, qres["nearest"], ws, nE, nV, Dl, Ul, wl 

689 ) 

690 transformed_vertices = D * X 

691 last_error = error 

692 error_vec = np.linalg.norm(qres["nearest"] - transformed_vertices, axis=-1) 

693 error = (error_vec * vertices_weight).mean() 

694 if return_records: 

695 records.append(transformed_vertices) 

696 cpt_iter += 1 

697 

698 if return_records: 

699 result = records 

700 else: 

701 result = transformed_vertices 

702 

703 result = _denormalize_by_source( 

704 source_mesh, target_geometry, target_positions, result, centroid, scale 

705 ) 

706 return result 

707 

708 

709def _from_mesh( 

710 mesh, 

711 input_points, 

712 from_vertices_only=False, 

713 return_barycentric_coordinates=False, 

714 return_normals=False, 

715 return_interpolated_normals=False, 

716 neighbors_count=10, 

717 **kwargs, 

718): 

719 """ 

720 Find the the closest points and associated attributes from a Trimesh. 

721 

722 Parameters 

723 ----------- 

724 mesh : Trimesh 

725 Trimesh from which the query is performed 

726 input_points : (m, 3) float 

727 Input query points 

728 from_vertices_only : bool 

729 If True, consider only the vertices and not the faces 

730 return_barycentric_coordinates : bool 

731 If True, return the barycentric coordinates 

732 return_normals : bool 

733 If True, compute the normals at each closest point 

734 return_interpolated_normals : bool 

735 If True, return the interpolated normal at each closest point 

736 neighbors_count : int 

737 The number of closest neighbors to query 

738 kwargs : dict 

739 Dict to accept other key word arguments (not used) 

740 Returns 

741 ---------- 

742 qres : dict 

743 Dictionary containing : 

744 - nearest points (m, 3) with key 'nearest' 

745 - distances to nearest point (m,) with key 'distances' 

746 - support triangle indices of the nearest points (m,) with key 'tids' 

747 - [optional] normals at nearest points (m,3) with key 'normals' 

748 - [optional] barycentric coordinates in support triangles (m,3) with key 

749 'barycentric_coordinates' 

750 - [optional] interpolated normals (m,3) with key 'interpolated_normals' 

751 """ 

752 input_points = np.asanyarray(input_points) 

753 neighbors_count = min(neighbors_count, len(mesh.vertices)) 

754 

755 if from_vertices_only or len(mesh.faces) == 0: 

756 # Consider only the vertices 

757 return _from_points( 

758 mesh.vertices, 

759 input_points, 

760 mesh.kdtree, 

761 return_normals=return_normals, 

762 neighbors_count=neighbors_count, 

763 ) 

764 # Else if we consider faces, use proximity.closest_point 

765 qres = {} 

766 from .proximity import closest_point 

767 from .triangles import points_to_barycentric 

768 

769 qres["nearest"], qres["distances"], qres["tids"] = closest_point(mesh, input_points) 

770 

771 if return_normals: 

772 qres["normals"] = mesh.face_normals[qres["tids"]] 

773 if return_barycentric_coordinates or return_interpolated_normals: 

774 qres["barycentric_coordinates"] = points_to_barycentric( 

775 mesh.vertices[mesh.faces[qres["tids"]]], qres["nearest"] 

776 ) 

777 

778 if return_interpolated_normals: 

779 # Interpolation from barycentric coordinates 

780 qres["interpolated_normals"] = np.einsum( 

781 "ij,ijk->ik", 

782 qres["barycentric_coordinates"], 

783 mesh.vertex_normals[mesh.faces[qres["tids"]]], 

784 ) 

785 return qres 

786 

787 

788def _from_points( 

789 target_points, 

790 input_points, 

791 kdtree=None, 

792 return_normals=False, 

793 neighbors_count=10, 

794 **kwargs, 

795): 

796 """ 

797 Find the the closest points and associated attributes 

798 from a set of 3D points. 

799 

800 Parameters 

801 ----------- 

802 target_points : (n, 3) float 

803 Points from which the query is performed 

804 input_points : (m, 3) float 

805 Input query points 

806 kdtree : scipy.cKDTree 

807 KDTree used for query. Computed if not provided 

808 return_normals : bool 

809 If True, compute the normals at each nearest point 

810 neighbors_count : int 

811 The number of closest neighbors to query 

812 kwargs : dict 

813 Dict to accept other key word arguments (not used) 

814 

815 Returns 

816 ---------- 

817 qres : dict 

818 Dictionary containing : 

819 - nearest points (m, 3) with key 'nearest' 

820 - distances to nearest point (m,) with key 'distances' 

821 - vertex indices of the nearest points (m,) with key 'vertex_indices' 

822 - [optional] normals at nearest points (m,3) with key 'normals' 

823 """ 

824 # Empty result 

825 target_points = np.asanyarray(target_points) 

826 input_points = np.asanyarray(input_points) 

827 neighbors_count = min(neighbors_count, len(target_points)) 

828 qres = {} 

829 if kdtree is None: 

830 kdtree = cKDTree(target_points) 

831 

832 if return_normals: 

833 assert neighbors_count >= 3 

834 distances, indices = kdtree.query(input_points, k=neighbors_count) 

835 nearest = target_points[indices, :] 

836 qres["normals"] = plane_fit(nearest)[1] 

837 qres["nearest"] = nearest[:, 0] 

838 qres["distances"] = distances[:, 0] 

839 qres["vertex_indices"] = indices[:, 0] 

840 else: 

841 qres["distances"], qres["vertex_indices"] = kdtree.query(input_points) 

842 qres["nearest"] = target_points[qres["vertex_indices"], :] 

843 

844 return qres 

845 

846 

847def nricp_sumner( 

848 source_mesh, 

849 target_geometry, 

850 source_landmarks=None, 

851 target_positions=None,