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 

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 

29def mesh_other( 

30 mesh, 

31 other, 

32 samples: Integer = 500, 

33 scale: bool = False, 

34 icp_first: Integer = 10, 

35 icp_final: Integer = 50, 

36 **kwargs, 

37): 

38 """ 

39 Align a mesh with another mesh or a PointCloud using 

40 the principal axes of inertia as a starting point which 

41 is refined by iterative closest point. 

42 

43 Parameters 

44 ------------ 

45 mesh : trimesh.Trimesh object 

46 Mesh to align with other 

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

48 Mesh or points in space 

49 samples : int 

50 Number of samples from mesh surface to align 

51 scale : bool 

52 Allow scaling in transform 

53 icp_first : int 

54 How many ICP iterations for the 9 possible 

55 combinations of sign flippage 

56 icp_final : int 

57 How many ICP iterations for the closest 

58 candidate from the wider search 

59 kwargs : dict 

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

61 

62 Returns 

63 ----------- 

64 mesh_to_other : (4, 4) float 

65 Transform to align mesh to the other object 

66 cost : float 

67 Average squared distance per point 

68 """ 

69 

70 def key_points(m, count): 

71 """ 

72 Return a combination of mesh vertices and surface samples 

73 with vertices chosen by likelihood to be important 

74 to registration. 

75 """ 

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

77 return np.vstack((m.vertices, m.sample(count - len(m.vertices)))) 

78 else: 

79 return m.sample(count) 

80 

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

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

83 

84 inverse = True 

85 search = mesh 

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

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

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

89 # do the expensive tree construction on the 

90 # smaller mesh and query the others points 

91 search = other 

92 inverse = False 

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

94 points_mesh = mesh 

95 else: 

96 points_mesh = other 

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

98 

99 if points_mesh.is_volume: 

100 points_PIT = points_mesh.principal_inertia_transform 

101 else: 

102 points_PIT = points_mesh.bounding_box_oriented.principal_inertia_transform 

103 

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

105 # case where other is just points 

106 points = other 

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

108 else: 

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

110 

111 # get the transform that aligns the search mesh principal 

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

113 if search.is_volume: 

114 search_PIT = search.principal_inertia_transform 

115 else: 

116 search_PIT = search.bounding_box_oriented.principal_inertia_transform 

117 

118 # transform that moves the principal axes of inertia 

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

120 # principal axes of the points 

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

122 

123 # permutations of cube rotations 

124 # the principal inertia transform has arbitrary sign 

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

126 # 180 degree rotations with a quick first ICP pass 

127 cubes = np.array( 

128 [ 

129 np.eye(4) * np.append(diag, 1) 

130 for diag in [ 

131 [1, 1, 1], 

132 [1, 1, -1], 

133 [1, -1, 1], 

134 [-1, 1, 1], 

135 [-1, -1, 1], 

136 [-1, 1, -1], 

137 [1, -1, -1], 

138 [-1, -1, -1], 

139 ] 

140 ] 

141 ) 

142 

143 # loop through permutations and run iterative closest point 

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

145 transforms = [None] * len(cubes) 

146 centroid = search.centroid 

147 

148 for i, flip in enumerate(cubes): 

149 # transform from points to search mesh 

150 # flipped around the centroid of search 

151 a_to_b = np.dot( 

152 transformations.transform_around(flip, centroid), 

153 np.linalg.inv(search_to_points), 

154 ) 

155 

156 # run first pass ICP 

157 matrix, _junk, cost = icp( 

158 a=points, 

159 b=search, 

160 initial=a_to_b, 

161 max_iterations=int(icp_first), 

162 scale=scale, 

163 **kwargs, 

164 ) 

165 

166 # save transform and costs from ICP 

167 transforms[i] = matrix 

168 costs[i] = cost 

169 

170 # run a final ICP refinement step 

171 matrix, _junk, cost = icp( 

172 a=points, 

173 b=search, 

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

175 max_iterations=int(icp_final), 

176 scale=scale, 

177 **kwargs, 

178 ) 

179 

180 # convert to per- point distance average 

181 cost /= len(points) 

