edit_anything / pose.py
memef4rmer's picture
Duplicate from jyseo/3DFuse
efe5745
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 * π # to rad
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)
# φ = v
φ = np.arccos(1 - 2 * (vs / π))
horz = hs
elev = np.rad2deg(π / 2 - φ)
eyes = np.zeros((n, 3))
eyes[:, 0] = r * sin(φ) * cos(π-θ) # x
eyes[:, 2] = r * sin(φ) * sin(π-θ) # z
eyes[:, 1] = r * cos(φ) # y
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 * ratio, 0.5)
_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 * ratio, 0.5)
_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