def get_body_name(self): mjcf_model = mjcf.from_path(self.source_xml_path) bodies = mjcf_model.worldbody.find_all('body') body_name_list = [] for body in bodies: body_name_list.append(body.name) return body_name_list
def get_act_name(self): mjcf_model = mjcf.from_path(self.source_xml_path) motors = mjcf_model.actuator.motor act_name_list = [] for motor in motors: act_name_list.append(motor.name) return act_name_list
def domain_randomizer( self, xml_name, out_dir, bounds={ "mass": [0.8, 1.2], "inertia": [0.8, 1.2], "friction": [0.8, 1.2], "damping": [0.8, 1.2] }): """ Change global joint damping and friction Assume that hardware shares same actuator """ mjcf_model = mjcf.from_path(self.source_xml_path) mjcf_model.default.joint.damping *= np.random.uniform( bounds["damping"][0], bounds["damping"][1]) mjcf_model.default.joint.frictionloss *= np.random.uniform( bounds["friction"][0], bounds["friction"][1]) #Assume mass and inertia of all boides change sameway bodies = mjcf_model.worldbody.find_all('body') mass_rand = np.random.uniform(bounds["mass"][0], bounds["mass"][1]) inertia_rand = np.random.uniform(bounds["inertia"][0], bounds["inertia"][1]) for body in bodies: body.inertial.mass *= mass_rand body.inertial.fullinertia *= inertia_rand mjcf.export_with_assets(mjcf_model, out_dir=out_dir, out_file_name=xml_name) mujoco_env.MujocoEnv.__init__(self, out_dir + xml_name, 5) utils.EzPickle.__init__(self)
def test_exceptions(self): out_dir = self.create_tempdir().full_path mjcf_model = mjcf.from_path(_TEST_MODEL_WITH_ASSETS) with six.assertRaisesRegex(self, ValueError, 'must end with \'.xml\''): mjcf.export_with_assets(mjcf_model, out_dir, out_file_name='invalid_extension.png')
def _build(self, name=None, use_pinch_site_as_tcp=False): """Initializes the JacoHand. Args: name: String, the name of this robot. Used as a prefix in the MJCF name name attributes. use_pinch_site_as_tcp: (optional) A boolean, if `True` the pinch site will be used as the tool center point. If `False` the grip site is used. """ self._mjcf_root = mjcf.from_path(_JACO_HAND_XML_PATH) if name: self._mjcf_root.model = name # Find MJCF elements that will be exposed as attributes. self._bodies = self.mjcf_model.find_all('body') self._tool_center_point = self._mjcf_root.find( 'site', _PINCH_SITE if use_pinch_site_as_tcp else _GRIP_SITE) self._joints = self._mjcf_root.find_all('joint') self._hand_geoms = list(self._mjcf_root.find('body', _HAND_BODY).geom) self._finger_geoms = [ geom for geom in self._mjcf_root.find_all('geom') if geom.name and geom.name.startswith('finger') ] self._grip_site = self._mjcf_root.find('site', _GRIP_SITE) self._pinch_site = self._mjcf_root.find('site', _PINCH_SITE) # Add actuators. self._finger_actuators = [ _add_velocity_actuator(joint) for joint in self._joints ]
def _build(self, params=None, name='walker', initializer=None): self.params = params self._mjcf_root = mjcf.from_path(_XML_PATH) if name: self._mjcf_root.model = name self.body_sites = [] super()._build(initializer=initializer)
def _build(self, name=None): """Initializes this arena. Args: name: (optional) A string, the name of this arena. If `None`, use the model name defined in the MJCF file. """ self._mjcf_root = mjcf.from_path(_ARENA_XML_PATH) if name: self._mjcf_root.model = name
def _build(self, name='walker', marker_rgba=None, camera_control=False, initializer=None, add_ears=False, camera_height=None): """Build a JumpingBallWithHead. Args: name: name of the walker. marker_rgba: RGBA value set to walker.marker_geoms to distinguish between walkers (in multi-agent setting). camera_control: If `True`, the walker exposes two additional actuated degrees of freedom to control the egocentric camera height and tilt. initializer: (Optional) A `WalkerInitializer` object. add_ears: a boolean. Same as the nose above but the red/blue balls are placed to the left/right of the agent. Better for egocentric vision. camera_height: A float specifying the height of the camera, or `None` if the camera height should be left as specified in the XML model. """ super()._build(initializer=initializer) self._mjcf_root = self._mjcf_root = mjcf.from_path(self._xml_path) if name: self._mjcf_root.model = name if camera_height is not None: self._mjcf_root.find('body', 'egocentric_camera').pos[2] = camera_height if add_ears: # Large ears head = self._mjcf_root.find('body', 'head_body') head.add('site', type='sphere', size=(.26, ), pos=(.22, 0, 0), rgba=(.7, 0, 0, 1)) head.add('site', type='sphere', size=(.26, ), pos=(-.22, 0, 0), rgba=(0, 0, .7, 1)) # Set corresponding marker color if specified. if marker_rgba is not None: for geom in self.marker_geoms: geom.set_attributes(rgba=marker_rgba) self._root_joints = None self._camera_control = camera_control if not camera_control: for name in ('camera_height', 'camera_tilt'): self._mjcf_root.find('actuator', name).remove() self._mjcf_root.find('joint', name).remove()
def _build(self, name='walker', marker_rgba=None, initializer=None): self._mjcf_root = mjcf.from_path(self._xml_path) if name: self._mjcf_root.model = name # Set corresponding marker color if specified. if marker_rgba is not None: for geom in self.marker_geoms: geom.set_attributes(rgba=marker_rgba) self._actuator_order = np.argsort(_CMU_MOCAP_JOINTS) self._inverse_order = np.argsort(self._actuator_order) super(_CMUHumanoidBase, self)._build(initializer=initializer)
def _build(self, name='walker', marker_rgba=None, include_face=False, initializer=None): self._mjcf_root = mjcf.from_path(self._xml_path) if name: self._mjcf_root.model = name # Set corresponding marker color if specified. if marker_rgba is not None: for geom in self.marker_geoms: geom.set_attributes(rgba=marker_rgba) self._actuator_order = np.argsort(_CMU_MOCAP_JOINTS) self._inverse_order = np.argsort(self._actuator_order) super()._build(initializer=initializer) if include_face: head = self._mjcf_root.find('body', 'head') head.add('geom', type='capsule', name='face', size=(0.065, 0.014), pos=(0.000341465, 0.048184, 0.01), quat=(0.717887, 0.696142, -0.00493334, 0), mass=0., contype=0, conaffinity=0) face_forwardness = head.pos[1] - .02 head_geom = self._mjcf_root.find('geom', 'head') nose_size = head_geom.size[0] / 4.75 face = head.add('body', name='face', pos=(0.0, 0.039, face_forwardness)) face.add('geom', type='capsule', name='nose', size=(nose_size, 0.01), pos=(0.0, 0.0, 0.0), quat=(1, 0.7, 0, 0), mass=0., contype=0, conaffinity=0, group=_WALKER_INVIS_GROUP)
def _build(self, name=None): """Initializes the JacoArm. Args: name: String, the name of this robot. Used as a prefix in the MJCF name name attributes. """ self._mjcf_root = mjcf.from_path(_JACO_ARM_XML_PATH) if name: self._mjcf_root.model = name # Find MJCF elements that will be exposed as attributes. self._joints = [self._mjcf_root.find('joint', name) for name in _ALL_JOINTS] self._wrist_site = self._mjcf_root.find('site', _WRIST_SITE) self._bodies = self.mjcf_model.find_all('body') # Add actuators. self._actuators = [_add_velocity_actuator(joint) for joint in self._joints] # Add torque sensors. self._joint_torque_sensors = [ _add_torque_sensor(joint) for joint in self._joints]
def test_export_model(self, xml_path, model_name): """Save processed MJCF model.""" out_dir = self.create_tempdir().full_path mjcf_model = mjcf.from_path(xml_path) mjcf.export_with_assets_as_zip(mjcf_model, out_dir=out_dir, model_name=model_name) # Read the .zip file in the output directory. # Check that the only directory is named `model_name`/, and put all the # contents under any directory in a dict a directory in a dict. zip_file_contents = {} zip_filename = os.path.join(out_dir, (model_name + '.zip')) self.assertTrue(zipfile.is_zipfile(zip_filename)) with zipfile.ZipFile(zip_filename, 'r') as zip_file: for zip_info in zip_file.infolist(): # Note: zip_info.is_dir() is not Python 2 compatible, but directories # inside a ZipFile are guaranteed to end with '/'. if not zip_info.filename.endswith(os.path.sep): with zip_file.open(zip_info.filename) as f: zip_file_contents[zip_info.filename] = f.read() else: self.assertEqual(os.path.join(model_name), zip_info.filename) # Check that the output directory contains an XML file of the correct name. xml_filename = os.path.join(model_name, model_name) + '.xml' self.assertIn(xml_filename, zip_file_contents) # Check that its contents match the output of `mjcf_model.to_xml_string()`. xml_contents = util.to_native_string( zip_file_contents.pop(xml_filename)) expected_xml_contents = mjcf_model.to_xml_string() self.assertEqual(xml_contents, expected_xml_contents) # Check that the other files in the directory match the contents of the # model's `assets` dict. assets = mjcf_model.get_assets() for asset_name, asset_contents in assets.items(): self.assertEqual( asset_contents, zip_file_contents[os.path.join(model_name, asset_name)])
def _build(self, name='walker', marker_rgba=None, initializer=None): """Build an PlanarWalker walker. Args: name: name of the walker. marker_rgba: (Optional) color the walker's legs with marker_rgba. initializer: (Optional) A `WalkerInitializer` object. """ super()._build(initializer=initializer) self._mjcf_root = mjcf.from_path(_XML_PATH) if name: self._mjcf_root.model = name # Set corresponding marker color if specified. if marker_rgba is not None: for geom in self.marker_geoms: geom.set_attributes(rgba=marker_rgba) # Initialize previous action. self._prev_action = np.zeros(shape=self.action_spec.shape, dtype=self.action_spec.dtype)
def test_export_model(self, xml_path, out_xml_name): """Save processed MJCF model.""" out_dir = self.create_tempdir().full_path mjcf_model = mjcf.from_path(xml_path) mjcf.export_with_assets(mjcf_model, out_dir=out_dir, out_file_name=out_xml_name) # Read the files in the output directory and put their contents in a dict. out_dir_contents = {} for filename in os.listdir(out_dir): with open(os.path.join(out_dir, filename), 'rb') as f: out_dir_contents[filename] = f.read() # Check that the output directory contains an XML file of the correct name. self.assertIn(out_xml_name, out_dir_contents) # Check that its contents match the output of `mjcf_model.to_xml_string()`. xml_contents = util.to_native_string( out_dir_contents.pop(out_xml_name)) expected_xml_contents = mjcf_model.to_xml_string() self.assertEqual(xml_contents, expected_xml_contents) # Check that the other files in the directory match the contents of the # model's `assets` dict. assets = mjcf_model.get_assets() self.assertDictEqual(out_dir_contents, assets) # Check that we can construct an MjModel instance using the path to the # output file. from_exported = wrapper.MjModel.from_xml_path( os.path.join(out_dir, out_xml_name)) # Check that this model is identical to one constructed directly from # `mjcf_model`. from_mjcf = wrapper.MjModel.from_xml_string(expected_xml_contents, assets=assets) self.assertEqual(from_exported.to_bytes(), from_mjcf.to_bytes())
def _build(self, name='j2s7'): """Initializes the JacoArm. Args: name: String, the name of this robot. Used as a prefix in the MJCF name name attributes. """ self._jaco_arm_xml_path = os.path.join(assets_path.KINOVA_ROOT, 'jaco_arm_{}.xml'.format(name)) self.num_joints = int(name[-1]) if self.num_joints == 6: self._large_joints = ('joint_1', 'joint_2', 'joint_3') self._small_joints = ('joint_4', 'joint_5', 'joint_6') elif self.num_joints == 7: self._large_joints = ('joint_1', 'joint_2', 'joint_3', 'joint_4') self._small_joints = ('joint_5', 'joint_6', 'joint_7') else: raise NotImplementedError self._all_joints = self._large_joints + self._small_joints self._mjcf_root = mjcf.from_path(self._jaco_arm_xml_path) if name: self._mjcf_root.model = name # Find MJCF elements that will be exposed as attributes. self._joints = [ self._mjcf_root.find('joint', name) for name in self._all_joints ] self._wrist_site = self._mjcf_root.find('site', _WRIST_SITE) self._bodies = self.mjcf_model.find_all('body') # Add actuators. self._actuators = [ self._add_velocity_actuator(joint) for joint in self._joints ] # Add torque sensors. self._joint_torque_sensors = [ _add_torque_sensor(joint) for joint in self._joints ]
def __init__(self, path="floor.xml"): self.mjcf_model = mjcf.from_path(path) pass
def test_default_model_filename(self): out_dir = self.create_tempdir().full_path mjcf_model = mjcf.from_path(_TEST_MODEL_WITH_ASSETS) mjcf.export_with_assets_as_zip(mjcf_model, out_dir, model_name=None) expected_name = mjcf_model.model + '.zip' self.assertTrue(os.path.isfile(os.path.join(out_dir, expected_name)))
def scene_to_npz(self, mesh_obj): integrated_xml = generate_integrated_xml(self.env_base_xml_file, mesh_obj.obj_xml_file, scale=mesh_obj.scale) abs_integrated_xml_path = os.path.join(self.env_xml_file_root, integrated_xml) visual_mjcf = mjcf.from_path( abs_integrated_xml_path ) #cannot have "include" in xml when loading with mjcf # camera pose is related the object pose? # for visualization if self.config.render: model = mujoco_py.load_model_from_path(abs_integrated_xml_path) sim = mujoco_py.MjSim(model) viewer = mujoco_py.MjViewer(sim) for _ in range(5000): sim.step() for _ in range(20): viewer.render() print(mesh_obj.name, self.config.camera_lookat_pos) # add cameras camera_positions, camera_quats = pcp_utils.cameras.generate_new_cameras_hemisphere( radius=self.config.camera_radius, lookat_point=self.config.camera_lookat_pos, pitch=self.config.camera_pitch, yaw=self.config.camera_yaw) # add all the cameras to the environtmnet for i, (pos, quat) in enumerate(zip(camera_positions, camera_quats)): visual_mjcf.worldbody.add('camera', name=f'vis_cam:{i}', pos=pos, quat=quat, fovy=self.config.camera_fov_y) visual_mjcf.worldbody.add('camera', name='ref_cam', pos=[0, -0.3, 0.5], zaxis=[0, -1, 0], fovy=self.config.camera_fov_y) # start simulate physics = mjcf.Physics.from_mjcf_model(visual_mjcf) physics.forward() for step_id in range(2000): physics.step() object_xpos = physics.data.xpos[physics.model.name2id( 'object0', 'body')] #print(f'object final location', object_xpos) object_xmat = physics.data.xmat[physics.model.name2id( 'object0', 'body')] #print(f'object final qpos', object_xmat) ep_imgs = list() ep_depths = list() ep_pix_T_cams = list() ep_camR_T_camXs = list() for i in range(len(camera_positions)): img = physics.render(self.config.camera_img_height, self.config.camera_img_width, camera_id=f'vis_cam:{i}') depth = physics.render(self.config.camera_img_height, self.config.camera_img_width, camera_id=f'vis_cam:{i}', depth=True) assert img.shape[ 0] == self.config.camera_img_height, "color img height is wrong" assert img.shape[ 1] == self.config.camera_img_width, "color img width is wrong" assert depth.shape[ 0] == self.config.camera_img_height, "depth img height is wrong" assert depth.shape[ 1] == self.config.camera_img_width, "depth img width is wrong" #if self.config.save_image: # imageio.imwrite(f"tmp/img_{i}.png", img) # imageio.imwrite(f"tmp/depth_{i}.png", np.minimum(depth, 5)) pix_T_cams = pcp_utils.cameras.get_intrinsics( self.config.camera_fov_y, self.config.camera_img_width, self.config.camera_img_height) origin_T_camX = pcp_utils.cameras.dm_get_extrinsics( physics, physics.model.name2id(f'vis_cam:{i}', 'camera')) camR_T_camX = np.dot(self.adam_T_origin, origin_T_camX) # camR should be the top of the table #camR_T_camX = ep_imgs.append(img) ep_depths.append(depth) ep_pix_T_cams.append(pix_T_cams) # this extrinsics is in mujoco frame ep_camR_T_camXs.append(camR_T_camX) img = physics.render(self.config.camera_img_height, self.config.camera_img_width, camera_id='ref_cam') depth = physics.render(self.config.camera_img_height, self.config.camera_img_width, camera_id='ref_cam', depth=True) pix_T_cams = pcp_utils.cameras.get_intrinsics( self.config.camera_fov_y, self.config.camera_img_width, self.config.camera_img_height) origin_T_camX = pcp_utils.cameras.dm_get_extrinsics( physics, physics.model.name2id('ref_cam', 'camera')) camR_T_camX = np.dot(self.adam_T_origin, origin_T_camX) if self.config.save_image: imageio.imwrite(f"tmp/img_ref.png", img) ep_imgs.append(img) ep_depths.append(depth) ep_pix_T_cams.append(pix_T_cams) ep_camR_T_camXs.append(camR_T_camX) composed_xmat = np.eye(4, dtype=np.float32) composed_xmat[:3, 3] = object_xpos composed_xmat[:3, :3] = np.reshape(object_xmat, [3, 3]) composed_xmat = np.dot(self.adam_T_origin, composed_xmat) object_xpos_adam = composed_xmat[:3, 3] #np.dot(self.adam_T_origin, pcp_utils.geom.pts_addone(np.reshape(object_xpos, [1, 3])).T)[:3, 0] object_xmat_adam = composed_xmat[:3, : 3] #np.dot(self.adam_T_origin[:3, :3], np.reshape(object_xmat, [3, 3])) bbox_points_from_mesh_adam = pcp_utils.np_vis.compute_bounding_box_from_obj_xml( mesh_obj.obj_xml_file, object_xpos_adam, object_xmat_adam, scale=mesh_obj.scale) if self.config.visualize_bbox: _, xyz_camRs, _ = pcp_utils.np_vis.unproject_depth( ep_depths, ep_pix_T_cams, ep_camR_T_camXs, camR_T_origin=None, #np.linalg.inv(self.origin_T_adam), clip_radius=5.0, do_vis=False) all_xyz_camR = np.concatenate(xyz_camRs, axis=0) object_pcd = pcp_utils.np_vis.get_pcd_object(all_xyz_camR, clip_radius=2.0) # transform object xpos and xmat to the adam coordinate (x right, y downs) bbox_lineset_from_mesh_adam = pcp_utils.np_vis.make_lineset( bbox_points_from_mesh_adam) o3d.visualization.draw_geometries( [object_pcd, bbox_lineset_from_mesh_adam]) #o3d.visualization.draw_geometries([object_pcd, bbox_lineset_from_mesh]) import ipdb ipdb.set_trace() # create a dictionary to save the data and ship it !! save_dict = AttrDict() """ last view is reference frame rgb_camXs: nviews x height x width x 3 depth_camXs: nviews x height x width pix_T_cams: nviews x 3 x 3 camR_T_camXs: nviews x 4 x 4 bbox_camR: 8 x 3 cluster_id: 'string' """ save_dict.rgb_camXs = np.stack(ep_imgs) save_dict.depth_camXs = np.stack(ep_depths) save_dict.pix_T_cams = np.stack(ep_pix_T_cams) save_dict.camR_T_camXs = np.stack(ep_camR_T_camXs) save_dict.bbox_camR = bbox_points_from_mesh_adam save_dict.cluster_id = mesh_obj.cluster_id scene_name = mesh_obj.name if self.config.save_image: imageio.imwrite(f"tmp/{scene_name}_nviews.png", np.concatenate(ep_imgs, axis=1)) vis_save_path = os.path.join(self.output_dir, f"visual_data_{scene_name}.npy") np.save(vis_save_path, save_dict) self.all_files.append( os.path.join(self.output_dir_exclud_root, f"visual_data_{scene_name}.npy")) print('---- done ----')
def create_env_and_collect_data(self, object_file, save_imgs, collect_visual_data, collect_touch_trajectories): """ Create the new environment with the specified object from xml_skeleton file object_file: path to the model file save_imgs : if True save the intermediate images from each view collect_visual_data : if True copies the data_mjcf to visual_mjcf and collects vis data collect_touch_data : if True copies the data_mjcf to touch_mjcf and collects touch data Function: saves the visual data, and the visual mjcf by converting it to xml saves the touch data, and the touch mjcf by converting it to xml """ self.okay = True self.object_name, stl_models_path = self._preprocess(object_file) stl_files_path = [osp.join(stl_models_path, stl_file) for stl_file in os.listdir(stl_models_path)] if len(stl_files_path) > 1: # that means the mesh is broken into different elements I am ignoring such meshes for now return data_mjcf = mjcf.from_path(self.xml_skeleton) for i, stl_file in enumerate(stl_files_path): data_mjcf.asset.add('mesh', name=f'mesh:{i}', file=stl_file) object_body = data_mjcf.worldbody.add('body', name=f'object:{i}', pos=[0, 0, 0], euler=self.object_angle) object_body.add('geom', type='mesh', name=f'object:{i}', mesh=f'mesh:{i}',\ group='1', condim='3', mass='1000') ## this is temporary for my heart satisfaction # temp_body = data_mjcf.worldbody.add('body', name=f'object:100', pos=[0.4, 0.0, 0.4], quat=[0, 0, 0, 1]) # temp_body.add('geom', type='sphere', name=f'object_geom:100', size='0.3', group='1', contype='0', conaffinity='0', mass='1000') # make the environment with empty file as of now self.modify_and_change_xml(data_mjcf.to_xml_string(), stl_files_path, xml_save_path='/home/bbgsp/anaconda3/envs/muj_exps/lib/python3.6/site-packages/dm_control/suite/changed.xml') task_kwargs = dict(xml='/home/bbgsp/anaconda3/envs/muj_exps/lib/python3.6/site-packages/dm_control/suite/changed.xml') # since I filtered out the invalid ones the next is not needed anymore and self.okay should start with true # def loader(): # env = empty.SUITE['empty'](**task_kwargs) # return env # try: # viewer.launch(loader) # self.okay = True # self.valid_cups.append(object_file) # except Exception as inst: # print('got an exception')\ ## .. not needed end .. ## # this creates the directory of the object based on pre-eliminary check of it being loadable if self.okay: # meaning I will want to collect data for this self.instance_storage_path = osp.join(self.obj_storage_path, self.object_name) if not osp.exists(self.instance_storage_path): os.makedirs(self.instance_storage_path) if self.okay and collect_touch_trajectories: print('---- COLLECTING TOUCH DATA ----') touch_mjcf = copy.deepcopy(data_mjcf) # touch_data, modified_touch_mjcf = self.touch_data_collector(stl_files_path, touch_mjcf, save_imgs) touch_data, modified_touch_mjcf, recon_imgs, chosen_cam_locs, chosen_cam_quats =\ self.touch_data_collector(stl_files_path, touch_mjcf, save_imgs) if len(chosen_cam_locs) == 1 and len(chosen_cam_quats) == 1: self.okay = False # if self.okay is false I will remove the above created directory if not self.okay: shutil.rmtree(self.instance_storage_path) if self.okay: touch_save_path = osp.join(self.instance_storage_path, "touch_data.npy") np.save(touch_save_path, touch_data) # save the modified xml too, NOTE not useful anymore since only one camera is modified over and over touch_xml_string = modified_touch_mjcf.to_xml_string() xml_save_path = osp.join(self.instance_storage_path, "touch_data_xml.xml") self.modify_and_change_xml(touch_xml_string, stl_files_path, xml_save_path) # save the reconstructed images # save the PIL images of the reconstruction # again for speed this would almost always be none if recon_imgs: for j, im in enumerate(recon_imgs): im = np.asarray(im) im = (im * 255.).astype(np.uint8) r_img = Image.fromarray(im) r_img.save(osp.join(self.instance_storage_path, f"recon_img_{j}.png")) # save the camera positions and quaternions cam_pos_save_path = osp.join(self.instance_storage_path, "cam_positions.npy") np.save(cam_pos_save_path, chosen_cam_locs) cam_quats_save_path = osp.join(self.instance_storage_path, "cam_quat.npy") np.save(cam_quats_save_path, chosen_cam_quats) if self.okay and collect_visual_data: # meaning the file loaded properly and everything is smooth till now # collect the visual data # now from the center of the object which is at 0,0,0 I want the camera to be 1m apart in the xy plane visual_mjcf = copy.deepcopy(data_mjcf) vis_data, modified_vis_mjcf = self.vis_data_collector(visual_mjcf, save_imgs) vis_save_path = osp.join(self.instance_storage_path, "visual_data.npy") np.save(vis_save_path, vis_data) # save the xml too vis_xml_save_path = osp.join(self.instance_storage_path, "object_mjcf.xml") self.modify_and_change_xml(modified_vis_mjcf.to_xml_string(), stl_files_path, vis_xml_save_path)
def main(xml_path, mesh_name=None, mesh_dir_path=None): current_folder = os.getcwd() xml_path = os.path.join(current_folder, xml_path) color_img_dir = os.path.join(current_folder, 'images') if not os.path.exists(color_img_dir): os.makedirs(color_img_dir) instance_storage_path = os.path.join(current_folder, 'instance') if not os.path.exists(instance_storage_path): os.makedirs(instance_storage_path) ## mujoco_py loading, do you really need this, actually this is ## required to render the environment, this is just for sanity checking model = mujoco_py.load_model_from_path(xml_path) sim = mujoco_py.MjSim(model) viewer = mujoco_py.MjViewer(sim) for _ in range(5000): sim.step() while True: viewer.render() ## ... End of mujoco py setup ... ## ## ... The code for collecting data actually begins here ... ## visual_mjcf = mjcf.from_path(xml_path) lookat_pos = visual_mjcf.worldbody.body['object0'].pos center_of_movement = lookat_pos camera_positions, camera_quats = utils.generate_new_cameras( 0.3, center_of_movement, lookat_pos, height=0.5, jitter_z=True, jitter_amount=0.08) ep_imgs = list() ep_depths = list() ep_intrinsics = list() ep_extrinsics = list() img_height = 128 img_width = 128 fov_y = 45 save_image = False recon_scene = True mesh_dir = os.path.join( "/Users/sfish0101/Documents/2020/Spring/quantized_policies/data_gen/meshes/03797390", mesh_name) mesh_dir = os.path.join(mesh_dir, "models/stl_models") # based on the mesh name also get the bounding box for the mesh and transform it to every # camera coordinate system physics = mjcf.Physics.from_mjcf_model(visual_mjcf) physics.forward() for _ in range(5000): physics.step() # now get the object position object_xpos = physics.data.xpos[physics.model.name2id('object0', 'body')] # now I need to call the file processing.py to get the bounding box corners # lengths, bbox_corners_origin= process_mesh(mesh_dir, object_xpos) # bbox_corners_origin *= 1.5 # now for each camera position and orientation get the image, depth, intrinsics and extrinsics for i, (pos, quat) in enumerate(zip(camera_positions, camera_quats)): print(f'generating for {i}/{len(camera_positions)} ...') visual_mjcf.worldbody.add('camera', name=f'vis_cam:{i}', pos=pos, quat=quat, fovy=fov_y) physics = mjcf.Physics.from_mjcf_model(visual_mjcf) # okay here the physics is initialized but since the object is placed at a location # not respecting the actual size fo the object I should simulate phyics for sometime physics.forward() for _ in range(5000): physics.step() img = physics.render(img_height, img_width, camera_id=f'vis_cam:{i}') depth = physics.render(img_height, img_width, camera_id=f'vis_cam:{i}', depth=True) assert img.shape[0] == img_height, "color img height is wrong" assert img.shape[1] == img_width, "color img width is wrong" assert depth.shape[0] == img_height, "depth img height is wrong" assert depth.shape[1] == img_width, "depth img width is wrong" if save_image: fig, ax = plt.subplots(2, sharex=True, sharey=True) ax[0].imshow(img) ax[1].imshow(depth) fig.savefig(f"{color_img_dir}/img_{i}.png") plt.close(fig=fig) # get the intrinsics and extrinsics and form the dict and store the data intrinsics = utils.dm_get_intrinsics(fov_y, img_width, img_height) extrinsics = utils.dm_get_extrinsics( physics, physics.model.name2id(f'vis_cam:{i}', 'camera')) ############# BOUNDING BOX DRAWING START ##################################################################### # now that you have the extrinsics which is from camera to origin, inverse of this will # give me from origin to camera coordinate system # cam_T_origin = np.linalg.inv(extrinsics) # bbox_t = np.c_[bbox_corners_origin, np.ones(len(bbox_corners_origin))] # bbox_corners_cam = np.dot(cam_T_origin, bbox_t.T).T # # so the bounding box is now in camera coordinate system, project it to image coordinate system # fx, fy, cx, cy = split_intrinsics(intrinsics) # X, Y, Z = bbox_corners_cam[:, 0], bbox_corners_cam[:, 1], bbox_corners_cam[:, 2] # x_pix = ((fx*X) / Z) + cx # y_pix = ((fy*Y) / Z) + cy # x_pix = 128 - x_pix # y_pix = 128 - y_pix # now I will use Adam's draw_box function # rgb = draw_boxes_on_image_py(img, corners_pix=np.c_[x_pix, y_pix][None], scores=np.ones(1,), tids=np.ones(1,)) # plt.imshow(rgb) # plt.show() ######## BOUNDING BOX DRAWING AND REGISTRATION COMPLETE ####################################################### ep_imgs.append(img) ep_depths.append(depth) ep_intrinsics.append(intrinsics) ep_extrinsics.append(extrinsics) visual_mjcf.worldbody.add('camera', name='ref_cam', pos=[0, -0.3, 0.5], zaxis=[0, -1, 0], fovy=fov_y) physics = mjcf.Physics.from_mjcf_model(visual_mjcf) physics.forward() # let it settle for _ in range(5000): physics.step() img = physics.render(img_height, img_width, camera_id='ref_cam') depth = physics.render(img_height, img_width, camera_id='ref_cam', depth=True) intrinsics = utils.dm_get_intrinsics(45, img_width, img_height) extrinsics = utils.dm_get_extrinsics( physics, physics.model.name2id('ref_cam', 'camera')) if save_image: fig, ax = plt.subplots(2, sharex=True, sharey=True) ax[0].imshow(img) ax[1].imshow(depth) fig.savefig(f'{color_img_dir}/img_ref.png') ep_imgs.append(img) ep_depths.append(depth) ep_intrinsics.append(intrinsics) ep_extrinsics.append(extrinsics) # next recreate the scene !! if recon_scene: # Here I can also return the points in the world and the points in ref_cam # this would remove the need for scaling rotation, translation etc recon_imgs, world_xyzs, camR_xyzs = utils.recreate_scene( ep_depths, ep_intrinsics, ep_extrinsics, camR_T_origin=np.linalg.inv(ep_extrinsics[-1]), clip_radius=5.0) frame = o3d.geometry.TriangleMesh.create_coordinate_frame( origin=np.zeros(3), size=0.8) # now here visualize the points first and then compute the bounding box using plane subtraction # collect all the points together world_xyzs_collated = np.concatenate(world_xyzs, axis=0) cup_points = segment_plane(world_xyzs_collated) # next compute the bounding box using these points bbox_points = compute_bounding_box_python(cup_points) bbox_lineset = make_lineset(bbox_points) # rotate the cup points, rotate the bbox points and see if they match camR_T_origin = np.linalg.inv(ep_extrinsics[-1]) h_cup = np.c_[cup_points, np.ones(len(cup_points))] r_cup = np.dot(camR_T_origin, h_cup.T).T[:, :3] # rotate the box h_bbox_points = np.c_[bbox_points, np.ones(len(bbox_points))] r_bbox_points = np.dot(camR_T_origin, h_bbox_points.T).T[:, :3] print(r_bbox_points) # make the rotated pcd r_cup_pcd = o3d.geometry.PointCloud() r_cup_pcd.points = o3d.utility.Vector3dVector(r_cup) r_bbox_lineset = make_lineset(r_bbox_points) # make the pcd and visualize cup_pcd = o3d.geometry.PointCloud() cup_pcd.points = o3d.utility.Vector3dVector(cup_points) # o3d.visualization.draw_geometries([cup_pcd, bbox_lineset, frame]) # o3d.visualization.draw_geometries([r_cup_pcd, r_bbox_lineset, frame]) # project the bounding box in the reference camera here ref_box_2d = get_box_corners_2d(r_bbox_points, int_mat=ep_intrinsics[-1]) ref_box_2d = ref_box_2d[None] ref_rgb = draw_boxes_on_image_py(ep_imgs[-1], ref_box_2d, scores=np.ones(1), tids=np.ones(1)) # plt.imshow(ref_rgb) # plt.show() for e, ext_mat in enumerate(ep_extrinsics): # also project the box onto the image of each camera temp_points = np.c_[bbox_points, np.ones(len(bbox_points))] box_in_camX = np.dot(np.linalg.inv(ext_mat), temp_points.T).T[:, :3] box_2d_camX = get_box_corners_2d(box_in_camX, int_mat=ep_intrinsics[e]) box_2d_camX = box_2d_camX[None] box_img = draw_boxes_on_image_py(ep_imgs[e], box_2d_camX, scores=np.ones(1), tids=np.ones(1)) # plt.imshow(box_img) # plt.show() for j, im in enumerate(recon_imgs): im = np.asarray(im) im = (im * 255.).astype(np.uint8) r_im = Image.fromarray(im) r_im.save( os.path.join(instance_storage_path, f'visual_recon_img_{j}.png')) # create a dictionary to save the data and ship it !! save_dict = AttrDict() save_dict.rgb_camXs = np.stack(ep_imgs) save_dict.depth_camXs = np.stack(ep_depths) save_dict.intrinsics = np.stack(ep_intrinsics) save_dict.extrinsics = np.stack(ep_extrinsics) save_dict.camR_T_origin = np.linalg.inv(ep_extrinsics[-1]) save_dict.bbox_in_ref_cam = r_bbox_points rgb_camRs = np.reshape(ep_imgs[-1], [1, img_height, img_width, 3]) rgb_camRs = np.tile(rgb_camRs, [len(ep_imgs), 1, 1, 1]) save_dict.rgb_camRs = rgb_camRs # everything should be len(51) for k in save_dict.keys(): if k == 'camR_T_origin' or k == 'bbox_in_ref_cam': continue assert len(save_dict[k]) == 51, "Data specific length is not right" vis_save_path = os.path.join(current_folder, f"visual_data_{mname}.npy") np.save(vis_save_path, save_dict) print('---- done ----')
def _build(self, easy_align=False, flanges=True, variation=0.0, color=(1., 0., 0.)): """Initializes a new `Duplo` instance. Args: easy_align: If True, the studs on the top of the brick will be capsules rather than cylinders. This makes alignment easier. flanges: Whether to use flanges on the bottom of the brick. These make the dynamics more expensive, but allow the bricks to be clicked together in partially overlapping configurations. variation: A float that controls the amount of variation in stud size (and therefore separation force). A value of 1.0 results in a distribution of separation forces that approximately matches the empirical distribution measured for real Duplo bricks. A value of 0.0 yields a deterministic separation force approximately equal to the mode of the empirical distribution. color: An optional tuple of (R, G, B) values specifying the color of the Duplo brick. These should be floats between 0 and 1. The default is red. Raises: ValueError: If `color` contains any value that is not between 0 and 1. """ self._mjcf_root = mjcf.from_path(_DUPLO_XML_PATH) stud = self._mjcf_root.default.find('default', 'stud') if easy_align: # Make cylindrical studs invisible and disable contacts. stud.geom.group = 3 stud.geom.contype = 9 stud.geom.conaffinity = 8 # Make capsule studs visible and enable contacts. stud_cap = self._mjcf_root.default.find('default', 'stud-capsule') stud_cap.geom.group = 0 stud_cap.geom.contype = 0 stud_cap.geom.conaffinity = 4 self._active_stud_dclass = stud_cap else: self._active_stud_dclass = stud if flanges: flange_dclass = self._mjcf_root.default.find('default', 'flange') flange_dclass.geom.contype = 4 # Enable contact with flanges. stud_size = _STUD_SIZE_PARAMS[(easy_align, flanges)] offset = (1 - variation) * stud_size.lower_quartile self._lower = offset + variation * stud_size.minimum self._upper = offset + variation * stud_size.maximum self._studs = np.ndarray((2, 4), dtype=object) self._holes = np.ndarray((2, 4), dtype=object) for row in range(2): for column in range(4): self._studs[row, column] = self._mjcf_root.find( 'site', 'stud_{}{}'.format(row, column)) self._holes[row, column] = self._mjcf_root.find( 'site', 'hole_{}{}'.format(row, column)) if not all(0 <= value <= 1 for value in color): raise ValueError(_COLOR_NOT_BETWEEN_0_AND_1.format(color)) self._mjcf_root.default.geom.rgba[:3] = color
def setUp(self): super(PhysicsTest, self).setUp() self.model = mjcf.from_path(ARM_MODEL) self.physics = mjcf.Physics.from_xml_string( self.model.to_xml_string(), assets=self.model.get_assets()) self.random = np.random.RandomState(0)
def test_default_model_filename(self): mjcf_model = mjcf.from_path(_TEST_MODEL_WITH_ASSETS) mjcf.export_with_assets(mjcf_model, _OUT_DIR, out_file_name=None) expected_name = mjcf_model.model + '.xml' self.assertTrue(os.path.isfile(os.path.join(_OUT_DIR, expected_name)))
def test_exceptions(self): mjcf_model = mjcf.from_path(_TEST_MODEL_WITH_ASSETS) with six.assertRaisesRegex(self, ValueError, 'must end with \'.xml\''): mjcf.export_with_assets(mjcf_model, _OUT_DIR, out_file_name='invalid_extension.png')
def __init__(self, path="rabbit_new_local.xml"): self.mjcf_model = mjcf.from_path(path) self.name = 'rabbit' self.com_name = 'torso'