182 

183 # we picked the smaller mesh to construct the tree 

184 # on so we may have calculated a transform backwards 

185 # to save computation, so just invert matrix here 

186 if inverse: 

187 mesh_to_other = np.linalg.inv(matrix) 

188 else: 

189 mesh_to_other = matrix 

190 

191 return mesh_to_other, cost 

192 

193 

194def procrustes( 

195 a: ArrayLike, 

196 b: ArrayLike, 

197 weights: ArrayLike | None = None, 

198 reflection: bool = True, 

199 translation: bool = True, 

200 scale: bool = True, 

201 return_cost: bool = True, 

202): 

203 """ 

204 Perform Procrustes' analysis to quickly align two corresponding 

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

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

207 correspond in order. 

208 

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

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

211 

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

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

214 

215 Parameters 

216 ---------- 

217 a : (n,3) float 

218 List of points in space 

219 b : (n,3) float 

220 List of points in space 

221 weights : (n,) float 

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

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

224 Translation and scaling are adjusted according to the weighting. 

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

226 non-uniform weighting, 

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

228 reflection : bool 

229 If the transformation is allowed reflections 

230 translation : bool 

231 If the transformation is allowed translation and rotation. 

232 scale : bool 

233 If the transformation is allowed scaling 

234 return_cost : bool 

235 Whether to return the cost and transformed a as well 

236 

237 Returns 

238 ---------- 

239 matrix : (4,4) float 

240 The transformation matrix sending a to b 

241 transformed : (n,3) float 

242 The image of a under the transformation 

243 cost : float 

244 The cost of the transformation 

245 """ 

246 

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

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

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

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

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

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

253 # weights are set to uniform if not provided. 

254 if weights is None: 

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

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

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

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

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

260 

261 # All zero entries are removed from further computations. 

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

263 # simply removing the zero entries. 

264 nonzero_weights = w_norm[:, 0] > 0.0 

265 a_nonzero = a_original[nonzero_weights] 

266 b_nonzero = b_original[nonzero_weights] 

267 w_norm = w_norm[nonzero_weights] 

268 

269 # Remove translation component 

270 if translation: 

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

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

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

274 else: 

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

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

277 

278 # Remove scale component 

279 if scale: 

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

281 # squared difference between each point and the center. 

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

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

284 else: 

285 ascale = 1 

286 bscale = 1 

287 

288 # Use SVD to find optimal orthogonal matrix R 

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

290 

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

292 

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

294 

295 if reflection: 

296 R = np.dot(u, vh) 

297 else: 

298 # no reflection allowed, so determinant must be 1.0 

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

300 

301 # Compute our 4D transformation matrix encoding 

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

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

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

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

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

307 

308 if return_cost: 

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

310 transformed = transform_points(a_original, matrix) 

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

312 # the transformed source points and the target points. 

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

314 return matrix, transformed, cost 

315 

316 return matrix 

317 

318 

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

320 """ 

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

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

323 initial transformation is roughly correct. Initial transformation can be 

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

325 points (often picked manually). 

326 

327 Parameters 

328 ---------- 

329 a : (n,3) float 

330 List of points in space. 

331 b : (m,3) float or Trimesh 

332 List of points in space or mesh. 

333 initial : (4,4) float 

334 Initial transformation. 

335 threshold : float 

336 Stop when change in cost is less than threshold 

337 max_iterations : int 

338 Maximum number of iterations 

339 kwargs : dict 

340 Args to pass to procrustes 

341 

342 Returns 

343 ---------- 

344 matrix : (4,4) float 

345 The transformation matrix sending a to b 

346 transformed : (n,3) float 

347 The image of a under the transformation 

348 cost : float 

349 The cost of the transformation 

350 """ 

351 

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

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

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

355 

356 if initial is None: 

357 initial = np.eye(4) 

358 

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

360 if not is_mesh: 

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

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

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

364 btree = cKDTree(b) 

365 

366 # transform a under initial_transformation 

367 a = transform_points(a, initial) 

368 total_matrix = initial 

369 

370 # start with infinite cost 

371 old_cost = np.inf 

372 

