File size: 4,486 Bytes
eafbe80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """
Loop Closure / Revisit metrics:
- View Recall PSNR: last frame vs first frame (revisit same view)
- Revisit SSIM: same
- Camera trajectory reference error (optional): expected pose vs GT pose when dataset provided
"""
from __future__ import annotations
import argparse
import json
import os
from typing import Any
import numpy as np
from .common import discover_loop_closure_videos, load_video_frames
try:
from skimage.metrics import peak_signal_noise_ratio as psnr_skimage
from skimage.metrics import structural_similarity as ssim_skimage
HAS_SKIMAGE = True
except ImportError:
HAS_SKIMAGE = False
try:
import cv2
HAS_CV2 = True
except ImportError:
HAS_CV2 = False
def _ensure_same_size(img1: np.ndarray, img2: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
if img1.shape == img2.shape:
return img1, img2
h, w = img1.shape[:2]
img2 = cv2.resize(img2, (w, h), interpolation=cv2.INTER_LINEAR)
return img1, img2
def _psnr(img1: np.ndarray, img2: np.ndarray, data_range: float = 255.0) -> float:
img1, img2 = _ensure_same_size(img1, img2)
if HAS_SKIMAGE:
return float(psnr_skimage(img1, img2, data_range=data_range))
mse = np.mean((img1.astype(np.float64) - img2.astype(np.float64)) ** 2)
return float(10 * np.log10((data_range ** 2) / mse)) if mse > 0 else 100.0
def _ssim(img1: np.ndarray, img2: np.ndarray, data_range: float = 255.0) -> float:
img1, img2 = _ensure_same_size(img1, img2)
if not HAS_SKIMAGE:
return 0.0
if img1.ndim == 3:
return float(ssim_skimage(img1, img2, data_range=data_range, channel_axis=2))
return float(ssim_skimage(img1, img2, data_range=data_range))
def view_recall_psnr_ssim(video_path: str) -> dict[str, Any] | None:
"""
For a loop video: first frame = start view, last frame = revisit. Compute PSNR/SSIM(last, first).
Returns dict with view_recall_psnr, view_recall_ssim, num_frames; or None if <2 frames.
"""
frames = load_video_frames(video_path)
n = frames.shape[0]
if n < 2:
return None
first = frames[0]
last = frames[-1]
return {
"view_recall_psnr": _psnr(first, last),
"view_recall_ssim": _ssim(first, last),
"num_frames": n,
}
def run_loop_closure(
evals_root: str,
dataset_base: str | None = None,
video_paths: list[tuple[str, str]] | None = None,
) -> dict[str, Any]:
"""
Compute loop closure metrics on 1_loop_4chunk and 3_multi_ctx_4chunk gen_only MP4s.
Optionally compute trajectory reference error when dataset_base is set (stub: aggregate empty).
"""
if video_paths is None:
video_paths = discover_loop_closure_videos(evals_root)
per_video = []
psnrs = []
ssims = []
for rel, absp in video_paths:
if not os.path.isfile(absp):
continue
res = view_recall_psnr_ssim(absp)
if res is None:
per_video.append({"rel": rel, "view_recall_psnr": None, "view_recall_ssim": None, "num_frames": 0})
continue
psnrs.append(res["view_recall_psnr"])
ssims.append(res["view_recall_ssim"])
per_video.append({"rel": rel, **res})
aggregate = {}
if psnrs:
aggregate["mean_view_recall_psnr"] = float(np.mean(psnrs))
aggregate["min_view_recall_psnr"] = float(np.min(psnrs))
aggregate["mean_view_recall_ssim"] = float(np.mean(ssims))
aggregate["min_view_recall_ssim"] = float(np.min(ssims))
if dataset_base:
aggregate["trajectory_ref_error_note"] = "Optional: set dataset_base and implement pose vs GT; currently not computed."
return {
"dimension": "loop_closure",
"params": {"dataset_base": dataset_base},
"per_video": per_video,
"aggregate": aggregate,
"num_videos": len(per_video),
}
def main():
p = argparse.ArgumentParser(description="Loop Closure / Revisit metrics")
p.add_argument("--evals_root", type=str, required=True, help="evals_ep0 root")
p.add_argument("--dataset_base", type=str, default=None, help="Dataset root for optional trajectory ref error")
p.add_argument("--output", type=str, default=None)
args = p.parse_args()
result = run_loop_closure(args.evals_root, dataset_base=args.dataset_base)
out = json.dumps(result, indent=2)
print(out)
if args.output:
with open(args.output, "w") as f:
f.write(out)
if __name__ == "__main__":
main()
|