74 lines
5.4 KiB
Python
74 lines
5.4 KiB
Python
"""Full-length contact registration; object poses are inferred targets, not physics."""
|
|
import json
|
|
from pathlib import Path
|
|
import numpy as np
|
|
from scipy.spatial.transform import Rotation,Slerp
|
|
from l20_contact import ContactRefiner
|
|
from l20_calibrated import ROOT
|
|
|
|
BASE=ROOT/'output/l20_2047635068';OUT=BASE/'bottle_grasp'
|
|
def main():
|
|
motion=np.load(BASE/'motion.npz');spec=json.loads((OUT/'object_spec.json').read_text())
|
|
solver=ContactRefiner(spec,OUT)
|
|
count=len(motion['qpos']);times=np.arange(count)
|
|
rotation=Rotation.from_quat(motion['wrist_quat_wxyz'][:,[1,2,3,0]])
|
|
reference=666
|
|
# Infer a constant bottle-to-palm registration over the holding phase.
|
|
rel_pos=rotation[reference].inv().apply(np.array(spec['position'])-motion['wrist_pos'][reference])
|
|
rel_rot=rotation[reference].inv()*Rotation.from_quat(np.array(spec['quat_wxyz'])[[1,2,3,0]])
|
|
bottle_pos=rotation.apply(np.tile(rel_pos,(count,1)))+motion['wrist_pos']
|
|
bottle_rot=rotation*rel_rot
|
|
# Contact onset bracket comes from visual inspection: open at 100, grasped by 200.
|
|
onset=160
|
|
bottle_pos[:onset]=bottle_pos[onset]
|
|
bottle_quat=bottle_rot.as_quat()[:,[3,0,1,2]]
|
|
bottle_quat[:onset]=bottle_quat[onset]
|
|
samples=np.unique(np.r_[np.arange(onset,count,8),count-1])
|
|
fits=[];reports=[]
|
|
for t in samples:
|
|
active=np.array([motion['qpos'][t,list(motion['joint_names']).index(n)] for n in solver.cal.active])
|
|
x,q,report=solver.solve(active,motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t])
|
|
fits.append(x);reports.append(report)
|
|
if len(fits)%20==0:print(f'contact keyframes {len(fits)}/{len(samples)}',flush=True)
|
|
fits=np.array(fits)
|
|
interp=np.stack([np.interp(times,samples,fits[:,j]) for j in range(22)],axis=1)
|
|
source_active=np.stack([motion['qpos'][:,list(motion['joint_names']).index(n)] for n in solver.cal.active],axis=1)
|
|
# Blend acquisition only; preserve pre-contact source motion.
|
|
weight=np.clip((times-(onset-30))/30,0,1);weight=weight*weight*(3-2*weight)
|
|
interp[:,:16]=source_active*(1-weight[:,None])+interp[:,:16]*weight[:,None]
|
|
interp[:,16:]*=weight[:,None]
|
|
repairs=[]
|
|
for t in range(onset,count):
|
|
solver.set_pose(interp[t],motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t])
|
|
if -solver.distances().min()>.0015:
|
|
x,_,r=solver.solve(source_active[t],motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t],initial=interp[t])
|
|
if r['max_penetration_after_mm']>1.5 and t>onset:
|
|
alternative,_,retry=solver.solve(source_active[t],motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t],initial=interp[t-1])
|
|
if retry['max_penetration_after_mm']<r['max_penetration_after_mm']:x,r=alternative,retry
|
|
interp[t]=x;repairs.append(dict(frame=t,**r))
|
|
if t%100==0:print(f'frame collision repair {t}/{count}: {len(repairs)} repaired',flush=True)
|
|
qpos=[];wp=[];wq=[];penetration=[]
|
|
for t,x in enumerate(interp):
|
|
solver.set_pose(x,motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t])
|
|
values=solver.expand(x[:16]);qpos.append([values[str(n)] for n in motion['joint_names']])
|
|
wp.append(solver.data.qpos[solver.wrist:solver.wrist+3].copy());wq.append(solver.data.qpos[solver.wrist+3:solver.wrist+7].copy())
|
|
penetration.append(max(0,-solver.distances().min()))
|
|
qpos=np.array(qpos);wp=np.array(wp);wq=np.array(wq)
|
|
commands=np.array([solver.cal.command(dict(zip(motion['joint_names'],q))) for q in qpos])
|
|
np.savez_compressed(OUT/'registered_motion.npz',qpos=qpos,joint_names=motion['joint_names'],wrist_pos=wp,wrist_quat_wxyz=wq,
|
|
bottle_pos=bottle_pos,bottle_quat_wxyz=bottle_quat,command_u8=commands,fps=30,
|
|
keyframes=samples,optimization_variables=interp,penetration_m=penetration,contact_onset_frame=onset)
|
|
np.savetxt(OUT/'registered_trajectory.csv',np.c_[times/30,wp,wq,qpos,bottle_pos,bottle_quat],delimiter=',',
|
|
header=','.join(['time_s','wrist_x','wrist_y','wrist_z','wrist_qw','wrist_qx','wrist_qy','wrist_qz']+list(motion['joint_names'])+['bottle_x','bottle_y','bottle_z','bottle_qw','bottle_qx','bottle_qy','bottle_qz']),comments='')
|
|
report=dict(frames=count,fps=30,contact_keyframes=len(samples),converged_keyframes=sum(r['converged'] for r in reports),
|
|
mean_keyframe_gap_before_mm=float(np.mean([r['mean_selected_surface_gap_before_mm'] for r in reports])),
|
|
mean_keyframe_gap_after_mm=float(np.mean([r['mean_selected_surface_gap_after_mm'] for r in reports])),
|
|
max_holding_penetration_mm=float(np.max(penetration[onset:])*1000),max_all_frame_penetration_mm=float(np.max(penetration)*1000),
|
|
contact_onset_frame=onset,source='current 2047635068 video + supplied tri-camera bottle mesh',
|
|
repaired_frames=len(repairs),repair_reports=repairs,
|
|
alignment='Inferred fixed bottle-to-palm transform during holding, with contact pose corrections; no measured object trajectory or camera extrinsic alignment',
|
|
boundary='Object poses in this file are prescribed registration targets. Dynamic grasp tests are separate and have not passed.',keyframe_reports=reports)
|
|
(OUT/'sequence_validation.json').write_text(json.dumps(report,indent=2));print({k:v for k,v in report.items() if not k.endswith('_reports')},flush=True)
|
|
|
|
if __name__=='__main__':main()
|