OpenMotionLab/MotionGPT
118
1"""Trackball class for 3D manipulation of viewpoints.2"""3import numpy as np4 5import trimesh.transformations as transformations6 7 8class Trackball(object):9 """A trackball class for creating camera transforms from mouse movements.10 """11 STATE_ROTATE = 012 STATE_PAN = 113 STATE_ROLL = 214 STATE_ZOOM = 315 16 def __init__(self, pose, size, scale,17 target=np.array([0.0, 0.0, 0.0])):18 """Initialize a trackball with an initial camera-to-world pose19 and the given parameters.20 21 Parameters22 ----------23 pose : [4,4]24 An initial camera-to-world pose for the trackball.25 26 size : (float, float)27 The width and height of the camera image in pixels.28 29 scale : float30 The diagonal of the scene's bounding box --31 used for ensuring translation motions are sufficiently32 fast for differently-sized scenes.33 34 target : (3,) float35 The center of the scene in world coordinates.36 The trackball will revolve around this point.37 """38 self._size = np.array(size)39 self._scale = float(scale)40 41 self._pose = pose42 self._n_pose = pose43 44 self._target = target45 self._n_target = target46 47 self._state = Trackball.STATE_ROTATE48 49 @property50 def pose(self):51 """autolab_core.RigidTransform : The current camera-to-world pose.52 """53 return self._n_pose54 55 def set_state(self, state):56 """Set the state of the trackball in order to change the effect of57 dragging motions.58 59 Parameters60 ----------61 state : int62 One of Trackball.STATE_ROTATE, Trackball.STATE_PAN,63 Trackball.STATE_ROLL, and Trackball.STATE_ZOOM.64 """65 self._state = state66 67 def resize(self, size):68 """Resize the window.69 70 Parameters71 ----------72 size : (float, float)73 The new width and height of the camera image in pixels.74 """75 self._size = np.array(size)76 77 def down(self, point):78 """Record an initial mouse press at a given point.79 80 Parameters81 ----------82 point : (2,) int83 The x and y pixel coordinates of the mouse press.84 """85 self._pdown = np.array(point, dtype=np.float32)86 self._pose = self._n_pose87 self._target = self._n_target88 89 def drag(self, point):90 """Update the tracball during a drag.91 92 Parameters93 ----------94 point : (2,) int95 The current x and y pixel coordinates of the mouse during a drag.96 This will compute a movement for the trackball with the relative97 motion between this point and the one marked by down().98 """99 point = np.array(point, dtype=np.float32)100 dx, dy = point - self._pdown101 mindim = 0.3 * np.min(self._size)102 103 target = self._target104 x_axis = self._pose[:3,0].flatten()105 y_axis = self._pose[:3,1].flatten()106 z_axis = self._pose[:3,2].flatten()107 eye = self._pose[:3,3].flatten()108 109 # Interpret drag as a rotation110 if self._state == Trackball.STATE_ROTATE:111 x_angle = -dx / mindim112 x_rot_mat = transformations.rotation_matrix(113 x_angle, y_axis, target114 )115 116 y_angle = dy / mindim117 y_rot_mat = transformations.rotation_matrix(118 y_angle, x_axis, target119 )120 121 self._n_pose = y_rot_mat.dot(x_rot_mat.dot(self._pose))122 123 # Interpret drag as a roll about the camera axis124 elif self._state == Trackball.STATE_ROLL:125 center = self._size / 2.0126 v_init = self._pdown - center127 v_curr = point - center128 v_init = v_init / np.linalg.norm(v_init)129 v_curr = v_curr / np.linalg.norm(v_curr)130 131 theta = (-np.arctan2(v_curr[1], v_curr[0]) +132 np.arctan2(v_init[1], v_init[0]))133 134 rot_mat = transformations.rotation_matrix(theta, z_axis, target)135 136 self._n_pose = rot_mat.dot(self._pose)137 138 # Interpret drag as a camera pan in view plane139 elif self._state == Trackball.STATE_PAN:140 dx = -dx / (5.0 * mindim) * self._scale141 dy = -dy / (5.0 * mindim) * self._scale142 143 translation = dx * x_axis + dy * y_axis144 self._n_target = self._target + translation145 t_tf = np.eye(4)146 t_tf[:3,3] = translation147 self._n_pose = t_tf.dot(self._pose)148 149 # Interpret drag as a zoom motion150 elif self._state == Trackball.STATE_ZOOM:151 radius = np.linalg.norm(eye - target)152 ratio = 0.0153 if dy > 0:154 ratio = np.exp(abs(dy) / (0.5 * self._size[1])) - 1.0155 elif dy < 0:156 ratio = 1.0 - np.exp(dy / (0.5 * (self._size[1])))157 translation = -np.sign(dy) * ratio * radius * z_axis158 t_tf = np.eye(4)159 t_tf[:3,3] = translation160 self._n_pose = t_tf.dot(self._pose)161 162 def scroll(self, clicks):163 """Zoom using a mouse scroll wheel motion.164 165 Parameters166 ----------167 clicks : int168 The number of clicks. Positive numbers indicate forward wheel169 movement.170 """171 target = self._target172 ratio = 0.90173 174 mult = 1.0175 if clicks > 0:176 mult = ratio**clicks177 elif clicks < 0:178 mult = (1.0 / ratio)**abs(clicks)179 180 z_axis = self._n_pose[:3,2].flatten()181 eye = self._n_pose[:3,3].flatten()182 radius = np.linalg.norm(eye - target)183 translation = (mult * radius - radius) * z_axis184 t_tf = np.eye(4)185 t_tf[:3,3] = translation186 self._n_pose = t_tf.dot(self._n_pose)187 188 z_axis = self._pose[:3,2].flatten()189 eye = self._pose[:3,3].flatten()190 radius = np.linalg.norm(eye - target)191 translation = (mult * radius - radius) * z_axis192 t_tf = np.eye(4)193 t_tf[:3,3] = translation194 self._pose = t_tf.dot(self._pose)195 196 def rotate(self, azimuth, axis=None):197 """Rotate the trackball about the "Up" axis by azimuth radians.198 199 Parameters200 ----------201 azimuth : float202 The number of radians to rotate.203 """204 target = self._target205 206 y_axis = self._n_pose[:3,1].flatten()207 if axis is not None:208 y_axis = axis209 x_rot_mat = transformations.rotation_matrix(azimuth, y_axis, target)210 self._n_pose = x_rot_mat.dot(self._n_pose)211 212 y_axis = self._pose[:3,1].flatten()213 if axis is not None:214 y_axis = axis215 x_rot_mat = transformations.rotation_matrix(azimuth, y_axis, target)216 self._pose = x_rot_mat.dot(self._pose)217