LuckyRobots/LuckyWorld
14
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects 1 Releases 13 Wiki Activity
LuckyWorld/SourceFiles/mujoco_menagerie-main/generate_gallery.py
Behron f571dc9459 Finished fixing merge conflicts and touched up camera settings. 
More merge conflicts. Taking main changes then will integrate mine.
2025-04-08 23:41:18 -06:00

482 lines
15 KiB
Python
Raw Blame History

# Copyright 2024 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Generate a markdown table with images of some of the models in Menagerie.
Requirements:
pip install absl-py dm_control pillow numpy tqdm mdutils opencv-python
Instructions:
`python generate_gallery.py` will create a markdown document called
`gallery.md` with a table of images. Copy this table into README.md to
display the images.
"""
import enum
import math
import pathlib
from absl import app
import cv2
from dm_control import mjcf
from mdutils import mdutils
import numpy as np
from PIL import Image
import tqdm.auto
DEFAULT_FOV = 40
class ModelType(int, enum.Enum):
ARM = 0
DUAL_ARM = 1
END_EFFECTOR = 2
MOBILE_MANIPULATOR = 3
QUADRUPED = 4
BIPED = 5
HUMANOID = 6
DRONE = 7
BIOMECHANICAL = 8
MISC = 9
NAME_MAP = {
"franka_emika_panda/panda": "panda",
"franka_emika_panda/hand": "panda gripper",
"franka_fr3/fr3": "franka fr3",
"ufactory_lite6/lite6": "lite6",
"flybody/fruitfly": "fruitfly",
"skydio_x2/x2": "skydio x2",
"unitree_h1/h1": "h1",
"bitcraze_crazyflie_2/cf2": "crazyflie 2",
"google_robot/robot": "google robot",
"unitree_a1/a1": "a1",
"google_barkour_v0/barkour_v0": "barkour v0",
"anybotics_anymal_b/anymal_b": "anymal b",
"unitree_go1/go1": "go1",
"unitree_z1/z1": "z1",
"anybotics_anymal_c/anymal_c": "anymal c",
"agility_cassie/cassie": "cassie",
"realsense_d435i/d435i": "d435i",
"universal_robots_ur5e/ur5e": "ur5e",
"aloha/aloha": "aloha 2",
"rethink_robotics_sawyer/sawyer": "sawyer",
"robotis_op3/op3": "op3",
"universal_robots_ur10e/ur10e": "ur10e",
"kuka_iiwa_14/iiwa14": "iiwa 14",
"trossen_vx300s/vx300s": "vx300s",
"unitree_g1/g1": "g1",
"robotiq_2f85/2f85": "2f85",
"ufactory_xarm7/hand": "xarm7 gripper",
"ufactory_xarm7/xarm7": "xarm7",
"hello_robot_stretch/stretch": "stretch 2",
"google_barkour_vb/barkour_vb": "barkour vb",
"unitree_go2/go2": "go2",
"boston_dynamics_spot/spot_arm": "spot",
"shadow_dexee/shadow_dexee": "dex-ee",
"pal_talos/talos": "talos",
"leap_hand/left_hand": "left leap",
"wonik_allegro/left_hand": "left allegro",
"shadow_hand/left_hand": "left shadow",
"kinova_gen3/gen3": "gen3",
"booster_t1/t1": "t1",
"agilex_piper/piper": "piper",
}
MODEL_MAP = {
"franka_emika_panda/panda": ModelType.ARM,
"franka_emika_panda/hand": ModelType.END_EFFECTOR,
"franka_fr3/fr3": ModelType.ARM,
"ufactory_lite6/lite6": ModelType.ARM,
"flybody/fruitfly": ModelType.BIOMECHANICAL,
"wonik_allegro/left_hand": ModelType.END_EFFECTOR,
"shadow_hand/left_hand": ModelType.END_EFFECTOR,
"skydio_x2/x2": ModelType.DRONE,
"unitree_h1/h1": ModelType.HUMANOID,
"bitcraze_crazyflie_2/cf2": ModelType.DRONE,
"google_robot/robot": ModelType.MOBILE_MANIPULATOR,
"unitree_a1/a1": ModelType.QUADRUPED,
"google_barkour_v0/barkour_v0": ModelType.QUADRUPED,
"anybotics_anymal_b/anymal_b": ModelType.QUADRUPED,
"unitree_go1/go1": ModelType.QUADRUPED,
"unitree_z1/z1": ModelType.ARM,
"anybotics_anymal_c/anymal_c": ModelType.QUADRUPED,
"agility_cassie/cassie": ModelType.BIPED,
"realsense_d435i/d435i": ModelType.MISC,
"universal_robots_ur5e/ur5e": ModelType.ARM,
"aloha/aloha": ModelType.DUAL_ARM,
"rethink_robotics_sawyer/sawyer": ModelType.ARM,
"robotis_op3/op3": ModelType.HUMANOID,
"universal_robots_ur10e/ur10e": ModelType.ARM,
"kuka_iiwa_14/iiwa14": ModelType.ARM,
"trossen_vx300s/vx300s": ModelType.ARM,
"unitree_g1/g1": ModelType.HUMANOID,
"robotiq_2f85/2f85": ModelType.END_EFFECTOR,
"ufactory_xarm7/hand": ModelType.END_EFFECTOR,
"ufactory_xarm7/xarm7": ModelType.ARM,
"hello_robot_stretch/stretch": ModelType.MOBILE_MANIPULATOR,
"google_barkour_vb/barkour_vb": ModelType.QUADRUPED,
"unitree_go2/go2": ModelType.QUADRUPED,
"boston_dynamics_spot/spot_arm": ModelType.QUADRUPED,
"shadow_dexee/shadow_dexee": ModelType.END_EFFECTOR,
"pal_talos/talos": ModelType.HUMANOID,
"leap_hand/left_hand": ModelType.END_EFFECTOR,
"kinova_gen3/gen3": ModelType.ARM,
"booster_t1/t1": ModelType.HUMANOID,
"agilex_piper/piper": ModelType.ARM,
}
DEFAULT_FOV = 40
CAMERA_MAP = {
"pal_talos/talos": dict(
pos="2.312 0.005 1.144",
xyaxes="-0.002 1.000 -0.000 -0.107 -0.000 0.994",
),
"skydio_x2/x2": dict(
pos="-0.580 -0.260 0.622",
xyaxes="0.442 -0.897 -0.000 0.428 0.211 0.879",
fovy=60,
),
"flybody/fruitfly": dict(
pos="0.430 -0.361 0.326", xyaxes="0.589 0.808 0.000 -0.486 0.354 0.799",
fovy=50
),
"wonik_allegro/left_hand": dict(
pos="0.002 0.043 0.432", xyaxes="0.052 -0.999 0.000 0.998 0.052 0.017"
),
"ufactory_xarm7/xarm7": dict(
pos="0.852 -0.383 0.860",
xyaxes="0.487 0.874 0.000 -0.354 0.197 0.914",
fovy=DEFAULT_FOV,
),
"ufactory_xarm7/hand": dict(
pos="-0.282 0.013 0.118",
xyaxes="-0.047 -0.999 0.000 0.160 -0.007 0.987",
fovy=45,
),
"shadow_hand/left_hand": dict(
pos="0.172 0.005 0.615",
xyaxes="-0.508 -0.861 -0.000 0.861 -0.508 0.017",
fovy=45,
),
"franka_emika_panda/panda": dict(
pos="0.412 1.106 0.849",
xyaxes="-0.994 0.108 0.000 -0.040 -0.369 0.928",
fovy=DEFAULT_FOV,
),
"franka_fr3/fr3": dict(
pos="0.412 1.106 0.849",
xyaxes="-0.994 0.108 0.000 -0.040 -0.369 0.928",
fovy=DEFAULT_FOV,
),
"franka_emika_panda/hand": dict(
pos="0.340 0.008 0.059",
xyaxes="-0.023 1.000 0.000 -0.084 -0.002 0.996",
fovy=DEFAULT_FOV,
),
"kuka_iiwa_14/iiwa14": dict(
pos="0.212 1.138 0.977",
xyaxes="-1.000 0.027 0.000 -0.012 -0.441 0.898",
fovy=45,
),
"ufactory_lite6/lite6": dict(
pos="0.077 -0.778 0.528",
xyaxes="1.000 -0.004 -0.000 0.001 0.274 0.962",
fovy=45,
),
"unitree_g1/g1": dict(
pos="1.466 -0.082 1.271",
xyaxes="0.072 0.997 -0.000 -0.377 0.027 0.926",
fovy=45,
),
"unitree_h1/h1": dict(
pos="2.098 0.006 1.893",
xyaxes="0.007 1.000 -0.000 -0.394 0.003 0.919",
fovy=45,
),
"robotis_op3/op3": dict(
pos="0.673 -0.024 0.447", xyaxes="0.035 0.999 0.000 -0.252 0.009 0.968",
fovy=45
),
"universal_robots_ur5e/ur5e": dict(
pos="0.603 1.012 0.595",
xyaxes="-0.932 0.363 -0.000 -0.080 -0.206 0.975",
fovy=DEFAULT_FOV,
),
"universal_robots_ur10e/ur10e": dict(
pos="1.286 -0.798 0.889",
xyaxes="0.696 0.718 -0.000 -0.224 0.218 0.950",
fovy=DEFAULT_FOV,
),
"unitree_z1/z1": dict(
pos="0.305 -0.400 0.552",
xyaxes="0.755 0.656 0.000 -0.359 0.413 0.837",
fovy=DEFAULT_FOV,
),
"unitree_go2/go2": dict(
pos="0.753 -0.427 0.433",
xyaxes="0.518 0.856 0.000 -0.284 0.172 0.943",
fovy=DEFAULT_FOV,
),
"unitree_go1/go1": dict(
pos="0.679 -0.553 0.530",
xyaxes="0.638 0.770 -0.000 -0.328 0.272 0.905",
fovy=DEFAULT_FOV,
),
"unitree_a1/a1": dict(
pos="0.654 -0.564 0.536",
xyaxes="0.676 0.737 -0.000 -0.327 0.299 0.896",
fovy=DEFAULT_FOV,
),
"trossen_vx300s/vx300s": dict(
pos="0.583 0.317 0.549",
xyaxes="-0.531 0.847 0.000 -0.434 -0.272 0.859",
fovy=DEFAULT_FOV,
),
"robotiq_2f85/2f85": dict(
pos="-0.009 -0.251 0.107",
xyaxes="0.999 -0.033 -0.000 0.005 0.150 0.989",
fovy=45,
),
"rethink_robotics_sawyer/sawyer": dict(
pos="1.014 -0.494 0.876",
xyaxes="0.555 0.832 -0.000 -0.372 0.248 0.895",
fovy=DEFAULT_FOV,
),
"realsense_d435i/d435i": dict(
pos="-0.000 -0.002 0.128",
xyaxes="1.000 -0.000 0.000 0.000 1.000 0.017",
fovy=DEFAULT_FOV,
),
"hello_robot_stretch/stretch": dict(
pos="1.464 -0.514 1.439",
xyaxes="0.271 0.963 -0.000 -0.362 0.102 0.927",
fovy=45,
),
"google_robot/robot": dict(
pos="1.753 -0.231 1.305",
xyaxes="0.025 1.000 0.000 -0.306 0.008 0.952",
fovy=50,
),
"google_barkour_vb/barkour_vb": dict(
pos="0.887 0.338 0.565",
xyaxes="-0.388 0.922 0.000 -0.460 -0.194 0.867",
fovy=DEFAULT_FOV,
),
"google_barkour_v0/barkour_v0": dict(
pos="0.733 0.320 0.558",
xyaxes="-0.416 0.909 -0.000 -0.441 -0.202 0.875",
fovy=DEFAULT_FOV,
),
"bitcraze_crazyflie_2/cf2": dict(
pos="0.037 -0.142 0.206",
xyaxes="0.963 0.268 0.000 -0.167 0.599 0.783",
fovy=DEFAULT_FOV,
),
"anybotics_anymal_b/anymal_b": dict(
pos="0.930 -1.239 1.221",
xyaxes="0.809 0.587 0.000 -0.308 0.424 0.852",
fovy=DEFAULT_FOV,
),
"anybotics_anymal_c/anymal_c": dict(
pos="1.547 -0.577 0.941",
xyaxes="0.378 0.926 -0.000 -0.358 0.146 0.922",
fovy=45,
),
"aloha/aloha": dict(
pos="0.484 1.158 0.836",
xyaxes="-0.939 0.345 -0.000 -0.162 -0.441 0.883",
fovy=DEFAULT_FOV,
),
"agility_cassie/cassie": dict(
pos="1.277 -1.122 1.053",
xyaxes="0.655 0.756 0.000 -0.196 0.170 0.966",
fovy=DEFAULT_FOV,
),
"boston_dynamics_spot/spot_arm": dict(
pos="1.326 -0.829 0.651",
xyaxes="0.573 0.820 -0.000 -0.154 0.108 0.982",
fovy=DEFAULT_FOV,
),
"shadow_dexee/shadow_dexee": dict(
pos="-0.003 -0.541 0.221",
xyaxes="1.000 0.000 -0.000 -0.000 0.131 0.991",
fovy=45,
),
"leap_hand/left_hand": dict(
pos="-0.123 0.096 0.512",
xyaxes="0.004 -1.000 -0.000 0.995 0.004 0.101",
),
"kinova_gen3/gen3": dict(
pos="0.252 -1.047 0.521",
xyaxes="0.988 0.156 -0.000 -0.024 0.153 0.988",
),
"booster_t1/t1": dict(
pos="1.499 -0.777 1.2",
xyaxes="0.453 0.892 0.000 -0.295 0.150 0.944",
fovy=DEFAULT_FOV,
),
"agilex_piper/piper": dict(
pos="0.288 -0.480 0.294",
xyaxes="0.866 0.500 0.000 -0.171 0.296 0.940",
fovy=50,
),
}
# pylint: disable=line-too-long
KEYFRAME_MAP = {
"pal_talos": "0 0 1.025 0 0 0 0 0 0.15 0 0 0.3 0.4 -0.5 -1.5 0 0 0 0 -0.4 0 0 0 0 0 -0.3 -0.4 0.5 -1.5 0 0 0 0 -0.4 0 0 0 0 0 0 0 -0.4 0.8 -0.4 0 0 0 -0.4 0.8 -0.4 0",
"robotis_op3": "0 0 0.2789 1 0 0 0 0.0 0.0 -0.0890 0.7931 -0.79 0.0874 -0.7946 0.7855 -0.0015 -0.0460 -0.1626 0.2316 0.1565 -0.0230 0.0 0.0445 0.1611 -0.2332 -0.1580 0.0215",
"google_barkour_vb": "0 0 0.21 1 0 0 0 0 0.5 1.0 0 0.5 1.0 0 0.5 1.0 0 0.5 1.0",
"hello_robot_stretch": "0 0 0 1 0 0 0 0 0 0.1325 0.07995 0.07995 0.07605 0.0702 1.585 0 0.198 0 0 0.126 0 0 0 0",
"google_robot": "-1.51699e-13 -1.16232e-12 -0.1444 2.9724 -0.146 -0.3759 1.15806e-12 0.5518 0.62275",
"aloha": "0.43988 -0.206468 1.08253 -0.443382 -1.084 -0.00397598 0.0084 0.00846495 -1.28822 -0.360594 0.717978 -0.000325086 -0.273415 6.76003e-05 0.0084 0.00839987",
"kuka_iiwa_14": "0 0 0 -1.5708 0 1.5708 0",
}
# pylint: enable=line-too-long
KEEP_LIGHT = ["go1", "a1", "op3", "aloha", "left_hand", "stretch", "piper"]
def create_arena():
arena = mjcf.RootElement()
arena.visual.quality.shadowsize = 8192
arena.visual.headlight.diffuse = (0.6,) * 3
arena.visual.headlight.ambient = (0.3,) * 3
arena.visual.headlight.specular = (0.2,) * 3
getattr(arena.visual, "global").offheight = 720
getattr(arena.visual, "global").offwidth = 1280
arena.asset.add(
"texture",
type="skybox",
builtin="gradient",
height=512,
width=512,
rgb1="1 1 1",
rgb2="1 1 1",
)
return arena
MODEL_XMLS = [pathlib.Path(f"../{k}.xml") for k in MODEL_MAP.keys()]
# Sort XML files.
def sort_func(xml):
name = f"{xml.parent.stem}/{xml.stem}"
return (MODEL_MAP[name], xml.stem)
MODEL_XMLS = sorted(MODEL_XMLS, key=sort_func)
def main(argv):
del argv
paths = []
pngs = []
for xml in tqdm.auto.tqdm(MODEL_XMLS):
try:
robot_maker = xml.parent.stem
robot_name = xml.stem
robot = f"{robot_maker}/{robot_name}"
if robot not in CAMERA_MAP:
continue
arena = create_arena()
if robot_maker in KEYFRAME_MAP:
arena.keyframe.add("key", qpos=KEYFRAME_MAP[robot_maker])
model_xml = mjcf.from_path(xml.as_posix(), escape_separators=True)
for light in model_xml.find_all("light"):
if robot_name not in KEEP_LIGHT:
light.remove()
if robot in CAMERA_MAP:
camera_kwargs = CAMERA_MAP[robot]
arena.worldbody.add("camera", name="thumbnail", **camera_kwargs)
if robot_maker == "aloha":
right_base = model_xml.find("body", "right\\base_link")
right_base.pos[0] = 0.3
left_base = model_xml.find("body", "left\\base_link")
left_base.pos[0] = -0.3
arena.include_copy(model_xml, override_attributes=True)
physics = mjcf.Physics.from_mjcf_model(arena)
try:
physics.reset(keyframe_id=0)
except:
physics.reset()
physics.forward()
if robot in CAMERA_MAP:
img = physics.render(height=500, width=500, camera_id="thumbnail")
else:
img = physics.render(height=500, width=500)
img = cv2.putText(
img.copy(),
NAME_MAP[robot],
(5, 480),
cv2.FONT_HERSHEY_SIMPLEX,
1.3,
(0, 0, 0),
1,
cv2.LINE_AA,
)
filename = f"assets/{robot_maker}-{robot_name}.png"
paths.append(filename)
png = np.zeros((500, 500, 4), dtype=np.uint8)
u, v = np.where(np.all(img == 255, axis=-1))
png[u, v, -1] = 0
png[u, v, :3] = 0
u, v = np.where(np.any(img != 255, axis=-1))
png[u, v, :3] = img[u, v]
png[u, v, -1] = 255
pngs.append(png.copy())
Image.fromarray(png).save(filename)
except Exception as e:
print(e)
print(f"failed to load {xml.as_posix()}")
n_models = len(paths)
n_cols = 5
n_rows = int(math.ceil(n_models / n_cols))
table = []
for r in range(n_rows):
row = []
for c in range(n_cols):
i = r * n_cols + c
if i >= n_models:
row.append("")
else:
row.append(f"<img src='{paths[i]}' width=100>")
table.extend(row)
mdfile = mdutils.MdUtils(file_name="gallery")
mdfile.new_table(columns=n_cols, rows=n_rows, text=table, text_align="center")
mdfile.create_md_file()
if __name__ == "__main__":
app.run(main)
Reference in New Issue View Git Blame Copy Permalink