Files
hand-motion-pipeline/scripts/fit_rgbd_objects.py
T
liyang ae28d55f81 Update to 2026-09-17 pipeline snapshot; add weights, L20 assets and recording via Git LFS
Source: RGB-D -> Dyn-HaMR -> L20 retargeting -> FoundationPose -> reference repair -> SPIDER,
documented in docs/PIPELINE_LATEST.md and docs/SETUP_AND_WEIGHTS.md. Adds FoundationPose and
nvdiffrast upstream snapshots, requirements/pipeline_venv.txt and the FoundationPose weight
manifest/downloader.

Assets (Git LFS): weights/ (WiLoR detector, HandFlow denoiser, UniDepth-L), FoundationPose
checkpoints, HaMeR checkpoint, Dyn-HaMR HMP model and BMC constraints, L20 URDF/meshes, the
20260915_171525 D405 recording and the two box CADs. MANO models are not redistributed
(third_party/hamer/_DATA/data/mano/README.txt). Environments, caches and run outputs excluded.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-17 11:43:37 +08:00

94 lines
6.4 KiB
Python

"""Estimate two colored object poses from RGB-D; partial-surface ICP, not pose truth."""
import os
os.environ.setdefault('OMP_NUM_THREADS','4')
from pathlib import Path
import json
import cv2
import numpy as np
import open3d as o3d
from scipy.spatial.transform import Rotation,Slerp
from scipy.ndimage import gaussian_filter1d
ROOT=Path(__file__).resolve().parents[1];SRC=ROOT/'docs/20260915_171525';BASE=ROOT/'output/20260915_171525'
meta=json.loads((SRC/'intrinsics.json').read_text());cam=np.load(BASE/'rgbd_camera.npz');N=len(cam['time'])
o3d.utility.random.seed(7)
models={}
for name in ['upper','lower']:
mesh=o3d.io.read_triangle_mesh(str(BASE/'object_assets'/f'{name}.stl'))
cloud=mesh.sample_points_uniformly(6000);cloud=cloud.voxel_down_sample(.003)
models[name]=(cloud,np.asarray(mesh.vertices).mean(0),mesh.get_axis_aligned_bounding_box().get_center())
cap=cv2.VideoCapture(str(SRC/'color.mp4'));clouds={n:[] for n in models};centers={n:[] for n in models}
for t in range(N):
ok,im=cap.read();assert ok
hsv=cv2.cvtColor(im,cv2.COLOR_BGR2HSV);z=cv2.imread(str(SRC/'depth'/f'{t:06d}.png'),-1)*meta['depth_scale_m']
for name in models:
hue=hsv[:,:,0];mask=((hue<12)|(hue>170)) if name=='upper' else ((hue>90)&(hue<135))
mask=(mask&(hsv[:,:,1]>100)&(hsv[:,:,2]>35)&(z>.12)&(z<.85)).astype(np.uint8)
# Restrict to the tabletop work area, then largest coherent colored component.
mask[:80]=0
count,labels,stats,_=cv2.connectedComponentsWithStats(mask,8)
if count<=1 or stats[1:,cv2.CC_STAT_AREA].max()<300:
clouds[name].append(None);centers[name].append(None);continue
selected=1+np.argmax(stats[1:,cv2.CC_STAT_AREA]);mask=(labels==selected).astype(np.uint8)
contours,_=cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
if name=='lower' and contours:
hull=cv2.convexHull(max(contours,key=cv2.contourArea));cv2.fillConvexPoly(mask,hull,1)
ycc=cv2.cvtColor(im,cv2.COLOR_BGR2YCrCb)
mask[cv2.inRange(ycc,np.array([0,133,77]),np.array([255,173,127]))>0]=0
mask[(z<=.12)|(z>=.85)]=0
mask=cv2.erode(mask,np.ones((3,3),np.uint8))>0
v,u=np.nonzero(mask);zz=z[v,u];points=np.c_[(u-meta['cx'])/meta['fx']*zz,(v-meta['cy'])/meta['fy']*zz,zz]
points=points@cam['c2w'][t,:3,:3].T+cam['c2w'][t,:3,3]
pc=o3d.geometry.PointCloud(o3d.utility.Vector3dVector(points)).voxel_down_sample(.004)
clouds[name].append(pc);centers[name].append(np.median(np.asarray(pc.points),axis=0))
cap.release()
arrays={};report={}
for name,(model,_,model_center) in models.items():
poses=[];good=[];records=[];previous=None;previous_center=None
frames=np.unique(np.r_[np.arange(0,N,3),N-1])
for t in frames:
target=clouds[name][t]
if target is None:
poses.append(np.eye(4) if previous is None else previous.copy());good.append(False);records.append(dict(frame=int(t),accepted=False));continue
points=np.asarray(target.points);center=centers[name][t]
seeds=[]
if previous is not None:
seed=previous.copy();seed[:3,3]+=center-previous_center;seeds.append(seed)
# Update observed face normal continuously, preserving the previous
# in-plane direction rather than reselecting a symmetric CAD yaw.
plane,inliers=target.segment_plane(.004,3,80)
normal=np.asarray(plane[:3]);normal/=np.linalg.norm(normal)
if normal@previous[:3,1]<0:normal=-normal
x=previous[:3,0]-normal*(normal@previous[:3,0]);x/=np.linalg.norm(x)
rr=np.column_stack([x,normal,np.cross(x,normal)])
normal_seed=np.eye(4);normal_seed[:3,:3]=rr;normal_seed[:3,3]=center-rr@model_center;seeds.append(normal_seed)
if previous is None:
_,axes=np.linalg.eigh(np.cov(points.T));normal=axes[:,0];x=axes[:,2];zaxis=np.cross(x,normal)
for sign in [1,-1]:
basis=np.column_stack([x,normal*sign,zaxis*sign])
for yaw in [0,np.pi/2,np.pi,3*np.pi/2]:
R=basis@Rotation.from_euler('y',yaw).as_matrix();T=np.eye(4);T[:3,:3]=R;T[:3,3]=center-R@model_center;seeds.append(T)
candidates=[]
for seed in seeds:
fit=o3d.pipelines.registration.registration_icp(target,model,.04,np.linalg.inv(seed),o3d.pipelines.registration.TransformationEstimationPointToPoint(),o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=35))
candidates.append((fit.inlier_rmse+(.1*(1-fit.fitness)),fit))
_,fit=min(candidates,key=lambda x:x[0]);T=np.linalg.inv(fit.transformation)
rotation_step=0. if previous is None else float(Rotation.from_matrix(previous[:3,:3].T@T[:3,:3]).magnitude())
accepted=bool(fit.fitness>.8 and fit.inlier_rmse<.015 and np.isfinite(T).all() and rotation_step<.7)
if accepted:previous=T;previous_center=center
poses.append(T if accepted or previous is None else previous.copy());good.append(accepted)
records.append(dict(frame=int(t),accepted=accepted,fitness=float(fit.fitness),rmse_m=float(fit.inlier_rmse),rotation_step_rad=rotation_step))
if t%60==0:print(name,t,'rmse',fit.inlier_rmse,'fitness',fit.fitness,flush=True)
poses=np.asarray(poses);good=np.asarray(good);ix=np.flatnonzero(good)
assert len(ix)>2,(name,len(ix))
knots=frames[ix];ts=np.clip(np.arange(N),knots[0],knots[-1]);p=np.stack([np.interp(ts,knots,poses[ix,k,3]) for k in range(3)],axis=1)
rot=Slerp(knots,Rotation.from_matrix(poses[ix,:3,:3]))(ts)
arrays[name+'_position_world']=gaussian_filter1d(p,1,axis=0)
arrays[name+'_quaternion_xyzw']=rot.as_quat()
arrays[name+'_keyframes']=frames;arrays[name+'_keyframe_valid']=good
# Interpolation is explicit; accepted ICP is geometric fit, not unique pose identification.
report[name]=dict(keyframes=len(frames),accepted_keyframes=int(good.sum()),median_rmse_m=float(np.median([r['rmse_m'] for r in records if r['accepted']])),frames=records)
np.savez_compressed(BASE/'object_poses.npz',time=cam['time'],**arrays)
report['boundary']='Red=upper, blue=lower assumed from recording. Fixed meter-scale STL, partial colored visible-surface ICP; pose symmetry and occlusion ambiguity unresolved. Interpolated between keyframes; not contact labels or ground truth.'
(BASE/'object_fit_validation.json').write_text(json.dumps(report,indent=2));print(json.dumps({n:{k:v for k,v in r.items() if k!='frames'} for n,r in report.items() if isinstance(r,dict)},indent=2))