(R, chunk_size)
| 55 | dists = rot_diff + lambda_t * trans_diff |
| 56 | return dists |
| 57 | def rotation_angle_batch_chunked(R, chunk_size): |
| 58 | N = R.shape[0] |
| 59 | rot_diff = np.empty((N, N), dtype=np.float32) |
| 60 | # Precompute R transpose once |
| 61 | R_t = R.transpose(0,2,1) |
| 62 | |
| 63 | for i_start in range(0, N, chunk_size): |
| 64 | i_end = min(N, i_start + chunk_size) |
| 65 | # Sub-block of R_t |
| 66 | R_i_t = R_t[i_start:i_end] # (B,3,3) |
| 67 | |
| 68 | for j_start in range(0, N, chunk_size): |
| 69 | j_end = min(N, j_start + chunk_size) |
| 70 | R_j = R[j_start:j_end] # (B,3,3) |
| 71 | # Compute R_i_t @ R_j for block |
| 72 | # R_i_t: (B,3,3) |
| 73 | # R_j: (B,3,3) but we need pairwise, so we expand dims |
| 74 | # This still can be large. If even BxB is too big, choose smaller chunks. |
| 75 | # shape (B,B,3,3) |
| 76 | R_mult = R_i_t[:, np.newaxis, :, :] @ R_j[np.newaxis, :, :, :] |
| 77 | # Compute trace |
| 78 | trace_vals = R_mult[...,0,0] + R_mult[...,1,1] + R_mult[...,2,2] |
| 79 | val = (trace_vals - 1.0) / 2.0 |
| 80 | np.clip(val, -1.0, 1.0, out=val) |
| 81 | angle_rad = np.arccos(val) |
| 82 | angle_deg = np.degrees(angle_rad) |
| 83 | block_rot_diff = angle_deg / 180.0 |
| 84 | rot_diff[i_start:i_end, j_start:j_end] = block_rot_diff.astype(np.float32) |
| 85 | return rot_diff |
| 86 | def extrinsic_distance_batch_chunked(extrinsics, lambda_t=1.0, chunk_size=1000): |
| 87 | R = extrinsics[:, :3, :3].astype(np.float32) |
| 88 | t = extrinsics[:, :3, 3].astype(np.float32) |
no outgoing calls
no test coverage detected