294 lines
10 KiB
Python
294 lines
10 KiB
Python
"""Shared helpers for optical-flow metrics.
|
|
|
|
Both ``optical_flow.gt_optical_flow`` and
|
|
``optical_flow.synthetic_optical_flow`` extract per-frame flow with
|
|
``ptlflow`` and reduce it through the same per-pixel / per-frame /
|
|
temporal aggregation pipeline. The pipeline lives here so the two
|
|
metrics stay byte-identical on the comparison side and only differ in
|
|
how they construct the *reference* flow field.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
_PER_FRAME_AGG_KEYS: tuple[str, ...] = (
|
|
"mf_epe",
|
|
"mf_angle_err",
|
|
"mf_cosine",
|
|
"mf_mag_ratio",
|
|
"pixel_epe_mean",
|
|
"pixel_epe_max",
|
|
"px_angle_rmse",
|
|
"grid_epe_mean",
|
|
"grid_epe_max",
|
|
"fl_all",
|
|
"foe_dist",
|
|
"flow_kl_2d",
|
|
)
|
|
|
|
|
|
def _trapezoid(vals: np.ndarray) -> float:
|
|
fn = getattr(np, "trapezoid", None) or np.trapz
|
|
return float(fn(vals))
|
|
|
|
|
|
def _estimate_foe(
|
|
flow: np.ndarray,
|
|
step: int = 8,
|
|
min_mag: float = 0.5,
|
|
) -> tuple[float, float]:
|
|
"""Least-squares Focus of Expansion. Returns (fx, fy)."""
|
|
H, W = flow.shape[:2]
|
|
ys = np.arange(step // 2, H, step)
|
|
xs = np.arange(step // 2, W, step)
|
|
yy, xx = np.meshgrid(ys, xs, indexing="ij")
|
|
yy = yy.ravel()
|
|
xx = xx.ravel()
|
|
uu = flow[yy, xx, 0]
|
|
vv = flow[yy, xx, 1]
|
|
|
|
mag = np.sqrt(uu**2 + vv**2)
|
|
valid = mag > min_mag
|
|
if valid.sum() < 10:
|
|
return W / 2.0, H / 2.0
|
|
|
|
xx = xx[valid].astype(np.float64)
|
|
yy = yy[valid].astype(np.float64)
|
|
uu = uu[valid].astype(np.float64)
|
|
vv = vv[valid].astype(np.float64)
|
|
|
|
# v * fx - u * fy = v * x - u * y
|
|
A = np.column_stack([vv, -uu])
|
|
b = vv * xx - uu * yy
|
|
result, _, _, _ = np.linalg.lstsq(A, b, rcond=None)
|
|
return float(result[0]), float(result[1])
|
|
|
|
|
|
def _flow_kl_2d(
|
|
flow_a: np.ndarray,
|
|
flow_b: np.ndarray,
|
|
n_angle_bins: int = 36,
|
|
n_mag_bins: int = 20,
|
|
min_mag: float = 0.5,
|
|
) -> float:
|
|
"""KL(P_a || P_b) over a joint (angle, log-magnitude) histogram."""
|
|
|
|
def _hist(flow: np.ndarray) -> np.ndarray | None:
|
|
u, v = flow[:, :, 0].ravel(), flow[:, :, 1].ravel()
|
|
mag = np.sqrt(u**2 + v**2)
|
|
angle = np.degrees(np.arctan2(v, u)) % 360
|
|
valid = mag >= min_mag
|
|
if valid.sum() < 10:
|
|
return None
|
|
mag = mag[valid]
|
|
angle = angle[valid]
|
|
mag_max = max(mag.max(), min_mag + 1.0)
|
|
mag_edges = np.logspace(np.log10(min_mag), np.log10(mag_max), n_mag_bins + 1)
|
|
angle_edges = np.linspace(0, 360, n_angle_bins + 1)
|
|
h, _, _ = np.histogram2d(angle, mag, bins=[angle_edges, mag_edges])
|
|
return h
|
|
|
|
ha, hb = _hist(flow_a), _hist(flow_b)
|
|
if ha is None or hb is None:
|
|
return 0.0
|
|
eps = 1.0
|
|
p = (ha + eps) / (ha + eps).sum()
|
|
q = (hb + eps) / (hb + eps).sum()
|
|
return float((p * np.log(p / q)).sum())
|
|
|
|
|
|
def compute_frame_metrics(
|
|
flow_gt: np.ndarray,
|
|
flow_gen: np.ndarray,
|
|
grid_size: int = 8,
|
|
min_mag: float = 0.5,
|
|
max_mag_pct: float = 80.0,
|
|
) -> dict[str, float]:
|
|
"""Per-frame comparison metrics between two HxWx2 flow fields.
|
|
|
|
Port of mhuo's compute_frame_metrics — see ptlflow_validation.py.
|
|
"""
|
|
metrics: dict[str, float] = {}
|
|
|
|
gt_mag_map = np.linalg.norm(flow_gt, axis=2)
|
|
gen_mag_map = np.linalg.norm(flow_gen, axis=2)
|
|
max_mag_map = np.maximum(gt_mag_map, gen_mag_map)
|
|
mag_hi = np.percentile(max_mag_map, max_mag_pct)
|
|
mag_mask = (max_mag_map >= min_mag) & (max_mag_map <= mag_hi)
|
|
n_valid = int(mag_mask.sum())
|
|
|
|
if n_valid > 0:
|
|
mean_gt = flow_gt[mag_mask].mean(axis=0)
|
|
mean_gen = flow_gen[mag_mask].mean(axis=0)
|
|
else:
|
|
mean_gt = flow_gt.reshape(-1, 2).mean(axis=0)
|
|
mean_gen = flow_gen.reshape(-1, 2).mean(axis=0)
|
|
|
|
metrics["mf_epe"] = float(np.linalg.norm(mean_gt - mean_gen))
|
|
|
|
mf_min_mag = 0.1
|
|
mag_gt = float(np.linalg.norm(mean_gt))
|
|
mag_gen = float(np.linalg.norm(mean_gen))
|
|
if mag_gt < mf_min_mag and mag_gen < mf_min_mag:
|
|
metrics["mf_angle_err"] = 0.0
|
|
metrics["mf_cosine"] = 1.0
|
|
elif mag_gt < mf_min_mag or mag_gen < mf_min_mag:
|
|
metrics["mf_angle_err"] = 90.0
|
|
metrics["mf_cosine"] = 0.0
|
|
elif mag_gt > 1e-6 and mag_gen > 1e-6:
|
|
cos_sim = float(np.dot(mean_gt, mean_gen) / (mag_gt * mag_gen))
|
|
cos_sim = float(np.clip(cos_sim, -1.0, 1.0))
|
|
metrics["mf_angle_err"] = float(np.degrees(np.arccos(cos_sim)))
|
|
metrics["mf_cosine"] = cos_sim
|
|
else:
|
|
metrics["mf_angle_err"] = 0.0
|
|
metrics["mf_cosine"] = 1.0
|
|
|
|
metrics["mf_mag_ratio"] = float(mag_gen / mag_gt) if mag_gt > 1e-6 else 1.0
|
|
|
|
epe_map = np.linalg.norm(flow_gt - flow_gen, axis=2)
|
|
if n_valid > 0:
|
|
metrics["pixel_epe_mean"] = float(epe_map[mag_mask].mean())
|
|
metrics["pixel_epe_max"] = float(epe_map[mag_mask].max())
|
|
else:
|
|
metrics["pixel_epe_mean"] = float(epe_map.mean())
|
|
metrics["pixel_epe_max"] = float(epe_map.max())
|
|
|
|
valid = mag_mask & (gt_mag_map > 0.5) & (gen_mag_map > 0.5)
|
|
if valid.sum() > 0:
|
|
dot = (flow_gt[:, :, 0] * flow_gen[:, :, 0] + flow_gt[:, :, 1] * flow_gen[:, :, 1])
|
|
cos_map = np.clip(dot / (gt_mag_map * gen_mag_map + 1e-8), -1.0, 1.0)
|
|
angle_map = np.degrees(np.arccos(cos_map))
|
|
metrics["px_angle_rmse"] = float(np.sqrt((angle_map[valid]**2).mean()))
|
|
else:
|
|
metrics["px_angle_rmse"] = 0.0
|
|
|
|
H, W = epe_map.shape
|
|
gh, gw = H // grid_size, W // grid_size
|
|
grid_vals = []
|
|
for gi in range(grid_size):
|
|
for gj in range(grid_size):
|
|
cell_mask = mag_mask[gi * gh:(gi + 1) * gh, gj * gw:(gj + 1) * gw]
|
|
cell_epe = epe_map[gi * gh:(gi + 1) * gh, gj * gw:(gj + 1) * gw]
|
|
if cell_mask.sum() > 0:
|
|
grid_vals.append(float(cell_epe[cell_mask].mean()))
|
|
else:
|
|
grid_vals.append(float(cell_epe.mean()))
|
|
metrics["grid_epe_mean"] = float(np.mean(grid_vals))
|
|
metrics["grid_epe_max"] = float(np.max(grid_vals))
|
|
|
|
if n_valid > 0:
|
|
outlier = (epe_map > 3.0) & (epe_map > 0.05 * gt_mag_map) & mag_mask
|
|
metrics["fl_all"] = float(outlier.sum() / n_valid)
|
|
else:
|
|
outlier = (epe_map > 3.0) & (epe_map > 0.05 * gt_mag_map)
|
|
metrics["fl_all"] = float(outlier.mean())
|
|
|
|
foe_gt_x, foe_gt_y = _estimate_foe(flow_gt)
|
|
foe_gen_x, foe_gen_y = _estimate_foe(flow_gen)
|
|
metrics["foe_dist"] = float(np.sqrt((foe_gt_x - foe_gen_x)**2 + (foe_gt_y - foe_gen_y)**2))
|
|
|
|
metrics["flow_kl_2d"] = _flow_kl_2d(flow_gt, flow_gen)
|
|
return metrics
|
|
|
|
|
|
def aggregate_temporal(per_frame: list[dict[str, float]], ) -> dict[str, float | int | None]:
|
|
"""Aggregate per-frame metric dicts into mean/std/max/auc/onset summaries.
|
|
|
|
Port of mhuo's compute_temporal_metrics.
|
|
"""
|
|
n = len(per_frame)
|
|
if n == 0:
|
|
return {"n_frames": 0}
|
|
|
|
summary: dict[str, float | int | None] = {"n_frames": n}
|
|
series: dict[str, np.ndarray] = {k: np.array([m[k] for m in per_frame]) for k in _PER_FRAME_AGG_KEYS}
|
|
for name, vals in series.items():
|
|
summary[f"{name}_mean"] = float(vals.mean())
|
|
summary[f"{name}_std"] = float(vals.std())
|
|
summary[f"{name}_max"] = float(vals.max())
|
|
summary[f"{name}_auc"] = _trapezoid(vals) / max(n - 1, 1)
|
|
|
|
epe_series = series["pixel_epe_mean"]
|
|
window = min(5, n)
|
|
if n >= window:
|
|
baseline = float(np.median(epe_series[:window]))
|
|
threshold = max(baseline * 2.0, 1.0)
|
|
kernel = np.ones(window) / window
|
|
smoothed = np.convolve(epe_series, kernel, mode="valid")
|
|
divergence_frame: int | None = None
|
|
for i, val in enumerate(smoothed):
|
|
if val > threshold:
|
|
divergence_frame = int(i)
|
|
break
|
|
summary["divergence_onset_frame"] = divergence_frame
|
|
summary["divergence_threshold"] = float(threshold)
|
|
else:
|
|
summary["divergence_onset_frame"] = None
|
|
summary["divergence_threshold"] = None
|
|
return summary
|
|
|
|
|
|
def tensor_to_bgr_list(video: torch.Tensor) -> list[np.ndarray]:
|
|
"""Convert ``(T, C, H, W)`` float [0,1] to a list of HWC BGR uint8 frames.
|
|
|
|
Casts + permutes + BGR-swaps on-device and transfers once, instead
|
|
of T per-frame ``.cpu()`` round-trips.
|
|
"""
|
|
bgr_u8 = (video.float() * 255.0).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).flip(-1).contiguous()
|
|
arr = bgr_u8.cpu().numpy()
|
|
return [arr[t] for t in range(arr.shape[0])]
|
|
|
|
|
|
def load_ptlflow_model(model_name: str, ckpt: str, device: torch.device):
|
|
"""Load a ``ptlflow`` model on *device* in eval mode."""
|
|
import ptlflow
|
|
model = ptlflow.get_model(model_name, ckpt_path=ckpt)
|
|
model.eval()
|
|
return model.to(device)
|
|
|
|
|
|
def extract_video_flows(
|
|
model,
|
|
video: torch.Tensor, # (T, C, H, W) float [0, 1]
|
|
*,
|
|
chunk: int,
|
|
device: torch.device,
|
|
) -> list[np.ndarray]:
|
|
"""Run *model* on every consecutive frame pair in *video*.
|
|
|
|
Returns a list of HxWx2 flow arrays of length ``T - 1``.
|
|
"""
|
|
from ptlflow.utils.io_adapter import IOAdapter
|
|
|
|
h, w = video.shape[2], video.shape[3]
|
|
io_adapter = IOAdapter(
|
|
output_stride=model.output_stride,
|
|
input_size=(h, w),
|
|
cuda=(device.type == "cuda"),
|
|
)
|
|
bgr_frames = tensor_to_bgr_list(video)
|
|
pairs = [(bgr_frames[i], bgr_frames[i + 1]) for i in range(len(bgr_frames) - 1)]
|
|
|
|
flows: list[np.ndarray] = []
|
|
for start in range(0, len(pairs), chunk):
|
|
end = min(start + chunk, len(pairs))
|
|
pair_tensors = []
|
|
for f1, f2 in pairs[start:end]:
|
|
inputs = io_adapter.prepare_inputs([f1, f2])
|
|
pair_tensors.append(inputs["images"])
|
|
batched_images = torch.cat(pair_tensors, dim=0)
|
|
with torch.no_grad():
|
|
preds = model({"images": batched_images})
|
|
preds["images"] = batched_images
|
|
preds = io_adapter.unscale(preds)
|
|
flows_tensor = preds["flows"]
|
|
if flows_tensor.dim() == 5:
|
|
flows_tensor = flows_tensor.squeeze(1)
|
|
for i in range(flows_tensor.shape[0]):
|
|
flow = flows_tensor[i].detach().cpu().permute(1, 2, 0).numpy()
|
|
flows.append(flow)
|
|
return flows
|