373 # avoid looping forever by capping iterations 

374 for _ in range(max_iterations): 

375 # Closest point in b to each point in a 

376 if is_mesh: 

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

378 else: 

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

380 closest = b[ix] 

381 

382 # align a with closest points 

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

384 

385 # update a with our new transformed points 

386 a = transformed 

387 total_matrix = np.dot(matrix, total_matrix) 

388 

389 if old_cost - cost < threshold: 

390 break 

391 else: 

392 old_cost = cost 

393 

394 return total_matrix, transformed, cost 

395 

396 

397def _normalize_by_source(source_mesh, target_geometry, target_positions): 

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

399 # target geometry accordingly 

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

401 target_geometry, PointCloud 

402 ): 

403 vertices = np.asanyarray(target_geometry) 

404 target_geometry = PointCloud(vertices) 

405 centroid, scale = source_mesh.centroid, source_mesh.scale 

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

407 # Dont forget to also transform the target positions 

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

409 if target_positions is not None: 

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

411 return target_geometry, target_positions, centroid, scale 

412 

413 

414def _denormalize_by_source( 

415 source_mesh, target_geometry, target_positions, result, centroid, scale 

416): 

417 # Utility function to transform source mesh from 

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

419 # and transform target geometry accordingly 

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

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

422 if target_positions is not None: 

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

424 if isinstance(result, list): 

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

426 else: 

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

428 return result 

429 

430 

431def nricp_amberg( 

432 source_mesh, 

433 target_geometry, 

434 source_landmarks=None, 

435 target_positions=None, 

436 steps=None, 

437 eps=0.0001, 

438 gamma=1, 

439 distance_threshold=0.1, 

440 return_records=False, 

441 use_faces=True, 

442 use_vertex_normals=True, 

443 neighbors_count=8, 

444): 

445 """ 

446 Non Rigid Iterative Closest Points 

447 

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

449 Nonrigid ICP Algorithms for Surface Registration." 

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

451 on a point cloud. The core algorithm is explained 

452 at the end of page 3 of the paper. 

453 

454 Comparison between nricp_amberg and nricp_sumner: 

455 * nricp_amberg fits to the target mesh in less steps 

456 * nricp_amberg can generate sharp edges 

457 * only vertices and their neighbors are considered 

458 * nricp_sumner tend to preserve more the original shape 

459 * nricp_sumner parameters are easier to tune 

460 * nricp_sumner solves for triangle positions whereas 

461 nricp_amberg solves for vertex transforms 

462 * nricp_sumner is less optimized when wn > 0 

463 

464 Parameters 

465 ---------- 

466 source_mesh : Trimesh 

467 Source mesh containing both vertices and faces. 

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

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

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

471 n landmarks on the the source mesh. 

472 Represented as vertex indices (n,) int. 

473 It can also be represented as a tuple of triangle 

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

475 target_positions : (n, 3) float 

476 Target positions assigned to source landmarks 

477 steps : Core parameters of the algorithm 

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

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

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

481 eps : float 

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

483 gamma : float 

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

485 Recommended value : 1. 

486 distance_threshold : float 

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

488 return_records : bool 

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

490 and tune the parameters to match a specific case. 

491 use_faces : bool 

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

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

494 use_vertex_normals : bool 

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

496 geometry matching points. 

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

498 neighbors_count : int 

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

500 no faces or if use_faces is False. 

501 

502 Returns 

503 ---------- 

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

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

506 onto the target geometry. 

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

508 iteration. 

509 """ 

510 

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

512 # Solve for Eq. 12 

513 U = nearest * vertices_weight[:, None] 

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

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

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

517 if use_landmarks: 

518 A_stack.append(wl * Dl) 

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

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

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

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

523 if use_landmarks: 

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

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

526 return X 

527 

528 def _node_arc_incidence(mesh, do_weight): 

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

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

531 nE = len(mesh.edges) 

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

533 cols = mesh.edges.flatten() 

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

535 data[1::2] = -1 

536 if do_weight: 

537 edge_lengths = np.linalg.norm( 

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

539 ) 

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

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

542 

