overwrite with WIP changes: src/evaluate.py

UCPE/src/evaluate.py

@@ -44,6 +44,10 @@ class Args(BaseSettings):
     frame_stride: Optional[int] = None
     pose_frames: Optional[int] = None
 
+    # Cap pred+GT to first N frames for video metrics AND force pose_frames=N.
+    # When set, output dir gets a `_<N>f` suffix so 49/81 runs don't collide.
+    eval_frames: Optional[int] = None
+
     model_config = SettingsConfigDict(
         env_prefix="EVAL_",
         cli_parse_args=True,
@@ -52,11 +56,12 @@ class Args(BaseSettings):
 
 
 def get_path(args):
+    suffix = f"_{args.eval_frames}f" if args.eval_frames is not None else ""
     if args.evaluate_gt:
         assert args.data == "PanShotDataset", "GT evaluation only supports PanShotDataset."
-        paths = {"i2v": (args.data_root / "PanShot" / "videos-test", args.data_root / "evaluate")}
+        paths = {"i2v": (args.data_root / "PanShot" / "videos-test", args.data_root / f"evaluate{suffix}")}
     elif args.test_res_path is not None:
-        paths = {args.test_res_path.name: (args.test_res_path, args.test_res_path.parent / f"evaluate_{args.test_res_path.name}")}
+        paths = {args.test_res_path.name: (args.test_res_path, args.test_res_path.parent / f"evaluate_{args.test_res_path.name}{suffix}")}
     else:
         run_id = os.environ.get('WANDB_RUN_ID', None)
         assert run_id is not None, "WANDB_RUN_ID environment variable must be set."
@@ -66,7 +71,7 @@ def get_path(args):
         for task in ["t2v", "i2v"]:
             task_path = predict_dir / task
             if task_path.exists():
-                paths[task] = (task_path, predict_dir / f"evaluate_{task}")
+                paths[task] = (task_path, predict_dir / f"evaluate_{task}{suffix}")
     print(f"Evaluation paths: {paths}")
     return paths
 
@@ -463,6 +468,27 @@ def video(args):
             video = rearrange(video, "C T H W -> T C H W")  # [T, C, H, W]
             video = video / 2. + 0.5  # to [0, 1]
 
+            if "video" in data:
+                # Baselines produce varying frame counts (seva/flexworld/viewcrafter=49,
+                # cf_ucpe/wan22_fun/bidir=81, gen3c=121). GT is 81. Truncate both to
+                # the common prefix so per-frame metrics (SSIM/LPIPS/PSNR/cs_image/FID)
+                # zip-align. FVD takes the truncated clip too. If args.eval_frames is
+                # set (e.g. EVAL_FRAMES=49), additionally cap to that length so 81-frame
+                # methods can be re-scored against shorter baselines on equal footing.
+                T_common = min(video.shape[0], gt_video.shape[0])
+                if args.eval_frames is not None:
+                    T_common = min(T_common, args.eval_frames)
+                video = video[:T_common]
+                gt_video = gt_video[:T_common]
+                # Resize pred to GT spatial dims when they differ (flexworld 1024-wide,
+                # viewcrafter 576x1024). I3D / Inception have their own internal resize
+                # for FVD/FID but per-frame metrics (SSIM/LPIPS/PSNR/geocalib) need
+                # spatial alignment.
+                if video.shape[-2:] != gt_video.shape[-2:]:
+                    import torch.nn.functional as F
+                    video = F.interpolate(video, size=gt_video.shape[-2:],
+                                          mode="bilinear", align_corners=False)
+
             for metric_name, metric in image_metrics.items():
                 if metric_name == "geocalib":
                     if "video" not in data:
@@ -698,11 +724,27 @@ def vipe(args):
             process.stdin.close()
             process.wait()
         
-        vipe_path = test_dir / "vipe"
+        # vipe SLAM output is shared across eval_frames runs (49f / 81f) because
+        # SLAM operates on the full rectified mp4 — frame-count truncation happens
+        # at pose-metric time, not at SLAM time. Put it outside the suffixed
+        # test_dir so both runs reuse it (saves ~30 min × 5 baselines).
+        vipe_path = test_res_path.parent / f"vipe_{test_res_path.name}"
+        # Invoke vipe env's python directly (skipping `conda run`) so the
+        # LD_LIBRARY_PATH we set below is actually inherited — `conda run
+        # --no-capture-output` re-applies its own env logic that drops our
+        # injection on nodes whose vipe env lacks the torch_libs.sh activate
+        # hook, leading to ImportError: libnccl.so.2.
+        vipe_env = Path("/home/tiger/miniconda3/envs/vipe")
+        vipe_site = vipe_env / "lib" / "python3.10" / "site-packages"
+        ld_paths = [str(vipe_site / "torch" / "lib")]
+        nvidia_dir = vipe_site / "nvidia"
+        if nvidia_dir.is_dir():
+            ld_paths.extend(str(p / "lib") for p in nvidia_dir.iterdir() if (p / "lib").is_dir())
+        env = os.environ.copy()
+        env["LD_LIBRARY_PATH"] = ":".join(ld_paths) + ":" + env.get("LD_LIBRARY_PATH", "")
         cmd = [
-            "conda", "run", "-n", "vipe",
-            "--no-capture-output",
-            "python", "/mnt/pfs/users/zhangchen/panshot/UCPE/thirdparty/vipe/run.py",
+            str(vipe_env / "bin" / "python"),
+            str(Path(__file__).resolve().parent.parent / "thirdparty" / "vipe" / "run.py"),
             "pipeline=default",
             "streams=raw_mp4_stream",
             f"streams.base_path={rectify_res_path}",
@@ -710,8 +752,8 @@ def vipe(args):
             "pipeline.output.save_artifacts=true",
             "pipeline.post.depth_align_model=null",
         ]
-        print(f"[CMD] {' '.join(cmd)}")
-        subprocess.run(cmd, check=True)
+        print(f"[CMD] LD_LIBRARY_PATH={env['LD_LIBRARY_PATH'][:80]}... {' '.join(cmd)}")
+        subprocess.run(cmd, check=True, env=env)
 
 
 def pose(args):
@@ -769,7 +811,9 @@ def pose(args):
 
     for task, (test_res_path, test_dir) in tasks.items():
         print(f"Evaluating task: {task}")
-        vipe_path = test_dir / "vipe"
+        # Read shared vipe SLAM output (see vipe() — single SLAM pass, reused by
+        # both eval_frames=49 and eval_frames=81 pose-metric runs).
+        vipe_path = test_res_path.parent / f"vipe_{test_res_path.name}"
         vipe_pose_path = vipe_path / "pose"
         vipe_video_ids = set(p.stem for p in vipe_pose_path.glob("*.npz"))
         # if valid_video_ids is not None:
@@ -789,9 +833,15 @@ def pose(args):
                 gt_c2w = gt_c2w[::args.frame_stride]
                 pred_c2w = pred_c2w[::args.frame_stride]
 
-            if args.pose_frames is not None:
-                gt_c2w = gt_c2w[:args.pose_frames]
-                pred_c2w = pred_c2w[:args.pose_frames]
+            # Truncate to the common length so pred (e.g. voyager=45) and GT (=81)
+            # match before the rot/trans/cammc matmul. Then optionally cap to
+            # pose_frames_eff (eval_frames overrides pose_frames).
+            T = min(len(gt_c2w), len(pred_c2w))
+            pose_frames_eff = args.eval_frames if args.eval_frames is not None else args.pose_frames
+            if pose_frames_eff is not None:
+                T = min(T, pose_frames_eff)
+            gt_c2w = gt_c2w[:T]
+            pred_c2w = pred_c2w[:T]
 
             # Relative + translation normalized
             gt_rel = normalize_t(relative_pose(gt_c2w.copy()))