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'
