71 lines
4.1 KiB
Python
71 lines
4.1 KiB
Python
"""Render fixed-world and wrist-following views of a moving L20 trajectory."""
|
|
import os
|
|
os.environ.setdefault('MUJOCO_GL','osmesa')
|
|
import argparse,json
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import mujoco,imageio.v2 as imageio
|
|
from PIL import Image,ImageDraw
|
|
from scipy.spatial.transform import Rotation
|
|
|
|
def main():
|
|
parser=argparse.ArgumentParser()
|
|
parser.add_argument('--directory',type=Path,default=Path(__file__).resolve().parents[1]/'output/l20_2047635068')
|
|
parser.add_argument('--frames',type=int,default=0)
|
|
args=parser.parse_args();out=args.directory
|
|
with np.load(out/'motion.npz') as archive:
|
|
motion={name:archive[name] for name in archive.files}
|
|
q=motion['qpos'][:args.frames or None];positions=motion['wrist_pos'][:len(q)]
|
|
model=mujoco.MjModel.from_xml_path(str(out/'l20_moving.xml'));data=mujoco.MjData(model)
|
|
addr=[model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,str(n))] for n in motion['joint_names']]
|
|
wrist=model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,'wrist_free')]
|
|
rotations=Rotation.from_quat(motion['wrist_quat_wxyz'][:,[1,2,3,0]]).as_matrix()
|
|
world=np.einsum('tij,tkj->tki',rotations,motion['actual'])+motion['wrist_pos'][:,None]
|
|
center=(world.min((0,1))+world.max((0,1)))/2
|
|
span=np.ptp(world,axis=(0,1)).max()
|
|
floor=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_GEOM,'floor')
|
|
model.geom_pos[floor,2]=world[:,:,2].min()-.035
|
|
model.vis.quality.shadowsize=1024;model.vis.quality.offsamples=1
|
|
model.vis.headlight.ambient[:]=.65;model.vis.headlight.diffuse[:]=.7
|
|
renderer=mujoco.Renderer(model,height=480,width=640)
|
|
cameras=[]
|
|
for azimuth in [135,25]:
|
|
cam=mujoco.MjvCamera();cam.lookat[:]=center;cam.distance=max(.65,span*2.3);cam.azimuth=azimuth;cam.elevation=10
|
|
cameras.append(cam)
|
|
option=mujoco.MjvOption();option.sitegroup[:]=0
|
|
fps=float(motion['fps'])
|
|
coupling=str(motion['coupling']) if 'coupling' in motion else 'JSON-calibrated'
|
|
writer=imageio.get_writer(out/'l20_moving_replay.mp4',fps=fps,codec='libx264',quality=8)
|
|
max_wrist_error=0.
|
|
try:
|
|
for t,pose in enumerate(q):
|
|
data.qpos[addr]=pose;data.qpos[wrist:wrist+3]=positions[t];data.qpos[wrist+3:wrist+7]=motion['wrist_quat_wxyz'][t]
|
|
mujoco.mj_forward(model,data)
|
|
base=mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_BODY,'hand_base_link')
|
|
max_wrist_error=max(max_wrist_error,float(np.abs(data.xpos[base]-positions[t]).max()))
|
|
panels=[]
|
|
cameras[1].lookat[:]=world[t].mean(0);cameras[1].distance=.58
|
|
for vi,cam in enumerate(cameras):
|
|
renderer.update_scene(data,camera=cam,scene_option=option)
|
|
if vi==0:
|
|
points=positions[np.unique(np.linspace(0,t,min(t+1,60),dtype=int))]
|
|
for a,b in zip(points[:-1],points[1:]):
|
|
scene=renderer.scene
|
|
mujoco.mjv_initGeom(scene.geoms[scene.ngeom],mujoco.mjtGeom.mjGEOM_CAPSULE,np.ones(3)*.001,a,np.eye(3).ravel(),np.array([.2,.8,1,1]))
|
|
mujoco.mjv_connector(scene.geoms[scene.ngeom],mujoco.mjtGeom.mjGEOM_CAPSULE,.0015,a,b)
|
|
scene.ngeom+=1
|
|
panels.append(renderer.render().copy())
|
|
img=Image.fromarray(np.concatenate(panels,axis=1));draw=ImageDraw.Draw(img)
|
|
draw.rectangle([0,0,1280,45],fill=(15,22,30))
|
|
draw.text((10,8),f'L20 right | frame {t+1}/{len(q)} | moving 6D wrist + {coupling} fingers',fill='white')
|
|
draw.text((10,27),'Fixed world camera / blue wrist trail',fill='white')
|
|
draw.text((650,27),'Following camera / finger detail',fill='white')
|
|
writer.append_data(np.asarray(img))
|
|
if t in [0,len(q)//2,len(q)-1]:img.save(out/f'preview_{t:04d}.png')
|
|
if t%200==0:print(f'render L20 {t}/{len(q)}',flush=True)
|
|
finally:writer.close();renderer.close()
|
|
(out/'render_validation.json').write_text(json.dumps(dict(frames=len(q),fps=fps,max_wrist_position_error_m=max_wrist_error),indent=2))
|
|
print('COMPLETE',flush=True)
|
|
|
|
if __name__=='__main__':main()
|