evalmde/utils/torch.py

from typing import List

import numpy as np
import torch
import torch.nn.functional as F


@torch.no_grad()
def get_angle_between(n1: torch.Tensor, n2: torch.Tensor) -> torch.Tensor:
    '''
    :param n1: shape (..., 3), norm > 0
    :param n2: shape (..., 3), norm > 0
    :return: shape (...), in radius
    '''
    return torch.acos((F.normalize(n1, dim=-1) * F.normalize(n2, dim=-1)).sum(dim=-1).clamp(-1, 1))


def reformat_as_torch_tensor(x, device=torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')):
    if isinstance(x, List):
        return torch.tensor(x, device=device)
    elif isinstance(x, np.ndarray):
        return torch.from_numpy(x).to(device=device)
    elif isinstance(x, torch.Tensor):
        return x.to(device=device)
    else:
        raise ValueError(f'Unsupported type: {type(x)}')