scripts/phase5_triangulate.py

import glob
import json
import os

import cv2
import numpy as np

ROOT = "/root/bdmc_pipeline"
H, W = 1920, 1080
fx = fy = H / (2 * np.tan(np.radians(72.0 / 2)))
K = np.array([[fx, 0, W / 2], [0, fy, H / 2], [0, 0, 1.0]])
Z_MIN, Z_MAX = 1.2, 9.0
RESID_TOL_M = 0.06

def unproj(xs, ys, z):
    return np.stack([(xs - W / 2) * z / fx, (ys - H / 2) * z / fy, z], 1).astype(np.float64)

def kabsch_rigid(A, B):
    ca, cb = A.mean(0), B.mean(0)
    Hm = (A - ca).T @ (B - cb)
    U, _, Vt = np.linalg.svd(Hm)
    d = np.sign(np.linalg.det(Vt.T @ U.T))
    D = np.diag([1, 1, d])
    R = Vt.T @ D @ U.T
    t = cb - R @ ca
    return R, t

def rigid_motion(Xa, Xb):
    for _ in range(3):
        R, t = kabsch_rigid(Xa, Xb)
        res = np.linalg.norm(Xa @ R.T + t - Xb, axis=1)
        keep = res < RESID_TOL_M
        if keep.sum() < 30:
            break
        Xa, Xb = Xa[keep], Xb[keep]
    return R, t, float(np.median(res[keep])) if keep.sum() else 9e9

def dlt_multiview(K, Rs, ts, pts2d):
    rows = []
    for R_i, t_i, x in zip(Rs, ts, pts2d):
        P = K @ np.hstack([R_i, t_i.reshape(3, 1)])
        u, v = x
        rows.append(u * P[2] - P[0])
        rows.append(v * P[2] - P[1])
    A = np.stack(rows)
    _, _, vt = np.linalg.svd(A, full_matrices=False)
    Xh = vt[-1]
    return Xh[:3] / Xh[3]

results = {}
all_stats = []
for path in sorted(glob.glob(f"{ROOT}/outputs/tracks/win_*.npz")):
    z = np.load(path)
    tracks, vis = z["tracks"], z["vis"]
    f0, n = int(z["start"]), int(z["n_frames"])
    N = tracks.shape[1]
    depths = []
    masks = []
    for i in range(f0, f0 + n):
        dep = np.load(f"{ROOT}/depth/f_{i:04d}.npy").astype(np.float32)
        m = cv2.imread(f"{ROOT}/masks/road_surface/f_{i:04d}.png", 0) > 127
        depths.append(dep)
        masks.append(m)
    clouds = []
    for i in range(n):
        x = tracks[i, :, 0]
        y = tracks[i, :, 1]
        xi = np.clip(x.round().astype(int), 0, W - 1)
        yi = np.clip(y.round().astype(int), 0, H - 1)
        ok = (vis[i].reshape(-1).astype(bool)) & masks[i][yi, xi] & (depths[i][yi, xi] > Z_MIN) & (depths[i][yi, xi] < Z_MAX)
        c = np.full((N, 3), np.nan)
        c[ok] = unproj(x[ok], y[ok], depths[i][yi[ok], xi[ok]])
        clouds.append(c)
    Rs, ts = [np.eye(3)], [np.zeros(3)]
    motion_res = []
    for i in range(n - 1):
        a, b = clouds[i], clouds[i + 1]
        both = np.isfinite(a).all(1) & np.isfinite(b).all(1)
        if both.sum() < 30:
            Rs.append(Rs[-1].copy())
            ts.append(ts[-1].copy())
            continue
        R, t, rmed = rigid_motion(a[both], b[both])
        motion_res.append(rmed)
        Rs.append(R @ Rs[-1])
        ts.append(R @ ts[-1] + t)
    world_pts, reproj_rmse, n_obs = [], [], []
    for j in range(N):
        obs_idx = [i for i in range(n) if vis[i, j] and np.isfinite(clouds[i][j]).all()]
        if len(obs_idx) < 6:
            continue
        try:
            Xw = dlt_multiview(K, [Rs[i] for i in obs_idx], [ts[i] for i in obs_idx],
                               [tracks[i, j] for i in obs_idx])
        except Exception:
            continue
        errs = []
        for i in obs_idx:
            p = K @ (Rs[i] @ Xw + ts[i])
            p = p[:2] / p[2]
            errs.append(np.hypot(*(p - tracks[i, j])))
        rmse = float(np.sqrt(np.mean(np.square(errs))))
        if rmse < 12.0:
            world_pts.append(Xw)
            reproj_rmse.append(rmse)
            n_obs.append(len(obs_idx))
    results[f"win_{f0:04d}"] = {
        "start": f0,
        "n_frames": n,
        "n_tracks_total": int(N),
        "n_triangulated": len(world_pts),
        "median_reproj_rmse_px": round(float(np.median(reproj_rmse)), 2) if reproj_rmse else None,
        "median_motion_residual_cm": round(float(np.median(motion_res)) * 100, 2) if motion_res else None,
    }
    if world_pts:
        np.savez_compressed(f"{ROOT}/outputs/tracks/triangulated_{f0:04d}.npz",
                            xyz=np.array(world_pts),
                            rmse=np.array(reproj_rmse),
                            n_obs=np.array(n_obs),
                            start=f0,
                            Rs=np.array(Rs), ts=np.array(ts))
    all_stats.append(results[f"win_{f0:04d}"])
    print(json.dumps(results[f"win_{f0:04d}"]), flush=True)

tot_t = sum(s["n_triangulated"] for s in all_stats)
tot_p = sum(s["n_tracks_total"] for s in all_stats)
summary = {"windows": all_stats,
           "triangulation_yield": round(tot_t / max(tot_p, 1), 3),
           "total_triangulated_points": tot_t}
with open(f"{ROOT}/outputs/phase5_triangulation.json", "w") as fj:
    json.dump(summary, fj, indent=1)
print("Phase 5b triangulation complete | yield", summary["triangulation_yield"], flush=True)