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>
95 lines
7.5 KiB
Python
95 lines
7.5 KiB
Python
"""Refit CAD poses from observed outlines and robust depth planes; fixed CAD scale."""
|
|
import os
|
|
os.environ.setdefault('OMP_NUM_THREADS','4')
|
|
from pathlib import Path
|
|
import json,cv2,numpy as np,open3d as o3d
|
|
from scipy.spatial.transform import Rotation,Slerp
|
|
ROOT=Path(__file__).resolve().parents[1];BASE=ROOT/'output/20260915_171525_dynhamr';OUT=BASE/'object_refit_tracked';OUT.mkdir(exist_ok=True)
|
|
meta=json.loads((ROOT/'docs/20260915_171525/intrinsics.json').read_text());cam=np.load(BASE/'rgbd_camera.npz');old=np.load(BASE/'object_poses.npz');N=352
|
|
K=np.array([[meta['fx'],0,meta['cx']],[0,meta['fy'],meta['cy']],[0,0,1.]])
|
|
meshes={n:np.asarray(o3d.io.read_triangle_mesh(str(BASE/'object_assets'/f'{n}.stl')).vertices) for n in ['upper','lower']}
|
|
poses={n:[] for n in meshes};records={n:[] for n in meshes};cap=cv2.VideoCapture(str(ROOT/'docs/20260915_171525/color.mp4'))
|
|
def project(v,T):
|
|
p=v@T[:3,:3].T+T[:3,3];return (p@K.T)[:,:2]/p[:,2,None]
|
|
def score(v,T,mask):
|
|
uv=project(v,T); hull=cv2.convexHull(np.rint(uv).astype(np.int32)); pred=np.zeros_like(mask);cv2.fillConvexPoly(pred,hull,1)
|
|
obs=np.zeros_like(mask);cs,_=cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE);cv2.fillConvexPoly(obs,cv2.convexHull(max(cs,key=cv2.contourArea)),1)
|
|
return float(np.sum((pred>0)&(obs>0))/max(1,np.sum((pred>0)|(obs>0))))
|
|
polygons={'upper':np.array([[491,268],[601,275],[631,399],[488,379]],dtype=np.float32),'lower':np.array([[309,254],[436,264],[442,382],[291,375]],dtype=np.float32)}
|
|
previous_gray=None
|
|
face_choice={}
|
|
for t in range(N):
|
|
ok,im=cap.read();assert ok
|
|
gray=cv2.cvtColor(im,cv2.COLOR_BGR2GRAY)
|
|
if previous_gray is not None:
|
|
for n,poly in polygons.items():
|
|
featuremask=np.zeros_like(gray);cv2.fillConvexPoly(featuremask,poly.astype('int32'),255)
|
|
points=cv2.goodFeaturesToTrack(previous_gray,300,.01,5,mask=featuremask)
|
|
if points is not None and len(points)>=8:
|
|
nxt,st,_=cv2.calcOpticalFlowPyrLK(previous_gray,gray,points,None,winSize=(21,21),maxLevel=3)
|
|
back,st2,_=cv2.calcOpticalFlowPyrLK(gray,previous_gray,nxt,None,winSize=(21,21),maxLevel=3)
|
|
good=(st.ravel()>0)&(st2.ravel()>0)&(np.linalg.norm(back-points,axis=(1,2))<1.)
|
|
if good.sum()>=8:
|
|
aff,inl=cv2.estimateAffinePartial2D(points[good],nxt[good],method=cv2.RANSAC,ransacReprojThreshold=2.)
|
|
H=None if aff is None else np.vstack([aff,[0,0,1]])
|
|
if H is not None and inl.sum()>=8:
|
|
proposed=cv2.perspectiveTransform(poly[None],H)[0]
|
|
if np.max(np.linalg.norm(proposed-poly,axis=1))<30:polygons[n]=proposed
|
|
previous_gray=gray
|
|
hsv=cv2.cvtColor(im,cv2.COLOR_BGR2HSV);dep=cv2.imread(str(ROOT/'docs/20260915_171525/depth'/f'{t:06d}.png'),-1)*meta['depth_scale_m']
|
|
for n,v in meshes.items():
|
|
hue=hsv[:,:,0]; mask=(((hue<12)|(hue>170)) if n=='upper' else ((hue>90)&(hue<135)))&(hsv[:,:,1]>90)&(hsv[:,:,2]>35)&(dep>.12)&(dep<.85);mask[:80]=False
|
|
mask &= ~(cv2.inRange(cv2.cvtColor(im,cv2.COLOR_BGR2YCrCb),np.array([0,133,77]),np.array([255,173,127]))>0)
|
|
num,labels,stats,_=cv2.connectedComponentsWithStats(mask.astype('uint8'),8);label=1+np.argmax(stats[1:,4]);mask=(labels==label).astype('uint8')
|
|
cs,_=cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE);hull=cv2.convexHull(max(cs,key=cv2.contourArea))
|
|
if cv2.contourArea(hull)>6500 and mask.sum()/cv2.contourArea(hull)>.55:
|
|
for eps in [.015,.025,.04,.06]:
|
|
quad=cv2.approxPolyDP(hull,eps*cv2.arcLength(hull,True),True)
|
|
if len(quad)==4:
|
|
poly=quad[:,0].astype('float32');center=poly.mean(0);poly=poly[np.argsort(np.arctan2(poly[:,1]-center[1],poly[:,0]-center[0]))];poly=np.roll(poly,-np.argmin(poly.sum(1)),axis=0)
|
|
poly=min([np.roll(poly,k,axis=0) for k in range(4)],key=lambda q:np.sum((q-polygons[n])**2))
|
|
polygons[n]=poly;break
|
|
poly=polygons[n].astype(float)
|
|
observed=np.zeros_like(mask);cv2.fillConvexPoly(observed,poly.astype('int32'),1)
|
|
mask=mask*observed
|
|
# Robust plane uses only colored measured pixels, never white insert/background hull fill.
|
|
yy,xx=np.nonzero(cv2.erode(mask,np.ones((3,3),np.uint8)));z=dep[yy,xx];pts=np.c_[(xx-meta['cx'])/meta['fx']*z,(yy-meta['cy'])/meta['fy']*z,z]
|
|
if len(pts)<30:pts=np.array([[0,0,.5],[.1,0,.5],[0,.1,.5]])
|
|
pc=o3d.geometry.PointCloud(o3d.utility.Vector3dVector(pts[::max(1,len(pts)//1500)]));plane,inliers=pc.segment_plane(.004,3,100);normal=np.array(plane[:3]);rays=np.c_[poly,np.ones(4)]@np.linalg.inv(K).T;corners=rays*(-plane[3]/(rays@normal))[:,None]
|
|
lo=v.min(0);hi=v.max(0)
|
|
if n=='upper':local=np.array([[lo[0],lo[1],hi[2]],[hi[0],lo[1],hi[2]],[hi[0],lo[1],lo[2]],[lo[0],lo[1],lo[2]]])
|
|
else:local=np.array([[lo[0],-.0172,lo[2]],[hi[0],-.0172,lo[2]],[hi[0],-.0172,hi[2]],[lo[0],-.0172,hi[2]]])
|
|
A=local-local.mean(0);B=corners-corners.mean(0);u,s,vt=np.linalg.svd(A.T@B);rr=vt.T@np.diag([1,1,np.linalg.det(vt.T@u.T)])@u.T
|
|
T=np.eye(4);T[:3,:3]=rr;T[:3,3]=corners.mean(0)-rr@local.mean(0)
|
|
before=np.eye(4);before[:3,:3]=Rotation.from_quat(old[n+'_quaternion_xyzw'][t]).as_matrix();before[:3,3]=old[n+'_position_world'][t];before=np.linalg.inv(cam['c2w'][t])@before
|
|
candidates=[]
|
|
for face in [lo[1],hi[1]]:
|
|
for flip in [False,True]:
|
|
loc=local.copy();loc[:,1]=face
|
|
if flip:loc=loc[[3,2,1,0]]
|
|
success,rv,tv=cv2.solvePnP(loc.astype(float),poly.astype(float),K,None,flags=cv2.SOLVEPNP_ITERATIVE)
|
|
if success:
|
|
candidate=np.eye(4);candidate[:3,:3]=Rotation.from_rotvec(rv.ravel()).as_matrix();candidate[:3,3]=tv.ravel();candidates.append(candidate)
|
|
previous=None if not poses[n] else np.linalg.inv(cam['c2w'][t])@poses[n][-1]
|
|
if n not in face_choice:face_choice[n]=int(np.argmax([score(v,candidate,observed) for candidate in candidates]))
|
|
T=candidates[face_choice[n]]
|
|
rr=T[:3,:3]
|
|
iou0=score(v,before,observed);iou1=score(v,T,observed);rmse=np.sqrt(np.mean(np.sum((local@rr.T+T[:3,3]-corners)**2,axis=1)))
|
|
center_step=0 if previous is None else np.linalg.norm((T[:3,:3]@v.mean(0)+T[:3,3])-(previous[:3,:3]@v.mean(0)+previous[:3,3]))
|
|
angle_step=0 if previous is None else Rotation.from_matrix(previous[:3,:3].T@T[:3,:3]).magnitude()
|
|
accepted=bool(iou1>.60 and np.isfinite(T).all() and center_step<.04 and angle_step<.3)
|
|
# Keep candidate for continuity diagnostics; rejected poses are excluded at interpolation.
|
|
poses[n].append(cam['c2w'][t]@T);records[n].append(dict(frame=t,accepted=accepted,old_outline_iou=iou0,new_outline_iou=iou1,corner_fit_rmse_m=float(rmse)))
|
|
if t in [0,176,351]:
|
|
pic=im.copy()
|
|
for transform,color in [(before,(0,200,255)),(T,(0,255,0))]:cv2.polylines(pic,[cv2.convexHull(np.rint(project(v,transform)).astype('int32'))],True,color,2)
|
|
cv2.imwrite(str(OUT/f'{n}_{t:04d}_overlay.jpg'),pic)
|
|
if t%80==0:print('refit',t,flush=True)
|
|
cap.release();(OUT/'all_candidates.json').write_text(json.dumps(records));arrays={'time':cam['time']};summary={}
|
|
for n in meshes:
|
|
P=np.asarray(poses[n]);good=np.array([r['accepted'] for r in records[n]]);idx=np.flatnonzero(good);assert len(idx)>N*.5,(n,len(idx))
|
|
ts=np.clip(np.arange(N),idx[0],idx[-1]);p=np.stack([np.interp(ts,idx,P[idx,k,3]) for k in range(3)],axis=1);rot=Slerp(idx,Rotation.from_matrix(P[idx,:3,:3]))(ts)
|
|
arrays.update({n+'_position_world':p,n+'_quaternion_xyzw':rot.as_quat(),n+'_keyframes':np.arange(N),n+'_keyframe_valid':good})
|
|
summary[n]=dict(accepted=int(good.sum()),old_iou_median=float(np.median([r['old_outline_iou'] for r in records[n]])),new_iou_median=float(np.median([r['new_outline_iou'] for r in records[n]])),frames=records[n])
|
|
np.savez_compressed(OUT/'object_poses.npz',**arrays);(OUT/'validation.json').write_text(json.dumps(summary,indent=2));print({n:{k:v for k,v in s.items() if k!='frames'} for n,s in summary.items()})
|