"""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))