543 def _create_D(vertex_3d_data): 

544 # Create Data matrix (Eq. 8) 

545 nV = len(vertex_3d_data) 

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

547 cols = np.arange(4 * nV) 

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

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

550 

551 def _create_X(nV): 

552 # Create Unknowns Matrix (Eq. 1) 

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

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

555 

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

557 # Create landmark terms (Eq. 11) 

558 Dl, Ul = None, None 

559 

560 if source_landmarks is None or target_positions is None: 

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

562 return Dl, Ul 

563 

564 if isinstance(source_landmarks, tuple): 

565 source_tids, source_barys = source_landmarks 

566 source_tri_vids = source_mesh.faces[source_tids] 

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

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

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

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

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

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

573 Ul0 = ( 

574 target_positions 

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

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

577 ) 

578 Ul1 = ( 

579 target_positions 

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

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

582 ) 

583 Ul2 = ( 

584 target_positions 

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

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

587 ) 

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

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

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

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

592 else: 

593 Dl = D[source_landmarks, :] 

594 Ul = target_positions 

595 return Dl, Ul 

596 

597 target_geometry, target_positions, centroid, scale = _normalize_by_source( 

598 source_mesh, target_geometry, target_positions 

599 ) 

600 

601 # Number of edges and vertices in source mesh 

602 nE = len(source_mesh.edges) 

603 nV = len(source_mesh.vertices) 

604 

605 # Initialize transformed vertices 

606 transformed_vertices = source_mesh.vertices.copy() 

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

608 M = _node_arc_incidence(source_mesh, True) 

609 # G (Eq. 10) 

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

611 # M kronecker G (Eq. 10) 

612 M_kron_G = sparse.kron(M, G) 

613 # D (Eq. 8) 

614 D = _create_D(source_mesh.vertices) 

615 # D but for normal computation from the transformations X 

616 DN = _create_D(source_mesh.vertex_normals) 

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

618 X = _create_X(nV) 

619 # Landmark related terms (Eq. 11) 

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

621 

622 # Parameters of the algorithm (Eq. 6) 

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

624 if steps is None: 

625 steps = [ 

626 [0.01, 10, 0.5, 10], 

627 [0.02, 5, 0.5, 10], 

628 [0.03, 2.5, 0.5, 10], 

629 [0.01, 0, 0.0, 10], 

630 ] 

631 if return_records: 

632 records = [transformed_vertices] 

633 

634 # Main loop 

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

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

637 # than 3 points per query, avoid normal estimation 

638 if not use_faces and neighbors_count < 3: 

639 wn = 0 

640 

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

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

643 cpt_iter = 0 

644 

645 # Current step iterations loop 

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

647 qres = _from_mesh( 

648 target_geometry, 

649 transformed_vertices, 

650 from_vertices_only=not use_faces, 

651 return_normals=wn > 0, 

652 return_interpolated_normals=wn > 0 and use_vertex_normals, 

653 neighbors_count=neighbors_count, 

654 ) 

655 

656 # Data weighting 

657 vertices_weight = np.ones(nV) 

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

659 

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

661 target_normals = qres["normals"] 

662 if use_vertex_normals and "interpolated_normals" in qres: 

663 target_normals = qres["interpolated_normals"] 

664 # Normal weighting = multiplying weights by cosines^wn 

665 source_normals = DN * X 

666 dot = util.diagonal_dot(source_normals, target_normals) 

667 # Normal orientation is only known for meshes as target 

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

669 vertices_weight = vertices_weight * dot**wn 

670 

671 # Actual system solve 

672 X = _solve_system( 

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

674 ) 

675 transformed_vertices = D * X 

676 last_error = error 

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

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

679 if return_records: 

680 records.append(transformed_vertices) 

681 cpt_iter += 1 

682 

683 if return_records: 

684 result = records 

685 else: 

686 result = transformed_vertices 

687 

688 result = _denormalize_by_source( 

689 source_mesh, target_geometry, target_positions, result, centroid, scale 

690 ) 

691 return result 

692 

693 

