File size: 2,780 Bytes
d6682b6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import numpy as np
import torch
from typing import Union, List
def lerp(
t: float, v0: Union[np.ndarray, torch.Tensor], v1: Union[np.ndarray, torch.Tensor]
) -> Union[np.ndarray, torch.Tensor]:
return (1 - t) * v0 + t * v1
def maybe_torch(v: np.ndarray, is_torch: bool):
if is_torch:
return torch.from_numpy(v)
return v
def normalize(v: np.ndarray, eps: float):
norm_v = np.linalg.norm(v)
if norm_v > eps:
v = v / norm_v
return v
class slerp:
def __init__(self):
pass
def execute(
self,
t: Union[float, List[float]],
v0: Union[List[torch.Tensor], torch.Tensor],
v1: Union[List[torch.Tensor], torch.Tensor],
DOT_THRESHOLD: float = 0.9995,
eps: float = 1e-8,
densities = None,
):
if type(v0) is list:
v0 = v0[0]
if type(v1) is list:
v1 = v1[0]
if type(t) is list:
t = t[0]
"""
Spherical linear interpolation
From: https://gist.github.com/dvschultz/3af50c40df002da3b751efab1daddf2c
Args:
t (float/np.ndarray): Float value between 0.0 and 1.0
v0 (np.ndarray): Starting vector
v1 (np.ndarray): Final vector
DOT_THRESHOLD (float): Threshold for considering the two vectors as
colinear. Not recommended to alter this.
Returns:
v2 (np.ndarray): Interpolation vector between v0 and v1
"""
is_torch = False
if not isinstance(v0, np.ndarray):
is_torch = True
v0 = v0.detach().cpu().float().numpy()
if not isinstance(v1, np.ndarray):
is_torch = True
v1 = v1.detach().cpu().float().numpy()
# Copy the vectors to reuse them later
v0_copy = np.copy(v0)
v1_copy = np.copy(v1)
# Normalize the vectors to get the directions and angles
v0 = normalize(v0, eps)
v1 = normalize(v1, eps)
# Dot product with the normalized vectors (can't use np.dot in W)
dot = np.sum(v0 * v1)
# If absolute value of dot product is almost 1, vectors are ~colinear, so use lerp
if np.abs(dot) > DOT_THRESHOLD:
res = lerp(t, v0_copy, v1_copy)
return maybe_torch(res, is_torch)
# Calculate initial angle between v0 and v1
theta_0 = np.arccos(dot)
sin_theta_0 = np.sin(theta_0)
# Angle at timestep t
theta_t = theta_0 * t
sin_theta_t = np.sin(theta_t)
# Finish the slerp algorithm
s0 = np.sin(theta_0 - theta_t) / sin_theta_0
s1 = sin_theta_t / sin_theta_0
res = s0 * v0_copy + s1 * v1_copy
return maybe_torch(res, is_torch) |