scripts/phase3_defects.py

import csv
import json
import os

import cv2
import numpy as np

ROOT = "/root/bdmc_pipeline"
DEPTH_DIR = f"{ROOT}/depth"
MASKS_DIR = f"{ROOT}/masks/road_surface"
FPS_A = 5.0
H, W = 1920, 1080
FOV_LONG_DEG = 72.0
fx = fy = H / (2 * np.tan(np.radians(FOV_LONG_DEG / 2)))
cx, cy = W / 2, H / 2

TILE = 160
STRIDE = 80
Z_FIT_MIN, Z_FIT_MAX = 1.5, 8.0
Z_EVAL_MAX = 7.0
PERSIST_FRAMES = 5
MIN_AREA_PX_FULLRES = 1000
MAX_TRACK_JUMP_PX = 60.0
SIGMA_K = 2.0
T_MIN, T_MAX = 0.04, 0.10

rng = np.random.default_rng(0)

def unproj(xs, ys, z):
    return np.stack([(xs - cx) * z / fx, (ys - cy) * z / fy, z], 1).astype(np.float32)

def frame_residual_map(f):
    dep = np.load(f"{DEPTH_DIR}/f_{f:04d}.npy").astype(np.float32)
    m = cv2.imread(f"{MASKS_DIR}/f_{f:04d}.png", 0) > 127
    valid = m & (dep > Z_FIT_MIN) & (dep < Z_FIT_MAX)
    dev = np.full((H, W), np.nan, np.float32)
    m_core = cv2.erode(m.astype(np.uint8), np.ones((48, 48), np.uint8)).astype(bool)
    sigmas = []
    for ty in range(0, H - TILE // 2, STRIDE):
        for tx in range(0, W - TILE // 2, STRIDE):
            tm = valid[ty:ty + TILE, tx:tx + TILE]
            ys_t, xs_t = np.nonzero(tm)
            if len(ys_t) < 800:
                continue
            ys = ys_t + ty
            xs = xs_t + tx
            if len(ys) > 4000:
                idx = rng.choice(len(ys), 4000, replace=False)
                ys, xs = ys[idx], xs[idx]
            pts = unproj(xs, ys, dep[ys, xs])
            cen = pts.mean(0)
            _, _, vt = np.linalg.svd(pts - cen, full_matrices=False)
            nrm = vt[2]
            d = -nrm @ cen
            res = pts @ nrm + d
            mad = float(np.median(np.abs(res - np.median(res))))
            sigmas.append(1.4826 * mad)
            sub = dep[ty:ty + TILE, tx:tx + TILE]
            sm = m_core[ty:ty + TILE, tx:tx + TILE] & (sub > Z_FIT_MIN) & (sub < Z_EVAL_MAX)
            ey, ex = np.nonzero(sm)
            if len(ey) == 0:
                continue
            epts = unproj(ex + tx, ey + ty, sub[ey, ex])
            s = epts @ nrm + d
            dev[ey + ty, ex + tx] = -(s - np.median(res))
    return dev, float(np.median(sigmas)) if sigmas else np.nan

os.makedirs(f"{ROOT}/outputs", exist_ok=True)

sig_all = []
devmaps = {}
print("Computing per-frame residuals...")
for f in range(1, 301):
    dev, sig = frame_residual_map(f)
    sig_all.append(sig)
    devmaps[f] = dev
    if f % 50 == 0 or f == 1:
        print(f"  Frame {f}: sigma={sig*100:.2f} cm" if not np.isnan(sig) else f"  Frame {f}: no sigma")

sig_arr = np.array(sig_all)
sigma_med = float(np.nanmedian(sig_arr))
thresh = float(np.clip(SIGMA_K * sigma_med, T_MIN, T_MAX))
print(f"\nMedian tile sigma: {sigma_med*100:.2f} cm -> threshold: {thresh*100:.1f} cm")

tracks, confirmed = [], []
for f in range(1, 301):
    dev = devmaps[f]
    t_f = float(np.clip(SIGMA_K * (sig_arr[f - 1] if not np.isnan(sig_arr[f - 1]) else sigma_med), T_MIN, T_MAX))
    hot = (dev < -t_f).astype(np.uint8)
    n, labels, st, cent = cv2.connectedComponentsWithStats(hot, 8)
    comps = []
    for ci in range(1, n):
        area = int(st[ci, cv2.CC_STAT_AREA])
        if area < MIN_AREA_PX_FULLRES:
            continue
        comps.append({"centroid": cent[ci], "area_px": area,
                      "peak_dev": float(-np.nanmax(dev[labels == ci]))})
    for tr in tracks:
        tr["matched"] = False
    for c in comps:
        best, bestd = None, 1e9
        for tr in tracks:
            if tr["matched"] or f - tr["last_f"] != 1:
                continue
            dist = float(np.hypot(*(c["centroid"] - tr["centroid"])))
            if dist < MAX_TRACK_JUMP_PX and dist < bestd:
                best, bestd = tr, dist
        if best is None:
            tracks.append({**c, "first_f": f, "last_f": f, "matched": True, "frames": [f],
                           "max_area": c["area_px"], "peak_dev": c["peak_dev"]})
        else:
            best.update(matched=True, last_f=f, centroid=c["centroid"],
                        max_area=max(best["max_area"], c["area_px"]),
                        peak_dev=max(best["peak_dev"], c["peak_dev"]))
            best["frames"].append(f)
    tracks[:] = [t for t in tracks if t.get("logged") or f - t["last_f"] <= 6]
    for tr in tracks:
        if not tr.get("logged") and len(tr["frames"]) >= PERSIST_FRAMES:
            tr["logged"] = True
            fi0 = tr["frames"][0]
            confirmed.append({
                "first_frame": fi0,
                "t_sec": round((fi0 - 1) / FPS_A, 2),
                "duration_frames": tr["frames"][-1] - fi0 + 1,
                "duration_s": round((tr["frames"][-1] - fi0 + 1) / FPS_A, 2),
                "peak_dev_cm": round(tr["peak_dev"] * 100, 1),
                "area_px": tr["max_area"],
                "centroid_px": [round(float(tr["centroid"][0]), 1), round(float(tr["centroid"][1]), 1)],
            })

cols = ["first_frame", "t_sec", "duration_frames", "duration_s", "peak_dev_cm", "area_px", "centroid_px"]
with open(f"{ROOT}/outputs/defects.csv", "w", newline="") as fh:
    wr = csv.DictWriter(fh, fieldnames=cols)
    wr.writeheader()
    wr.writerows([{k: r[k] for k in cols} for r in confirmed])

ds = np.zeros((300, H // 4, W // 4), np.uint8)
for f in range(1, 301):
    d4 = cv2.resize(np.nan_to_num(devmaps[f], nan=0.0), (W // 4, H // 4))
    ds[f - 1] = np.clip((d4 * 1000 + 128), 0, 255).astype(np.uint8)
np.savez_compressed(f"{ROOT}/outputs/devmaps_small.npz", dev=ds)

summary = {
    "method": "tile-local RANSAC-free SVD plane fit, 160px tiles stride 80, temporal window folded via persistence filter",
    "z_fit_band_m": [Z_FIT_MIN, Z_FIT_MAX],
    "tile_sigma_median_cm": round(sigma_med * 100, 2),
    "threshold_rule": f"clamp({SIGMA_K}*sigma, {T_MIN*100:.0f}cm, {T_MAX*100:.0f}cm)",
    "threshold_used_cm": round(thresh * 100, 1),
    "persist_frames_required": PERSIST_FRAMES,
    "min_area_px_fullres": MIN_AREA_PX_FULLRES,
    "fov_long_axis_deg_assumed": FOV_LONG_DEG,
    "n_confirmed_defects": len(confirmed),
}
with open(f"{ROOT}/outputs/phase3_summary.json", "w") as fj:
    json.dump(summary, fj, indent=1)
with open(f"{ROOT}/outputs/tile_sigma_per_frame.json", "w") as fj:
    json.dump({"sigma_cm": [round(s * 100, 2) for s in sig_all]}, fj)

print(f"\nPhase 3 complete: {len(confirmed)} confirmed defects")
print(json.dumps(summary, indent=1))