learning3d 0.0.1__py3-none-any.whl

learning3d/ops/quaternion.py

@@ -0,0 +1,218 @@
+# Copyright (c) 2018-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+import torch
+import numpy as np
+
+# PyTorch-backed implementations
+
+
+def qmul(q, r):
+    """
+    Multiply quaternion(s) q with quaternion(s) r.
+    Expects two equally-sized tensors of shape (*, 4), where * denotes any number of dimensions.
+    Returns q*r as a tensor of shape (*, 4).
+    """
+    assert q.shape[-1] == 4
+    assert r.shape[-1] == 4
+
+    original_shape = q.shape
+
+    # Compute outer product
+    terms = torch.bmm(r.view(-1, 4, 1), q.view(-1, 1, 4))
+
+    w = terms[:, 0, 0] - terms[:, 1, 1] - terms[:, 2, 2] - terms[:, 3, 3]
+    x = terms[:, 0, 1] + terms[:, 1, 0] - terms[:, 2, 3] + terms[:, 3, 2]
+    y = terms[:, 0, 2] + terms[:, 1, 3] + terms[:, 2, 0] - terms[:, 3, 1]
+    z = terms[:, 0, 3] - terms[:, 1, 2] + terms[:, 2, 1] + terms[:, 3, 0]
+    return torch.stack((w, x, y, z), dim=1).view(original_shape)
+
+
+def qrot(q, v):
+    """
+    Rotate vector(s) v about the rotation described by quaternion(s) q.
+    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
+    where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 3).
+    """
+    assert q.shape[-1] == 4
+    assert v.shape[-1] == 3
+    assert q.shape[:-1] == v.shape[:-1]
+
+    original_shape = list(v.shape)
+    q = q.view(-1, 4)
+    v = v.view(-1, 3)
+
+    qvec = q[:, 1:]
+    uv = torch.cross(qvec, v, dim=1)
+    uuv = torch.cross(qvec, uv, dim=1)
+    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)
+
+
+def qeuler(q, order, epsilon=0):
+    """
+    Convert quaternion(s) q to Euler angles.
+    Expects a tensor of shape (*, 4), where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 3).
+    """
+    assert q.shape[-1] == 4
+
+    original_shape = list(q.shape)
+    original_shape[-1] = 3
+    q = q.view(-1, 4)
+
+    q0 = q[:, 0]
+    q1 = q[:, 1]
+    q2 = q[:, 2]
+    q3 = q[:, 3]
+
+    if order == "xyz":
+        x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        y = torch.asin(torch.clamp(2 * (q1 * q3 + q0 * q2), -1 + epsilon, 1 - epsilon))
+        z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
+    elif order == "yzx":
+        x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+        y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3))
+        z = torch.asin(torch.clamp(2 * (q1 * q2 + q0 * q3), -1 + epsilon, 1 - epsilon))
+    elif order == "zxy":
+        x = torch.asin(torch.clamp(2 * (q0 * q1 + q2 * q3), -1 + epsilon, 1 - epsilon))
+        y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q1 * q1 + q3 * q3))
+    elif order == "xzy":
+        x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+        y = torch.atan2(2 * (q0 * q2 + q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3))
+        z = torch.asin(torch.clamp(2 * (q0 * q3 - q1 * q2), -1 + epsilon, 1 - epsilon))
+    elif order == "yxz":
+        x = torch.asin(torch.clamp(2 * (q0 * q1 - q2 * q3), -1 + epsilon, 1 - epsilon))
+        y = torch.atan2(2 * (q1 * q3 + q0 * q2), 1 - 2 * (q1 * q1 + q2 * q2))
+        z = torch.atan2(2 * (q1 * q2 + q0 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+    elif order == "zyx":
+        x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        y = torch.asin(torch.clamp(2 * (q0 * q2 - q1 * q3), -1 + epsilon, 1 - epsilon))
+        z = torch.atan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
+    else:
+        raise
+
+    return torch.stack((x, y, z), dim=1).view(original_shape)
+
+
+# Numpy-backed implementations
+
+
+def qmul_np(q, r):
+    q = torch.from_numpy(q).contiguous()
+    r = torch.from_numpy(r).contiguous()
+    return qmul(q, r).numpy()
+
+
+def qrot_np(q, v):
+    q = torch.from_numpy(q).contiguous()
+    v = torch.from_numpy(v).contiguous()
+    return qrot(q, v).numpy()
+
+
+def qeuler_np(q, order, epsilon=0, use_gpu=False):
+    if use_gpu:
+        q = torch.from_numpy(q).cuda()
+        return qeuler(q, order, epsilon).cpu().numpy()
+    else:
+        q = torch.from_numpy(q).contiguous()
+        return qeuler(q, order, epsilon).numpy()
+
+
+def qfix(q):
+    """
+    Enforce quaternion continuity across the time dimension by selecting
+    the representation (q or -q) with minimal distance (or, equivalently, maximal dot product)
+    between two consecutive frames.
+
+    Expects a tensor of shape (L, J, 4), where L is the sequence length and J is the number of joints.
+    Returns a tensor of the same shape.
+    """
+    assert len(q.shape) == 3
+    assert q.shape[-1] == 4
+
+    result = q.copy()
+    dot_products = np.sum(q[1:] * q[:-1], axis=2)
+    mask = dot_products < 0
+    mask = (np.cumsum(mask, axis=0) % 2).astype(bool)
+    result[1:][mask] *= -1
+    return result
+
+
+def expmap_to_quaternion(e):
+    """
+    Convert axis-angle rotations (aka exponential maps) to quaternions.
+    Stable formula from "Practical Parameterization of Rotations Using the Exponential Map".
+    Expects a tensor of shape (*, 3), where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 4).
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+    e = e.reshape(-1, 3)
+
+    theta = np.linalg.norm(e, axis=1).reshape(-1, 1)
+    w = np.cos(0.5 * theta).reshape(-1, 1)
+    xyz = 0.5 * np.sinc(0.5 * theta / np.pi) * e
+    return np.concatenate((w, xyz), axis=1).reshape(original_shape)
+
+
+def euler_to_quaternion(e, order):
+    """
+    Convert Euler angles to quaternions.
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+
+    e = e.reshape(-1, 3)
+
+    x = e[:, 0]
+    y = e[:, 1]
+    z = e[:, 2]
+
+    rx = np.stack(
+        (np.cos(x / 2), np.sin(x / 2), np.zeros_like(x), np.zeros_like(x)), axis=1
+    )
+    ry = np.stack(
+        (np.cos(y / 2), np.zeros_like(y), np.sin(y / 2), np.zeros_like(y)), axis=1
+    )
+    rz = np.stack(
+        (np.cos(z / 2), np.zeros_like(z), np.zeros_like(z), np.sin(z / 2)), axis=1
+    )
+
+    result = None
+    for coord in order:
+        if coord == "x":
+            r = rx
+        elif coord == "y":
+            r = ry
+        elif coord == "z":
+            r = rz
+        else:
+            raise
+        if result is None:
+            result = r
+        else:
+            result = qmul_np(result, r)
+
+    # Reverse antipodal representation to have a non-negative "w"
+    if order in ["xyz", "yzx", "zxy"]:
+        result *= -1
+
+    return result.reshape(original_shape)
+
+
+def qinv(q):
+    # expectes q in (w,x,y,z) format
+    w = q[:, 0:1]
+    v = q[:, 1:]
+    inv = torch.cat([w, -v], dim=1)
+    return inv

learning3d/ops/se3.py

@@ -0,0 +1,157 @@
+""" 3-d rigid body transfomation group and corresponding Lie algebra. """
+import torch
+from .sinc import sinc1, sinc2, sinc3
+from . import so3
+
+def twist_prod(x, y):
+    x_ = x.view(-1, 6)
+    y_ = y.view(-1, 6)
+
+    xw, xv = x_[:, 0:3], x_[:, 3:6]
+    yw, yv = y_[:, 0:3], y_[:, 3:6]
+
+    zw = so3.cross_prod(xw, yw)
+    zv = so3.cross_prod(xw, yv) + so3.cross_prod(xv, yw)
+
+    z = torch.cat((zw, zv), dim=1)
+
+    return z.view_as(x)
+
+def liebracket(x, y):
+    return twist_prod(x, y)
+
+
+def mat(x):
+    # size: [*, 6] -> [*, 4, 4]
+    x_ = x.view(-1, 6)
+    w1, w2, w3 = x_[:, 0], x_[:, 1], x_[:, 2]
+    v1, v2, v3 = x_[:, 3], x_[:, 4], x_[:, 5]
+    O = torch.zeros_like(w1)
+
+    X = torch.stack((
+        torch.stack((  O, -w3,  w2, v1), dim=1),
+        torch.stack(( w3,   O, -w1, v2), dim=1),
+        torch.stack((-w2,  w1,   O, v3), dim=1),
+        torch.stack((  O,   O,   O,  O), dim=1)), dim=1)
+    return X.view(*(x.size()[0:-1]), 4, 4)
+
+def vec(X):
+    X_ = X.view(-1, 4, 4)
+    w1, w2, w3 = X_[:, 2, 1], X_[:, 0, 2], X_[:, 1, 0]
+    v1, v2, v3 = X_[:, 0, 3], X_[:, 1, 3], X_[:, 2, 3]
+    x = torch.stack((w1, w2, w3, v1, v2, v3), dim=1)
+    return x.view(*X.size()[0:-2], 6)
+
+def genvec():
+    return torch.eye(6)
+
+def genmat():
+    return mat(genvec())
+
+def exp(x):
+    x_ = x.view(-1, 6)
+    w, v = x_[:, 0:3], x_[:, 3:6]
+    t = w.norm(p=2, dim=1).view(-1, 1, 1)
+    W = so3.mat(w)
+    S = W.bmm(W)
+    I = torch.eye(3).to(w)
+
+    # Rodrigues' rotation formula.
+    #R = cos(t)*eye(3) + sinc1(t)*W + sinc2(t)*(w*w');
+    #  = eye(3) + sinc1(t)*W + sinc2(t)*S
+    R = I + sinc1(t)*W + sinc2(t)*S
+
+    #V = sinc1(t)*eye(3) + sinc2(t)*W + sinc3(t)*(w*w')
+    #  = eye(3) + sinc2(t)*W + sinc3(t)*S
+    V = I + sinc2(t)*W + sinc3(t)*S
+
+    p = V.bmm(v.contiguous().view(-1, 3, 1))
+
+    z = torch.Tensor([0, 0, 0, 1]).view(1, 1, 4).repeat(x_.size(0), 1, 1).to(x)
+    Rp = torch.cat((R, p), dim=2)
+    g = torch.cat((Rp, z), dim=1)
+
+    return g.view(*(x.size()[0:-1]), 4, 4)
+
+def inverse(g):
+    g_ = g.view(-1, 4, 4)
+    R = g_[:, 0:3, 0:3]
+    p = g_[:, 0:3, 3]
+    Q = R.transpose(1, 2)
+    q = -Q.matmul(p.unsqueeze(-1))
+
+    z = torch.Tensor([0, 0, 0, 1]).view(1, 1, 4).repeat(g_.size(0), 1, 1).to(g)
+    Qq = torch.cat((Q, q), dim=2)
+    ig = torch.cat((Qq, z), dim=1)
+
+    return ig.view(*(g.size()[0:-2]), 4, 4)
+
+
+def log(g):
+    g_ = g.view(-1, 4, 4)
+    R = g_[:, 0:3, 0:3]
+    p = g_[:, 0:3, 3]
+
+    w = so3.log(R)
+    H = so3.inv_vecs_Xg_ig(w)
+    v = H.bmm(p.contiguous().view(-1, 3, 1)).view(-1, 3)
+
+    x = torch.cat((w, v), dim=1)
+    return x.view(*(g.size()[0:-2]), 6)
+
+def transform(g, a):
+    # g : SE(3),  * x 4 x 4
+    # a : R^3,    * x 3[x N]
+    g_ = g.view(-1, 4, 4)
+    R = g_[:, 0:3, 0:3].contiguous().view(*(g.size()[0:-2]), 3, 3)
+    p = g_[:, 0:3, 3].contiguous().view(*(g.size()[0:-2]), 3)
+    if len(g.size()) == len(a.size()):
+        b = R.matmul(a) + p.unsqueeze(-1)
+    else:
+        b = R.matmul(a.unsqueeze(-1)).squeeze(-1) + p
+    return b
+
+def group_prod(g, h):
+    # g, h : SE(3)
+    g1 = g.matmul(h)
+    return g1
+
+
+class ExpMap(torch.autograd.Function):
+    """ Exp: se(3) -> SE(3)
+    """
+    @staticmethod
+    def forward(ctx, x):
+        """ Exp: R^6 -> M(4),
+            size: [B, 6] -> [B, 4, 4],
+              or  [B, 1, 6] -> [B, 1, 4, 4]
+        """
+        ctx.save_for_backward(x)
+        g = exp(x)
+        return g
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, = ctx.saved_tensors
+        g = exp(x)
+        gen_k = genmat().to(x)
+
+        # Let z = f(g) = f(exp(x))
+        # dz = df/dgij * dgij/dxk * dxk
+        #    = df/dgij * (d/dxk)[exp(x)]_ij * dxk
+        #    = df/dgij * [gen_k*g]_ij * dxk
+
+        dg = gen_k.matmul(g.view(-1, 1, 4, 4))
+        # (k, i, j)
+        dg = dg.to(grad_output)
+
+        go = grad_output.contiguous().view(-1, 1, 4, 4)
+        dd = go * dg
+        grad_input = dd.sum(-1).sum(-1)
+
+        return grad_input
+
+Exp = ExpMap.apply
+
+
+#EOF

learning3d/ops/sinc.py

@@ -0,0 +1,229 @@
+""" sinc(t) := sin(t) / t """
+import torch
+from torch import sin, cos
+
+def sinc1(t):
+    """ sinc1: t -> sin(t)/t """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t[s] ** 2
+    r[s] = 1 - t2/6*(1 - t2/20*(1 - t2/42))  # Taylor series O(t^8)
+    r[c] = sin(t[c]) / t[c]
+
+    return r
+
+def sinc1_dt(t):
+    """ d/dt(sinc1) """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t ** 2
+    r[s] = -t[s]/3*(1 - t2[s]/10*(1 - t2[s]/28*(1 - t2[s]/54)))  # Taylor series O(t^8)
+    r[c] = cos(t[c])/t[c] - sin(t[c])/t2[c]
+
+    return r
+
+def sinc1_dt_rt(t):
+    """ d/dt(sinc1) / t """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t ** 2
+    r[s] = -1/3*(1 - t2[s]/10*(1 - t2[s]/28*(1 - t2[s]/54)))  # Taylor series O(t^8)
+    r[c] = (cos(t[c]) / t[c] - sin(t[c]) / t2[c]) / t[c]
+
+    return r
+
+
+def rsinc1(t):
+    """ rsinc1: t -> t/sinc1(t) """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t[s] ** 2
+    r[s] = (((31*t2)/42 + 7)*t2/60 + 1)*t2/6 + 1  # Taylor series O(t^8)
+    r[c] = t[c] / sin(t[c])
+
+    return r
+
+def rsinc1_dt(t):
+    """ d/dt(rsinc1) """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t[s] ** 2
+    r[s] = ((((127*t2)/30 + 31)*t2/28 + 7)*t2/30 + 1)*t[s]/3  # Taylor series O(t^8)
+    r[c] = 1/sin(t[c]) - (t[c]*cos(t[c]))/(sin(t[c])*sin(t[c]))
+
+    return r
+
+def rsinc1_dt_csc(t):
+    """ d/dt(rsinc1) / sin(t) """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t[s] ** 2
+    r[s] = t2*(t2*((4*t2)/675 + 2/63) + 2/15) + 1/3  # Taylor series O(t^8)
+    r[c] = (1/sin(t[c]) - (t[c]*cos(t[c]))/(sin(t[c])*sin(t[c]))) / sin(t[c])
+
+    return r
+
+
+def sinc2(t):
+    """ sinc2: t -> (1 - cos(t)) / (t**2) """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t ** 2
+    r[s] = 1/2*(1-t2[s]/12*(1-t2[s]/30*(1-t2[s]/56)))  # Taylor series O(t^8)
+    r[c] = (1-cos(t[c]))/t2[c]
+
+    return r
+
+def sinc2_dt(t):
+    """ d/dt(sinc2) """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t ** 2
+    r[s] = -t[s]/12*(1 - t2[s]/5*(1.0/3 - t2[s]/56*(1.0/2 - t2[s]/135)))  # Taylor series O(t^8)
+    r[c] = sin(t[c])/t2[c] - 2*(1-cos(t[c]))/(t2[c]*t[c])
+
+    return r
+
+
+def sinc3(t):
+    """ sinc3: t -> (t - sin(t)) / (t**3) """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t[s] ** 2
+    r[s] = 1/6*(1-t2/20*(1-t2/42*(1-t2/72)))  # Taylor series O(t^8)
+    r[c] = (t[c]-sin(t[c]))/(t[c]**3)
+
+    return r
+
+def sinc3_dt(t):
+    """ d/dt(sinc3) """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t[s] ** 2
+    r[s] = -t[s]/60*(1 - t2/21*(1 - t2/24*(1.0/2 - t2/165)))  # Taylor series O(t^8)
+    r[c] = (3*sin(t[c]) - t[c]*(cos(t[c]) + 2))/(t[c]**4)
+
+    return r
+
+
+def sinc4(t):
+    """ sinc4: t -> 1/t^2 * (1/2 - sinc2(t))
+                  = 1/t^2 * (1/2 - (1 - cos(t))/t^2)
+    """
+    e = 0.01
+    r = torch.zeros_like(t)
+    a = torch.abs(t)
+
+    s = a < e
+    c = (s == 0)
+    t2 = t ** 2
+    r[s] = 1/24*(1-t2/30*(1-t2/56*(1-t2/90)))  # Taylor series O(t^8)
+    r[c] = (0.5 - (1 - cos(t))/t2) / t2
+
+
+class Sinc1_autograd(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, theta):
+        ctx.save_for_backward(theta)
+        return sinc1(theta)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        theta, = ctx.saved_tensors
+        grad_theta = None
+        if ctx.needs_input_grad[0]:
+            grad_theta = grad_output * sinc1_dt(theta).to(grad_output)
+        return grad_theta
+
+Sinc1 = Sinc1_autograd.apply
+
+class RSinc1_autograd(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, theta):
+        ctx.save_for_backward(theta)
+        return rsinc1(theta)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        theta, = ctx.saved_tensors
+        grad_theta = None
+        if ctx.needs_input_grad[0]:
+            grad_theta = grad_output * rsinc1_dt(theta).to(grad_output)
+        return grad_theta
+
+RSinc1 = RSinc1_autograd.apply
+
+class Sinc2_autograd(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, theta):
+        ctx.save_for_backward(theta)
+        return sinc2(theta)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        theta, = ctx.saved_tensors
+        grad_theta = None
+        if ctx.needs_input_grad[0]:
+            grad_theta = grad_output * sinc2_dt(theta).to(grad_output)
+        return grad_theta
+
+Sinc2 = Sinc2_autograd.apply
+
+class Sinc3_autograd(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, theta):
+        ctx.save_for_backward(theta)
+        return sinc3(theta)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        theta, = ctx.saved_tensors
+        grad_theta = None
+        if ctx.needs_input_grad[0]:
+            grad_theta = grad_output * sinc3_dt(theta).to(grad_output)
+        return grad_theta
+
+Sinc3 = Sinc3_autograd.apply
+
+
+#EOF