import open3d as o3d
import os
import numpy as np
import open3d.visualization.gui as gui
import open3d.visualization.rendering as rendering
[docs]
class Visualizer:
"""Class to help with custom rendering"""
def __init__(self, client=None):
"""
:param client: The client to be used for the visualizer.
:type client: int, optional, default=None
"""
self.client = client
self.meshes = {}
self.meshes = {}
self.R_urdf = {}
if self.client.config.skin.use:
self.fpath_to_skin = {}
# Init GUI, windows, etc.
self.gui = gui.Application.instance
self.gui.initialize()
self.window = self.gui.create_window("pyCub", 640, 480, 0, 0)
self.gui.menubar = gui.Menu()
self.menu = self.gui.menubar
self.scene = gui.SceneWidget()
self.scene.scene = rendering.Open3DScene(self.window.renderer)
self.vis = self.scene.scene
self.scene.background_color = o3d.visualization.gui.Color(1, 1, 1, 1)
self.vis.set_background([1, 1, 1, 1])
self.vis.set_lighting(self.vis.NO_SHADOWS, [0, 0, 0])
# prepare default material
self.mat = rendering.MaterialRecord()
self.mat.shader = 'defaultLit'
# Prepare the menu
self.file_menu = gui.Menu()
self.file_menu.add_item("RGB image", 0)
self.file_menu.add_item("depth image", 1)
menu_counter = 2
for eye in ["l_eye", "r_eye"]:
for image_type in ["RGB", "RGB_D"]:
self.file_menu.add_item(f"{eye} {image_type} image", menu_counter)
self.file_menu.set_enabled(menu_counter, False)
menu_counter += 1
self.menu.add_menu("Save", self.file_menu)
self.show_menu = gui.Menu()
self.show_menu.add_item("l_eye", menu_counter)
self.show_menu.add_item("r_eye", menu_counter + 1)
self.menu.add_menu("Show", self.show_menu)
self.file_menu.add_item("Stream l_eye to file", menu_counter + 2)
self.file_menu.set_enabled(menu_counter + 2, False)
self.file_menu.add_item("Stream r_eye to file", menu_counter + 3)
self.file_menu.set_enabled(menu_counter + 3, False)
show_callbacks = []
for menu_id in range(menu_counter + 4):
c = self.MenuCallback(menu_id, self)
if menu_id in [6, 7]:
show_callbacks.append(c)
self.window.set_on_menu_item_activated(menu_id, c)
self.window.add_child(self.scene)
self.file_dir = os.path.dirname(os.path.abspath(__file__))
self.orientations = []
self.positions = []
self.colors = []
self.paths = []
self.eye_windows = {}
# Run this, to load the initial meshes, compute bounding boxes and init camera and center of rotation
self.read_info(self.client.robot)
self.show_first()
self.show_mesh() # this is here only for bounding box computation
# compute center of all objects
scene_mesh = o3d.geometry.TriangleMesh()
for f_path, m in self.meshes.items():
R = self.vis.get_geometry_transform(f_path)
m.transform(R)
scene_mesh += m
m.transform(np.linalg.inv(R))
bbox = scene_mesh.get_axis_aligned_bounding_box()
center = bbox.get_center()
# look at the center
self.scene.look_at(center, center+[-1, 0, 0], [0, 0, 1])
self.scene.center_of_rotation = center
for obj_id, obj_name, _ in self.client.other_objects:
self.read_info(obj_id)
self.show_first(urdf_name=obj_name)
self.show_mesh()
self.is_alive = True
if self.client.config.eyes.l_eye:
show_callbacks[0]()
if self.client.config.eyes.r_eye:
show_callbacks[1]()
[docs]
def show_first(self, urdf_name="robot"):
"""
Show the first batch of meshes in the visualizer. It loads the meshes and saves the to dict for quicker use later
:param urdf_name: The name of the urdf to be used.
:type urdf_name: str, optional, default="robot"
"""
for mesh_id in range(len(self.positions) // 3):
# get correct values for given mesh
col = self.colors[mesh_id * 3:(mesh_id + 1) * 3]
f_path = self.paths[mesh_id]
self.meshes[f_path] = o3d.io.read_triangle_mesh(f_path)
# Just for visualization
if not self.meshes[f_path].has_triangle_normals():
self.meshes[f_path].compute_triangle_normals()
if not self.meshes[f_path].has_vertex_normals():
self.meshes[f_path].compute_vertex_normals()
self.meshes[f_path].paint_uniform_color(col[:3])
# URDF rotations and translations
init_xyz, init_rpy, scale, link_name = self.find_xyz_rpy(os.path.basename(f_path), urdf_name=urdf_name)
R_urdf = np.eye(4)
R_urdf[:3, :3] = np.reshape(self.client.getMatrixFromQuaternion(self.client.getQuaternionFromEuler(init_rpy)), (3, 3))
R_urdf[:3, 3] = init_xyz
self.R_urdf[f_path] = R_urdf
self.meshes[f_path].scale(scale, self.meshes[f_path].get_center())
self.meshes[f_path].translate(self.meshes[f_path].get_center()*scale, relative=False)
# Add the mesh
self.vis.add_geometry(f_path, geometry=self.meshes[f_path], material=self.mat)
if self.client.config.skin.use:
if link_name in self.client.skin_point_clouds:
self.fpath_to_skin[f_path] = link_name
self.vis.add_geometry(f_path+"_skin", geometry=self.client.skin_point_clouds[link_name], material=self.mat)
else:
self.fpath_to_skin[f_path] = None
[docs]
def show_mesh(self):
"""
Function to parse info about meshes from PyBullet
"""
for mesh_id in range(len(self.positions) // 3):
# get correct values for given mesh
pos = self.positions[mesh_id * 3:(mesh_id + 1) * 3]
ori = self.orientations[mesh_id * 4:(mesh_id + 1) * 4]
f_path = self.paths[mesh_id]
R_urdf = self.R_urdf[f_path]
# get ori and position as 4x4 transformation matrix
R = np.eye(4)
R[:3, :3] = np.reshape(self.client.getMatrixFromQuaternion(ori), (3, 3))
R[:3, 3] = pos
self.vis.set_geometry_transform(f_path, R @ R_urdf)
for ew in self.eye_windows.values():
ew.vis.set_geometry_transform(f_path, R @ R_urdf)
if ew.link_name in f_path:
R_ = R @ R_urdf
ew.center = (R_ @ np.hstack((self.meshes[f_path].get_center(), 1)))[:3]
R_[:3, 3] = [0, 0, 0]
ew.dir = (R_ @ [0, 0, 1, 1])[:3]
ew.dir = 0.1*ew.dir/np.linalg.norm(ew.dir)
if self.client.config.skin.use and self.fpath_to_skin[f_path] is not None:
self.client.skin_point_clouds[self.fpath_to_skin[f_path]].paint_uniform_color([1, 0, 0])
colors = np.asarray(self.client.skin_point_clouds[self.fpath_to_skin[f_path]].colors)
colors[self.client.skin_activations[self.fpath_to_skin[f_path]] > 0] = [0, 1, 0]
self.client.skin_point_clouds[self.fpath_to_skin[f_path]].colors = o3d.utility.Vector3dVector(colors)
self.vis.remove_geometry(f_path+"_skin")
self.vis.add_geometry(f_path+"_skin", geometry=self.client.skin_point_clouds[self.fpath_to_skin[f_path]], material=self.mat)
self.vis.set_geometry_transform(f_path+"_skin", R)
[docs]
def read_info(self, obj_id):
"""
Read info from PyBullet
:param obj_id: id of the object; given by pybullet
:type obj_id: int
:return: 0 for success
:rtype: int
"""
# All meshes from the current object
visualData = self.client.getVisualShapeData(obj_id)
self.positions = []
self.orientations = []
self.colors = []
self.paths = []
for m in visualData:
# Get information about individual parts of the object
f_path = m[self.client.visualShapeData["FILE"]].decode("utf-8")
if f_path == "":
continue
col = m[self.client.visualShapeData["COLOR"]]
link = m[self.client.visualShapeData["LINK"]]
# non-base links
if link != -1:
# get link info
linkState = self.client.getLinkState(obj_id, link, computeLinkVelocity=0,
computeForwardKinematics=0)
# get orientation and position wrt URDF - better than in world
ori = linkState[self.client.linkInfo["URDFORI"]]
pos = linkState[self.client.linkInfo["URDFPOS"]]
# link == -1 is base. For that, getBasePosition... needs to be used. This joint must not contain URDF visual xyz and rpy
else:
pos, ori = self.client.getBasePositionAndOrientation(obj_id)
self.positions += pos
self.orientations += ori
self.colors += col[:-1]
self.paths.append(f_path)
return 0
[docs]
def render(self):
"""
Render all the things
"""
# read info
self.read_info(self.client.robot)
# add new
self.show_mesh()
for obj_id, _, fixed in self.client.other_objects:
if not fixed:
self.read_info(obj_id)
self.show_mesh()
self.window.post_redraw()
for ew in self.eye_windows.values():
ew.scene.look_at(ew.center + ew.dir, ew.center, [0, 0, 1])
ew.window.post_redraw()
ew.get_image()
if ew.save_images_bool:
ew.save_images()
if not self.gui.run_one_tick():
self.is_alive = False
self.gui.quit()
[docs]
def find_xyz_rpy(self, mesh_name, urdf_name="robot"):
"""
Find the xyz, rpy and scales values.
:param mesh_name: The name of the mesh.
:type mesh_name: str
:param urdf_name: The name of the urdf.
:type urdf_name: str, optional, default="robot"
:return: The xyz, rpy, and scales, link_name
"""
for link in self.client.urdfs[urdf_name].links:
if hasattr(link, "visual"):
for idx in range(len(link.visual.geometry.mesh.filename)):
if os.path.basename(link.visual.geometry.mesh.filename) == mesh_name:
xyz = link.visual.origin.xyz
rpy = link.visual.origin.rpy
scale = link.visual.geometry.mesh.scale[0]
link_name = link.name
return xyz, rpy, scale, link_name
[docs]
class EyeWindow:
MENU_IDS = {"l_eye": [2, 3, 8], "r_eye": [4, 5, 9]}
POSITIONS = {"l_eye": [320, 560], "r_eye": [0, 560]}
def __init__(self, eye, parent):
"""
Class to handle windows for eye rendering
:param eye: name of the eye
:type eye: str
:param parent: The parent class (Visualizer).
:type parent: int
"""
self.eye = eye
self.link_name = self.eye + "_pupil"
self.parent = parent
self.window = self.parent.gui.create_window(self.eye, 320, 240,
self.POSITIONS[self.eye][0], self.POSITIONS[self.eye][1])
self.window.set_on_close(self.on_close)
self.scene = gui.SceneWidget()
self.scene.set_on_mouse(self.on_mouse)
self.scene.scene = rendering.Open3DScene(self.window.renderer)
self.vis = self.scene.scene
self.scene.background_color = o3d.visualization.gui.Color(1, 1, 1, 1)
self.vis.set_background([1, 1, 1, 1])
self.vis.set_lighting(self.vis.NO_SHADOWS, [0, 0, 0])
self.window.show_menu(False)
self.window.add_child(self.scene)
self.center = None
self.dir = None
self.img_counter = 0
self.save_images_bool = False
self.last_image = None
for menu_id in self.MENU_IDS[self.eye]:
self.parent.file_menu.set_enabled(menu_id, True)
# Add all existing meshes to the visualizer and transform them to current position
for mesh in self.parent.meshes:
self.vis.add_geometry(mesh, geometry=self.parent.meshes[mesh], material=self.parent.mat)
# This is needed basically only for stationary objects, but it does not really matter if done for all
self.vis.set_geometry_transform(mesh, self.parent.vis.get_geometry_transform(mesh))
self.parent.eye_windows[self.eye] = self
[docs]
def on_close(self):
"""
Small function to delete the window from the parent class
"""
for menu_id in self.MENU_IDS[self.eye]:
self.parent.file_menu.set_enabled(menu_id, False)
del self.parent.eye_windows[self.eye]
return True
[docs]
def on_mouse(self, event):
"""
Small function to ignore mouse events
:param event: Mouse event
:type event: gui.MouseEvent
"""
return self.scene.EventCallbackResult.CONSUMED
[docs]
def get_image(self):
self.vis.scene.render_to_image(self.save_image)
[docs]
def save_images(self):
"""
Function to save stream of images to file
"""
p = os.path.join(self.parent.file_dir, "..", "images", self.eye, "RGB")
if not os.path.exists(p):
os.makedirs(p)
p = os.path.join(p, str(self.img_counter)+".png")
self.img_counter += 1
o3d.io.write_image(p, self.last_image)
[docs]
def save_image(self, im):
"""
Callback to get images from open3d
:param im: the image to be saves
:type im: o3d.geometry.Image
"""
self.last_image = im