hand-motion-pipeline/scripts/track_red_box_foundationpose.py

"""Offline RGB-D red-box tracking, with explicit partial-view and quality flags."""
import os
os.environ.setdefault('OMP_NUM_THREADS','4')
import sys,json,time,argparse,logging
from pathlib import Path
import cv2,numpy as np
ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT/'third_party/FoundationPose'))
from estimater import FoundationPose,ScorePredictor,PoseRefinePredictor
from Utils import nvdiffrast_render,make_mesh_tensors
import torch,trimesh,nvdiffrast.torch as dr
from scipy.spatial.transform import Rotation

def redmask(b,z):
 h=cv2.cvtColor(b,cv2.COLOR_BGR2HSV)
 m=(((h[:,:,0]<10)|(h[:,:,0]>170))&(h[:,:,1]>115)&(h[:,:,2]>65)&(z>.1)&(z<.85)).astype('uint8')
 m[:150]=0
 n,l,st,c=cv2.connectedComponentsWithStats(m,8)
 if n<2:return m*0
 j=1+np.argmax(st[1:,4]);return (l==j).astype('uint8')

def main():
 a=argparse.ArgumentParser();a.add_argument('--end',type=int,default=893);a.add_argument('--init',type=int,default=30);a.add_argument('--output',default='output/foundationpose_red_20260916');args=a.parse_args()
 out=ROOT/args.output;out.mkdir(parents=True,exist_ok=True);(out/'snapshots').mkdir(exist_ok=True)
 logging.basicConfig(level=logging.WARNING)
 src=ROOT/'docs/20260916_104026';meta=json.loads((src/'intrinsics.json').read_text());K=np.array([[meta['fx'],0,meta['cx']],[0,meta['fy'],meta['cy']],[0,0,1.]])
 cap=cv2.VideoCapture(str(src/'color.mp4'));frames=[]
 while True:
  ok,b=cap.read()
  if not ok:break
  frames.append(b)
 cap.release();N=min(len(frames),args.end);assert len(frames)==meta['frames']
 mesh=trimesh.load(ROOT/'docs/上半.stl');mesh.visual.vertex_colors=np.tile([200,35,35,255],(len(mesh.vertices),1))
 ctx=dr.RasterizeCudaContext();est=FoundationPose(model_pts=mesh.vertices,model_normals=mesh.vertex_normals,mesh=mesh,glctx=ctx,debug=0,debug_dir=str(out/'debug'))
 original_tensors=make_mesh_tensors(mesh);H,W=frames[0].shape[:2]
 poses=np.full((N,4,4),np.nan);records=[None]*N;init_state=None
 sequence=list(range(args.init,N))+list(range(args.init-1,-1,-1));start=time.time()
 for i in sequence:
  b=frames[i];z=cv2.imread(str(src/'depth'/f'{i:06d}.png'),-1).astype(np.float32)*meta['depth_scale_m'];z[(z<.1)|(z>.85)]=0;m=redmask(b,z);yy,xx=np.nonzero(m)
  row={'frame':i,'red_pixels':len(xx),'status':'not_observable','accepted':False};pic=b.copy()
  # After the red box is put aside it is severely cropped; do not invent its full pose.
  observable=i<=335  # Manually reviewed red manipulation interval for this recording.
  if i==args.init-1:est.pose_last=init_state.clone()
  if observable:
   rgb=cv2.cvtColor(b,cv2.COLOR_BGR2RGB)
   if i==args.init:
    cv2.imwrite(str(out/'initial_mask.png'),m*255);cv2.imwrite(str(out/'initial_frame.png'),b)
    pose=est.register(K=K,rgb=rgb,depth=z,ob_mask=m,iteration=5);init_state=est.pose_last.clone()
   else:pose=est.track_one(rgb=rgb,depth=z,K=K,iteration=2)
   poses[i]=pose
   with torch.inference_mode():
    _,rd,_=nvdiffrast_render(K=K,H=H,W=W,ob_in_cams=torch.as_tensor(pose[None],device='cuda'),glctx=ctx,mesh_tensors=original_tensors)
   rd=rd[0].cpu().numpy();render=rd>0;overlap=render&(m>0)&(z>0)
   coverage=float((render&(m>0)).sum()/max(1,m.sum()));err=float(np.median(np.abs(rd[overlap]-z[overlap]))) if overlap.any() else None
   partial_left=bool(i>=315 or (m[:,:3]>0).any())
   accepted=coverage>.65 and err is not None and err<.02 and not partial_left
   row.update(status='partial_view' if partial_left else ('tracked' if accepted else 'low_confidence'),accepted=accepted,red_coverage=coverage,depth_median_error_m=err)
   # Show all predicted surfaces, including those behind the hand, for honest overlay QA.
   pic[render]=(pic[render]*.65+np.array([0,220,0])*.35).astype('uint8')
   contours,_=cv2.findContours(render.astype('uint8'),cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE);cv2.drawContours(pic,contours,-1,(0,255,0),1)
  records[i]=row
  cv2.putText(pic,f'{i:04d} {row["status"]}',(15,30),0,.7,(0,255,255),2)
  if i%30==0 or i in [args.init,89,119,149,179,209,239,269,299,329]:cv2.imwrite(str(out/'snapshots'/f'{i:06d}.jpg'),pic)
  if i%30==0:print(json.dumps({**row,'elapsed_s':round(time.time()-start,1)}),flush=True);np.save(out/'poses_checkpoint.npy',poses)
 np.savez_compressed(out/'poses.npz',T_camera_from_object=poses,time_s=(np.array(meta['timestamps_ms'][:N])-meta['timestamps_ms'][0])/1000,accepted=np.array([r['accepted'] for r in records]),frame_index=np.arange(N),K=K)
 (out/'frame_metrics.json').write_text(json.dumps(records,indent=2))
 # Export original STL origin, and its geometric bounding-box center, separately.
 rows=[]
 for i,T in enumerate(poses):
  q=Rotation.from_matrix(T[:3,:3]).as_quat() if np.isfinite(T).all() else np.full(4,np.nan)
  center=T[:3,:3]@mesh.bounds.mean(0)+T[:3,3]
  rows.append([i,(meta['timestamps_ms'][i]-meta['timestamps_ms'][0])/1000,*T[:3,3],*q,*center,int(records[i]['accepted'])])
 np.savetxt(out/'poses.csv',rows,delimiter=',',header='frame,time_s,origin_x_m,origin_y_m,origin_z_m,qx,qy,qz,qw,center_x_m,center_y_m,center_z_m,accepted',comments='')
 summary={'frames':N,'estimated':int(np.isfinite(poses).all((1,2)).sum()),'accepted_by_heuristics':sum(r['accepted'] for r in records),'init_frame':args.init,'offline_backward_frames':args.init,'coordinate':'original STL object coordinates to fixed OpenCV camera: x right, y down, z forward; meters assumed and checked by depth residual','manual_reviewed_tracking_interval_inclusive':[0,335],'mesh':'docs/上半.stl','mesh_extents_m':mesh.extents.tolist(),'boundary':'Predicted poses, not ground truth. Color/depth quality gates are heuristics; no contact or dynamics validation. Unobservable frames are NaN, not held.','elapsed_s':time.time()-start}
 (out/'summary.json').write_text(json.dumps(summary,indent=2));print(json.dumps(summary),flush=True)
if __name__=='__main__':main()