VistaDream / ops /trajs /__init__.py
hpwang's picture
[Init]
fd5e0f7
import numpy as np
from ops.sky import Sky_Seg_Tool
from ops.utils import dpt2xyz
from .spiral import spiral_camera_poses
class Trajectory_Generation():
def __init__(self,
scene = None,
method = 'spiral') -> None:
'''
method = 'spiral'/ rot' / 'spin'
'''
self._method = method
self.forward_ratio = scene.traj_forward_ratio
self.backward_ratio = scene.traj_backward_ratio
self.min_percentage = scene.traj_min_percentage
self.max_percentage = scene.traj_max_percentage
def _radius(self, xyz):
# get range
_min = np.percentile(xyz,self.min_percentage,axis=0)
_max = np.percentile(xyz,self.max_percentage,axis=0)
_range = _max - _min
# set radius to mean range of three axes
self.radius = np.mean(_range)
def _traj_spiral(self, nframe):
trajs = spiral_camera_poses(nframe, self.radius, self.forward_ratio, self.backward_ratio)
return trajs
def __call__(self, nframe, xyz):
if xyz.ndim > 2:
xyz = xyz.reshape(-1,3)
self._radius(xyz)
if self._method == 'rot':
trajs = self._traj_rot(nframe)
elif self._method == 'spin':
trajs = self._traj_spin(nframe)
elif self._method == 'spiral':
trajs = self._traj_spiral(nframe)
else:
raise TypeError('method = rot / spiral')
return trajs
def _generate_trajectory(cfg, scene, nframes=None):
method = scene.traj_type
nframe = cfg.scene.traj.n_sample*6 if nframes is None else nframes
sky,dpt,intrinsic = scene.frames[0].sky,scene.frames[0].dpt,scene.frames[0].intrinsic
xyz = dpt2xyz(dpt,intrinsic)
init_xyz = xyz[~sky]
generator = Trajectory_Generation(scene=scene,method=method)
traj = generator(nframe,init_xyz)
return traj