File size: 3,353 Bytes
90fd8f8 |
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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import roma
from kiui.op import safe_normalize
def get_rays(pose, h, w, fovy, opengl=True):
x, y = torch.meshgrid(
torch.arange(w, device=pose.device),
torch.arange(h, device=pose.device),
indexing="xy",
)
x = x.flatten()
y = y.flatten()
cx = w * 0.5
cy = h * 0.5
focal = h * 0.5 / np.tan(0.5 * np.deg2rad(fovy))
camera_dirs = F.pad(
torch.stack(
[
(x - cx + 0.5) / focal,
(y - cy + 0.5) / focal * (-1.0 if opengl else 1.0),
],
dim=-1,
),
(0, 1),
value=(-1.0 if opengl else 1.0),
) # [hw, 3]
rays_d = camera_dirs @ pose[:3, :3].transpose(0, 1) # [hw, 3]
rays_o = pose[:3, 3].unsqueeze(0).expand_as(rays_d) # [hw, 3]
rays_o = rays_o.view(h, w, 3)
rays_d = safe_normalize(rays_d).view(h, w, 3)
return rays_o, rays_d
def orbit_camera_jitter(poses, strength=0.1):
# poses: [B, 4, 4], assume orbit camera in opengl format
# random orbital rotate
B = poses.shape[0]
rotvec_x = poses[:, :3, 1] * strength * np.pi * (torch.rand(B, 1, device=poses.device) * 2 - 1)
rotvec_y = poses[:, :3, 0] * strength * np.pi / 2 * (torch.rand(B, 1, device=poses.device) * 2 - 1)
rot = roma.rotvec_to_rotmat(rotvec_x) @ roma.rotvec_to_rotmat(rotvec_y)
R = rot @ poses[:, :3, :3]
T = rot @ poses[:, :3, 3:]
new_poses = poses.clone()
new_poses[:, :3, :3] = R
new_poses[:, :3, 3:] = T
return new_poses
def grid_distortion(images, strength=0.5):
# images: [B, C, H, W]
# num_steps: int, grid resolution for distortion
# strength: float in [0, 1], strength of distortion
B, C, H, W = images.shape
num_steps = np.random.randint(8, 17)
grid_steps = torch.linspace(-1, 1, num_steps)
# have to loop batch...
grids = []
for b in range(B):
# construct displacement
x_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive
x_steps = (x_steps + strength * (torch.rand_like(x_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb
x_steps = (x_steps * W).long() # [num_steps]
x_steps[0] = 0
x_steps[-1] = W
xs = []
for i in range(num_steps - 1):
xs.append(torch.linspace(grid_steps[i], grid_steps[i + 1], x_steps[i + 1] - x_steps[i]))
xs = torch.cat(xs, dim=0) # [W]
y_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive
y_steps = (y_steps + strength * (torch.rand_like(y_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb
y_steps = (y_steps * H).long() # [num_steps]
y_steps[0] = 0
y_steps[-1] = H
ys = []
for i in range(num_steps - 1):
ys.append(torch.linspace(grid_steps[i], grid_steps[i + 1], y_steps[i + 1] - y_steps[i]))
ys = torch.cat(ys, dim=0) # [H]
# construct grid
grid_x, grid_y = torch.meshgrid(xs, ys, indexing='xy') # [H, W]
grid = torch.stack([grid_x, grid_y], dim=-1) # [H, W, 2]
grids.append(grid)
grids = torch.stack(grids, dim=0).to(images.device) # [B, H, W, 2]
# grid sample
images = F.grid_sample(images, grids, align_corners=False)
return images
|