Fix DeMemWM revisit metric import

algorithms/dememwm/df_video.py

@@ -1,5 +1,6 @@
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 import csv
+from dataclasses import dataclass
 import os
 import random
 import time
@@ -27,7 +28,6 @@ from datasets.video.memory_selection import (
     _select_revisit,
 )
 from utils.logging_utils import log_video, get_validation_metrics_for_videos
-from algorithms.worldmem.revisit_consistency import RevisitPair, compute_frame_metric_rows
 from .df_base import DiffusionForcingBase
 from .models.vae import VAE_models
 from .models.diffusion import Diffusion
@@ -42,6 +42,160 @@ _DEMEMWM_GEOMETRY_PROJ_MARKERS = (".query_pose_proj.", ".key_pose_proj.", ".time
 _DEMEMWM_ADALN_MLP_MARKERS = (".r_adaLN_modulation.", ".r_mlp.")
 
 
+@dataclass(frozen=True)
+class RevisitPair:
+    clip_id: str
+    source_index: int
+    target_index: int
+    source_frame: int
+    target_frame: int
+    gap: int
+    fov_overlap: float
+    plucker_overlap: float
+    position_distance: float
+    yaw_delta_deg: float
+    pitch_delta_deg: float
+    positional: bool
+
+
+def _best_revisit_pair_by_target(pairs: Sequence[RevisitPair]) -> dict[int, RevisitPair]:
+    best: dict[int, RevisitPair] = {}
+    for pair in pairs:
+        rank = (float(pair.fov_overlap), float(pair.plucker_overlap), int(pair.gap), -int(pair.source_index))
+        current = best.get(int(pair.target_index))
+        if current is None:
+            best[int(pair.target_index)] = pair
+            continue
+        current_rank = (
+            float(current.fov_overlap),
+            float(current.plucker_overlap),
+            int(current.gap),
+            -int(current.source_index),
+        )
+        if rank > current_rank:
+            best[int(pair.target_index)] = pair
+    return best
+
+
+def _metric_psnr_from_mse(mse: torch.Tensor) -> torch.Tensor:
+    return 10.0 * torch.log10(1.0 / mse.clamp_min(torch.finfo(mse.dtype).eps))
+
+
+def _metric_central_crop(tensor: torch.Tensor, crop_fraction: float) -> torch.Tensor:
+    crop_fraction = float(crop_fraction)
+    if crop_fraction == 1.0:
+        return tensor
+    h = int(tensor.shape[-2])
+    w = int(tensor.shape[-1])
+    crop_h = max(1, int(round(h * crop_fraction)))
+    crop_w = max(1, int(round(w * crop_fraction)))
+    top = max(0, (h - crop_h) // 2)
+    left = max(0, (w - crop_w) // 2)
+    return tensor[..., top : top + crop_h, left : left + crop_w]
+
+
+def _metric_lpips_scalar(lpips_model, pred: torch.Tensor, target: torch.Tensor) -> float:
+    if lpips_model is None:
+        return float("nan")
+    pred = torch.clamp(pred.float(), 0.0, 1.0)
+    target = torch.clamp(target.float(), 0.0, 1.0).to(device=pred.device)
+    if hasattr(lpips_model, "reset"):
+        lpips_model.reset()
+    with torch.no_grad():
+        try:
+            value = lpips_model(pred, target)
+        except TypeError:
+            lpips_model.update(pred, target)
+            value = lpips_model.compute()
+    if hasattr(lpips_model, "reset"):
+        lpips_model.reset()
+    if torch.is_tensor(value):
+        return float(value.detach().float().mean().cpu().item())
+    return float(value)
+
+
+def compute_frame_metric_rows(
+    *,
+    branch: str,
+    clip_id: str,
+    pred: torch.Tensor,
+    gt: torch.Tensor | None,
+    poses: torch.Tensor | None,
+    pairs: Sequence[RevisitPair],
+    context_frames: int,
+    lpips_model=None,
+    self_crop_fraction: float = 0.50,
+    focal_length: float = 0.35,
+    compute_plucker_similarity: bool = True,
+) -> list[dict[str, object]]:
+    if pred.ndim == 5:
+        pred = pred[:, 0]
+    if gt is not None and gt.ndim == 5:
+        gt = gt[:, 0]
+    pred = pred.detach()
+    gt = None if gt is None else gt.detach().to(device=pred.device)
+
+    best_pairs = _best_revisit_pair_by_target(pairs)
+    gt_psnr_values = None
+    if gt is not None:
+        gt_mse = (torch.clamp(pred.float(), 0.0, 1.0) - torch.clamp(gt.float(), 0.0, 1.0)).square()
+        gt_psnr_values = _metric_psnr_from_mse(gt_mse.flatten(start_dim=1).mean(dim=1)).detach().cpu().tolist()
+
+    rows: list[dict[str, object]] = []
+    for out_index in range(int(pred.shape[0])):
+        frame_index = int(context_frames) + out_index
+        pair = best_pairs.get(frame_index)
+        gt_lpips = float("nan")
+        gt_psnr = float("nan")
+        if gt is not None and gt_psnr_values is not None:
+            gt_psnr = float(gt_psnr_values[out_index])
+            gt_lpips = _metric_lpips_scalar(lpips_model, pred[out_index : out_index + 1], gt[out_index : out_index + 1])
+
+        self_psnr = float("nan")
+        self_lpips = float("nan")
+        source_index = ""
+        gap = ""
+        fov_overlap = ""
+        plucker_overlap = ""
+        self_status = "not_revisit"
+        if pair is not None:
+            source_index = int(pair.source_index)
+            gap = int(pair.gap)
+            fov_overlap = float(pair.fov_overlap)
+            plucker_overlap = float(pair.plucker_overlap)
+            source_out_index = int(pair.source_index) - int(context_frames)
+            if 0 <= source_out_index < int(pred.shape[0]):
+                pred_source = _metric_central_crop(pred[source_out_index : source_out_index + 1], self_crop_fraction)
+                pred_target = _metric_central_crop(pred[out_index : out_index + 1], self_crop_fraction)
+                self_mse = (torch.clamp(pred_source.float(), 0.0, 1.0) - torch.clamp(pred_target.float(), 0.0, 1.0)).square().mean()
+                self_psnr = float(_metric_psnr_from_mse(self_mse.reshape(1))[0].detach().cpu().item())
+                self_lpips = _metric_lpips_scalar(lpips_model, pred_target, pred_source)
+                self_status = "central_crop"
+            else:
+                self_status = "source_before_prediction_horizon"
+
+        rows.append(
+            {
+                "branch": str(branch),
+                "clip_id": str(clip_id),
+                "frame_index": int(frame_index),
+                "output_index": int(out_index),
+                "is_revisit": pair is not None,
+                "source_index": source_index,
+                "gap": gap,
+                "fov_overlap": fov_overlap,
+                "plucker_overlap": plucker_overlap,
+                "gt_lpips": gt_lpips,
+                "gt_psnr": gt_psnr,
+                "self_lpips": self_lpips,
+                "self_psnr": self_psnr,
+                "self_plucker_similarity": float("nan"),
+                "self_consistency_status": self_status,
+            }
+        )
+    return rows
+
+
 def _cfg_get(cfg, key: str, default=None):
     if cfg is None:
         return default