74 lines
4.5 KiB
Python
74 lines
4.5 KiB
Python
"""Render contact registration separately from the free-object physics test."""
|
|
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
|
|
|
|
OUT=Path(__file__).resolve().parents[1]/'output/l20_2047635068/bottle_grasp'
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('--physics',action='store_true');p.add_argument('--frames',type=int,default=0);args=p.parse_args()
|
|
physics=args.physics
|
|
path=OUT/('physics_scene.xml' if physics else 'contact_scene.xml')
|
|
motion=np.load(OUT/('physics_motion.npz' if physics else 'registered_motion.npz'))
|
|
model=mujoco.MjModel.from_xml_path(str(path));data=mujoco.MjData(model)
|
|
model.vis.headlight.ambient[:]=.65;model.vis.headlight.diffuse[:]=.7;model.vis.quality.offsamples=1
|
|
renderer=mujoco.Renderer(model,height=480,width=640)
|
|
option=mujoco.MjvOption();option.geomgroup[3]=0;option.sitegroup[:]=0
|
|
count=len(motion['qpos']);count=min(count,args.frames) if args.frames else count
|
|
fps=float(motion['fps'])
|
|
obj=model.jnt_qposadr[model.joint('object_free').id]
|
|
if not physics:
|
|
addresses=[model.jnt_qposadr[model.joint(str(n)).id] for n in motion['joint_names']]
|
|
wrist=model.jnt_qposadr[model.joint('wrist_free').id]
|
|
reference=int(motion['contact_onset_frame'])
|
|
rotation=Rotation.from_quat(motion['bottle_quat_wxyz'][reference,[1,2,3,0]]).inv()
|
|
translation=np.array([0,0,.3])-rotation.apply(motion['bottle_pos'][reference])
|
|
bottle_pos=rotation.apply(motion['bottle_pos'])+translation
|
|
wrist_pos=rotation.apply(motion['wrist_pos'])+translation
|
|
bottle_quat=(rotation*Rotation.from_quat(motion['bottle_quat_wxyz'][:,[1,2,3,0]])).as_quat()[:,[3,0,1,2]]
|
|
wrist_quat=(rotation*Rotation.from_quat(motion['wrist_quat_wxyz'][:,[1,2,3,0]])).as_quat()[:,[3,0,1,2]]
|
|
both=np.concatenate([bottle_pos+[0,0,.08],wrist_pos])
|
|
center=(both.min(0)+both.max(0))/2;distance=max(.75,np.ptp(both,axis=0).max()*2.5)
|
|
floor=model.geom('floor').id;model.geom_pos[floor,2]=min(both[:,2])-.06
|
|
else:center=np.array([0,0,.38]);distance=.65
|
|
cameras=[]
|
|
for azimuth in [45,225]:
|
|
cam=mujoco.MjvCamera();cam.lookat[:]=center;cam.distance=distance;cam.azimuth=azimuth;cam.elevation=12;cameras.append(cam)
|
|
filename='physics_test.mp4' if physics else 'contact_registration.mp4'
|
|
writer=imageio.get_writer(OUT/filename,fps=fps,codec='libx264',quality=8)
|
|
try:
|
|
for t in range(count):
|
|
if physics:
|
|
data.qpos[:]=motion['qpos'][t];data.mocap_pos[0]=motion['wrist'][t,:3];data.mocap_quat[0]=motion['wrist'][t,3:]
|
|
else:
|
|
data.qpos[addresses]=motion['qpos'][t]
|
|
data.qpos[wrist:wrist+7]=np.r_[wrist_pos[t],wrist_quat[t]]
|
|
data.qpos[obj:obj+7]=np.r_[bottle_pos[t],bottle_quat[t]]
|
|
cameras[1].lookat[:]=(bottle_pos[t]+wrist_pos[t])/2+[0,0,.04];cameras[1].distance=.65
|
|
mujoco.mj_forward(model,data)
|
|
panels=[]
|
|
for camera in cameras:
|
|
renderer.update_scene(data,camera=camera,scene_option=option);panels.append(renderer.render().copy())
|
|
img=Image.fromarray(np.concatenate(panels,axis=1));draw=ImageDraw.Draw(img);draw.rectangle([0,0,1280,48],fill=(15,22,30))
|
|
if physics:
|
|
title='FREE OBJECT PHYSICS TEST | 16 actuators + nonlinear coupling | moving wrist'
|
|
note=f"t={t/fps:.2f}s | slip={motion['relative_drift'][t]*1000:.1f} mm | contacts={motion['contact_count'][t]} | result: NOT RETAINED"
|
|
color=(255,120,100)
|
|
else:
|
|
title='L20 + tri-camera bottle | CONTACT REGISTRATION | object poses prescribed, not a physical grasp'
|
|
note=f"frame {t+1}/{count} | moving wrist | maximum penetration {motion['penetration_m'][t]*1000:.2f} mm | estimated object-to-hand alignment"
|
|
color='white'
|
|
draw.text((10,8),title,fill='white');draw.text((10,28),note,fill=color)
|
|
writer.append_data(np.asarray(img))
|
|
if t in [0,count//2,count-1]:img.save(OUT/(('physics' if physics else 'registration')+f'_{t:04d}.png'))
|
|
if t%200==0:print(f'{filename} {t}/{count}',flush=True)
|
|
finally:writer.close();renderer.close()
|
|
(OUT/(filename+'.json')).write_text(json.dumps({'frames':count,'fps':fps,'physics':physics},indent=2))
|
|
print('COMPLETE '+filename,flush=True)
|
|
|
|
if __name__=='__main__':main()
|