liyang
ae28d55f81
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>
94 lines
6.4 KiB
Python
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))
|