import torch
import torch.nn.functional as F
class ONB(object):
def __init__(self, batch_size : int):
"""
uvw - xyz - tbn
"""
self.batch_size = batch_size
self.axis = torch.zeros(batch_size, 3, 3) # (batch, 3, 3)
self.axis[:, 0, 0] = 1
self.axis[:, 1, 1] = 1
self.axis[:, 2, 2] = 1
def u(self) -> torch.Tensor:
"""
Returns:
u axis of shape (batch, 3)
"""
return self.axis[:, 0, :]
def v(self) -> torch.Tensor:
return self.axis[:, 1, :]
def w(self) -> torch.Tensor:
return self.axis[:, 2, :]
def inverse_transform(self, p : torch.Tensor) -> torch.Tensor:
"""
Convert local coordinate(in onb) back to global
coordinate(onb in).
Args:
p: local coordinates of shape (batch, 3) or (batch, N, 3)
Returns:
global coordinates of shape (batch, 3) or (batch, N, 3)
"""
assert(self.batch_size == p.size(0))
assert(len(p.size()) in [2, 3])
if len(p.size()) == 2:
return p[:, 0:1] * self.u() + p[:, 1:2] * self.v() + p[:, 2:3] * self.w()
elif len(p.size()) == 3:
u = self.u().unsqueeze(1)
v = self.v().unsqueeze(1)
w = self.w().unsqueeze(1)
return p[:, :, [0]] * u + p[:, :, [1]] * v + p[:, :, [2]] * w
def transform(self, p : torch.Tensor) -> torch.Tensor:
"""
Convert global coordinate(onb in) to local
coordinate(in onb).
Args:
p: global coordinates of shape (batch, 3) or (batch, lightnum, 3)
Returns:
local coordinates of shape (batch, 3) or (batch, lightnum, 3)
"""
assert(self.batch_size == p.size(0))
assert(len(p.size()) in [2, 3])
if len(p.size()) == 2:
x = torch.sum(p * self.u(), dim=1, keepdim=True)
y = torch.sum(p * self.v(), dim=1, keepdim=True)
z = torch.sum(p * self.w(), dim=1, keepdim=True)
return torch.cat([x, y, z], dim=1)
elif len(p.size()) == 3:
lightnum = p.size(1)
u = self.u().unsqueeze(1).repeat(1, lightnum, 1)
v = self.v().unsqueeze(1).repeat(1, lightnum, 1)
w = self.w().unsqueeze(1).repeat(1, lightnum, 1)
x = torch.sum(p * u, dim=2, keepdim=True)
y = torch.sum(p * v, dim=2, keepdim=True)
z = torch.sum(p * w, dim=2, keepdim=True)
return torch.cat([x, y, z], dim=2)
def build_from_ntb(
self,
n : torch.Tensor,
t : torch.Tensor,
b : torch.Tensor,
) -> None:
"""
Args:
n, t, b: The local frame, of shape (batch, 3)
"""
batch_size = n.size(0)
self.axis = torch.zeros((batch_size, 3, 3)).to(n.device)
self.axis[:, 2, :] = F.normalize(n, dim=1)
self.axis[:, 1, :] = F.normalize(b, dim=1)
self.axis[:, 0, :] = F.normalize(t, dim=1)
def build_from_w(self, normal : torch.Tensor) -> None:
"""
Build the local frame based on the normal.
Args:
normal: The normal coordinates of shape (batch, 3)
"""
assert(self.batch_size == normal.size(0))
device = normal.device
n = F.normalize(normal, dim=1)
nz = n[:, [2]]
batch_size = n.shape[0]
constant_001 = torch.zeros_like(normal).to(device)
constant_001[:, 2] = 1.0
constant_100 = torch.zeros_like(normal).to(device)
constant_100[:, 0] = 1.0
nz_notequal_1 = torch.gt(torch.abs(nz - 1.0), 1e-6)
nz_notequal_m1 = torch.gt(torch.abs(nz + 1.0), 1e-6)
t = torch.where(nz_notequal_1 & nz_notequal_m1, constant_001, constant_100)
# Optix version
# b = F.normalize(torch.cross(normal, t), dim=1)
# t = torch.cross(b, normal)
# Original pytorch version
t = F.normalize(torch.cross(t, normal), dim=1)
b = torch.cross(n, t)
self.axis = torch.zeros((batch_size, 3, 3)).to(device)
self.axis[:, 2, :] = n
self.axis[:, 1, :] = b
self.axis[:, 0, :] = t
def rotate_frame(self, theta : torch.Tensor) -> None:
"""
Rotate local frame along the normal axis
Args:
theta: the degrees of counterclockwise rotation, of shape (batch, 1)
"""
assert(self.batch_size == theta.size(0))
n = self.w()
t = self.u()
b = self.v()
t = F.normalize(t * torch.cos(theta) + b * torch.sin(theta), dim=1)
b = F.normalize(torch.cross(n, t), dim=1)
self.axis = torch.zeros((self.batch_size, 3, 3)).to(theta.device)
self.axis[:, 0, :] = t
self.axis[:, 1, :] = b
self.axis[:, 2, :] = n
def _back_hemi_octa_map(self, n_2d : torch.Tensor) -> torch.Tensor:
"""
The original normal is (0, 0, 1), we should use this method to
perturb the original normal to get a new normal and then build
a new local frame based on the new normal.
Args:
n_2d: shape (batch, 2)
Returns:
local_n: shape (batch, 3), which is define in geometry local
frame.
"""
p = (n_2d - 0.5) * 2.0
resultx = (p[:, [0]] + p[:, [1]]) * 0.5
resulty = (p[:, [1]] - p[:, [0]]) * 0.5
resultz = 1.0 - torch.abs(resultx) - torch.abs(resulty)
result = torch.cat([resultx, resulty, resultz], dim=1)
return F.normalize(result, dim=1)
def hemi_octa_map(self, dir : torch.Tensor) -> torch.Tensor:
"""
Args:
dir: shape (batch, 3)
Returns:
n2d: shape (batch, 2), which is define in circle coordinate
"""
high_dim = False
batch_size = dir.shape[0]
if len(dir.shape) > 2:
high_dim = True
dir = dir.reshape(-1, 3)
p = dir/torch.sum(torch.abs(dir), dim=1, keepdim=True) # (batch,3)
n_2d = torch.cat([p[:,[0]] - p[:,[1]],p[:,[0]] + p[:, [1]]],dim=1) * 0.5 + 0.5
if high_dim:
n_2d = n_2d.reshape(batch_size, -1, 2)
return n_2d
def build_from_n2d(self, n_2d : torch.Tensor, theta : torch.Tensor) -> None:
"""
Args:
n_2d: tensor of shape (batch, 2). the param defines how
to perturb local normal.
theta: the degrees of rotation of tangent.
"""
assert(self.batch_size == n_2d.size(0))
local_n = self._back_hemi_octa_map(n_2d)
self.build_from_w(local_n)
self.rotate_frame(theta)
