Compute the angle spanned by the unit vectors in arr and ref. Args: arr: (*b_shape, *d_shape, 3) ref: (*b_shape, *d_shape, 3) ndim_b: number of dimension of b_shape. If None, = 0. normalized: whetehr arr and ref are unit vectors
(
arr: torch.Tensor,
ref: torch.Tensor,
ndim_b: int = None,
normalized: bool = True,
valid_mask: torch.Tensor = None,
)
| 388 | |
| 389 | |
| 390 | def compute_diff_angle( |
| 391 | arr: torch.Tensor, |
| 392 | ref: torch.Tensor, |
| 393 | ndim_b: int = None, |
| 394 | normalized: bool = True, |
| 395 | valid_mask: torch.Tensor = None, |
| 396 | ): |
| 397 | """ |
| 398 | Compute the angle spanned by the unit vectors in arr and ref. |
| 399 | |
| 400 | Args: |
| 401 | arr: (*b_shape, *d_shape, 3) |
| 402 | ref: (*b_shape, *d_shape, 3) |
| 403 | ndim_b: |
| 404 | number of dimension of b_shape. If None, = 0. |
| 405 | normalized: |
| 406 | whetehr arr and ref are unit vectors |
| 407 | valid_mask: (*b_shape, *d_shape,) |
| 408 | |
| 409 | Returns: |
| 410 | err: (*b_shape,) angle in degree |
| 411 | """ |
| 412 | if ndim_b is None: |
| 413 | ndim_b = 0 |
| 414 | |
| 415 | if not normalized: |
| 416 | arr = torch.nn.functional.normalize(arr, p=2, dim=-1) |
| 417 | ref = torch.nn.functional.normalize(ref, p=2, dim=-1) |
| 418 | |
| 419 | # make sure arr and ref points to the same direction |
| 420 | out = torch.sum(arr * ref, dim=-1) # (*b, *d) |
| 421 | arr = arr * out.sign().unsqueeze(-1) |
| 422 | |
| 423 | # recompute inner product |
| 424 | out = torch.sum(arr * ref, dim=-1) # (*b, *d) |
| 425 | |
| 426 | angle = torch.arccos(out.clamp(min=-1 + 1e-9, max=1 - 1e-9)) * (180. / torch.pi) # (*b, *d) in degree |
| 427 | if valid_mask is None: |
| 428 | angle = angle.reshape(*(arr.shape[:ndim_b]), -1) # (*b, numel_d) |
| 429 | angle = angle.mean(dim=-1) # (*b,) |
| 430 | else: |
| 431 | valid_mask = valid_mask.view( |
| 432 | *(valid_mask.shape), *([1] * (angle.ndim - valid_mask.ndim))).expand_as(angle) |
| 433 | valid_mask = valid_mask.reshape(*(arr.shape[:ndim_b]), -1) # (*b, numel_d) |
| 434 | angle = angle.reshape(*(arr.shape[:ndim_b]), -1) # (*b, numel_d) |
| 435 | angle = (angle * valid_mask).sum(dim=-1) / valid_mask.sum(-1) |
| 436 | return angle |
| 437 | |
| 438 | |
| 439 | def compute_accuracy( |