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>
201 lines
17 KiB
Python
201 lines
17 KiB
Python
"""Physically-consistent object reference: rest segments frozen (and snapped to the table),
|
|
held segments rigidly bound to the holding hand, free segments lightly smoothed, camera odometry smoothed.
|
|
|
|
Input : output/registered_spider_20260915/foundationpose_objects.npz (raw FoundationPose, raw odometry)
|
|
output/20260915_171525_dynhamr/registered/ (hands, raw-odometry world)
|
|
Output: output/final_spider_20260915/foundationpose_objects.npz (same schema; *_valid all True where a pose is
|
|
defined, *_observed_raw keeps the original acceptance flags), rgbd_camera_smooth.npz, object_segments.npz,
|
|
object_smoothing.json. Validation renders the CAD into the sensor depth before/after.
|
|
"""
|
|
import os, sys, json
|
|
os.environ.setdefault('OMP_NUM_THREADS', '4')
|
|
from pathlib import Path
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT / 'third_party/FoundationPose')); sys.path.insert(0, str(ROOT / 'scripts'))
|
|
import numpy as np, cv2, torch, trimesh, open3d as o3d, nvdiffrast.torch as dr
|
|
from scipy.ndimage import gaussian_filter1d
|
|
from scipy.spatial.transform import Rotation as R, Slerp
|
|
from Utils import nvdiffrast_render
|
|
from red_box_distance import signed_distance
|
|
|
|
SRC = ROOT / 'docs/20260915_171525'; OLD = ROOT / 'output/registered_spider_20260915'; HAND = ROOT / 'output/20260915_171525_dynhamr/registered'
|
|
OUT = ROOT / 'output/final_spider_20260915'; OUT.mkdir(exist_ok=True)
|
|
meta = json.loads((SRC / 'intrinsics.json').read_text()); fp = dict(np.load(OLD / 'foundationpose_objects.npz')); N = 352
|
|
K = fp['K']; c2w_raw = fp['c2w']; w2cR = np.transpose(c2w_raw[:, :3, :3], (0, 2, 1)); tO = -np.einsum('tij,tj->ti', w2cR, c2w_raw[:, :3, 3])
|
|
SIG_CAM, SIG_FREE, SIG_HELD, BLEND = 2., 2., 4., 4
|
|
REST_WIN, REST_POS, REST_ROT, REST_MIN = 15, .006, 2., 10 # 0.5 s window, 6 mm, 2 deg, min 10 frames
|
|
ENGAGE, RELEASE, MIN_TIPS = .025, .035, 2
|
|
|
|
def smooth_rot(rot, sigma, idx=None):
|
|
n = len(rot); out = []
|
|
for t in range(n):
|
|
ix = np.arange(max(0, t - int(3 * sigma)), min(n, t + int(3 * sigma) + 1)); w = np.exp(-.5 * ((ix - t) / sigma) ** 2)
|
|
out.append(rot[ix].mean(weights=w))
|
|
return R.concatenate(out)
|
|
def smooth_T(T, sigma):
|
|
out = T.copy(); out[:, :3, 3] = gaussian_filter1d(T[:, :3, 3], sigma, axis=0, mode='nearest'); out[:, :3, :3] = smooth_rot(R.from_matrix(T[:, :3, :3]), sigma).as_matrix(); return out
|
|
def interp_T(T, valid):
|
|
idx = np.flatnonzero(valid); t = np.arange(len(T)); out = T.copy(); clip = np.clip(t, idx[0], idx[-1])
|
|
out[:, :3, 3] = np.stack([np.interp(t, idx, T[idx, k, 3]) for k in range(3)], 1); out[:, :3, :3] = Slerp(idx, R.from_matrix(T[idx, :3, :3]))(clip).as_matrix(); out[:, 3, :] = [0, 0, 0, 1]; return out
|
|
def inv(T): o = np.eye(4); o[:3, :3] = T[:3, :3].T; o[:3, 3] = -T[:3, :3].T @ T[:3, 3]; return o
|
|
|
|
# --- camera: smoothed odometry ---
|
|
c2w_s = smooth_T(c2w_raw, SIG_CAM)
|
|
np.savez_compressed(OUT / 'rgbd_camera_smooth.npz', c2w=c2w_s, intrinsics=np.array([meta[k] for k in ['fx', 'fy', 'cx', 'cy']]), time=fp['time'], valid=fp['camera_valid'], method='registered_spider c2w, Gaussian sigma 2 frames on position and SO(3)')
|
|
# --- hands in camera frame (independent of odometry) ---
|
|
hands = {}
|
|
for side in ['right', 'left']:
|
|
hj = np.load(HAND / f'human_joints_{side}.npz', allow_pickle=True); mo = np.load(HAND / f'l20_{side}_stable/motion.npz', allow_pickle=True)
|
|
tips_w = hj['joints'][:, [4, 8, 12, 16, 20]] + hj['wrist_world'][:, None]; tips_c = np.einsum('tij,tkj->tki', w2cR, tips_w) + tO[:, None]
|
|
Th = np.tile(np.eye(4), (N, 1, 1)); Th[:, :3, :3] = np.einsum('tij,tjk->tik', w2cR, mo['wrist_world_R']); Th[:, :3, 3] = np.einsum('tij,tj->ti', w2cR, mo['wrist_world']) + tO
|
|
hands[side] = dict(tips_cam=tips_c, T_cam=Th)
|
|
# --- table plane (frame 0, camera-0 frame), as in prepare ---
|
|
cap = cv2.VideoCapture(str(SRC / 'color.mp4')); ok, img0 = cap.read(); cap.release(); hsv0 = cv2.cvtColor(img0, cv2.COLOR_BGR2HSV)
|
|
dep0 = cv2.imread(str(SRC / 'depth/000000.png'), -1) * meta['depth_scale_m']; yy, xx = np.indices(dep0.shape)
|
|
mask = (xx > 100) & (xx < 780) & (yy > 200) & (yy < 460) & (hsv0[:, :, 1] < 50) & (dep0 > .2) & (dep0 < 1)
|
|
pts = np.stack([(xx - K[0, 2]) * dep0 / K[0, 0], (yy - K[1, 2]) * dep0 / K[1, 1], dep0], -1)[mask][::3]
|
|
pc = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(pts)); plane, _ = pc.segment_plane(.004, 3, 1000); n_cam = np.array(plane[:3]); off = plane[3]
|
|
if n_cam[2] > 0: n_cam, off = -n_cam, -off # normal points up (toward the camera side)
|
|
# --- meshes ---
|
|
scenes = {}
|
|
for name, stl in [('upper', '上半.stl'), ('lower', '下半.stl')]:
|
|
closed = trimesh.load(ROOT / f'output/collision_fix_20260915/collision_v2/{name}_closed.ply', process=False); assert closed.is_watertight
|
|
sc = o3d.t.geometry.RaycastingScene(); sc.add_triangles(o3d.t.geometry.TriangleMesh.from_legacy(o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(closed.vertices), o3d.utility.Vector3iVector(closed.faces))))
|
|
scenes[name] = dict(scene=sc, closed=closed, render=trimesh.load(ROOT / 'docs' / stl))
|
|
def tip_dist(name, T_cam, tips_cam):
|
|
local = (tips_cam - T_cam[:3, 3]) @ T_cam[:3, :3]; return np.abs(signed_distance(scenes[name]['scene'], scenes[name]['closed'], local))
|
|
|
|
ctx = dr.RasterizeCudaContext(); mt = {}
|
|
for name in scenes:
|
|
m_ = scenes[name]['render']; mt[name] = {'pos': torch.as_tensor(m_.vertices, device='cuda', dtype=torch.float32), 'faces': torch.as_tensor(m_.faces, device='cuda', dtype=torch.int32), 'vnormals': torch.zeros((len(m_.vertices), 3), device='cuda'), 'vertex_color': torch.ones((len(m_.vertices), 3), device='cuda')}
|
|
def render(name, T):
|
|
with torch.inference_mode(): _, d, _ = nvdiffrast_render(K=K, H=480, W=848, ob_in_cams=torch.as_tensor(T[None], device='cuda', dtype=torch.float32), glctx=ctx, mesh_tensors=mt[name])
|
|
return d[0].cpu().numpy()
|
|
_cap = cv2.VideoCapture(str(SRC / 'color.mp4')); _frames = {}
|
|
def frame_data(f):
|
|
if f not in _frames:
|
|
_cap.set(cv2.CAP_PROP_POS_FRAMES, f); ok, b = _cap.read(); assert ok
|
|
z = cv2.imread(str(SRC / 'depth' / f'{f:06d}.png'), -1).astype('float32') * meta['depth_scale_m']; h = cv2.cvtColor(b, cv2.COLOR_BGR2HSV); _frames[f] = (z, h)
|
|
return _frames[f]
|
|
def depth_residual(name, f, T):
|
|
z, h = frame_data(f); hue = h[:, :, 0]
|
|
cm = (((hue < 12) | (hue > 170)) if name == 'upper' else ((hue > 95) & (hue < 135))) & (h[:, :, 1] > 90) & (h[:, :, 2] > 40) & (z > .1) & (z < 1.); cm[:200] = False
|
|
d = render(name, T); over = (d > 0) & cm
|
|
return float(np.median(np.abs(d[over] - z[over])) * 1000) if over.sum() > 200 else None
|
|
report = {'parameters': dict(camera_sigma=SIG_CAM, free_sigma=SIG_FREE, held_relative_sigma=SIG_HELD, blend_frames=BLEND, rest_window=REST_WIN, rest_pos_m=REST_POS, rest_rot_deg=REST_ROT, engage_m=ENGAGE, release_m=RELEASE, min_tips=MIN_TIPS), 'objects': {}}
|
|
out = dict(fp); labels = {}
|
|
for name in ['upper', 'lower']:
|
|
T_raw = fp[name + '_T_camera'].copy(); valid = fp[name + '_valid'].astype(bool).copy(); obs = valid.copy()
|
|
last = N if name == 'upper' else 177 # lower unobservable after 176: hold
|
|
T = interp_T(T_raw, valid); T[last:] = T[last - 1]
|
|
# holder per frame (camera frame, hysteresis per hand)
|
|
holder = np.full(N, '', dtype=object); engaged = {s: np.zeros(5, bool) for s in hands}
|
|
for t in range(N):
|
|
best, best_n = '', 0
|
|
for s in hands:
|
|
dist = tip_dist(name, T[t], hands[s]['tips_cam'][t]); engaged[s] = np.where(engaged[s], dist < RELEASE, dist < ENGAGE)
|
|
n_on = int(engaged[s].sum())
|
|
if n_on >= MIN_TIPS and (n_on > best_n or (n_on == best_n and s == holder[t - 1] if t else False)): best, best_n = s, n_on
|
|
holder[t] = best
|
|
# rest test on lightly pre-smoothed poses, over a 0.5 s window
|
|
pre = smooth_T(T, 2.); rest = np.zeros(N, bool); h = REST_WIN // 2
|
|
for t in range(N):
|
|
a, b = max(0, t - h), min(N - 1, t + h); dp = np.linalg.norm(pre[b, :3, 3] - pre[a, :3, 3]); dr_ = (R.from_matrix(pre[a, :3, :3]).inv() * R.from_matrix(pre[b, :3, :3])).magnitude() * 180 / np.pi
|
|
rest[t] = dp < REST_POS and dr_ < REST_ROT
|
|
rest[last:] = True
|
|
lab = np.array(['free'] * N, dtype=object); lab[rest] = 'rest'
|
|
for t in range(N):
|
|
if not rest[t] and holder[t]: lab[t] = 'held:' + holder[t]
|
|
# drop short rest runs; then merge rest-(free with no hand)-rest, since a static object cannot move by itself
|
|
t = 0
|
|
while t < N:
|
|
u = t
|
|
while u < N and lab[u] == lab[t]: u += 1
|
|
if lab[t] == 'rest' and u - t < REST_MIN: lab[t:u] = 'free'
|
|
t = u
|
|
changed = True
|
|
while changed:
|
|
changed = False; t = 0
|
|
while t < N:
|
|
u = t
|
|
while u < N and lab[u] == lab[t]: u += 1
|
|
if lab[t] == 'free' and t > 0 and u < N and lab[t - 1] == 'rest' and lab[u] == 'rest' and not any(holder[t:u]):
|
|
lab[t:u] = 'rest'; changed = True
|
|
t = u
|
|
# per-segment poses
|
|
S = T.copy(); alt = T.copy(); segs = []; t = 0; depth_split = 0
|
|
while t < N:
|
|
u = t
|
|
while u < N and lab[u] == lab[t]: u += 1
|
|
kind = lab[t]; sl = slice(t, u)
|
|
if kind == 'rest':
|
|
# consensus pose: among the median and sampled observed poses, take the one the sensor depth agrees with best
|
|
fr = [f for f in range(t, u, 3) if obs[f]]
|
|
Tm = np.eye(4); Tm[:3, :3] = R.from_matrix(T[sl, :3, :3]).mean().as_matrix(); Tm[:3, 3] = np.median(T[sl, :3, 3], axis=0)
|
|
cands = [('median', Tm)] + [(f'obs{f}', T[f]) for f in range(t, u, max(1, (u - t) // 12)) if obs[f]]
|
|
def score(Tc): rs = [depth_residual(name, f, Tc) for f in fr]; rs = [x for x in rs if x is not None]; return float(np.median(rs)) if rs else 1e9
|
|
scored = [(score(Tc), lab_, Tc) for lab_, Tc in cands]; best_score, best_lab, Tr = min(scored, key=lambda x: x[0]); Tr = Tr.copy()
|
|
cl = scenes[name]['closed']; hv = (cl.vertices @ Tr[:3, :3].T + Tr[:3, 3]) @ n_cam + off; gap = float(hv.min()) # signed height of lowest vertex above the table
|
|
snapped = False
|
|
if abs(gap) < .03:
|
|
Ts = Tr.copy(); Ts[:3, 3] -= n_cam * gap; snap_score = score(Ts)
|
|
if snap_score <= best_score + 1.0: Tr, snapped, best_score = Ts, True, snap_score
|
|
# split test: if the consensus pose disagrees with the sensor for a sustained run, the object actually moved
|
|
per = [(f, depth_residual(name, f, Tr), depth_residual(name, f, T[f])) for f in fr]; bad = [f for f, a, b in per if a is not None and b is not None and a > b + 3.0]
|
|
runs = [];
|
|
for f in bad:
|
|
if runs and f - runs[-1][-1] <= 3: runs[-1].append(f)
|
|
else: runs.append([f])
|
|
long_runs = [r for r in runs if len(r) >= 3]
|
|
if long_runs and u - t >= 2 * REST_MIN and depth_split < 3:
|
|
cut = long_runs[0][0] if long_runs[0][0] - t >= REST_MIN else long_runs[0][-1] + 1
|
|
if t + REST_MIN <= cut <= u - REST_MIN:
|
|
lab[cut:cut + 3] = 'free'; depth_split += 1; continue # re-segment from t with the new boundary
|
|
S[sl] = Tr; segs.append(dict(kind='rest', start=t, end=u - 1, table_gap_mm=gap * 1000, snapped=snapped, pose_source=best_lab, depth_residual_mm=best_score, median_pose_residual_mm=scored[0][0], split_runs=[(r[0], r[-1]) for r in long_runs]))
|
|
elif kind.startswith('held'):
|
|
s = kind.split(':')[1]; Th = hands[s]['T_cam'][sl]; rel = np.array([inv(Th[i]) @ T[t + i] for i in range(u - t)])
|
|
rel_s = smooth_T(rel, SIG_HELD) if u - t > 2 else rel; S[sl] = np.array([Th[i] @ rel_s[i] for i in range(u - t)])
|
|
alt[sl] = smooth_T(T[sl], SIG_FREE) if u - t > 2 else T[sl]
|
|
segs.append(dict(kind=kind, start=t, end=u - 1, relative_position_2nd_diff_rms_mm_before=float(np.sqrt((np.linalg.norm(np.diff(rel[:, :3, 3], n=2, axis=0), axis=1) ** 2).mean()) * 1000) if u - t > 2 else 0., after=float(np.sqrt((np.linalg.norm(np.diff(rel_s[:, :3, 3], n=2, axis=0), axis=1) ** 2).mean()) * 1000) if u - t > 2 else 0.))
|
|
else:
|
|
S[sl] = smooth_T(T[sl], SIG_FREE) if u - t > 2 else T[sl]; segs.append(dict(kind='free', start=t, end=u - 1))
|
|
t = u
|
|
# held segments: keep the hand-bound version only if the sensor depth agrees at least as well as plain smoothing
|
|
for sg in segs:
|
|
if not sg['kind'].startswith('held') or sg['end'] - sg['start'] < 3: continue
|
|
fr = [f for f in range(sg['start'], sg['end'] + 1, 3) if obs[f]]
|
|
rb = [depth_residual(name, f, S[f]) for f in fr]; ra = [depth_residual(name, f, alt[f]) for f in fr]; rb = [x for x in rb if x is not None]; ra = [x for x in ra if x is not None]
|
|
sg['depth_residual_bound_mm'] = float(np.median(rb)) if rb else None; sg['depth_residual_smoothed_mm'] = float(np.median(ra)) if ra else None
|
|
if rb and ra and np.median(rb) > np.median(ra) + 1.0: S[sg['start']:sg['end'] + 1] = alt[sg['start']:sg['end'] + 1]; sg['chosen'] = 'own-trajectory smoothing'
|
|
else: sg['chosen'] = 'hand-bound'
|
|
# blend across segment boundaries
|
|
B = S.copy()
|
|
for i in range(1, len(segs)):
|
|
b = segs[i]['start']; a0, a1 = max(0, b - BLEND), min(N, b + BLEND); rot = R.from_matrix(np.stack([S[a0, :3, :3], S[a1 - 1, :3, :3]])); sl = Slerp([a0, a1 - 1], rot)
|
|
for k in range(a0, a1):
|
|
w = (k - a0) / max(1, a1 - 1 - a0); B[k, :3, 3] = (1 - w) * S[a0, :3, 3] + w * S[a1 - 1, :3, 3]; B[k, :3, :3] = sl(k).as_matrix()
|
|
# keep observed frames close to observation: report deviation
|
|
dev = np.linalg.norm(B[obs, :3, 3] - T_raw[obs, :3, 3], axis=1) * 1000; rdev = (R.from_matrix(B[obs, :3, :3]).inv() * R.from_matrix(T_raw[obs, :3, :3])).magnitude() * 180 / np.pi
|
|
def jit(X): p = X[:, :3, 3]; rr = R.from_matrix(X[:, :3, :3]); return dict(pos_2nd_diff_rms_mm=float(np.sqrt((np.linalg.norm(np.diff(p, n=2, axis=0), axis=1) ** 2).mean()) * 1000), pos_step_p95_mm=float(np.percentile(np.linalg.norm(np.diff(p, axis=0), axis=1), 95) * 1000), rot_step_p95_deg=float(np.percentile((rr[:-1].inv() * rr[1:]).magnitude() * 180 / np.pi, 95)))
|
|
out[name + '_T_camera'] = B; out[name + '_T_world'] = c2w_s @ B; out[name + '_valid'] = np.r_[np.ones(last, bool), np.zeros(N - last, bool)]; out[name + '_observed_raw'] = obs
|
|
labels[name] = lab.astype(str)
|
|
report['objects'][name] = dict(segments=segs, frames_rest=int((lab == 'rest').sum()), frames_held=int(np.char.startswith(lab.astype(str), 'held').sum()), frames_free=int((lab == 'free').sum()), jitter_before=jit(T[:last]), jitter_after=jit(B[:last]), deviation_from_observed_mm=dict(median=float(np.median(dev)), p95=float(np.percentile(dev, 95)), max=float(dev.max())), deviation_rot_deg=dict(median=float(np.median(rdev)), p95=float(np.percentile(rdev, 95))))
|
|
out['c2w'] = c2w_s
|
|
np.savez_compressed(OUT / 'foundationpose_objects.npz', **out); np.savez_compressed(OUT / 'object_segments.npz', **{k: v for k, v in labels.items()})
|
|
# --- validation against the sensor depth (every 3rd frame) ---
|
|
val = {n: {'raw': [], 'smoothed': []} for n in scenes}
|
|
for t in range(0, N, 3):
|
|
for name in scenes:
|
|
if not fp[name + '_valid'][t]: continue
|
|
for key, T in [('raw', fp[name + '_T_camera'][t]), ('smoothed', out[name + '_T_camera'][t])]:
|
|
r = depth_residual(name, t, T)
|
|
if r is not None: val[name][key].append(r)
|
|
for name in scenes: report['objects'][name]['depth_residual_mm'] = {k: dict(median=float(np.median(v)), p95=float(np.percentile(v, 95)), frames=len(v)) for k, v in val[name].items()}
|
|
report['camera'] = dict(pos_2nd_diff_rms_mm_before=float(np.sqrt((np.linalg.norm(np.diff(c2w_raw[:, :3, 3], n=2, axis=0), axis=1) ** 2).mean()) * 1000), after=float(np.sqrt((np.linalg.norm(np.diff(c2w_s[:, :3, 3], n=2, axis=0), axis=1) ** 2).mean()) * 1000), max_change_mm=float(np.linalg.norm(c2w_s[:, :3, 3] - c2w_raw[:, :3, 3], axis=1).max() * 1000))
|
|
report['table_plane_camera0'] = dict(normal=n_cam.tolist(), offset=float(off))
|
|
(OUT / 'object_smoothing.json').write_text(json.dumps(report, indent=2, default=str)); print(json.dumps({k: v for k, v in report.items() if k != 'objects'}, indent=1))
|
|
for name in scenes:
|
|
r = report['objects'][name]; print(name, 'rest/held/free', r['frames_rest'], r['frames_held'], r['frames_free'], '| jitter before', {k: round(v, 2) for k, v in r['jitter_before'].items()}, '| after', {k: round(v, 2) for k, v in r['jitter_after'].items()})
|
|
print(' deviation from observed', {k: round(v, 1) for k, v in r['deviation_from_observed_mm'].items()}, '| depth residual', r['depth_residual_mm'])
|
|
print(' segments:', [(s['kind'], s['start'], s['end'], ('snap' if s.get('snapped') else '') + (f"gap{s['table_gap_mm']:.0f}" if 'table_gap_mm' in s else '') + (f" {s['chosen']} b{s['depth_residual_bound_mm']:.1f}/s{s['depth_residual_smoothed_mm']:.1f}" if s.get('chosen') and s.get('depth_residual_bound_mm') is not None else '')) for s in r['segments']])
|