hand-motion-pipeline/scripts/export_camera_delivery.py

"""Package the 1333-frame replay cameras, source cameras and frame mapping."""
import os
os.environ.setdefault('MUJOCO_GL', 'osmesa')
import csv
import hashlib
import json
from pathlib import Path
import shutil
import subprocess
import tarfile

import cv2
import mujoco
import numpy as np
from scipy.spatial.transform import Rotation

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT/'output/camera_delivery_1333'


def save_json(path, value):
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2)+'\n')


def pose(R, t):
    result = np.broadcast_to(np.eye(4), (*np.asarray(R).shape[:-2],4,4)).copy()
    result[..., :3,:3] = R
    result[..., :3,3] = t
    return result


def probe(path):
    return json.loads(subprocess.check_output([
        'ffprobe','-v','error','-select_streams','v:0','-show_streams',
        '-show_frames','-show_entries',
        'stream=width,height,avg_frame_rate,time_base,nb_frames,start_time,duration:frame=best_effort_timestamp,best_effort_timestamp_time',
        '-of','json',str(path)]))


def main():
    OUT.mkdir(exist_ok=True)
    replay = ROOT/'output/l20_full_replay'
    original = Path('/home/timessage/下载/2047635068.mp4')
    final = np.load(replay/'motion.npz')
    temporal = np.load(ROOT/'output/l20_2047635068/jitter_audit/temporal/motion.npz')
    reconstruction = np.load(ROOT/'output/results_2047635068/world_results.npz')
    vipe_k = np.load(ROOT/'output/vipe_2047635068/intrinsics/2047635068.npz')
    vipe_pose = np.load(ROOT/'output/vipe_2047635068/pose/2047635068.npz')
    N = len(final['qpos'])
    assert N == 1333 and np.array_equal(vipe_k['inds'],np.arange(N))
    assert np.array_equal(vipe_pose['inds'],np.arange(N))
    video_meta = {name:probe(path) for name,path in {
        'original':original,'full_replay':replay/'full_replay.mp4',
        'comparison':replay/'original_vs_mujoco.mp4'}.items()}
    for value in video_meta.values():
        assert len(value['frames']) == N
    timestamps = {name:np.array([float(x['best_effort_timestamp_time']) for x in value['frames']])
                  for name,value in video_meta.items()}
    with (OUT/'frame_mapping.csv').open('w') as f:
        writer=csv.writer(f)
        writer.writerow(['hdf5_frame_0based','source_video_frame_0based','source_extracted_jpg',
                         'source_pts_ticks','source_pts_s','hdf5_time_s','full_replay_pts_s','comparison_pts_s'])
        for i in range(N):
            writer.writerow([i,i,f'{i+1:06d}.jpg',video_meta['original']['frames'][i]['best_effort_timestamp'],
                             timestamps['original'][i],final['time'][i],timestamps['full_replay'][i],timestamps['comparison'][i]])
    # Reconstruct the fixed world transforms from the actual build script.
    source_R = Rotation.from_quat(temporal['wrist_quat_wxyz'][:,[1,2,3,0]])
    bottle_R = source_R * Rotation.from_euler('x',-np.pi/2)
    bottle_p = temporal['wrist_pos'] + source_R.apply(np.tile([.12,-.065,.15],(N,1)))
    align = bottle_R[600].inv()
    shift = np.array([0,0,.3])-align.apply(bottle_p[600])
    before_floor = align.apply(temporal['wrist_pos'])+shift+final['correction'][:,16:19]
    residual = final['wrist_pos']-before_floor
    floor_shift = residual.mean(axis=0)
    assert np.max(np.abs(residual-floor_shift)) < 1e-9
    assert np.max(np.abs(floor_shift[:2])) < 1e-9
    final_from_temporal = pose(align.as_matrix(),shift+floor_shift)
    temporal_from_dyn = pose(temporal['scene_rotation'],temporal['scene_translation'])
    final_from_dyn = final_from_temporal @ temporal_from_dyn
    track = int(np.load(ROOT/'output/dex_2047635068/human_joints.npz')['source_track'])
    camera_from_dyn = pose(reconstruction['cam_R'][track],reconstruction['cam_t'][track])
    camera_from_final = camera_from_dyn @ np.linalg.inv(final_from_dyn)
    final_from_camera = np.linalg.inv(camera_from_final)
    # world_results cam_t already contains world_scale; do not scale twice.
    K4 = reconstruction['intrins'].astype(float)
    K = np.array([[K4[0],0,K4[2]],[0,K4[1],K4[3]],[0,0,1.]])
    np.savez_compressed(OUT/'source_camera.npz',
                        frame_index=np.arange(N), time=final['time'], intrinsics_fx_fy_cx_cy=vipe_k['data'],
                        K=K, camera_from_final_world=camera_from_final,
                        final_world_from_camera=final_from_camera,
                        camera_from_dynhamr_world=camera_from_dyn,
                        final_world_from_dynhamr_world=final_from_dyn,
                        vipe_world_from_camera_raw=vipe_pose['data'])
    with (OUT/'source_intrinsics.csv').open('w') as f:
        w=csv.writer(f);w.writerow(['frame_0based','fx','fy','cx','cy','width','height'])
        for i,row in enumerate(vipe_k['data']):w.writerow([i,*row,1280,720])
    with (OUT/'source_extrinsics.csv').open('w') as f:
        w=csv.writer(f);w.writerow(['frame_0based']+[f'camera_from_final_world_{i}{j}' for i in range(4) for j in range(4)])
        for i,T in enumerate(camera_from_final):w.writerow([i,*T.ravel()])
    source_config = dict(width=1280,height=720,fps=30,K=K.tolist(),
                         model='estimated pinhole',calibration='ViPE/Dyn-HaMR estimate, not measured calibration',
                         distortion_coefficients=None,undistortion='No explicit lens undistortion found in this inference path; capture-device processing unknown.',
                         preprocessing='Input video and extracted JPG are 1280x720; no full-image crop/resize observed in those inputs. Network-internal crops/resizes do not redefine this full-image K.',
                         camera_axes='OpenCV: +X right, +Y down, +Z forward; pixel coordinates u right, v down',
                         world_axes='Final scene.xml right-handed display world, Z display up; physical gravity not calibrated',
                         transform_convention='column vectors: p_camera = camera_from_final_world @ p_final_world; inverse is final_world_from_camera',
                         world_scale_already_applied=float(reconstruction['world_scale'].item()),
                         intrinsics_max_variation=float(np.ptp(vipe_k['data'],axis=0).max()),
                         warning='Contact registration changes hand motion per frame. This camera transform preserves the fixed world change, not inverse contact fitting. Exact original-video hand overlap is not guaranteed.')
    save_json(OUT/'source_camera.json',source_config)
    save_json(OUT/'world_transforms.json',dict(final_world_from_dynhamr_world=final_from_dyn.tolist(),
                  temporal_world_from_dynhamr_world=temporal_from_dyn.tolist(),
                  final_world_from_temporal_world=final_from_temporal.tolist(),
                  recovered_floor_translation=floor_shift.tolist(),
                  per_frame_contact_correction_included_in_camera=False,
                  derivation='scripts/build_l20_full_replay.py: alignment using bottle orientation at frame 600, shift to [0,0,.3], then floor clearance. Floor shift recovered from saved wrist minus fixed transform minus saved contact correction, verified constant across all frames.'))
    # Recover actual mono renderer cameras, including the moving close-up lookat.
    m=mujoco.MjModel.from_xml_path(str(replay/'scene.xml'));d=mujoco.MjData(m)
    option=mujoco.MjvOption();option.sitegroup[:]=0;option.geomgroup[3:]=0
    scene=mujoco.MjvScene(m,maxgeom=10000)
    cameras=[]
    for distance,azimuth,elevation in [(1.05,135,-23),(.62,35,-12)]:
        c=mujoco.MjvCamera();c.lookat[:]=final['camera_center'];c.distance=distance;c.azimuth=azimuth;c.elevation=elevation;cameras.append(c)
    arrays={name:[] for name in ['wide_world_from_camera_cv','close_world_from_camera_cv','close_lookat','wide_gl_pos','close_gl_pos']}
    for row in final['qpos']:
        d.qpos[:]=row;mujoco.mj_forward(m,d)
        cameras[1].lookat[:]=(d.xpos[m.body('hand_base_link').id]+d.xpos[m.body('right_object').id])/2
        cameras[1].lookat[2]+=.06
        arrays['close_lookat'].append(cameras[1].lookat.copy())
        for name,c in zip(['wide','close'],cameras):
            mujoco.mjv_updateScene(m,d,option,None,c,mujoco.mjtCatBit.mjCAT_ALL,scene)
            gl=mujoco.mjv_averageCamera(scene.camera[0],scene.camera[1])
            forward=gl.forward.astype(float);forward/=np.linalg.norm(forward)
            right=np.cross(forward,gl.up);right/=np.linalg.norm(right)
            up=np.cross(right,forward)
            arrays[name+'_world_from_camera_cv'].append(pose(np.column_stack([right,-up,forward]),gl.pos))
            arrays[name+'_gl_pos'].append(gl.pos.copy())
    arrays={k:np.asarray(v) for k,v in arrays.items()}
    fy=720/(2*np.tan(np.deg2rad(float(m.vis.global_.fovy))/2))
    render_K=np.array([[fy,0,320],[0,fy,360],[0,0,1.]])
    np.savez_compressed(OUT/'render_cameras.npz',frame_index=np.arange(N),time=final['time'],
                        K_panel=render_K,**arrays,
                        wide_camera_from_world_cv=np.linalg.inv(arrays['wide_world_from_camera_cv']),
                        close_camera_from_world_cv=np.linalg.inv(arrays['close_world_from_camera_cv']))
    save_json(OUT/'render_camera_config.json',dict(mujoco_version=mujoco.__version__,
        panel_width=640,panel_height=720,full_video_width=1280,full_video_height=720,
        vertical_fov_degrees=float(m.vis.global_.fovy),K_panel=render_K.tolist(),
        pixel_convention='K uses continuous image-edge coordinates, pixel centers at (i+0.5,j+0.5); subtract 0.5 from cx/cy for integer pixel-center indexing.',
        wide=dict(lookat=final['camera_center'].tolist(),distance=1.05,azimuth=135,elevation=-23,panel_x=0),
        close=dict(lookat='(hand_base_link.xpos + right_object.xpos)/2 + [0,0,.06], per frame',distance=.62,azimuth=35,elevation=-12,panel_x=640),
        stereo='monoscopic, average of MjvScene left/right GL cameras',
        option=dict(sitegroup='all zero',geomgroup_3_and_above='zero'),
        banners=dict(top_y=[0,64],bottom_y=[690,720]),
        comparison_video=dict(width=1600,height=720,
             original_panel='source scaled 1280x720 -> 960x540, offset [0,90]',
             render_panel='full_replay crop x=640,y=64,w=640,h=626, then pad y=64 and place x=960',
             source_image_to_comparison=[[.75,0,0],[0,.75,90],[0,0,1]],
             close_panel_to_comparison=[[1,0,960],[0,1,0],[0,0,1]])))
    # Check time correspondence, coordinate algebra and a rendered frame.
    sample_frames=[0,160,600,1332]
    pixel_errors=[]
    cap=cv2.VideoCapture(str(original))
    for i in sample_frames:
        cap.set(cv2.CAP_PROP_POS_FRAMES,i);ok,im=cap.read();assert ok
        jpg=cv2.imread(str(ROOT/f'third_party/Dyn-HaMR/test/images/2047635068/{i+1:06d}.jpg'))
        assert jpg.shape==im.shape
        pixel_errors.append(float(np.mean(np.abs(im.astype(float)-jpg.astype(float)))))
    cap.release()
    renderer=mujoco.Renderer(m,height=720,width=640)
    render_errors=[]
    cap=cv2.VideoCapture(str(replay/'full_replay.mp4'))
    for i in [0,600,1332]:
        d.qpos[:]=final['qpos'][i];mujoco.mj_forward(m,d)
        cameras[1].lookat[:]=arrays['close_lookat'][i]
        panels=[]
        for c in cameras:
            renderer.update_scene(d,camera=c,scene_option=option);panels.append(renderer.render().copy())
        generated=np.concatenate(panels,axis=1)
        cap.set(cv2.CAP_PROP_POS_FRAMES,i);ok,video=cap.read();assert ok
        video=cv2.cvtColor(video,cv2.COLOR_BGR2RGB)
        error=float(np.mean(np.abs(generated[65:689].astype(float)-video[65:689].astype(float))))
        render_errors.append(error)
        if i==0:cv2.imwrite(str(OUT/'render_check_frame_0000.png'),cv2.cvtColor(generated,cv2.COLOR_RGB2BGR))
    renderer.close();cap.release()
    assert max(render_errors)<8,render_errors
    assert max(pixel_errors)<8,pixel_errors
    points=np.column_stack([np.linspace(0,.1,N),np.ones(N)*.2,np.ones(N),np.ones(N)])
    lhs=np.einsum('tij,tj->ti',camera_from_dyn,points)
    rhs=np.einsum('tij,tj->ti',camera_from_final,points@final_from_dyn.T)
    validation=dict(frames=N,source_to_hdf5_time_max_error_s=float(np.max(np.abs(timestamps['original']-final['time']))),
                    world_transform_roundtrip_max_error=float(np.max(np.abs(lhs-rhs))),
                    floor_shift_consistency_max_error_m=float(np.max(np.abs(residual-floor_shift))),
                    source_jpg_sample_frames=sample_frames,source_jpg_mean_absolute_pixel_error=pixel_errors,
                    rerender_checked_frames=[0,600,1332],rerender_mean_absolute_pixel_error=render_errors,
                    cameras_finite=bool(np.isfinite(camera_from_final).all()),
                    calibration_is_measured=False,source_distortion_calibration_available=False)
    assert validation['source_to_hdf5_time_max_error_s']<1e-6
    assert validation['world_transform_roundtrip_max_error']<1e-8
    save_json(OUT/'validation.json',validation)
    save_json(OUT/'video_metadata.json',{k:v['streams'][0] for k,v in video_meta.items()})
    target=OUT/'replay';target.mkdir(exist_ok=True)
    for name in ['scene.xml','motion.npz','full_replay.mp4','original_vs_mujoco.mp4']:
        shutil.copy2(replay/name,target/name)
    shutil.copytree(replay/'assets',target/'assets',dirs_exist_ok=True)
    shutil.copy2(original,OUT/'original_2047635068.mp4')
    shutil.copy2(ROOT/'scripts/play_l20_full.py',OUT/'play_l20_full.py')
    shutil.copy2(replay/'make_comparison.sh',OUT/'make_comparison_original.sh')
    shutil.copy2(ROOT/'scripts/build_l20_full_replay.py',OUT/'build_l20_full_replay_reference.py')
    shutil.copy2(ROOT/'output/hdf5_delivery/right_1333/demonstrations_right.hdf5',OUT/'demonstrations_right.hdf5')
    print(json.dumps(validation,indent=2))


if __name__ == '__main__':
    main()