61 lines
2.9 KiB
Python
61 lines
2.9 KiB
Python
"""Compare two trajectories with exactly the same fixed world camera."""
|
|
import os
|
|
os.environ.setdefault('MUJOCO_GL','osmesa')
|
|
from pathlib import Path
|
|
import json
|
|
import mujoco
|
|
import numpy as np
|
|
import imageio.v2 as imageio
|
|
from PIL import Image,ImageDraw,ImageFont
|
|
from scipy.spatial.transform import Rotation
|
|
|
|
ROOT=Path(__file__).resolve().parents[1]
|
|
out=ROOT/'output/l20_15886123_stable'
|
|
motions=[]
|
|
for folder in [ROOT/'output/l20_15886123',out]:
|
|
with np.load(folder/'motion.npz') as f:
|
|
motions.append({k:f[k] for k in f.files})
|
|
count=len(motions[0]['qpos']);fps=float(motions[0]['fps'])
|
|
assert count==len(motions[1]['qpos'])
|
|
model=mujoco.MjModel.from_xml_path(str(out/'l20_moving.xml'))
|
|
data=mujoco.MjData(model)
|
|
addr=[model.jnt_qposadr[model.joint(str(n)).id] for n in motions[0]['joint_names']]
|
|
wrist=model.jnt_qposadr[model.joint('wrist_free').id]
|
|
world=[]
|
|
for m in motions:
|
|
r=Rotation.from_quat(m['wrist_quat_wxyz'][:,[1,2,3,0]]).as_matrix()
|
|
world.append(np.einsum('tij,tkj->tki',r,m['actual'])+m['wrist_pos'][:,None])
|
|
points=np.concatenate(world).reshape(-1,3)
|
|
center=(points.min(0)+points.max(0))/2
|
|
model.geom_pos[model.geom('floor').id,2]=points[:,2].min()-.035
|
|
model.vis.quality.offsamples=1
|
|
model.vis.headlight.ambient[:]=.65
|
|
model.vis.headlight.diffuse[:]=.7
|
|
cam=mujoco.MjvCamera();cam.lookat[:]=center;cam.distance=max(.65,float(np.ptp(points,axis=0).max())*2.1);cam.azimuth=100;cam.elevation=5
|
|
option=mujoco.MjvOption();option.sitegroup[:]=0
|
|
renderer=mujoco.Renderer(model,height=480,width=640)
|
|
font=ImageFont.truetype('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',18)
|
|
writer=imageio.get_writer(out/'comparison.mp4',fps=fps,codec='libx264',quality=8)
|
|
try:
|
|
for t in range(count):
|
|
panels=[]
|
|
for m in motions:
|
|
data.qpos[addr]=m['qpos'][t]
|
|
data.qpos[wrist:wrist+3]=m['wrist_pos'][t]
|
|
data.qpos[wrist+3:wrist+7]=m['wrist_quat_wxyz'][t]
|
|
mujoco.mj_forward(model,data)
|
|
renderer.update_scene(data,camera=cam,scene_option=option)
|
|
panels.append(renderer.render().copy())
|
|
im=Image.fromarray(np.concatenate(panels,axis=1));d=ImageDraw.Draw(im)
|
|
d.rectangle([0,0,1280,50],fill=(15,22,30))
|
|
d.text((12,8),'Before - previous smoothing',font=font,fill='white')
|
|
d.text((652,8),'After - temporal regularization',font=font,fill='white')
|
|
d.text((12,29),f'Same fixed camera | {t/fps:.2f} s | moving wrist preserved',font=font,fill=(190,210,225))
|
|
writer.append_data(np.asarray(im))
|
|
if t==280:im.save(out/'comparison_preview.png')
|
|
if t%200==0:print('AB render',t,count,flush=True)
|
|
finally:
|
|
writer.close();renderer.close()
|
|
(out/'comparison_camera.json').write_text(json.dumps(dict(same_camera=True,fixed_world_camera=True,frames=count,fps=fps,lookat=center.tolist(),distance=cam.distance,azimuth=cam.azimuth,elevation=cam.elevation),indent=2))
|
|
print('COMPLETE',flush=True)
|