def get_image_depth_and_mask(scene: pyrender.Scene, scene_setup_loader: DiceScene.SceneSetupLoader, width: int, height: int, keep_nodes_in_scene: bool): """Renders an image ggiven a scene and seetup, along with the depth and segmentation mask labelling each die.""" r = pyrender.OffscreenRenderer(width, height) color_bg, depth_bg = r.render(scene) depth_nodes = [] for node in scene_setup_loader.dice_nodes: scene.add_node(node) color_node, depth_node = r.render(scene) depth_nodes.append(depth_node) scene.remove_node(node) scene_setup_loader.add_loaded_to_scene(scene) color_final, depth_final = r.render(scene) if not keep_nodes_in_scene: scene_setup_loader.remove_nodes_from_scene(scene) #Initialize labels of pixels to -1 (for background) labels_mask = np.ones((height, width), dtype=np.int8) * -1 for index, depth_for_node in enumerate(depth_nodes): depth_not_background = np.not_equal(depth_bg, depth_for_node) depth_at_foreground = np.equal(depth_final, depth_for_node) depth_at_dice = np.logical_and(depth_not_background, depth_at_foreground) labels_mask[depth_at_dice] = index return color_final, depth_final, labels_mask
def render_big_gallery(results_dir, nb=30, pts_colors=[0.5, 0.5, 0.5], draw_text=False): ''' pts_colors: [0,0,0] return np array of a big image ''' cam = PerspectiveCamera(yfov=(YFOV)) cam_pose = CAM_POSE point_l = PointLight(color=np.ones(3), intensity=POINT_LIGHT_INTENSITY) scene = Scene(bg_color=np.array([1, 1, 1, 0])) # cam and light _ = scene.add(cam, pose=cam_pose) _ = scene.add(point_l, pose=cam_pose) input_ply_filenames = get_all_filnames(results_dir, nb) r = OffscreenRenderer(viewport_width=640 * 2, viewport_height=480 * 2, point_size=POINT_SIZE) pc_pose = PC_POSE images = [] for _, input_pf in enumerate(input_ply_filenames): input_pc = read_ply_xyz(input_pf) colors = np.array(pts_colors) colors = np.tile(colors, (input_pc.shape[0], 1)) input_pc_node = add_point_cloud_mesh_to_scene(input_pc, scene, pc_pose, colors) renderred_color, _ = r.render(scene) scene.remove_node(input_pc_node) if draw_text: im_here = Image.fromarray(renderred_color) d = ImageDraw.Draw(im_here) fnt = ImageFont.truetype( font='/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf', size=100) d.text((0, 0), input_pf.split('/')[-1], fill=(0, 0, 0, 255), font=fnt) renderred_color = np.array(im_here) images.append(renderred_color) big_gallery = np.concatenate(images, axis=0) r.delete() return big_gallery
def render_big_gallery_overlay(dir_1, dir_2, pts_color_1=[0.5, 0.5, 0.5], pts_color_2=[0.5, 0.5, 0.5], nb=30): ''' return np array of a big image ''' cam = PerspectiveCamera(yfov=(YFOV)) cam_pose = CAM_POSE point_l = PointLight(color=np.ones(3), intensity=POINT_LIGHT_INTENSITY) scene = Scene(bg_color=np.array([1, 1, 1, 0])) # cam and light _ = scene.add(cam, pose=cam_pose) _ = scene.add(point_l, pose=cam_pose) input_ply_filenames_1 = get_all_filnames(dir_1, nb) input_ply_filenames_2 = get_all_filnames(dir_2, nb) r = OffscreenRenderer(viewport_width=640 * 2, viewport_height=480 * 2, point_size=POINT_SIZE) pc_pose = PC_POSE images = [] for idx, input_pf in enumerate(input_ply_filenames_1): input_pc_1 = read_ply_xyz(input_pf) input_pc_2 = read_ply_xyz(input_ply_filenames_2[idx]) color_1 = np.array(pts_color_1) color_1 = np.tile(color_1, (input_pc_1.shape[0], 1)) color_2 = np.array(pts_color_2) color_2 = np.tile(color_2, (input_pc_2.shape[0], 1)) input_pc_node_1 = add_point_cloud_mesh_to_scene( input_pc_1, scene, pc_pose, color_1) input_pc_node_2 = add_point_cloud_mesh_to_scene( input_pc_2, scene, pc_pose, color_2) renderred_color, _ = r.render(scene) scene.remove_node(input_pc_node_1) scene.remove_node(input_pc_node_2) images.append(renderred_color) big_gallery = np.concatenate(images, axis=0) r.delete() return big_gallery
class HVTPClient: def __init__(self): """ """ self.scene = Scene(ambient_light=np.array([0.02, 0.02, 0.02, 1.0])) self.scene.add(PointLight(color=[0.5, 0.2, 0.3], intensity=2.0)) self.scene.add( SpotLight(color=[0.1, 0.6, 0.3], intensity=2.0, innerConeAngle=0.05, outerConeAngle=0.5)) self.scene.add( DirectionalLight(color=[0.33, 0.33, 0.33], intensity=2.0)) self.root_node = None # self.scene.add_node(self.root_node) def create(self, ip, port): """ """ self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.client_socket.connect((ip, port)) self.client_socket.setblocking(True) def destroy(self): """ """ self.client_socket.shutdown(socket.SHUT_RDWR) self.client_socket.close() def run(self): """ """ self.viewer = Viewer(self.scene, use_raymond_lighting=False, cull_faces=False, run_in_thread=True) threading.Thread(target=self.accept_messages).start() def accept_messages(self): """ Polls our socket for changes that have been sent by the server. """ print("Starting message acceptor") while True: try: packet_header = self.client_socket.recv( HVTPConstants.HVTP_HEADER_SIZE_IN_BYTES) print("Received!") if len(packet_header) == 0: print("Connection closed by server") sys.exit() # Read the header i = 0 offs = 4 hvtp_magic = packet_header[(i * offs):(offs + i * offs)].decode( HVTPConstants.HVTP_ENCODING) i += 1 hvtp_version = int.from_bytes( packet_header[(i * offs):(offs + i * offs)], byteorder='big', signed=False) i += 1 hvtp_length = int.from_bytes( packet_header[(i * offs):(offs + i * offs)], byteorder='big', signed=False) i += 1 hvtp_type = packet_header[(i * offs):(offs + i * offs)].decode( HVTPConstants.HVTP_ENCODING) print(hvtp_magic) print(hvtp_version) print(hvtp_length) print(hvtp_type) # Read the rest of the packet payload = self.client_socket.recv(hvtp_length) loaded_payload = trimesh.load(BytesIO(payload), file_type='glb') mesh = Mesh.from_trimesh(list(loaded_payload.dump())) # Add to the scene self.add_to_scene(mesh) except IOError as e: # This is normal on non blocking connections - when there are no incoming data error is going to be raised # Some operating systems will indicate that using AGAIN, and some using WOULDBLOCK error code # We are going to check for both - if one of them - that's expected, means no incoming data, continue as normal # If we got different error code - something happened if e.errno != errno.EAGAIN and e.errno != errno.EWOULDBLOCK: print('Reading error: {}'.format(str(e))) sys.exit() except Exception as e: # Any other exception - something happened, exit print('Reading error: {}'.format(str(e))) sys.exit() def add_to_scene(self, mesh): """ """ # Grab the viewer mutex self.viewer.render_lock.acquire() # Remove anything before and insert the mesh into the scene-graph if self.root_node != None: self.scene.remove_node(self.root_node) self.root_node = self.scene.add(mesh) # Release the viewer mutex self.viewer.render_lock.release()
def remove_nodes_from_scene(self, scene: pyrender.Scene): """Remove nodes in scene corresponding to internal list of dice nodes.""" for mesh_node in self.dice_nodes: if scene.has_node(mesh_node): scene.remove_node(mesh_node)
def test_scenes(): # Basics s = Scene() assert np.allclose(s.bg_color, np.ones(4)) assert np.allclose(s.ambient_light, np.zeros(3)) assert len(s.nodes) == 0 assert s.name is None s.name = 'asdf' s.bg_color = None s.ambient_light = None assert np.allclose(s.bg_color, np.ones(4)) assert np.allclose(s.ambient_light, np.zeros(3)) assert s.nodes == set() assert s.cameras == set() assert s.lights == set() assert s.point_lights == set() assert s.spot_lights == set() assert s.directional_lights == set() assert s.meshes == set() assert s.camera_nodes == set() assert s.light_nodes == set() assert s.point_light_nodes == set() assert s.spot_light_nodes == set() assert s.directional_light_nodes == set() assert s.mesh_nodes == set() assert s.main_camera_node is None assert np.all(s.bounds == 0) assert np.all(s.centroid == 0) assert np.all(s.extents == 0) assert np.all(s.scale == 0) # From trimesh scene tms = trimesh.load('tests/data/WaterBottle.glb') s = Scene.from_trimesh_scene(tms) assert len(s.meshes) == 1 assert len(s.mesh_nodes) == 1 # Test bg color formatting s = Scene(bg_color=[0, 1.0, 0]) assert np.allclose(s.bg_color, np.array([0.0, 1.0, 0.0, 1.0])) # Test constructor for nodes n1 = Node() n2 = Node() n3 = Node() nodes = [n1, n2, n3] s = Scene(nodes=nodes) n1.children.append(n2) s = Scene(nodes=nodes) n3.children.append(n2) with pytest.raises(ValueError): s = Scene(nodes=nodes) n3.children = [] n2.children.append(n3) n3.children.append(n2) with pytest.raises(ValueError): s = Scene(nodes=nodes) # Test node accessors n1 = Node() n2 = Node() n3 = Node() nodes = [n1, n2] s = Scene(nodes=nodes) assert s.has_node(n1) assert s.has_node(n2) assert not s.has_node(n3) # Test node poses for n in nodes: assert np.allclose(s.get_pose(n), np.eye(4)) with pytest.raises(ValueError): s.get_pose(n3) with pytest.raises(ValueError): s.set_pose(n3, np.eye(4)) tf = np.eye(4) tf[:3, 3] = np.ones(3) s.set_pose(n1, tf) assert np.allclose(s.get_pose(n1), tf) assert np.allclose(s.get_pose(n2), np.eye(4)) nodes = [n1, n2, n3] tf2 = np.eye(4) tf2[:3, :3] = np.diag([-1, -1, 1]) n1.children.append(n2) n1.matrix = tf n2.matrix = tf2 s = Scene(nodes=nodes) assert np.allclose(s.get_pose(n1), tf) assert np.allclose(s.get_pose(n2), tf.dot(tf2)) assert np.allclose(s.get_pose(n3), np.eye(4)) n1 = Node() n2 = Node() n3 = Node() n1.children.append(n2) s = Scene() s.add_node(n1) with pytest.raises(ValueError): s.add_node(n2) s.set_pose(n1, tf) assert np.allclose(s.get_pose(n1), tf) assert np.allclose(s.get_pose(n2), tf) s.set_pose(n2, tf2) assert np.allclose(s.get_pose(n2), tf.dot(tf2)) # Test node removal n1 = Node() n2 = Node() n3 = Node() n1.children.append(n2) n2.children.append(n3) s = Scene(nodes=[n1, n2, n3]) s.remove_node(n2) assert len(s.nodes) == 1 assert n1 in s.nodes assert len(n1.children) == 0 assert len(n2.children) == 1 s.add_node(n2, parent_node=n1) assert len(n1.children) == 1 n1.matrix = tf n3.matrix = tf2 assert np.allclose(s.get_pose(n3), tf.dot(tf2)) # Now test ADD function s = Scene() m = Mesh([], name='m') cp = PerspectiveCamera(yfov=2.0) co = OrthographicCamera(xmag=1.0, ymag=1.0) dl = DirectionalLight() pl = PointLight() sl = SpotLight() n1 = s.add(m, name='mn') assert n1.mesh == m assert len(s.nodes) == 1 assert len(s.mesh_nodes) == 1 assert n1 in s.mesh_nodes assert len(s.meshes) == 1 assert m in s.meshes assert len(s.get_nodes(node=n2)) == 0 n2 = s.add(m, pose=tf) assert len(s.nodes) == len(s.mesh_nodes) == 2 assert len(s.meshes) == 1 assert len(s.get_nodes(node=n1)) == 1 assert len(s.get_nodes(node=n1, name='mn')) == 1 assert len(s.get_nodes(name='mn')) == 1 assert len(s.get_nodes(obj=m)) == 2 assert len(s.get_nodes(obj=m, obj_name='m')) == 2 assert len(s.get_nodes(obj=co)) == 0 nsl = s.add(sl, name='sln') npl = s.add(pl, parent_name='sln') assert nsl.children[0] == npl ndl = s.add(dl, parent_node=npl) assert npl.children[0] == ndl nco = s.add(co) ncp = s.add(cp) assert len(s.light_nodes) == len(s.lights) == 3 assert len(s.point_light_nodes) == len(s.point_lights) == 1 assert npl in s.point_light_nodes assert len(s.spot_light_nodes) == len(s.spot_lights) == 1 assert nsl in s.spot_light_nodes assert len(s.directional_light_nodes) == len(s.directional_lights) == 1 assert ndl in s.directional_light_nodes assert len(s.cameras) == len(s.camera_nodes) == 2 assert s.main_camera_node == nco s.main_camera_node = ncp s.remove_node(ncp) assert len(s.cameras) == len(s.camera_nodes) == 1 assert s.main_camera_node == nco s.remove_node(n2) assert len(s.meshes) == 1 s.remove_node(n1) assert len(s.meshes) == 0 s.remove_node(nsl) assert len(s.lights) == 0 s.remove_node(nco) assert s.main_camera_node is None s.add_node(n1) s.clear() assert len(s.nodes) == 0 # Trigger final errors with pytest.raises(ValueError): s.main_camera_node = None with pytest.raises(ValueError): s.main_camera_node = ncp with pytest.raises(ValueError): s.add(m, parent_node=n1) with pytest.raises(ValueError): s.add(m, name='asdf') s.add(m, name='asdf') s.add(m, parent_name='asdf') with pytest.raises(ValueError): s.add(m, parent_name='asfd') with pytest.raises(TypeError): s.add(None) s.clear() # Test bounds m1 = Mesh.from_trimesh(trimesh.creation.box()) m2 = Mesh.from_trimesh(trimesh.creation.box()) m3 = Mesh.from_trimesh(trimesh.creation.box()) n1 = Node(mesh=m1) n2 = Node(mesh=m2, translation=[1.0, 0.0, 0.0]) n3 = Node(mesh=m3, translation=[0.5, 0.0, 1.0]) s.add_node(n1) s.add_node(n2) s.add_node(n3) assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [1.5, 0.5, 1.5]]) s.clear() s.add_node(n1) s.add_node(n2, parent_node=n1) s.add_node(n3, parent_node=n2) assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [2.0, 0.5, 1.5]]) tf = np.eye(4) tf[:3, 3] = np.ones(3) s.set_pose(n3, tf) assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [2.5, 1.5, 1.5]]) s.remove_node(n2) assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]]) s.clear() assert np.allclose(s.bounds, 0.0)
choice = input("Generate dataset? [y/n]: ") iterations = 5 mode = CameraPose.NONE csvfile = './outputs/flowers_dataset.csv' csvMode = 'w' # 'a' => append / 'w' => write (overwrites existing csv file) meshes = [(flower_mesh, 'flower'), (stem_mesh, 'stem'), (center_mesh, 'center')] skip_default_view = False salt = '001' offset = 0 skip_sub_mesh = {name: False for (mesh, name) in meshes} if flower_type in no_stem_flowers: skip_sub_mesh['stem'] = True if choice == 'y': # Remove camera that was used as preview from the scene scene.remove_node(cam_node) generate( scene=scene, renderer=r, flags=flags, alpha=alpha, beta=beta, camera=camera, iterations=iterations, flower_pose=flower_pose, meshes=meshes, light_intensitites = varying_light_intensities, skip_sub_mesh= skip_sub_mesh, at_x=at_x, at_y=at_y, at_z=at_z, skip_default_view=skip_default_view, offset=offset, distance=distance, mode=mode, csvfile=csvfile, csvMode=csvMode, flower_name=file_name, salt=salt ) print("Finished")
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()
except Exception as e: write_log("Fatal: Render error!") write_log(e, verbose=False) # Drop current row row_mat = None break assert color.shape[1] == g_single_viewport_size[0] and color.shape[0] == g_single_viewport_size[1], "Fatal: Inconsistent color map size" if row_mat is None: row_mat = color else: assert row_mat.shape[0] == color.shape[0], "Fatal: Inconsistent cell size" row_mat = np.hstack((row_mat, color)) for node in mesh_nodes: scene.remove_node(node) if row_mat is None: write_log("Fatal: Failed to render the current row!") continue if image_mat is None: image_mat = row_mat else: assert image_mat.shape[1] == row_mat.shape[1], "Fatal: Inconsistent row size" image_mat = np.vstack((image_mat, row_mat)) if image_mat is None: write_log("Fatal: Failed to render the current image!") continue
viewer_flags = { 'mouse_pressed': False, 'rotate': False, 'rotate_rate': np.pi / 6.0, 'rotate_axis': np.array([0.0, 1.0, 0.0]), 'view_center': np.array([0.0, 0.0, 0.0]), 'record': False, 'use_raymond_lighting': False, 'use_direct_lighting': False, 'lighting_intensity': 3.0, 'use_perspective_cam': True, 'save_directory': '/home/wanglz14/Desktop/video/', 'window_title': 'FOV Human', 'refresh_rate': 10.0, 'fullscreen': False, 'show_world_axis': False, 'show_mesh_axes': False, 'caption': None } v = Viewer(scene, viewport_size=(800, 800), render_flags=render_flags, viewer_flags=viewer_flags, run_in_thread=True) mesh = trimesh.load('../results/test/01.ply') points_mesh = Mesh.from_trimesh(mesh, smooth=True) v._render_lock.acquire() scene.remove_node(human_node) human_node = scene.add(points_mesh) v._render_lock.release()