method_temporal2.py

#!/usr/bin/env python3
"""Refinement of temporal smoothing: graceful SOFT forward-backward (HMM posterior) instead of hard
Viterbi, + a confidence SNAP. Soft posterior degrades to the mass-mean when uncertain (fixes the
hard-Viterbi blow-ups on noisy venues) while concentrating on the right mode when the trajectory is
consistent. Snap: when the posterior is confident, commit to the argmax mode (exact sub-metre hits)."""
import json
from pathlib import Path
import numpy as np, torch
exec(open('/workspace/ego6d_rag/method_temporal.py').read().split('\nfor sigma')[0])  # val_modes, recover_t0, viterbi, fnet, VENUES, AN


def forward_backward(centers, mass, valid, sigma):
    T, K, _ = centers.shape
    e = np.where(valid, mass, 0.0) + 1e-9; e = e / e.sum(1, keepdims=True)
    A = []; a = np.zeros((T, K)); a[0] = e[0]
    for t in range(1, T):
        d2 = ((centers[t-1][:, None] - centers[t][None]) ** 2).sum(-1)
        At = np.exp(-d2 / (2 * sigma * sigma)); At[~valid[t-1]] = 0; At = At / (At.sum(1, keepdims=True) + 1e-9)
        A.append(At); a[t] = e[t] * (a[t-1] @ At); a[t] /= (a[t].sum() + 1e-9)
    b = np.ones((T, K))
    for t in range(T-2, -1, -1):
        b[t] = (A[t] * (e[t+1] * b[t+1])[None, :]).sum(1); b[t] /= (b[t].sum() + 1e-9)
    g = a * b; g = g / (g.sum(1, keepdims=True) + 1e-9)
    return (g[..., None] * centers).sum(1), g


def run(sigma, snap_thr):
    rows = {}
    for v in VENUES:
        V = Venue(v); C, MA, VA, GT, SID = val_modes(V); t0 = recover_t0(V, GT, SID)
        e_mass, e_soft, e_snap, e_oracle = [], [], [], []
        for sid in np.unique(SID):
            idx = np.where(SID == sid)[0]; order = idx[np.argsort(t0[idx])]
            c, ma, va, gt = C[order], MA[order], VA[order], GT[order]; ar = np.arange(len(order))
            soft, g = forward_backward(c, ma, va, sigma)
            conf = g.max(1); amode = g.argmax(1)
            massw = ma / ma.sum(1, keepdims=True).clip(min=1e-9); pm = (massw[..., None] * c).sum(1)
            snap = np.where((conf > snap_thr)[:, None], c[ar, amode], soft)            # snap to mode when confident
            e_mass += list(np.linalg.norm(pm - gt, axis=1)); e_soft += list(np.linalg.norm(soft - gt, axis=1))
            e_snap += list(np.linalg.norm(snap - gt, axis=1))
            dm = np.linalg.norm(c - gt[:, None], axis=-1); dm[~va] = 1e9; e_oracle += list(dm.min(1))
        rows[v] = dict(mass=float(np.median(e_mass)), soft=float(np.median(e_soft)), snap=float(np.median(e_snap)),
                       oracle=float(np.median(e_oracle)), snap_r1=float(np.mean(np.array(e_snap) < 1.0)),
                       mass_r1=float(np.mean(np.array(e_mass) < 1.0)))
    return rows


best = None
for sigma in [0.75, 1.0, 1.5]:
    for thr in [0.5, 0.7]:
        rows = run(sigma, thr); mac = lambda k: float(np.mean([rows[v][k] for v in VENUES]))
        print(f"sigma={sigma} snap>{thr}:  mass {mac('mass'):.2f}  soft {mac('soft'):.2f}  snap {mac('snap'):.2f}  "
              f"oracle {mac('oracle'):.2f}  snapR@1 {mac('snap_r1'):.2f} (massR@1 {mac('mass_r1'):.2f})", flush=True)
        if best is None or mac('snap') < best[0]: best = (mac('snap'), sigma, thr, rows)

_, sigma, thr, rows = best; mac = lambda k: float(np.mean([rows[v][k] for v in VENUES]))
print(f"\n=== BEST: sigma={sigma}, snap>{thr} ===  {'venue':6}{'mass':>6}{'soft':>6}{'snap':>6}{'oracle':>7}{'snapR@1':>8}")
for v in VENUES:
    r = rows[v]; print(f"  {v:6} {r['mass']:6.2f}{r['soft']:6.2f}{r['snap']:6.2f}{r['oracle']:7.2f}{r['snap_r1']:8.2f}")
print(f"  {'MACRO':6} {mac('mass'):6.2f}{mac('soft'):6.2f}{mac('snap'):6.2f}{mac('oracle'):7.2f}{mac('snap_r1'):8.2f}")
print(f"\n[refs] FCx top-1 3.31 | mass-mean 3.27 | hard-Viterbi 3.29 (blows up Loc_28) | Method1/1b 3.4-3.7 | oracle ~1.0")
json.dump({'sigma': sigma, 'thr': thr, 'macro': {k: mac(k) for k in rows[VENUES[0]]}, 'per_venue': rows},
          open(AN / 'temporal_soft_best.json', 'w'), indent=1)