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>
53 lines
4.4 KiB
Python
53 lines
4.4 KiB
Python
"""FoundationPose RGB-D tracking for both CAD parts, preserving camera/world transforms."""
|
|
import os
|
|
os.environ.setdefault('OMP_NUM_THREADS','4')
|
|
import sys,json,time,subprocess
|
|
from pathlib import Path
|
|
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT/'third_party/FoundationPose'))
|
|
import cv2,numpy as np,torch,trimesh,nvdiffrast.torch as dr
|
|
from estimater import FoundationPose,ScorePredictor,PoseRefinePredictor
|
|
from Utils import nvdiffrast_render,make_mesh_tensors
|
|
OUT=ROOT/'output/foundationpose_spider_20260915';SRC=ROOT/'docs/20260915_171525';m=json.loads((SRC/'intrinsics.json').read_text());N=m['frames'];K=np.array([[m['fx'],0,m['cx']],[0,m['fy'],m['cy']],[0,0,1.]])
|
|
cam=np.load(ROOT/'output/20260915_171525_dynhamr/rgbd_camera.npz');c=cv2.VideoCapture(str(SRC/'color.mp4'));frames=[]
|
|
while True:
|
|
ok,b=c.read()
|
|
if not ok:break
|
|
frames.append(b)
|
|
assert len(frames)==N
|
|
ctx=dr.RasterizeCudaContext();score=ScorePredictor();refine=PoseRefinePredictor();arrays={'K':K,'time':np.array(m['timestamps_ms'])/1000-m['timestamps_ms'][0]/1000,'c2w':cam['c2w'],'camera_valid':cam['valid']};meshes={};allmetrics={}
|
|
for name,file,init,color in [('upper','上半.stl',0,[210,35,35,255]),('lower','下半.stl',80,[35,65,210,255])]:
|
|
mesh=trimesh.load(ROOT/'docs'/file);mesh.visual.vertex_colors=np.tile(color,(len(mesh.vertices),1));meshes[name]=mesh
|
|
est=FoundationPose(model_pts=mesh.vertices,model_normals=mesh.vertex_normals,mesh=mesh,scorer=score,refiner=refine,glctx=ctx,debug=0,debug_dir=str(OUT/name/'debug'));mt=make_mesh_tensors(mesh);P=np.full((N,4,4),np.nan);valid=np.zeros(N,bool);metrics=[None]*N;start_state=None
|
|
for i in list(range(init,N))+list(range(init-1,-1,-1)):
|
|
b=frames[i];z=cv2.imread(str(SRC/'depth'/f'{i:06d}.png'),-1).astype('float32')*m['depth_scale_m'];z[(z<.1)|(z>1.)]=0
|
|
h=cv2.cvtColor(b,cv2.COLOR_BGR2HSV);hue=h[:,:,0];mask=((((hue<12)|(hue>170)) if name=='upper' else ((hue>95)&(hue<135)))&(h[:,:,1]>90)&(h[:,:,2]>40)&(z>.1)&(z<1.)).astype('uint8');mask[:200]=0
|
|
num,lab,st,_=cv2.connectedComponentsWithStats(mask,8)
|
|
if num>1:mask=(lab==1+np.argmax(st[1:,4])).astype('uint8')
|
|
if i==init-1:est.pose_last=start_state.clone()
|
|
rgb=cv2.cvtColor(b,cv2.COLOR_BGR2RGB)
|
|
if i==init:
|
|
cv2.imwrite(str(OUT/name/'initial_mask.png'),mask*255);p=est.register(K=K,rgb=rgb,depth=z,ob_mask=mask,iteration=5);start_state=est.pose_last.clone()
|
|
else:p=est.track_one(rgb=rgb,depth=z,K=K,iteration=2)
|
|
P[i]=p
|
|
with torch.inference_mode():_,depth,_=nvdiffrast_render(K=K,H=480,W=848,ob_in_cams=torch.as_tensor(p[None],device='cuda'),glctx=ctx,mesh_tensors=mt)
|
|
depth=depth[0].cpu().numpy();over=(depth>0)&(mask>0)&(z>0);coverage=float(over.sum()/max(1,mask.sum()));err=float(np.median(np.abs(depth[over]-z[over]))) if over.any() else 1.
|
|
valid[i]=mask.sum()>500 and coverage>.65 and err<.02
|
|
# The lower part becomes occluded during assembly; independent lower tracking is not reliable there.
|
|
if name=='lower' and i>=177:valid[i]=False
|
|
metrics[i]={'frame':i,'valid':bool(valid[i]),'visible_color_pixels':int(mask.sum()),'coverage':coverage,'depth_residual_m':err}
|
|
if i%60==0:print(name,metrics[i],flush=True)
|
|
arrays[name+'_T_camera']=P;arrays[name+'_T_world']=cam['c2w']@P;arrays[name+'_valid']=valid;allmetrics[name]=metrics
|
|
np.savez_compressed(OUT/'foundationpose_objects.npz',**arrays);(OUT/'object_metrics.json').write_text(json.dumps(allmetrics,indent=2))
|
|
writer=subprocess.Popen(['ffmpeg','-y','-v','error','-f','rawvideo','-pix_fmt','bgr24','-s','848x480','-r','30','-i','-','-an','-c:v','libx264','-crf','20','-pix_fmt','yuv420p',str(OUT/'foundationpose_overlay.mp4')],stdin=subprocess.PIPE)
|
|
for i,b in enumerate(frames):
|
|
pic=b.copy()
|
|
for name,col in [('upper',[0,220,0]),('lower',[220,180,0])]:
|
|
with torch.inference_mode():_,d,_=nvdiffrast_render(K=K,H=480,W=848,ob_in_cams=torch.tensor(arrays[name+'_T_camera'][i:i+1],device='cuda',dtype=torch.float32),glctx=ctx,mesh_tensors=make_mesh_tensors(meshes[name]))
|
|
mask=d[0].cpu().numpy()>0
|
|
if arrays[name+'_valid'][i]:pic[mask]=(pic[mask]*.65+np.array(col)*.35).astype('uint8')
|
|
cv2.putText(pic,f'{i:03d} red:{arrays["upper_valid"][i]} blue:{arrays["lower_valid"][i]}',(8,25),0,.55,(0,255,255),2)
|
|
writer.stdin.write(pic.tobytes())
|
|
if i in [0,40,80,120,160,200,240,280,351]:cv2.imwrite(str(OUT/f'objects_{i:04d}.jpg'),pic)
|
|
writer.stdin.close();assert writer.wait()==0
|
|
print('FOUNDATIONPOSE_DONE', {n:int(arrays[n+'_valid'].sum()) for n in meshes},flush=True)
|