694def _from_mesh( 

695 mesh, 

696 input_points, 

697 from_vertices_only=False, 

698 return_barycentric_coordinates=False, 

699 return_normals=False, 

700 return_interpolated_normals=False, 

701 neighbors_count=10, 

702 **kwargs, 

703): 

704 """ 

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

706 

707 Parameters 

708 ----------- 

709 mesh : Trimesh 

710 Trimesh from which the query is performed 

711 input_points : (m, 3) float 

712 Input query points 

713 from_vertices_only : bool 

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

715 return_barycentric_coordinates : bool 

716 If True, return the barycentric coordinates 

717 return_normals : bool 

718 If True, compute the normals at each closest point 

719 return_interpolated_normals : bool 

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

721 neighbors_count : int 

722 The number of closest neighbors to query 

723 kwargs : dict 

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

725 Returns 

726 ---------- 

727 qres : dict 

728 Dictionary containing : 

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

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

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

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

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

734 'barycentric_coordinates' 

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

736 """ 

737 input_points = np.asanyarray(input_points) 

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

739 

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

741 # Consider only the vertices 

742 return _from_points( 

743 mesh.vertices, 

744 input_points, 

745 mesh.kdtree, 

746 return_normals=return_normals, 

747 neighbors_count=neighbors_count, 

748 ) 

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

750 qres = {} 

751 from .proximity import closest_point 

752 from .triangles import points_to_barycentric 

753 

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

755 

756 if return_normals: 

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

758 if return_barycentric_coordinates or return_interpolated_normals: 

759 qres["barycentric_coordinates"] = points_to_barycentric( 

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

761 ) 

762 

763 if return_interpolated_normals: 

764 # Interpolation from barycentric coordinates 

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

766 "ij,ijk->ik", 

767 qres["barycentric_coordinates"], 

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

769 ) 

770 return qres 

771 

772 

773def _from_points( 

774 target_points, 

775 input_points, 

776 kdtree=None, 

777 return_normals=False, 

778 neighbors_count=10, 

779 **kwargs, 

780): 

781 """ 

782 Find the the closest points and associated attributes 

783 from a set of 3D points. 

784 

785 Parameters 

786 ----------- 

787 target_points : (n, 3) float 

788 Points from which the query is performed 

789 input_points : (m, 3) float 

790 Input query points 

791 kdtree : scipy.cKDTree 

792 KDTree used for query. Computed if not provided 

793 return_normals : bool 

794 If True, compute the normals at each nearest point 

795 neighbors_count : int 

796 The number of closest neighbors to query 

797 kwargs : dict 

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

799 

800 Returns 

801 ---------- 

802 qres : dict 

803 Dictionary containing : 

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

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

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

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

808 """ 

809 # Empty result 

810 target_points = np.asanyarray(target_points) 

811 input_points = np.asanyarray(input_points) 

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

813 qres = {} 

814 if kdtree is None: 

815 kdtree = cKDTree(target_points) 

816 

817 if return_normals: 

818 assert neighbors_count >= 3 

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

820 nearest = target_points[indices, :] 

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

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

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

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

825 else: 

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

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

828 

829 return qres 

830 

831 

832def nricp_sumner( 

833 source_mesh, 

834 target_geometry, 

835 source_landmarks=None, 

836 target_positions=None, 

837 steps=None, 

838 distance_threshold=0.1, 

839 return_records=False, 

840 use_faces=True, 

841 use_vertex_normals=True, 

842 neighbors_count=8, 

843 face_pairs_type="vertex", 

844): 

845 """ 

846 Non Rigid Iterative Closest Points 

847 

848 Implementation of the correspondence computation part of 

849 "Sumner and Popovic 2004: Deformation Transfer for Triangle Meshes" 

850 Allows to register non-rigidly a mesh on another geometry. 

851 

852 Comparison between nricp_amberg and nricp_sumner: 

853 * nricp_amberg fits to the target mesh in less steps 

854 * nricp_amberg can generate sharp edges 

855 * only vertices and their neighbors are considered 

856 * nricp_sumner tend to preserve more the original shape 

857 * nricp_sumner parameters are easier to tune 

858 * nricp_sumner solves for triangle positions whereas