|
import numpy as np |
|
from numpy import sin, cos |
|
from math import pi as π |
|
from my3d import camera_pose |
|
from my.config import BaseConf |
|
|
|
|
|
def get_K(H, W, FoV_x): |
|
FoV_x = FoV_x / 180 * π |
|
f = 1 / np.tan(FoV_x / 2) * (W / 2) |
|
|
|
K = np.array([ |
|
[f, 0, -(W/2 - 0.5)], |
|
[0, -f, -(H/2 - 0.5)], |
|
[0, 0, -1] |
|
]) |
|
|
|
|
|
return K |
|
|
|
|
|
SIDEVIEW_PROMPTS = [ |
|
"front view of", "side view of", "backside view of", "side view of" |
|
] |
|
|
|
TOPVIEW_PROMPT = "overhead view of" |
|
|
|
|
|
def train_eye_with_prompts(r, n): |
|
hs = np.random.rand(n) * 360 |
|
vs = np.random.rand(n) * np.deg2rad(100) |
|
vs = np.clip(vs, 1e-2, π-1e-2) |
|
|
|
prompts = [] |
|
v_thresh = np.deg2rad(30) |
|
for i in range(n): |
|
_p = "" |
|
if vs[i] < v_thresh: |
|
_p = TOPVIEW_PROMPT |
|
else: |
|
_a = hs[i] |
|
_a = (_a + 45) % 360 |
|
_quad = int(_a // 90) |
|
_p = SIDEVIEW_PROMPTS[_quad] |
|
prompts.append(_p) |
|
|
|
θ = np.deg2rad(hs) |
|
|
|
φ = np.arccos(1 - 2 * (vs / π)) |
|
|
|
horz = hs |
|
elev = np.rad2deg(π / 2 - φ) |
|
|
|
eyes = np.zeros((n, 3)) |
|
|
|
eyes[:, 0] = r * sin(φ) * cos(π-θ) |
|
eyes[:, 2] = r * sin(φ) * sin(π-θ) |
|
eyes[:, 1] = r * cos(φ) |
|
|
|
return eyes, prompts, horz, elev |
|
|
|
|
|
def spiral_poses( |
|
radius, height, |
|
num_steps=20, num_rounds=1, |
|
center=np.array([0, 0, 0]), up=np.array([0, 1, 0]), |
|
): |
|
eyes = [] |
|
for i in range(num_steps): |
|
ratio = (i + 1) / num_steps |
|
Δz = height * (1 - ratio) |
|
|
|
θ = ratio * (360 * num_rounds) |
|
θ = θ / 180 * π |
|
|
|
_r = max(radius * sin(ratio * π / 2), 0.5) |
|
Δx, Δy = _r * np.array([np.cos(θ), np.sin(θ)]) |
|
eyes.append(center + [Δx, Δz, Δy]) |
|
|
|
poses = [ |
|
camera_pose(e, center - e, up) for e in eyes |
|
] |
|
return poses |
|
|
|
def circular_poses( |
|
radius, height, |
|
num_steps=36, num_rounds=1, |
|
center=np.array([0, 0, 0]), up=np.array([0, 1, 0]), |
|
): |
|
eyes = [] |
|
horz = [] |
|
for i in range(num_steps): |
|
ratio = (i + 1) / num_steps |
|
Δz = radius * (np.cos(np.deg2rad(60))) |
|
|
|
θ = ratio * (360 * num_rounds) |
|
θ = θ / 180 * π |
|
|
|
_r = max(radius * sin(ratio * π / 2), 0.5) |
|
Δx, Δy = radius * np.array([np.sin(np.deg2rad(60))*np.cos(θ), np.sin(np.deg2rad(60))*np.sin(θ)]) |
|
eyes.append(center + [Δx, Δz, Δy]) |
|
horz.append(np.rad2deg(θ)) |
|
poses = [ |
|
camera_pose(e, center - e, up) for e in eyes |
|
] |
|
horz=np.array(horz) |
|
return poses |
|
|
|
class PoseConfig(BaseConf): |
|
rend_hw: int = 64 |
|
FoV: float = 60.0 |
|
R: float = 1.5 |
|
|
|
def make(self): |
|
cfgs = self.dict() |
|
hw = cfgs.pop("rend_hw") |
|
cfgs["H"] = hw |
|
cfgs["W"] = hw |
|
return Poser(**cfgs) |
|
|
|
|
|
class Poser(): |
|
def __init__(self, H, W, FoV, R): |
|
self.H, self.W = H, W |
|
self.R = R |
|
self.K = get_K(H, W, FoV) |
|
self.FoV = FoV |
|
|
|
def sample_train(self, n, device): |
|
eyes, prompts, horz, elev = train_eye_with_prompts(r=self.R, n=n) |
|
up = np.array([0, 1, 0]) |
|
|
|
|
|
poses = [ |
|
camera_pose(e, -e, up) for e in eyes |
|
] |
|
poses = np.stack(poses, 0) |
|
|
|
|
|
|
|
FoV = np.random.rand(n) * 30 + 40 |
|
|
|
|
|
random_Ks = [ |
|
get_K(self.H, self.W, FoV[i]) |
|
for i in range(len(poses)) |
|
|
|
] |
|
|
|
|
|
angles_list = [] |
|
|
|
for horizontal, elevation, fov in zip(horz, elev, FoV): |
|
angles_list.append([horizontal, elevation, fov]) |
|
|
|
return random_Ks, poses, prompts, angles_list |
|
|
|
def sample_test(self, n): |
|
poses = circular_poses(self.R, self.R, n, num_rounds=3) |
|
poses = np.stack(poses, axis=0) |
|
return self.K, poses |
