class PybulletPhysicsEngine(PhysicsEngine):
"""Wrapper for pybullet physics engine that is tied to a single ID"""
def __init__(self, urdf_cache_dir, debug=False):
PhysicsEngine.__init__(self)
self._physics_client = None
self._debug = debug
self._urdf_cache_dir = urdf_cache_dir
if not os.path.isabs(self._urdf_cache_dir):
self._urdf_cache_dir = os.path.join(
os.getcwd(), self._urdf_cache_dir
)
if not os.path.exists(os.path.join(self._urdf_cache_dir, "plane")):
os.makedirs(os.path.join(self._urdf_cache_dir, "plane"))
shutil.copy(
pkg_resources.resource_filename(
"sd_maskrcnn", "data/plane/plane.urdf"
),
os.path.join(self._urdf_cache_dir, "plane", "plane.urdf"),
)
shutil.copy(
pkg_resources.resource_filename(
"sd_maskrcnn", "data/plane/plane_convex_piece_0.obj"
),
os.path.join(
self._urdf_cache_dir, "plane", "plane_convex_piece_0.obj"
),
)
def add(self, obj, static=False):
# create URDF
urdf_filename = os.path.join(
self._urdf_cache_dir,
KEY_SEP_TOKEN.join(obj.key.split(KEY_SEP_TOKEN)[:-1]),
"{}.urdf".format(
KEY_SEP_TOKEN.join(obj.key.split(KEY_SEP_TOKEN)[:-1])
),
)
urdf_dir = os.path.dirname(urdf_filename)
if not os.path.exists(urdf_filename):
try:
os.makedirs(urdf_dir)
except:
self._logger.warning(
"Failed to create dir %s. The object may have been created simultaneously by another process"
% (urdf_dir)
)
self._logger.info(
"Exporting URDF for object {}".format(
KEY_SEP_TOKEN.join(obj.key.split(KEY_SEP_TOKEN)[:-1])
)
)
# Fix center of mass (for rendering) and density and export
geometry = obj.mesh.copy()
geometry.apply_translation(-obj.mesh.center_mass)
trimesh.exchange.export.export_urdf(geometry, urdf_dir)
com = obj.mesh.center_mass
pose = self._convert_pose(obj.pose, com)
obj_t = pose.translation
obj_q_wxyz = pose.quaternion
obj_q_xyzw = np.roll(obj_q_wxyz, -1)
try:
obj_id = pybullet.loadURDF(
urdf_filename,
obj_t,
obj_q_xyzw,
useFixedBase=static,
physicsClientId=self._physics_client,
)
except:
raise Exception("Failed to load %s" % (urdf_filename))
if self._debug:
self._add_to_scene(obj)
self._key_to_id[obj.key] = obj_id
self._key_to_com[obj.key] = com
def get_velocity(self, key):
obj_id = self._key_to_id[key]
return pybullet.getBaseVelocity(
obj_id, physicsClientId=self._physics_client
)
def get_pose(self, key):
obj_id = self._key_to_id[key]
obj_t, obj_q_xyzw = pybullet.getBasePositionAndOrientation(
obj_id, physicsClientId=self._physics_client
)
obj_q_wxyz = np.roll(obj_q_xyzw, 1)
pose = RigidTransform(
rotation=obj_q_wxyz,
translation=obj_t,
from_frame="obj",
to_frame="world",
)
pose = self._deconvert_pose(pose, self._key_to_com[key])
return pose
def remove(self, key):
obj_id = self._key_to_id[key]
pybullet.removeBody(obj_id, physicsClientId=self._physics_client)
self._key_to_id.pop(key)
self._key_to_com.pop(key)
if self._debug:
self._remove_from_scene(key)
def step(self):
pybullet.stepSimulation(physicsClientId=self._physics_client)
if self._debug:
time.sleep(0.04)
self._update_scene()
def reset(self):
if self._physics_client is not None:
self.stop()
self.start()
def start(self):
if self._physics_client is None:
self._physics_client = pybullet.connect(pybullet.DIRECT)
pybullet.setGravity(
0, 0, -GRAVITY_ACCEL, physicsClientId=self._physics_client
)
self._key_to_id = {}
self._key_to_com = {}
if self._debug:
self._create_scene()
self._viewer = Viewer(
self._scene, use_raymond_lighting=True, run_in_thread=True
)
def stop(self):
if self._physics_client is not None:
pybullet.disconnect(self._physics_client)
self._physics_client = None
if self._debug:
self._scene = None
self._viewer.close_external()
while self._viewer.is_active:
pass
def __del__(self):
self.stop()
del self
def _convert_pose(self, pose, com):
new_pose = pose.copy()
new_pose.translation = pose.rotation.dot(com) + pose.translation
return new_pose
def _deconvert_pose(self, pose, com):
new_pose = pose.copy()
new_pose.translation = pose.rotation.dot(-com) + pose.translation
return new_pose
def _create_scene(self):
self._scene = Scene()
camera = PerspectiveCamera(
yfov=0.833, znear=0.05, zfar=3.0, aspectRatio=1.0
)
cn = Node()
cn.camera = camera
pose_m = np.array(
[
[0.0, 1.0, 0.0, 0.0],
[1.0, 0.0, 0.0, 0.0],
[0.0, 0.0, -1.0, 0.88],
[0.0, 0.0, 0.0, 1.0],
]
)
pose_m[:, 1:3] *= -1.0
cn.matrix = pose_m
self._scene.add_node(cn)
self._scene.main_camera_node = cn
def _add_to_scene(self, obj):
self._viewer.render_lock.acquire()
n = Node(
mesh=Mesh.from_trimesh(obj.mesh),
matrix=obj.pose.matrix,
name=obj.key,
)
self._scene.add_node(n)
self._viewer.render_lock.release()
def _remove_from_scene(self, key):
self._viewer.render_lock.acquire()
if self._scene.get_nodes(name=key):
self._scene.remove_node(
next(iter(self._scene.get_nodes(name=key)))
)
self._viewer.render_lock.release()
def _update_scene(self):
self._viewer.render_lock.acquire()
for key in self._key_to_id.keys():
obj_pose = self.get_pose(key).matrix
if self._scene.get_nodes(name=key):
next(iter(self._scene.get_nodes(name=key))).matrix = obj_pose
self._viewer.render_lock.release()
class MeshViewer(object):
def __init__(self, width=1200, height=800, use_offscreen=True):
#super().__init__()
self.width, self.height = width, height
self.use_offscreen = use_offscreen
self.render_wireframe = False
self.mat_constructor = pyrender.MetallicRoughnessMaterial
self.trimesh_to_pymesh = pyrender.Mesh.from_trimesh
self.scene = pyrender.Scene(bg_color=colors['white'],
ambient_light=(0.3, 0.3, 0.3))
pc = pyrender.PerspectiveCamera(yfov=np.pi / 3.0,
aspectRatio=float(width) / height)
camera_pose = np.eye(4)
camera_pose[:3, 3] = np.array([0, 0, 3.0])
self.camera_node = self.scene.add(pc,
pose=camera_pose,
name='pc-camera')
self.figsize = (width, height)
if self.use_offscreen:
self.viewer = pyrender.OffscreenRenderer(*self.figsize)
self.use_raymond_lighting(4.)
else:
self.viewer = Viewer(self.scene,
use_raymond_lighting=True,
viewport_size=self.figsize,
cull_faces=False,
run_in_thread=True)
def set_background_color(self, color=colors['white']):
self.scene.bg_color = color
def set_cam_trans(self, trans=[0, 0, 3.0]):
if isinstance(trans, list): trans = np.array(trans)
camera_pose = np.eye(4)
camera_pose[:3, 3] = trans
self.scene.set_pose(self.camera_node, pose=camera_pose)
def update_camera_pose(self, camera_pose):
self.scene.set_pose(self.camera_node, pose=camera_pose)
def close_viewer(self):
if self.viewer.is_active:
self.viewer.close_external()
def set_meshes(self, meshes, group_name='static', poses=[]):
for node in self.scene.get_nodes():
if node.name is not None and '%s-mesh' % group_name in node.name:
self.scene.remove_node(node)
if len(poses) < 1:
for mid, mesh in enumerate(meshes):
if isinstance(mesh, trimesh.Trimesh):
mesh = pyrender.Mesh.from_trimesh(mesh)
self.scene.add(mesh, '%s-mesh-%2d' % (group_name, mid))
else:
for mid, iter_value in enumerate(zip(meshes, poses)):
mesh, pose = iter_value
if isinstance(mesh, trimesh.Trimesh):
mesh = pyrender.Mesh.from_trimesh(mesh)
self.scene.add(mesh, '%s-mesh-%2d' % (group_name, mid), pose)
def set_static_meshes(self, meshes, poses=[]):
self.set_meshes(meshes, group_name='static', poses=poses)
def set_dynamic_meshes(self, meshes, poses=[]):
self.set_meshes(meshes, group_name='dynamic', poses=poses)
def _add_raymond_light(self):
from pyrender.light import DirectionalLight
from pyrender.node import Node
thetas = np.pi * np.array([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0])
phis = np.pi * np.array([0.0, 2.0 / 3.0, 4.0 / 3.0])
nodes = []
for phi, theta in zip(phis, thetas):
xp = np.sin(theta) * np.cos(phi)
yp = np.sin(theta) * np.sin(phi)
zp = np.cos(theta)
z = np.array([xp, yp, zp])
z = z / np.linalg.norm(z)
x = np.array([-z[1], z[0], 0.0])
if np.linalg.norm(x) == 0:
x = np.array([1.0, 0.0, 0.0])
x = x / np.linalg.norm(x)
y = np.cross(z, x)
matrix = np.eye(4)
matrix[:3, :3] = np.c_[x, y, z]
nodes.append(
Node(light=DirectionalLight(color=np.ones(3), intensity=1.0),
matrix=matrix))
return nodes
def use_raymond_lighting(self, intensity=1.0):
if not self.use_offscreen:
sys.stderr.write(
'Interactive viewer already uses raymond lighting!\n')
return
for n in self._add_raymond_light():
n.light.intensity = intensity / 3.0
if not self.scene.has_node(n):
self.scene.add_node(n) #, parent_node=pc)
def render(self, render_wireframe=None, RGBA=False):
from pyrender.constants import RenderFlags
flags = RenderFlags.SHADOWS_DIRECTIONAL
if RGBA: flags |= RenderFlags.RGBA
if render_wireframe is not None and render_wireframe == True:
flags |= RenderFlags.ALL_WIREFRAME
elif self.render_wireframe:
flags |= RenderFlags.ALL_WIREFRAME
color_img, depth_img = self.viewer.render(self.scene, flags=flags)
return color_img
def save_snapshot(self, fname):
if not self.use_offscreen:
sys.stderr.write(
'Currently saving snapshots only works with off-screen renderer!\n'
)
return
color_img = self.render()
cv2.imwrite(fname, color_img)
