def rotate(obj_path, out_path, angle_list, base=False, random=False):
mesh = load_mesh(obj_path)
mesh.vertices -= mesh.center_mass
# Calculate the momentum and principal axes
inertia = mesh.moment_inertia
if base == True:
p_axis = Principal_Axes(inertia)
p_R = np.concatenate((np.concatenate(p_axis, axis=0).reshape(
-1, 3), np.array([[0., 0., 0.]])),
axis=0)
R = np.concatenate((p_R, np.array([[0.], [0.], [0.], [1]])), axis=1)
else:
if random == False:
alpha, beta, gamma = np.radians(angle_list[0]), np.radians(angle_list[1]),\
np.radians(angle_list[2])
Rx = transform.rotation_matrix(alpha, [1, 0, 0])
Ry = transform.rotation_matrix(beta, [0, 1, 0])
Rz = transform.rotation_matrix(gamma, [0, 0, 1])
R = transform.concatenate_matrices(Rx, Ry, Rz)
else:
# Random Rotation Matrix
R = transform.random_rotation_matrix()
# Rotation
mesh.apply_transform(R)
mesh.export(out_path)
def test_matrix_to_rpy():
xr45 = tfs.rotation_matrix(np.pi / 2, [1, 0, 0])[:3,:3]
yr45 = tfs.rotation_matrix(np.pi / 2, [0, 1, 0])[:3,:3]
zr45 = tfs.rotation_matrix(np.pi / 2, [0, 0, 1])[:3,:3]
assert np.allclose([np.pi / 2, 0, 0], matrix_to_rpy(xr45))
assert np.allclose([0, np.pi / 2, 0], matrix_to_rpy(yr45))
assert np.allclose([0, 0, np.pi / 2], matrix_to_rpy(zr45))
def test_rpy_to_matrix():
xr45 = tfs.rotation_matrix(np.pi / 2, [1, 0, 0])[:3,:3]
yr45 = tfs.rotation_matrix(np.pi / 2, [0, 1, 0])[:3,:3]
zr45 = tfs.rotation_matrix(np.pi / 2, [0, 0, 1])[:3,:3]
c = zr45.dot(yr45.dot(xr45))
assert np.allclose(rpy_to_matrix([np.pi / 2, 0, 0]), xr45)
assert np.allclose(rpy_to_matrix([0, np.pi / 2, 0]), yr45)
assert np.allclose(rpy_to_matrix([0, 0, np.pi / 2]), zr45)
assert np.allclose(rpy_to_matrix(np.pi / 2 * np.ones(3)), c)
def test_xyz_rpy_to_matrix():
xr45 = tfs.rotation_matrix(np.pi / 2, [1, 0, 0])
yr45 = tfs.rotation_matrix(np.pi / 2, [0, 1, 0])
zr45 = tfs.rotation_matrix(np.pi / 2, [0, 0, 1])
xr45[:3,3] = np.array([1,2,3])
yr45[:3,3] = np.array([2,3,1])
zr45[:3,3] = np.array([3,1,2])
assert np.allclose(xyz_rpy_to_matrix([1, 2, 3, np.pi / 2, 0, 0]), xr45)
assert np.allclose(xyz_rpy_to_matrix([2, 3, 1, 0, np.pi / 2, 0]), yr45)
assert np.allclose(xyz_rpy_to_matrix([3, 1, 2, 0, 0, np.pi / 2]), zr45)
def test_matrix_to_xyz_rpy():
xr45 = tfs.rotation_matrix(np.pi / 2, [1, 0, 0])
yr45 = tfs.rotation_matrix(np.pi / 2, [0, 1, 0])
zr45 = tfs.rotation_matrix(np.pi / 2, [0, 0, 1])
xr45[:3,3] = np.array([1,2,3])
yr45[:3,3] = np.array([2,3,1])
zr45[:3,3] = np.array([3,1,2])
assert np.allclose([1, 2, 3, np.pi / 2, 0, 0], matrix_to_xyz_rpy(xr45))
assert np.allclose([2, 3, 1, 0, np.pi / 2, 0], matrix_to_xyz_rpy(yr45))
assert np.allclose([3, 1, 2, 0, 0, np.pi / 2], matrix_to_xyz_rpy(zr45))
def callback(dt):
if window.rotate:
for widget in widgets.values():
if isinstance(widget, trimesh.viewer.SceneWidget):
scene = widget.scene
camera = scene.camera
axis = tf.transform_points([[0, 1, 0]],
camera.transform,
translate=False)[0]
camera.transform = tf.rotation_matrix(
np.deg2rad(window.rotate),
axis,
point=scene.centroid,
) @ camera.transform
widget.view['ball']._n_pose = camera.transform
return
if window.scenes_group and (window.next or window.play):
try:
scenes = next(window.scenes_group)
for key, widget in widgets.items():
if isinstance(widget, trimesh.viewer.SceneWidget):
widget.scene.geometry.update(scenes[key].geometry)
widget.scene.graph.load(
scenes[key].graph.to_edgelist())
widget._draw()
elif isinstance(widget, glooey.Image):
widget.set_image(numpy_to_image(scenes[key]))
except StopIteration:
window.play = False
window.next = False
def _get_random_stable_pose(self, stable_poses, stable_poses_probs, thres=0.005):
"""Return a stable pose according to their likelihood.
Args:
stable_poses (list[np.ndarray]): List of stable poses as 4x4 matrices.
stable_poses_probs (list[float]): List of probabilities.
thres (float): Threshold of pose stability to include for sampling
Returns:
np.ndarray: homogeneous 4x4 matrix
"""
# Random pose with unique (avoid symmetric poses) stability prob > thres
_,unique_idcs = np.unique(stable_poses_probs.round(decimals=3), return_index=True)
unique_idcs = unique_idcs[::-1]
unique_stable_poses_probs = stable_poses_probs[unique_idcs]
n = len(unique_stable_poses_probs[unique_stable_poses_probs>thres])
index = unique_idcs[np.random.randint(n)]
# index = np.random.choice(len(stable_poses), p=stable_poses_probs)
inplane_rot = tra.rotation_matrix(
angle=np.random.uniform(0, 2.0 * np.pi), direction=[0, 0, 1]
)
return inplane_rot.dot(stable_poses[index])
def rxyz(thetax, thetay, thetaz):
"""Calculates rotation matrices by multiplying in the order x,y,z.
Parameters
----------
thetax : rotation around x in degrees.
thetay : rotation around y in degrees.
thetaz : rotation around z in degrees.
"""
# Convert radians to degrees
rx = tf.rotation_matrix(thetax, [1,0,0])
ry = tf.rotation_matrix(thetay, [0,1,0])
rz = tf.rotation_matrix(thetaz, [0,0,1])
rxyz = tf.concatenate_matrices(rx,ry,rz)
return rxyz
def mesh(self, verts):
mesh = trimesh.Trimesh(vertices=verts, faces=self.faces)
# this transform is necessary to get correct image
# because z axis is other way around in trimesh
transform = trans.rotation_matrix(np.deg2rad(-180), [1, 0, 0],
mesh.centroid)
mesh.apply_transform(transform)
return mesh
def rotate(self, azimuth, axis=None):
"""Rotate the trackball about the "Up" axis by azimuth radians.
Parameters
----------
azimuth : float
The number of radians to rotate.
"""
target = self._target
y_axis = self._n_pose[:3, 1].flatten()
if axis is not None:
y_axis = axis
x_rot_mat = transformations.rotation_matrix(azimuth, y_axis, target)
self._n_pose = x_rot_mat.dot(self._n_pose)
y_axis = self._pose[:3, 1].flatten()
if axis is not None:
y_axis = axis
x_rot_mat = transformations.rotation_matrix(azimuth, y_axis, target)
self._pose = x_rot_mat.dot(self._pose)
def rxyz(thetax, thetay, thetaz, as_degrees=False):
"""Calculates rotation matrices by multiplying in the order x,y,z.
Parameters
----------
thetax : rotation around x in degrees.
thetay : rotation around y in degrees.
thetaz : rotation around z in degrees.
"""
if as_degrees is True:
thetax = thetax * math.pi / 180.
thetay = thetay * math.pi / 180.
thetaz = thetaz * math.pi / 180.
# Convert radians to degrees
rx = tf.rotation_matrix(thetax, [1, 0, 0])
ry = tf.rotation_matrix(thetay, [0, 1, 0])
rz = tf.rotation_matrix(thetaz, [0, 0, 1])
rxyz = tf.concatenate_matrices(rx, ry, rz)
return rxyz
def _get_cylinder(radius, normal, center, height=0.1, sections=10):
"""Get a thin disk at a given location and orientation."""
rotation = rotation_matrix(
angle=_angle_between([0, 0, 1], normal),
direction=np.cross([0, 0, 1], normal),
)
cyl = cylinder(radius=radius,
height=height,
transform=rotation,
sections=sections)
cyl.vertices += center
return cyl
def rotate(self, azimuth, axis=None):
"""Rotate the trackball about the "Up" axis by azimuth radians.
Parameters
----------
azimuth : float
The number of radians to rotate.
"""
target = self._target
y_axis = self._n_pose[:3, 1].flatten()
if axis is not None:
y_axis = axis
x_rot_mat = transformations.rotation_matrix(azimuth, y_axis, target)
self._n_pose = x_rot_mat.dot(self._n_pose)
y_axis = self._pose[:3, 1].flatten()
if axis is not None:
y_axis = axis
x_rot_mat = transformations.rotation_matrix(azimuth, y_axis, target)
self._pose = x_rot_mat.dot(self._pose)
def mesh(self, verts):
mesh = trimesh.Trimesh(vertices=verts, faces=self.faces,
vertex_colors=[200, 255, 255, 255],
face_colors=[0, 0, 0, 0],
use_embree=False,
process=False)
# this transform is necessary to get correct image
# because z axis is other way around in trimesh
transform = trans.rotation_matrix(np.deg2rad(-180), [1, 0, 0], mesh.centroid)
mesh.apply_transform(transform)
return mesh
def _create_mesh(self, vertices):
print("vertices.shape", vertices.shape)
mesh = trimesh.Trimesh(vertices=vertices,
faces=self._smpl_faces,
vertex_colors=[200, 255, 255, 255],
face_colors=[0, 0, 0, 0],
use_embree=False,
process=False)
transform = trans.rotation_matrix(np.deg2rad(-180), [1, 0, 0],
mesh.centroid)
mesh.apply_transform(transform)
return mesh
def _get_random_stable_pose(self, stable_poses, stable_poses_probs):
"""Return a stable pose according to their likelihood.
Args:
stable_poses (list[np.ndarray]): List of stable poses as 4x4 matrices.
stable_poses_probs (list[float]): List of probabilities.
Returns:
np.ndarray: homogeneous 4x4 matrix
"""
index = np.random.choice(len(stable_poses), p=stable_poses_probs)
inplane_rot = tra.rotation_matrix(angle=np.random.uniform(
0, 2.0 * np.pi),
direction=[0, 0, 1])
return inplane_rot.dot(stable_poses[index])
def callback(dt):
if window.rotate:
for widget in widgets.values():
if isinstance(widget, trimesh.viewer.SceneWidget):
axis = tf.transform_points(
[[0, 1, 0]],
widget.scene.camera_transform,
translate=False,
)[0]
widget.scene.camera_transform[...] = (tf.rotation_matrix(
np.deg2rad(window.rotate * rotation_scaling),
axis,
point=widget.scene.centroid,
) @ widget.scene.camera_transform)
widget.view["ball"]._n_pose = widget.scene.camera_transform
return
if window.scenes_group and (window.next or window.play):
try:
scenes = next(window.scenes_group)
clear = scenes.get("__clear__", False) or window._clear
window._clear = False
for key, widget in widgets.items():
scene = scenes[key]
if isinstance(widget, trimesh.viewer.SceneWidget):
assert isinstance(scene, trimesh.Scene)
if clear:
widget.clear()
widget.scene = scene
else:
widget.scene.geometry.update(scene.geometry)
widget.scene.graph.load(scene.graph.to_edgelist())
widget.scene.camera_transform[
...] = scene.camera_transform
widget.view[
"ball"]._n_pose = widget.scene.camera_transform
widget._draw()
elif isinstance(widget, glooey.Image):
widget.set_image(numpy_to_image(scene))
except StopIteration:
print("Reached the end of the scenes")
window.play = False
window.next = False
def mesh(self, at, groups=None):
"""Generate mesh for gear pair at position given by `at`.
.. note: `at` is/must be edited by reference
"""
sign = np.sign(self.disp) if self.disp != 0 else 1
if abs(self.disp) < self.stretch_margin:
stretch = self.disp + sign * self.height / 2
else:
stretch = sign * (self.height + self.gears.p.stretch +
2 * self.stretch_margin) / 2
at.dot(translation_matrix([0, 0, stretch]), out=at)
at.dot(rotation_matrix(self.angle, [0, 0, 1]), out=at)
p_mesh = self.gears.p.mesh(at)
at.dot(translation_matrix(
[self.ap, 0, sign * self.gears.p.stretch / 2]),
out=at)
g_mesh = self.gears.g.mesh(at)
if groups is not None:
groups[-1].append(p_mesh)
groups.append([g_mesh])
return (p_mesh, g_mesh)
def _get_class_specific_orientations_and_extents(self, cannonincal_Rs,
cannonical_extent_idxs):
self._canonical_quaternions = dict()
self._canonical_extents = dict()
for class_id in self._class_ids:
self._canonical_quaternions[class_id] = dict()
self._canonical_extents[class_id] = dict()
R_z_fix = ttf.rotation_matrix(self._z_rotations[class_id],
[0.0, 0.0, 1.0])
extent_orig = self._object_models.get_cad(
class_id=class_id).extents
for up_axis_key in self._class_up_axis_keys[class_id]:
q_s = list()
for R in cannonincal_Rs[up_axis_key]:
R_tot = np.matmul(R, R_z_fix)
q = ttf.quaternion_from_matrix(R_tot)
q_s.append(np.concatenate((q[1:], q[0:1])))
self._canonical_quaternions[class_id][up_axis_key] = q_s
extents = list()
extent_idxs = cannonical_extent_idxs[up_axis_key]
for extent_idx in extent_idxs:
extent = [
extent_orig[extent_idx[0]],
extent_orig[extent_idx[1]],
extent_orig[extent_idx[2]],
]
extents.append(extent)
self._canonical_extents[class_id][up_axis_key] = extents
def drag(self, point):
"""Update the tracball during a drag.
Parameters
----------
point : (2,) int
The current x and y pixel coordinates of the mouse during a drag.
This will compute a movement for the trackball with the relative
motion between this point and the one marked by down().
"""
point = np.array(point, dtype=np.float32)
dx, dy = point - self._pdown
# dy = -dy
mindim = 0.3 * np.min(self._size)
target = self._target
x_axis = self._pose[:3, 0].flatten()
y_axis = self._pose[:3, 1].flatten()
z_axis = self._pose[:3, 2].flatten()
eye = self._pose[:3, 3].flatten()
direction = -self._pose[:3, 2]
# Interpret drag as a rotation
if self._state == Trackball.STATE_ROTATE:
if self.rotation_mode == 'trackball':
intersection_point = linePlaneCollision(
planeNormal=np.array([0, 0, 1.0]),
planePoint=np.array([0, 0, 0.0]),
rayDirection=direction,
rayPoint=eye)
# x_angle = -dx / mindim
# x_rot_mat = transformations.rotation_matrix(
# x_angle, np.array([0, 0, 1.]), intersection_point
# )
# print(intersection_point)
y_angle = dy / mindim
y_rot_mat = transformations.rotation_matrix(
y_angle, project_onto_plane(-x_axis, [0., 0., 1.0]),
intersection_point)
x_angle = -dx / mindim
x_rot_mat_ = transformations.rotation_matrix(
x_angle, y_axis, point=intersection_point)
# self._n_pose = x_rot_mat_.dot(self._pose)
self._n_pose = x_rot_mat_.dot(y_rot_mat.dot(self._pose))
else:
self._n_pose = self._pose.copy()
self._n_pose = self.tilt(-dy / 10,
self._n_pose,
cm_orig=self._n_pose)
self._n_pose = self.yaw(dx / 10, self._n_pose)
# Interpret drag as a roll about the camera axis
elif self._state == Trackball.STATE_ROLL:
center = self._size / 2.0
v_init = self._pdown - center
v_curr = point - center
v_init = v_init / np.linalg.norm(v_init)
v_curr = v_curr / np.linalg.norm(v_curr)
theta = (-np.arctan2(v_curr[1], v_curr[0]) +
np.arctan2(v_init[1], v_init[0]))
rot_mat = transformations.rotation_matrix(theta, z_axis, target)
self._n_pose = rot_mat.dot(self._pose)
# Interpret drag as a camera pan in view plane
elif self._state == Trackball.STATE_PAN:
dx = -dx / (2.0 * mindim) * self._scale
dy = dy / (2.0 * mindim) * self._scale
translation = dx * x_axis + dy * y_axis
self._n_target = self._target + translation
t_tf = np.eye(4)
t_tf[:3, 3] = translation
self._n_pose = t_tf.dot(self._pose)
# Interpret drag as a zoom motion
elif self._state == Trackball.STATE_ZOOM:
radius = np.linalg.norm(eye - target)
ratio = 0.0
if dy > 0:
ratio = np.exp(abs(dy) / (0.5 * self._size[1])) - 1.0
elif dy < 0:
ratio = 1.0 - np.exp(dy / (0.5 * (self._size[1])))
translation = -np.sign(dy) * ratio * radius * z_axis
t_tf = np.eye(4)
t_tf[:3, 3] = translation
self._n_pose = t_tf.dot(self._pose)
def _get_cannonical_object_rotation_mats_and_extent_idxs(self):
# final data container
all_R_s = dict()
all_extents = dict()
# visualization
scene = trimesh.Scene()
scene.add_geometry(trimesh.creation.axis())
# axes
axes = dict()
axes["x+"] = [1.0, 0.0, 0.0]
axes["x-"] = [-1.0, 0.0, 0.0]
axes["y+"] = [0.0, 1.0, 0.0]
axes["y-"] = [0.0, -1.0, 0.0]
axes["z+"] = [0.0, 0.0, 1.0]
axes["z-"] = [0.0, 0.0, -1.0]
# axis indices
axis_indices = dict()
axis_indices["x+"] = 0
axis_indices["x-"] = 0
axis_indices["y+"] = 1
axis_indices["y-"] = 1
axis_indices["z+"] = 2
axis_indices["z-"] = 2
# axis up transformations
transformations = dict()
transformations["x+"] = ("y+", -math.pi / 2)
transformations["x-"] = ("y+", math.pi / 2)
transformations["y+"] = ("x+", math.pi / 2)
transformations["y-"] = ("x+", -math.pi / 2)
transformations["z+"] = ("x+", 0.0)
transformations["z-"] = ("x+", math.pi)
for up_axis_key, transformation in transformations.items():
rotation_axis_key = transformation[0]
remaining_axis_key = self._get_remaining_axis_key(
up_axis_key, rotation_axis_key)
rotation_angle = transformation[1]
rotation_axis = axes[rotation_axis_key]
R_0 = ttf.rotation_matrix(rotation_angle, rotation_axis)
R_90 = np.matmul(
R_0, ttf.rotation_matrix(math.pi / 2, axes[up_axis_key]))
R_180 = np.matmul(
R_90, ttf.rotation_matrix(math.pi / 2, axes[up_axis_key]))
R_270 = np.matmul(
R_180, ttf.rotation_matrix(math.pi / 2, axes[up_axis_key]))
# store extents
# ToDo: change so that unaligned objects have correct extents, now that some of the objects are z rotated
extents = list()
for i in range(2):
extents.append([
axis_indices[rotation_axis_key],
axis_indices[remaining_axis_key],
axis_indices[up_axis_key],
])
extents.append([
axis_indices[remaining_axis_key],
axis_indices[rotation_axis_key],
axis_indices[up_axis_key],
])
# store quaternions
R_s = list()
for R in [R_0, R_90, R_180, R_270]:
R_s.append(R)
all_R_s[up_axis_key] = R_s
all_extents[up_axis_key] = extents
return all_R_s, all_extents
def rotation_z(angle):
"""create matrix for rotation around z axis"""
return transformations.rotation_matrix(angle, Z_HAT)
def test_nodes():
x = Node()
assert x.name is None
assert x.camera is None
assert x.children == []
assert x.skin is None
assert np.allclose(x.matrix, np.eye(4))
assert x.mesh is None
assert np.allclose(x.rotation, [0, 0, 0, 1])
assert np.allclose(x.scale, np.ones(3))
assert np.allclose(x.translation, np.zeros(3))
assert x.weights is None
assert x.light is None
x.name = 'node'
# Test node light/camera/mesh tests
c = PerspectiveCamera(yfov=2.0)
m = Mesh([])
d = DirectionalLight()
x.camera = c
assert x.camera == c
with pytest.raises(TypeError):
x.camera = m
x.camera = d
x.camera = None
x.mesh = m
assert x.mesh == m
with pytest.raises(TypeError):
x.mesh = c
x.mesh = d
x.light = d
assert x.light == d
with pytest.raises(TypeError):
x.light = m
x.light = c
# Test transformations getters/setters/etc...
# Set up test values
x = np.array([1.0, 0.0, 0.0])
y = np.array([0.0, 1.0, 0.0])
t = np.array([1.0, 2.0, 3.0])
s = np.array([0.5, 2.0, 1.0])
Mx = transformations.rotation_matrix(np.pi / 2.0, x)
qx = np.roll(transformations.quaternion_about_axis(np.pi / 2.0, x), -1)
Mxt = Mx.copy()
Mxt[:3, 3] = t
S = np.eye(4)
S[:3, :3] = np.diag(s)
Mxts = Mxt.dot(S)
My = transformations.rotation_matrix(np.pi / 2.0, y)
qy = np.roll(transformations.quaternion_about_axis(np.pi / 2.0, y), -1)
Myt = My.copy()
Myt[:3, 3] = t
x = Node(matrix=Mx)
assert np.allclose(x.matrix, Mx)
assert np.allclose(x.rotation, qx)
assert np.allclose(x.translation, np.zeros(3))
assert np.allclose(x.scale, np.ones(3))
x.matrix = My
assert np.allclose(x.matrix, My)
assert np.allclose(x.rotation, qy)
assert np.allclose(x.translation, np.zeros(3))
assert np.allclose(x.scale, np.ones(3))
x.translation = t
assert np.allclose(x.matrix, Myt)
assert np.allclose(x.rotation, qy)
x.rotation = qx
assert np.allclose(x.matrix, Mxt)
x.scale = s
assert np.allclose(x.matrix, Mxts)
x = Node(matrix=Mxt)
assert np.allclose(x.matrix, Mxt)
assert np.allclose(x.rotation, qx)
assert np.allclose(x.translation, t)
assert np.allclose(x.scale, np.ones(3))
x = Node(matrix=Mxts)
assert np.allclose(x.matrix, Mxts)
assert np.allclose(x.rotation, qx)
assert np.allclose(x.translation, t)
assert np.allclose(x.scale, s)
# Individual element getters
x.scale[0] = 0
assert np.allclose(x.scale[0], 0)
x.translation[0] = 0
assert np.allclose(x.translation[0], 0)
x.matrix = np.eye(4)
x.matrix[0, 0] = 500
assert x.matrix[0, 0] == 1.0
# Failures
with pytest.raises(ValueError):
x.matrix = 5 * np.eye(4)
with pytest.raises(ValueError):
x.matrix = np.eye(5)
with pytest.raises(ValueError):
x.matrix = np.eye(4).dot([5, 1, 1, 1])
with pytest.raises(ValueError):
x.rotation = np.array([1, 2])
with pytest.raises(ValueError):
x.rotation = np.array([1, 2, 3])
with pytest.raises(ValueError):
x.rotation = np.array([1, 2, 3, 4])
with pytest.raises(ValueError):
x.translation = np.array([1, 2, 3, 4])
with pytest.raises(ValueError):
x.scale = np.array([1, 2, 3, 4])
def drag(self, point):
"""Update the tracball during a drag.
Parameters
----------
point : (2,) int
The current x and y pixel coordinates of the mouse during a drag.
This will compute a movement for the trackball with the relative
motion between this point and the one marked by down().
"""
point = np.array(point, dtype=np.float32)
dx, dy = point - self._pdown
mindim = 0.3 * np.min(self._size)
target = self._target
x_axis = self._pose[:3, 0].flatten()
y_axis = self._pose[:3, 1].flatten()
z_axis = self._pose[:3, 2].flatten()
eye = self._pose[:3, 3].flatten()
# Interpret drag as a rotation
if self._state == Trackball.STATE_ROTATE:
x_angle = -dx / mindim
x_rot_mat = transformations.rotation_matrix(
x_angle, y_axis, target)
y_angle = dy / mindim
y_rot_mat = transformations.rotation_matrix(
y_angle, x_axis, target)
self._n_pose = y_rot_mat.dot(x_rot_mat.dot(self._pose))
# Interpret drag as a roll about the camera axis
elif self._state == Trackball.STATE_ROLL:
center = self._size / 2.0
v_init = self._pdown - center
v_curr = point - center
v_init = v_init / np.linalg.norm(v_init)
v_curr = v_curr / np.linalg.norm(v_curr)
theta = -np.arctan2(v_curr[1], v_curr[0]) + np.arctan2(
v_init[1], v_init[0])
rot_mat = transformations.rotation_matrix(theta, z_axis, target)
self._n_pose = rot_mat.dot(self._pose)
# Interpret drag as a camera pan in view plane
elif self._state == Trackball.STATE_PAN:
dx = -dx / (5.0 * mindim) * self._scale
dy = -dy / (5.0 * mindim) * self._scale
translation = dx * x_axis + dy * y_axis
self._n_target = self._target + translation
t_tf = np.eye(4)
t_tf[:3, 3] = translation
self._n_pose = t_tf.dot(self._pose)
# Interpret drag as a zoom motion
elif self._state == Trackball.STATE_ZOOM:
radius = np.linalg.norm(eye - target)
ratio = 0.0
if dy > 0:
ratio = np.exp(abs(dy) / (0.5 * self._size[1])) - 1.0
elif dy < 0:
ratio = 1.0 - np.exp(dy / (0.5 * (self._size[1])))
translation = -np.sign(dy) * ratio * radius * z_axis
t_tf = np.eye(4)
t_tf[:3, 3] = translation
self._n_pose = t_tf.dot(self._pose)
def opengl_camera_transform(transform=None):
if transform is None:
transform = np.eye(4)
return transform @ tf.rotation_matrix(np.deg2rad(-180), [1, 0, 0])
def generate(
n,
method=None,
align_pai=False,
density=(2200, 2600),
directory='.',
name=None,
seed=None,
show=False,
start_index=0,
max_index=None,
make_log=True,
**kwargs,
):
"""Generates n pairs of files (.obj and .urdf) from meshes created
with a given method.
Args:
n: number of objects to be generated.
method: method used to generate the seed. Either a callable or the
string identifier of one of the implemented methods.
align_pia: whether to align the mesh's principal axes of inertia with
the frame axes. If False, the mesh's oriented bounding box is
aligned instead.
density: used for the model mass and inertia. Either a scalar
or a tuple with the limits of a uniform random distribution
directory: path to the location where to save the models.
name: name of the models (to be suffixed with an index). If None, only
the index is used to name the files.
seed: seed of the random generator.
show: whether to show a visualization of each model.
start_index: index of the first generated object. Used to generate
objects on a directory without overwriting previous ones.
max_index: maximum index expected in the directory. Used to correctly
set the number of leading zeros of the index in the name of the files
if more files are expected to be generated outside this function call.
If None, the maximum index generated in this call is used.
"""
# Set method
if method is None:
method = box
if isinstance(method, str):
method = methods[method]
elif not callable(method):
raise TypeError("method must be callable or a string.")
# Load the template urdf
with open(os.path.join(os.path.dirname(__file__), 'template.urdf')) as f:
urdf = f.read()
# Create directory if needed
if not os.path.isdir(directory):
os.makedirs(directory)
# Set logging
if make_log:
log_name = os.path.join(directory,
name + '.csv' if name else 'log.csv')
if start_index and os.path.isfile(log_name):
logf = open(log_name, 'a')
else:
logf = open(log_name, 'w')
logf.write('Name,Volume,Rectangularity,AspectRatio,NumVertices\n')
else:
logf = None
max_index = max(max_index or n + start_index - 1, 1)
name_format = '{:0' + str(int(np.log10(max_index)) + 1) + '}'
if isinstance(name, str):
name_format = '{}_{}'.format(name, name_format)
name = lambda i: name_format.format(i)
# Create random number generator.
random = np.random.default_rng(seed)
for i in range(start_index, start_index + n):
# Create mesh
mesh = method(seed=random, **kwargs)
# Align the principal axes with (Z,Y,X).
if align_pai:
mesh.apply_transform(mesh.principal_inertia_transform)
else:
mesh.apply_obb()
mesh.apply_transform(
transformations.rotation_matrix(angle=np.pi / 2,
direction=[0, 1, 0]))
# # Correct residual translation due to rotation
# mesh.apply_translation(-mesh.center_mass)
mesh.process()
assert mesh.is_watertight
# Set density
if np.isscalar(density):
mesh.density = density
grayscale = 0.5
else:
mesh.density = random.uniform(density[0], density[1])
# Set color according to density
grayscale = 0.6 - 0.2 * (mesh.density - density[0]) / (density[1] -
density[0])
if show:
mesh.show()
name_i = name(i)
# Log shape metrics
if logf is not None:
extents = mesh.bounding_box_oriented.extents
logf.write('{},{},{},{},{}\n'.format(
name_i,
mesh.volume,
mesh.volume / mesh.bounding_box_oriented.volume,
max(extents) / min(extents),
len(mesh.vertices),
))
# Export mesh to .obj file
fname = os.path.join(directory, name_i)
with open(fname + '.obj', 'w') as f:
mesh.export(file_obj=f, file_type='obj')
# Create .urdf file from template with mesh specs
with open(fname + '.urdf', 'w') as f:
f.write(
urdf.format(
name=name_i,
friction=0.6,
mass=mesh.mass,
x=mesh.center_mass[0],
y=mesh.center_mass[1],
z=mesh.center_mass[2],
ixx=mesh.moment_inertia[0, 0],
ixy=mesh.moment_inertia[0, 1],
ixz=mesh.moment_inertia[0, 2],
iyy=mesh.moment_inertia[1, 1],
iyz=mesh.moment_inertia[1, 2],
izz=mesh.moment_inertia[2, 2],
mesh=name_i + '.obj',
r=grayscale,
g=grayscale,
b=grayscale,
a=1.,
))
def drag(self, point):
"""Update the tracball during a drag.
Parameters
----------
point : (2,) int
The current x and y pixel coordinates of the mouse during a drag.
This will compute a movement for the trackball with the relative
motion between this point and the one marked by down().
"""
point = np.array(point, dtype=np.float32)
dx, dy = point - self._pdown
mindim = 0.3 * np.min(self._size)
target = self._target
x_axis = self._pose[:3, 0].flatten()
y_axis = self._pose[:3, 1].flatten()
z_axis = self._pose[:3, 2].flatten()
eye = self._pose[:3, 3].flatten()
# Interpret drag as a rotation
if self._state == Trackball.STATE_ROTATE:
x_angle = -dx / mindim
x_rot_mat = transformations.rotation_matrix(
x_angle, y_axis, target
)
y_angle = dy / mindim
y_rot_mat = transformations.rotation_matrix(
y_angle, x_axis, target
)
self._n_pose = y_rot_mat.dot(x_rot_mat.dot(self._pose))
# Interpret drag as a roll about the camera axis
elif self._state == Trackball.STATE_ROLL:
center = self._size / 2.0
v_init = self._pdown - center
v_curr = point - center
v_init = v_init / np.linalg.norm(v_init)
v_curr = v_curr / np.linalg.norm(v_curr)
theta = (-np.arctan2(v_curr[1], v_curr[0]) +
np.arctan2(v_init[1], v_init[0]))
rot_mat = transformations.rotation_matrix(theta, z_axis, target)
self._n_pose = rot_mat.dot(self._pose)
# Interpret drag as a camera pan in view plane
elif self._state == Trackball.STATE_PAN:
dx = -dx / (5.0 * mindim) * self._scale
dy = -dy / (5.0 * mindim) * self._scale
translation = dx * x_axis + dy * y_axis
self._n_target = self._target + translation
t_tf = np.eye(4)
t_tf[:3, 3] = translation
self._n_pose = t_tf.dot(self._pose)
# Interpret drag as a zoom motion
elif self._state == Trackball.STATE_ZOOM:
radius = np.linalg.norm(eye - target)
ratio = 0.0
if dy > 0:
ratio = np.exp(abs(dy) / (0.5 * self._size[1])) - 1.0
elif dy < 0:
ratio = 1.0 - np.exp(dy / (0.5 * (self._size[1])))
translation = -np.sign(dy) * ratio * radius * z_axis
t_tf = np.eye(4)
t_tf[:3, 3] = translation
self._n_pose = t_tf.dot(self._pose)
