Add files using upload-large-folder tool

UCPE/.gitignore

@@ -0,0 +1,182 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# PyPI configuration file
+.pypirc
+
+# Project-specific files
+/data
+/debug
+/.vscode
+pyrightconfig.json
+/logs
+wandb/
+*.pth
+/outputs
+/models

UCPE/CLAUDE.md

@@ -0,0 +1,139 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+UCPE (Unified Camera Positional Encoding) is a CVPR 2026 paper implementing geometry-consistent camera control for text-to-video generation. It introduces Relative Ray Encoding + Absolute Orientation Encoding, adding only 0.5% parameters (35.5M) over a 7.3B base model (Wan2.1-T2V-1.3B). The primary pretrained run ID is `6wodf04s`. The installable package name is `diffsynth` (via `pip install -e .`).
+
+## Common Commands
+
+### Environment Setup
+```bash
+conda create -n UCPE python=3.11 -y
+conda activate UCPE
+conda install -c conda-forge "ffmpeg<8" libiconv libgl -y
+pip install -r requirements.txt
+pip install --no-build-isolation --no-cache-dir flash-attn==2.8.0.post2
+pip install -e .
+cd thirdparty/equilib && pip install -e . && cd ../..
+wandb login
+```
+
+### Training
+```bash
+# Full pipeline (train + predict on both datasets):
+bash scripts/train.sh
+# Or manually with custom config:
+source scripts/set_Wan2.1-T2V-1.3B.sh
+export WANDB_NAME="relray_absmap_comp8"
+python src/main.py fit --model.camera_condition="relray_absmap" --model.attn_compress=8
+```
+
+Note: `scripts/train.sh` runs fit + predict sequentially and contains many commented-out ablation/baseline configs. The `set_*.sh` scripts set `CKPT_PATH="last"` by default, so training auto-resumes from the last checkpoint. To start fresh, override or delete existing checkpoints.
+
+### Prediction
+```bash
+# Manually (scripts/predict.sh contains mostly commented-out historical runs):
+source scripts/set_Wan2.1-T2V-1.3B.sh
+WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.attn_compress=8
+# On RealEstate10k:
+WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.attn_compress=8 --data=Re10kDataModule --trainer.limit_predict_batches=13
+```
+
+### Demo Inference
+```bash
+# Using convenience script (uses pretrained run 6wodf04s, runs all three demo configs: lens, pose, teaser):
+bash scripts/demo.sh
+# Or manually (note: use `python -m src.main` not `python src/main.py` for demo):
+source scripts/set_Wan2.1-T2V-1.3B.sh
+export WANDB_MODE=disabled
+export WANDB_RUN_ID=<run_id>
+export PL_PREDICT__DATA="DemoDataModule"
+export PL_PREDICT__MODEL__CKPT_PATH="logs/<run_id>/checkpoints/pytorch_model.bin"
+python -m src.main predict --data.input_file="demo/teaser.json"
+```
+
+### Evaluation
+```bash
+bash scripts/evaluate.sh  # uses hardcoded WANDB_RUN_ID=6wodf04s
+# Or manually (set WANDB_RUN_ID to your trained model):
+export EVAL_DATA_ROOT="data/UCPE"
+export EVAL_NUM_FRAMES=81
+WANDB_RUN_ID=<run_id> python src/evaluate.py
+# For RealEstate10k evaluation:
+export EVAL_DATA="Re10kDataset"
+export EVAL_DATA_ROOT="data/RealEstate10k"
+export EVAL_POSE_FRAMES=16
+export EVAL_FRAME_STRIDE=4
+export EVAL_LIMIT_EVAL_VIDEOS=100
+WANDB_RUN_ID=<run_id> python src/evaluate.py
+```
+Evaluation requires the FVD model `i3d_pretrained_400.pt` in `models/FVD/`. Note: `evaluate.sh` sets `HF_HUB_OFFLINE=1` independently and does NOT source `set_*.sh`.
+
+### Latent Caching (pre-training)
+```bash
+python src/cache.py
+```
+
+### Data Setup
+- **PanShot**: Download from [Hugging Face](https://huggingface.co/datasets/chengzhag/PanShot) to `data/UCPE/PanShot-7z`, then `cd data/UCPE/PanShot-7z && bash extract_panshot.sh`
+- **RealEstate10k**: Download poses from the [official site](https://google.github.io/realestate10k/) to `data/RealEstate10k/`, plus captions from CameraCtrl
+- **Pretrained weights**: Download from OneDrive link in README, place in `logs/` folder
+
+### Testing & Linting
+There is no test suite and no linting toolchain configured for this project.
+
+## Architecture
+
+### Core Source (`src/`)
+
+- **main.py** - Lightning training harness using `LightningCLI`. `PanShotTrainModule` loads the base WanVideoPipeline, calls `patch_dit()` to inject camera control modules, and `enable_grad()` to selectively train camera parameters. Data modules: `PanShotDataModule`, `Re10kDataModule`, `DemoDataModule`.
+
+- **camera_control.py** - Core UCPE implementation. Key function `patch_dit(pipe, method, ...)` injects `UcpeSelfAttention` modules into DiT blocks. Supports multiple encoding methods: `relray_absmap` (primary), `relray`, `plucker`, `recammaster`, `prope`, `gta`. Camera math uses the Unified Camera Model (UCM) with parameters `x_fov` (field of view) and `xi` (mirror parameter, 0=pinhole, >0=fisheye).
+
+- **dataset.py** - Dataset classes loading video + camera pose (`[T, 3, 4]` extrinsics) + intrinsics (`x_fov`, `xi`) + captions. Includes trajectory normalization and optional yaw zeroing. Note: `Re10kDataModule` has `overwrite_xfov=100.0` default since Re10k doesn't store intrinsics.
+
+- **evaluate.py** - Multi-metric evaluation: video quality (FVD, FID, CLIP-score), camera accuracy (FOV error, distortion, pitch/roll), and pose accuracy (rotation/translation error, CAMMC).
+
+- **cache.py** - Pre-computes VAE latent embeddings for training efficiency.
+
+### Pipeline (`diffsynth/`)
+
+- **pipelines/wan_video_panshot.py** - `WanVideoPipeline` with processing units chained sequentially. The `UCPECameraControl` unit generates camera embeddings from poses/intrinsics that feed into `UcpeSelfAttention` modules in the DiT. Training loss via `training_loss()` using flow-matching diffusion.
+
+- **models/wan_video_dit.py** - `WanModel` DiT (Diffusion Transformer) with spatial-temporal attention blocks. Camera conditioning is integrated via the injected `UcpeSelfAttention` modules.
+
+### Key Data Flow
+
+1. Video frames encoded to latents via VAE (cached by `cache.py`)
+2. Camera poses + intrinsics processed into ray embeddings by `UCPECameraControl`
+3. Ray embeddings injected into DiT self-attention via `UcpeSelfAttention` (PRoPE-style)
+4. Diffusion training: predict noise at random timestep, MSE loss, gradients only flow through camera modules
+
+### Configuration
+
+Training is configured via environment variables and CLI args to `src/main.py` (Lightning CLI). The env var convention uses `PL_{STAGE}__{SECTION}__{PARAM}` format — e.g., `PL_FIT__MODEL__FPS=16` maps to `--model.fps=16` for the `fit` subcommand. One of the `scripts/set_*.sh` scripts must be sourced before any run to set defaults for all stages (FIT, VALIDATE, TEST, PREDICT) and `HF_HUB_OFFLINE=1`:
+- `scripts/set_Wan2.1-T2V-1.3B.sh` — primary text-to-video model (model_id: `Wan-AI/Wan2.1-T2V-1.3B`)
+- `scripts/set_Wan2.2-TI2V-5B.sh` — text+image-to-video 5B model (model_id: `Wan2.2-TI2V-5B`, data at `/tmp/data/UCPE`)
+
+Key model params: `camera_condition` (encoding method), `attn_compress` (attention compression factor, default 8), `adaptation_method` ("parallel"/"before"/"after"), `ti2v_input_image_prob` (probability of conditioning on input image for TI2V, default 0.5), `num_predict` (number of predictions per input, controllable via `PL_PREDICT__MODEL__NUM_PREDICT`). Logging via Weights & Biases; checkpoints saved to `logs/{WANDB_RUN_ID}/checkpoints/`.
+
+### Demo Input Format
+
+JSON array of objects:
+```json
+{"pose_path": "path/to/pose.npy", "x_fov": 100.0, "xi": 0.0, "caption": "text"}
+```
+Where `pose.npy` contains `[T, 3, 4]` camera extrinsics.
+
+## Third-party Dependencies (`thirdparty/`)
+
+- **equilib** - Omnidirectional image projection (must be pip installed separately)
+- **prope** - Position-based Relative Positional Encoding for attention
+- **PanFlow/PanFlowAPI** - Source dataset processing for PanShot curation
+- **GeoCalib, UniK3D, Q-Align, vipe** - Evaluation tools (separate conda envs recommended)
+
+## Data Tools (`tools/`)
+
+20 scripts for dataset curation: PanShot processing (`process_panshot.py`, `caption_panshot.py`), PanFlow pipeline (`filter_panflow.py` -> `align_panflow.py` -> `match_panflow.py` -> `normalize_panflow.py`), visualization (`visualize_pose.py`), and export (`export_figure.py`, `export_table.py`).

UCPE/README.md

@@ -0,0 +1,425 @@
+# 📷 UCPE
+
+<p align="center">
+<h1 align="center">Unified Camera Positional Encoding for Controlled Video Generation</h1>
+<p align="center">
+  <p align="center">
+    <a href="https://chengzhag.github.io/">Cheng Zhang</a><sup>1</sup><sup>,2</sup>
+    ·
+    <a href="https://leeby68.github.io/">Boying Li</a><sup>1</sup>
+    ·
+    <a href="https://www.linkedin.com/in/meng-wei-66687a105/?originalSubdomain=au">Meng Wei</a><sup>1</sup>
+    ·
+    <a href="https://yanpei.me/">Yan-Pei Cao</a><sup>3</sup>
+    ·
+    <a href="https://www.monash.edu/mada/architecture/people/camilo-cruz-gambardella/">Camilo Cruz Gambardella</a><sup>1,2</sup>
+    ·
+    <a href="https://research.monash.edu/en/persons/dinh-phung/">Dinh Phung</a><sup>1</sup>
+    ·
+    <a href="https://jianfei-cai.github.io/">Jianfei Cai</a><sup>1</sup><br>
+    <sup>1</sup>Monash University <sup>2</sup>Building 4.0 CRC <sup>3</sup>VAST
+  </p>
+  <h2 align="center"><a href="https://arxiv.org/abs/2512.07237">Paper</a> | <a href="https://chengzhag.github.io/publication/ucpe/">Project Page</a> | <a href="https://youtu.be/rMX7gxH8jBM">Video</a> | <a href="https://huggingface.co/datasets/chengzhag/PanShot">Hugging Face</a></h2>
+</p>
+
+[![Watch the video](images/thumbnail.png)](https://youtu.be/rMX7gxH8jBM)
+*Our UCPE introduces a geometry-consistent alternative to Plücker rays as one of the core contributions, enabling better generalization in Transformers. We hope to inspire future research on camera-aware architectures.
+
+## 📢 Updates
+- \[2026.03.19\] 🔧 Fixed a bug in Plücker encoding (thanks to [@fengq1a0](https://github.com/fengq1a0)'s [issue #5](https://github.com/chengzhag/UCPE/issues/5)).
+- \[2026.02.21\] 🎉 **UCPE accepted to CVPR 2026**
+- \[2026.02.04\] 📁 **PanShot Dataset And Curation Code** (controllable camera data synthesized from [PanFlow](https://github.com/chengzhag/PanFlow))
+- \[2026.02.04\] 🎯 **Full Training, Evaluation, Visualization Code**
+- \[2025.12.07\] ⚡ **Quick Demo** code released
+
+## 🚀 TLDR
+
+🔥 **Camera-controlled text-to-video generation**, now with **intrinsics**, **distortion** and **orientation** control!
+
+<p align="center">
+  <img src="images/cameras.png" alt="Camera lenses" height="120px">
+  &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp;
+  <img src="images/orientation.png" alt="Orientation control" height="140px">
+</p>
+
+📷 UCPE integrates **Relative Ray Encoding**—which delivers significantly better generalization than Plücker across diverse camera motion, intrinsics and lens distortions—with **Absolute Orientation Encoding** for controllable pitch and roll, enabling a unified camera representation for Transformers and state-of-the-art camera-controlled video generation with just **0.5% extra parameters** (35.5M over the 7.3B parameters of the base model)
+
+<p align="center">
+  <img src="images/video-ucpe.gif"
+       alt="UCPE"
+       style="max-height:480px; width:auto;">
+</p>
+
+## 🛠️ Installation
+
+```bash
+conda create -n UCPE python=3.11 -y
+conda activate UCPE
+conda install -c conda-forge "ffmpeg<8" libiconv libgl -y
+pip install -r requirements.txt
+pip install --no-build-isolation --no-cache-dir flash-attn==2.8.0.post2
+pip install -e .
+
+cd thirdparty/equilib
+pip install -e .
+```
+
+We use wandb to log and visualize the training process. You can create an account then login to wandb by running the following command:
+
+```bash
+wandb login
+```
+
+<details>
+<summary>Below are installations for tools used in evaluation and dataset processing
+that can be skipped if you do not need these tools.</summary>
+
+```bash
+cd ../GeoCalib
+pip install -e .
+pip install -e siclib
+
+cd ../UniK3D
+pip install -e . --extra-index-url https://download.pytorch.org/whl/cu121
+
+cd ../Q-Align
+conda create -n qalign python=3.9 -y
+conda activate qalign
+pip install -e .
+pip install jsonlines "numpy<2" protobuf pydantic-settings
+
+cd ../vipe
+conda env create -f envs/base.yml
+conda activate vipe
+pip install -r envs/requirements.txt
+pip install --no-build-isolation -e .
+```
+</details>
+<br>
+
+## ⚡ Quick Demo
+
+Download our finetuned weights from [OneDrive](https://monashuni-my.sharepoint.com/:f:/g/personal/cheng_zhang_monash_edu/IgCoTNrYOJRJRKtk5A6I1yiCAR9c64-BOrsId5GYsUxE9y4?e=hD26qU) and put it in `logs/` folder. Then run:
+
+```bash
+bash scripts/demo.sh
+```
+
+The generated videos will be saved in `logs/6wodf04s/demo`, examples shown below:
+
+* `demo/lens.json`: Our **Relative Ray Encoding** not only generalizes to but also enables controllability over a wide range of camera intrinsics and lens distortions.
+
+<p align="center">
+  <img src="images/video-lens.gif"
+       alt="Lens control"
+       style="max-height:480px; width:auto;">
+</p>
+
+* `demo/pose.json`: The geometry-consistent design of **Relative Ray Encoding** further allows strong generalization and controllability over diverse camera motions.
+
+<p align="center">
+  <img src="images/video-pose.gif"
+       alt="Pose control"
+       style="max-height:480px; width:auto;">
+</p>
+
+* `demo/teaser.json`: Our **Absolute Orientation Encoding** further eliminate the ambiguity in pitch and roll in previous T2V methods, enabling precise control over initial camera orientation.
+
+<p align="center">
+  <img src="images/video-orientation.gif"
+       alt="Orientation control"
+       style="max-height:480px; width:auto;">
+</p>
+
+
+## 🌏 PanShot Dataset
+
+Please download the PanShot dataset from [Hugging Face](https://huggingface.co/datasets/chengzhag/PanShot) to `data/UCPE/PanShot-7z` by:
+
+```bash
+huggingface-cli download chengzhag/PanShot --repo-type dataset --local-dir data/UCPE/PanShot-7z
+```
+
+Then extract the dataset by:
+```bash
+cd data/UCPE/PanShot-7z
+bash extract_panshot.sh
+cd ../../..
+```
+The extracted dataset will be saved in `data/UCPE/PanShot`.
+Please then copy the other files to form the following folder structure:
+
+```
+├── captioned-test.jsonl
+├── captioned-train.jsonl
+├── max_rotation-test.json
+├── meta-test
+├── meta-train
+├── pose-test
+├── pose-train
+├── videos-test
+└── videos-train
+```
+
+<details>
+<summary>If you want to go through the dataset curation process, Please follow these three steps.</summary>
+
+> **Shortcut for steps 1 & 2:** You can skip the CameraBench and PanFlow curation steps by downloading our pre-processed data directly:
+> ```bash
+> huggingface-cli download --repo-type dataset chengzhag/UCPE --local-dir data/UCPE
+> cd data/UCPE && bash unpack_hf.sh && cd ../..
+> ```
+> Then proceed directly to step 3 (PanShot).
+
+### CameraBench
+
+Download the dataset from multiple sources:
+
+```bash
+cd data
+huggingface-cli download --repo-type dataset syCen/CameraBench --local-dir CameraBench
+cd CameraBench
+huggingface-cli download --repo-type dataset syCen/Videos4CameraBnech --local-dir data/videos
+wget https://huggingface.co/datasets/chancharikm/cambench_train_videos/resolve/main/videos.zip
+unzip videos.zip -d videos
+cd ../..
+```
+
+Process the dataset:
+
+```bash
+conda activate UCPE
+python tools/process_camerabench.py  # set split = "train" and split = "test"
+
+conda activate vipe
+cd thirdparty/vipe
+python thirdparty/vipe/run.py pipeline=default streams=raw_mp4_stream streams.base_path=data/UCPE/CameraBench/videos/ pipeline.output.path=data/UCPE/CameraBench/vipe/ pipeline.output.save_artifacts=true pipeline.post.depth_align_model=null
+
+conda activate UCPE
+python tools/geocalib_camerabench.py
+python tools/filter_camerabench.py
+```
+
+Processed dataset will be saved in `data/UCPE/CameraBench`.
+
+### PanFlow
+
+Download the pretrained model `PanoFlow(RAFT)-wo-CFE.pth` of Panoflow at [weiyun](https://share.weiyun.com/SIpeQTNE), then put it in `models/PanoFlow` folder.
+
+Our PanShot dataset is built upon [PanFlow](https://github.com/chengzhag/PanFlow) dataset's videos and slam_poses. Please download follow their [instructions](https://github.com/chengzhag/PanFlow/tree/main/curation#download-data) on how to download the full videos and download their meta and slam_poses files following [Full Dataset](https://github.com/chengzhag/PanFlow/tree/main#-full-dataset).
+
+Then process the dataset with:
+
+```bash
+conda activate UCPE
+python tools/filter_panflow.py
+
+conda activate qalign
+python tools/score_panflow.py
+
+conda activate UCPE
+python tools/align_panflow.py  # set split = "train" and split = "test"
+python tools/match_panflow.py  # set split = "train" and split = "test"
+python tools/normalize_panflow.py  # set split = "train" and split = "test"
+```
+
+
+### PanShot
+
+Export your YouTube cookies to `~/.config/cookies.txt` in Netscape format for 4k download. Then download and process the dataset:
+
+```bash
+conda activate UCPE
+python tools/process_panshot.py  # set split = "train" and split = "test"
+python tools/caption_panshot.py  # set split = "train" and split = "test"
+```
+
+</details>
+<br>
+
+## 🏡 RealEstate10k Dataset
+
+We use RealEstate10k Dataset for evaluation, so only poses and captions are needed. Plesae download the RealEstate10k poses from the official [website](https://google.github.io/realestate10k/) ([RealEstate10K.tgz](https://storage.cloud.google.com/realestate10k-public-files/RealEstate10K.tar.gz)) and unpack it to `data/RealEstate10k` folder. Then download the captions from [CameraCtrl](https://github.com/hehao13/CameraCtrl) ([train](https://drive.google.com/file/d/1nytBYjTa0bJ-8AMJWVCtKT2XwkJR3Jra/view) and [test](https://drive.google.com/file/d/1AGEJYbfip0jcp-ymgU9uCjUHzqETivYP/view))
+
+The final folder structure should be like this:
+```
+├── captions
+│   ├── test.json
+│   └── train.json
+├── pose_files
+│   ├── test
+│   └── train
+└── traj_normalization.txt
+```
+
+## 🎯 Training and Evaluation
+
+Prepare the latent cache and train the model with:
+
+```bash
+python src/cache.py
+bash scripts/train.sh
+```
+
+We used 8 A800 GPUs for training, which takes about 1 day. You'll get a WANDB_RUN_ID (e.g., `6wodf04s`) after starting the training. The logs will be synced to your wandb account and the checkpoints will be saved in `logs/<WANDB_RUN_ID>/checkpoints/`. You can use other commented settings in `scripts/train.sh` for ablation studies and baselines.
+
+For evaluation, first download the pretrained model `i3d_pretrained_400.pt` in [common_metrics_on_video_quality](https://github.com/JunyaoHu/common_metrics_on_video_quality/blob/main/fvd/videogpt/i3d_pretrained_400.pt), then put it in `models/FVD` folder. Evaluate results with:
+
+```bash
+bash scripts/evaluate.sh
+```
+
+Please change the `WANDB_RUN_ID` in `scripts/evaluate.sh` on your own trained model and check other commented settings for ablation studies and baselines.
+We note that there are some jitters in the synthesized videos due to inaccurate ViPE pose estimation. Therefore, our evaluation script uses the filtered RealEstate10k test set to avoid those cases.
+
+
+## 🔧 Tools
+
+<details>
+<summary>We also provide tools for visualizing camera trajectories, exporting figures and tables for paper, and visualizing camera statistics.</summary>
+
+Visualize camera trajectories:
+
+```bash
+# Export static camera trajectory visualizations
+python -m tools.visualize_panshot --out_path=data/UCPE/PanShot/pose_vis-test/ --zero_first_yaw
+python -m tools.visualize_re10k --pose_file_path=data/RealEstate10k/pose_files/test/ --filter_file=data/RealEstate10k/filter_files/filter_test_81.txt --relative_c2w --num_videos=150 --out_path=data/RealEstate10k/pose_vis/test/
+
+# Export animated camera trajectory visualizations
+python -m tools.visualize_panshot --out_path=data/UCPE/PanShot/pose_anim-test/ --zero_first_yaw --animate_camera
+python -m tools.visualize_re10k --pose_file_path=data/RealEstate10k/pose_files/test/ --filter_file=data/RealEstate10k/filter_files/filter_test_81.txt --relative_c2w --num_videos=150 --out_path=data/RealEstate10k/pose_anim/test/ --animate_camera
+```
+
+Export figures for paper:
+
+```bash
+# Teaser figure
+python -m tools.export_figure \
+    --methods \
+    "UCPE" "logs/6wodf04s/demo/t2v" \
+    --input_file \
+    "demo/teaser.json" \
+    --output_dir \
+    "outputs/figures/teaser" \
+    --animate_latup
+
+# Try other demo configs
+    # "demo/pose.json" \
+    # "demo/lens.json" \
+
+# Comparison on PanShot dataset
+python -m tools.export_figure \
+    --data=PanShotDataset \
+    --data_root="data/UCPE" \
+    --methods \
+    "ReCamMaster" "logs/khnmur4b/predict/t2v" \
+    "Wan CameraCtrl" "logs/9hjx47bc/predict/t2v" \
+    "UCPE" "logs/6wodf04s/predict/t2v" \
+    --output_dir \
+    "outputs/figures/panshot" \
+    --sample_frames=3 \
+    --animate_latup
+
+# Comparison on RealEstate10k dataset
+python -m tools.export_figure \
+    --data=Re10kDataset \
+    --data_root="data/RealEstate10k" \
+    --methods \
+    "ReCamMaster" "logs/lg1mxf9u/RealEstate10k/t2v" \
+    "Wan CameraCtrl" "logs/3yf7psvi/RealEstate10k/t2v" \
+    "CameraCtrl" "/mnt/pfs/users/zhangchen/panshot/CameraCtrl/out/re10k" \
+    "AC3D" "/mnt/pfs/users/zhangchen/panshot/ac3d/out/5B/test/10000" \
+    "UCPE" "logs/coo9rjaq/RealEstate10k/t2v" \
+    --output_dir \
+    "outputs/figures/re10k" \
+    --sample_frames=3
+```
+
+Export table for paper:
+
+```bash
+# Comparison on PanShot (w/o Absolute Orientation Control)
+python -m tools.export_table \
+    --pad_cols 1 \
+    --methods \
+    "ReCamMaster" "logs/lg1mxf9u/predict/evaluate_t2v/overall/last.json" \
+    "Wan CameraCtrl" "logs/3yf7psvi/predict/evaluate_t2v/overall/last.json" \
+    "UCPE" "logs/coo9rjaq/predict/evaluate_t2v/overall/last.json" \
+    --metrics \
+    "video_metrics/vfov_err" "video_metrics/k1_err" "video_metrics/k2_err" \
+    "video_metrics/pitch_err" "video_metrics/roll_err" \
+    "pose/rot_err" "pose/trans_err" "pose/cammc" \
+    "video_metrics/fvd" "video_metrics/fid" \
+    "video_metrics/cs_text"
+
+# Comparison on PanShot (w/ Absolute Orientation Control)
+python -m tools.export_table \
+    --pad_cols 1 \
+    --methods \
+    "ReCamMaster" "logs/khnmur4b/predict/evaluate_t2v/overall/last.json" \
+    "Wan CameraCtrl" "logs/9hjx47bc/predict/evaluate_t2v/overall/last.json" \
+    "UCPE" "logs/6wodf04s/predict/evaluate_t2v/overall/last.json" \
+    --metrics \
+    "video_metrics/vfov_err" "video_metrics/k1_err" "video_metrics/k2_err" \
+    "video_metrics/pitch_err" "video_metrics/roll_err" \
+    "pose/rot_err" "pose/trans_err" "pose/cammc" \
+    "video_metrics/fvd" "video_metrics/fid" \
+    "video_metrics/cs_text"
+
+# Ablation Study on PanShot
+python -m tools.export_table \
+    --pad_cols 1 \
+    --methods \
+    "1/2-dim (\$128 \times 6\$)" "logs/r0hmwcag/predict/evaluate_t2v/overall/last.json" \
+    "1/4-dim (\$128 \times 3\$)" "logs/nv4al3mj/predict/evaluate_t2v/overall/last.json" \
+    "1/8-dim (\$192 \times 1\$)" "logs/6wodf04s/predict/evaluate_t2v/overall/last.json" \
+    "1/12-dim (\$128 \times 1\$)" "logs/lkxh4srz/predict/evaluate_t2v/overall/last.json" \
+    "Pre-Attn" "logs/p03o7rqy/predict/evaluate_t2v/overall/last.json" \
+    "Post-Attn" "logs/82awngqn/predict/evaluate_t2v/overall/last.json" \
+    "PRoPE" "logs/wekc4yx6/predict/evaluate_t2v/overall/last.json" \
+    "GTA" "logs/z0cfx65s/predict/evaluate_t2v/overall/last.json" \
+    --metrics \
+    "video_metrics/vfov_err" "video_metrics/k1_err" "video_metrics/k2_err" \
+    "video_metrics/pitch_err" "video_metrics/roll_err" \
+    "pose/rot_err" "pose/trans_err" "pose/cammc" \
+    "video_metrics/fvd" "video_metrics/fid" \
+    "video_metrics/cs_text"
+
+# Comparison on RealEstate10k
+python -m tools.export_table \
+    --methods \
+    "ReCamMaster" "logs/lg1mxf9u/RealEstate10k/evaluate_t2v/overall/last.json" \
+    "Wan CameraCtrl" "logs/3yf7psvi/RealEstate10k/evaluate_t2v/overall/last.json" \
+    "CameraCtrl" "../CameraCtrl/out/evaluate_re10k/overall/last.json" \
+    "AC3D" "../ac3d/out/5B/test/evaluate_10000/overall/last.json" \
+    "UCPE" "logs/coo9rjaq/RealEstate10k/evaluate_t2v/overall/last.json" \
+    --metrics \
+    "pose/rot_err" "pose/trans_err" "pose/cammc" \
+    "qalign/image_quality" "qalign/image_aesthetic" "qalign/video_quality"
+```
+
+Visualize camera statistics:
+```bash
+# PanShot
+python -m tools.dataset_statistics \
+    --data=PanShotDataset \
+    --data_root=data/UCPE \
+    --output_dir=outputs/suppl/panshot \
+    --color=C0
+
+# RE10K
+python -m tools.dataset_statistics \
+    --data=Re10kDataset \
+    --data_root=data/RealEstate10k \
+    --output_dir=outputs/suppl/re10k \
+    --color=C1
+```
+
+</details>
+<br>
+
+## 💡 Acknowledgements
+
+Our paper cannot be completed without the amazing open-source projects [Wan2.1](https://github.com/Wan-Video/Wan2.1), [AC3D](https://github.com/snap-research/ac3d), [ReCamMaster](https://github.com/KlingTeam/ReCamMaster), [CameraCtrl](https://github.com/hehao13/CameraCtrl), [prope](https://github.com/liruilong940607/prope), [vllm](https://github.com/vllm-project/vllm), [stella_vslam](https://github.com/stella-cv/stella_vslam)...
+
+Also check out our Pan-Series works [PanFlow](https://github.com/chengzhag/PanFlow), [PanFusion](https://github.com/chengzhag/PanFusion) and [PanSplat](https://github.com/chengzhag/PanSplat) towards 3D scene generation with panoramic images!

UCPE/commands.sh

@@ -0,0 +1,45 @@
+conda activate UCPE
+source scripts/set_Wan2.2-TI2V-5B.sh
+export WANDB_NAME="relray_absmap_comp8_ti2v"
+export WANDB_RUN_ID="ti2v_relray_absmap_comp8"
+python src/main.py fit --model.camera_condition="relray_absmap" --model.attn_compress=8 --data.data_root=/tmp/data/UCPE
+
+
+
+conda activate vipe
+python thirdparty/vipe/run.py pipeline=default \
+    streams=raw_mp4_stream \
+    streams.base_path=data/UCPE/CameraBench/videos/ \
+    pipeline.output.path=data/UCPE/CameraBench/vipe/ \
+    pipeline.output.save_artifacts=true \
+    pipeline.post.depth_align_model=null
+
+
+
+# === Download and prepare datasets ===
+DATA_DIR="data/UCPE"
+mkdir -p "$DATA_DIR/PanShot"
+
+# Download PanShot dataset (raw 7z archives with poses, captions, meta)
+hf download chengzhag/PanShot --repo-type dataset --local-dir "$DATA_DIR/PanShot-7z"
+
+# Extract PanShot archives (pose-train/, meta-train/, etc.)
+python scripts/extract_panshot_poses.py --panshot_root "$DATA_DIR/PanShot-7z"
+
+
+# latent generating
+conda activate UCPE
+source scripts/set_Wan2.2-TI2V-5B.sh
+python src/cache.py --model.model_id=Wan2.2-TI2V-5B --data.video_subdir=videos_704 --data.data_root=data/UCPE
+
+
+conda activate UCPE
+source scripts/set_Wan2.2-TI2V-5B.sh
+export WANDB_NAME="relray_absmap_comp8_ti2v_704"
+export WANDB_RUN_ID="ti2v_relray_absmap_comp8_704"
+
+python src/main.py fit \
+  --model.camera_condition="relray_absmap" \
+  --model.attn_compress=8 \
+  --data.video_subdir=videos_704 \
+  --data.data_root=data/UCPE 

UCPE/demo/lens.json

@@ -0,0 +1,74 @@
+[
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 120.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 140.0,
+        "xi": 0.8,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 180.0,
+        "xi": 2.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 200.0,
+        "xi": 2.3,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 120.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 140.0,
+        "xi": 0.8,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 180.0,
+        "xi": 2.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 200.0,
+        "xi": 2.3,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    }
+]

UCPE/demo/pose.json

@@ -0,0 +1,134 @@
+[
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/aQVLCO-OOqM-11-1-no_rot_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/5451cefde53f06f1.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/Msw3A5IPPD0-96-0-yaw_pitch_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/ecc2b1b9cb1fa9ad.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/N99hmkAJwlQ-3-0-linear_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/38f7ba7fd9a83069.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/20e7a3651ec30386.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/1jVPI3bNCBc-5-0-no_rot_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/HTdKkiu771g-3-0-no_rot_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/YUPoNo1j420-5-0-no_rot_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/aQVLCO-OOqM-11-1-no_rot_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/5451cefde53f06f1.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/Msw3A5IPPD0-96-0-yaw_pitch_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/ecc2b1b9cb1fa9ad.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/N99hmkAJwlQ-3-0-linear_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/38f7ba7fd9a83069.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/20e7a3651ec30386.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/1jVPI3bNCBc-5-0-no_rot_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/HTdKkiu771g-3-0-no_rot_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/YUPoNo1j420-5-0-no_rot_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    }
+]

UCPE/demo/teaser.json

@@ -0,0 +1,164 @@
+[
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/aQVLCO-OOqM-11-1-no_rot_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/aQVLCO-OOqM-11-1-no_rot_aug_0.npy",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy golden dog resting on a garden stone wall surrounded by blooming flowers, gazing curiously toward the garden with its ears gently perked and fur glowing in the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy golden dog lying on a garden stone wall surrounded by blooming flowers, calmly watching the garden with gentle eyes as warm golden sunlight filters through the leaves and soft shadows move across its fur."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A fluffy gray cat lounging on a garden stone wall surrounded by blooming flowers, as the camera moves in a smooth clockwise arc around it, capturing the cat’s curious gaze, swaying tail, and the warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful gray cat leaping from a garden stone wall, its eyes focused ahead and body stretched midair as petals and dust swirl around in warm golden sunlight filtering through the leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A playful golden dog leaping from a garden stone path, its gaze focused ahead as petals swirl through the air and warm sunlight glows on its fur amid softly illuminated leaves."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A curious orange cat resting on a sunlit wooden table beside a window, with warm afternoon light streaming across its fur and soft shadows on the polished surface, surrounded by a few open books and green houseplants that gently frame the cozy, golden-lit room."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/f74139ac48f19b3c.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "A gentle brown dog resting on a sunlit wooden floor beside a window, with warm afternoon light streaming across its fur and soft shadows stretching over the polished surface, surrounded by a few open books and green houseplants that frame the cozy, golden-lit room."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/5451cefde53f06f1.txt",
+        "x_fov": 100.0,
+        "xi": 0.0,
+        "caption": "Aerial view ascending through a dense cityscape of tall glass skyscrapers and steel towers, with sunlight glinting off reflective windows and rooftop details gradually unfolding, revealing layered streets, terraces, and the vast geometry of the urban skyline ahead."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/Msw3A5IPPD0-96-0-yaw_pitch_aug_0.npy",
+        "x_fov": 120.0,
+        "xi": 0.0,
+        "caption": "Aerial drone view gliding over a historic European-style town with stone buildings, terracotta rooftops, and narrow winding streets, where soft golden light reveals weathered walls, old clock towers, and quiet courtyards, conveying the timeless charm and layered history of the place."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/Msw3A5IPPD0-96-0-yaw_pitch_aug_0.npy",
+        "x_fov": 120.0,
+        "xi": 0.0,
+        "caption": "Aerial drone view soaring above a dense green forest with tall trees of varying heights and winding trails, as a clear river weaves through the landscape reflecting sunlight between the branches, revealing the depth, texture, and gentle rhythm of the forest stretching across rolling hills."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/Msw3A5IPPD0-96-0-yaw_pitch_aug_0.npy",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A low aerial drone view gliding just above the treetops of a dense forest, weaving through layers of tall and short trees as a winding river glimmers between them, with sunlight flickering across leaves and water, revealing shifting textures, flowing shadows, and the serene rhythm of the landscape over rolling green hills."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/ecc2b1b9cb1fa9ad.txt",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A close-up ground-level view following a cat walking slowly through tall green grass, its tail swaying as it occasionally turns its head toward the camera with a curious gaze, while sunlight glimmers through the leaves and soft shadows ripple across the ground."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/ecc2b1b9cb1fa9ad.txt",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A close-up ground-level view following a cat walking along a narrow forest path covered with fallen leaves and dappled sunlight, its tail swaying as it glances back toward the camera, while soft beams of light filter through the trees and gentle shadows move across the mossy ground."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/ecc2b1b9cb1fa9ad.txt",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A close-up ground-level view following a dog walking leisurely through tall green grass, its ears gently perked and tail wagging as it occasionally glances back toward the camera with a friendly gaze, while sunlight filters through the blades and soft shadows dance across the ground."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/ecc2b1b9cb1fa9ad.txt",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A close-up ground-level view following a dog trotting along a sandy beach near the shoreline, its paws leaving prints in the wet sand as gentle waves roll in, sunlight glinting off the water and sea breeze rustling its fur while it occasionally looks back with a playful expression."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/N99hmkAJwlQ-3-0-linear_aug_0.npy",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A dramatic first-person drone view racing through a waterfall gorge where torrents crash between moss-covered cliffs and overhanging trees, with mist swirling, sunlight glinting off wet rocks and flowing water, and glimpses of ferns and boulders adding cinematic intensity."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/N99hmkAJwlQ-3-0-linear_aug_0.npy",
+        "x_fov": 180.0,
+        "xi": 2.0,
+        "caption": "A high-speed first-person drone view diving through a narrow gorge where a powerful waterfall crashes between moss-covered rocks and tall trees cling to the cliffs, with mist swirling through the air and sunlight scattering off droplets as the roaring water fills the scene with energy and depth."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/N99hmkAJwlQ-3-0-yaw_aug_0.npy",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A daytime first-person dashcam view driving along a city street with moderate traffic, showing cars and trucks ahead moving between lane markings, with roadside buildings, power lines, and trees under a partly cloudy sky as soft sunlight and shadows drift across the road."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/38f7ba7fd9a83069.txt",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A first-person dashcam view from inside a car looking through the windshield while approaching an intersection on a sunny afternoon, smoothly turning right onto a quiet tree-lined street with parked cars and small shops, as sunlight filters through the leaves and reflections glide across the glass."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/38f7ba7fd9a83069.txt",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A daytime dashcam view driving along a suburban street under a partly cloudy sky, with rows of cars ahead moving slowly in traffic, white pickup trucks and sedans lining both lanes, power lines stretching overhead, and trees and low buildings framing the road in soft sunlight."
+    },
+    {
+        "pose_path": "data/RealEstate10k/pose_files/test/20e7a3651ec30386.txt",
+        "x_fov": 200.0,
+        "xi": 2.3,
+        "caption": "A first-person view moving through a large warehouse with tall shelves and stacked boxes, illuminated by cool fluorescent lights, before coming to a steady stop in front of a metal rack filled with labeled packages, as faint echoes, reflections, and mechanical ambience fill the quiet industrial space."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/1jVPI3bNCBc-5-0-no_rot_aug_0.npy",
+        "x_fov": 200.0,
+        "xi": 2.3,
+        "caption": "A first-person view moving through a large warehouse with tall shelves and stacked boxes, illuminated by cool fluorescent lights, before coming to a steady stop in front of a metal rack filled with labeled packages, as faint echoes, reflections, and mechanical ambience fill the quiet industrial space."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/HTdKkiu771g-3-0-no_rot_aug_0.npy",
+        "x_fov": 200.0,
+        "xi": 2.3,
+        "caption": "A time-lapse view from an airplane window at night, showing the wing silhouetted against the glowing Milky Way that slowly rotates above the horizon, while a vast sea of clouds drifts below like rolling waves, illuminated by faint starlight and distant lightning flashes that shimmer across the metallic surface of the wing."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/YUPoNo1j420-5-0-no_rot_aug_0.npy",
+        "x_fov": 200.0,
+        "xi": 2.3,
+        "caption": "A time-lapse view from an airplane window at night, showing the wing silhouetted against the glowing Milky Way that slowly rotates above the horizon, while a vast sea of clouds drifts below like rolling waves, illuminated by faint starlight and distant lightning flashes that shimmer across the metallic surface of the wing."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/HTdKkiu771g-3-0-no_rot_aug_0.npy",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A time-lapse night view from the foot of a snow-covered mountain, with towering pine trees in the foreground and the majestic peak centered in the frame, as the Milky Way arches slowly across the sky, casting faint light over the snowy slopes and shimmering through the crisp, clear mountain air."
+    },
+    {
+        "pose_path": "data/UCPE/PanShot/pose-test/YUPoNo1j420-5-0-no_rot_aug_0.npy",
+        "x_fov": 160.0,
+        "xi": 1.5,
+        "caption": "A time-lapse night view from the foot of a snow-covered mountain, with towering pine trees in the foreground and the majestic peak centered in the frame, as the Milky Way arches slowly across the sky, casting faint light over the snowy slopes and shimmering through the crisp, clear mountain air."
+    }
+]

UCPE/diffsynth/__init__.py

@@ -0,0 +1,6 @@
+from .data import *
+from .models import *
+from .prompters import *
+from .schedulers import *
+from .pipelines import *
+from .controlnets import *

UCPE/requirements.txt

@@ -0,0 +1,36 @@
+torch==2.7.1
+torchvision
+transformers
+imageio
+imageio[ffmpeg]
+safetensors
+einops
+sentencepiece
+protobuf
+modelscope
+ftfy
+pynvml
+pandas
+accelerate
+peft
+datasets
+qwen-vl-utils[decord]
+jsonlines
+matplotlib
+ffmpeg-python
+vllm>0.7.2
+viser
+plotly
+yt-dlp
+pyarrow
+lightning
+websockets
+jsonargparse[signatures]
+deepspeed
+wandb
+omegaconf
+pydantic
+pydantic-settings
+torchmetrics[image]
+tyro
+seaborn

UCPE/scripts/compare_panshot.sh

@@ -0,0 +1,145 @@
+#!/bin/bash
+# Run UCPE inference on the PanShot test set and stitch each result with its
+# ground-truth reference into a side-by-side comparison mp4.
+#
+# Usage:
+#   bash scripts/compare_panshot.sh [N]
+#
+# Args:
+#   N   number of samples to generate (default: 8). Pass 0 for the full test set.
+#
+# Env overrides:
+#   WANDB_RUN_ID   default: 6wodf04s
+#   CKPT_PATH      default: logs/$WANDB_RUN_ID/checkpoints/pytorch_model.bin
+#   VIDEO_SUBDIR   default: videos_704  (PanShot videos-test is empty on this machine,
+#                                        videos_704-test is the available variant)
+
+set -euo pipefail
+
+cd "$(dirname "$0")/.."
+
+source scripts/set_Wan2.1-T2V-1.3B.sh
+
+export WANDB_MODE=disabled
+export WANDB_RUN_ID="${WANDB_RUN_ID:-6wodf04s}"
+
+CKPT_PATH="${CKPT_PATH:-logs/${WANDB_RUN_ID}/checkpoints/pytorch_model.bin}"
+VIDEO_SUBDIR="${VIDEO_SUBDIR:-videos_704}"
+NUM_GPUS="${NUM_GPUS:-8}"
+N="${1:-8}"
+
+# Use the local Wan2.1-T2V-1.3B symlink (-> HF cache snapshot) instead of the
+# default "Wan-AI/Wan2.1-T2V-1.3B" model_id, which would trigger a slow
+# ModelScope re-download (HF_HUB_OFFLINE=1 doesn't suppress ModelScope).
+if [[ -d "Wan2.1-T2V-1.3B" ]]; then
+    export PL_PREDICT__MODEL__MODEL_ID="Wan2.1-T2V-1.3B"
+fi
+
+if [[ ! -f "$CKPT_PATH" ]]; then
+    echo "ERROR: UCPE adapter not found at $CKPT_PATH" >&2
+    exit 1
+fi
+
+export PL_PREDICT__MODEL__CKPT_PATH="$CKPT_PATH"
+# set_*.sh sets PL_PREDICT__CKPT_PATH=last (Lightning resume), but we don't have a
+# last.ckpt for 6wodf04s — clear it so Lightning uses the in-memory model that
+# already loaded pytorch_model.bin in PanShotTrainModule.__init__.
+unset PL_PREDICT__CKPT_PATH
+
+LIMIT_ARG=()
+if [[ "$N" != "0" ]]; then
+    # Lightning's limit_predict_batches is per-rank under DDP, so ceil(N / NUM_GPUS)
+    # gives ~N total samples across all ranks.
+    PER_RANK=$(( (N + NUM_GPUS - 1) / NUM_GPUS ))
+    LIMIT_ARG=(--trainer.limit_predict_batches="$PER_RANK")
+    EFFECTIVE_N=$(( PER_RANK * NUM_GPUS ))
+else
+    EFFECTIVE_N=0  # full test set
+fi
+
+# Map the requested GPU count to CUDA_VISIBLE_DEVICES if not already set.
+if [[ -z "${CUDA_VISIBLE_DEVICES:-}" ]]; then
+    CVD=$(seq -s, 0 $((NUM_GPUS - 1)))
+    export CUDA_VISIBLE_DEVICES="$CVD"
+fi
+
+# DDP strategy + multi-GPU. For NUM_GPUS=1 use auto strategy (no DDP overhead).
+STRATEGY_ARG=()
+if [[ "$NUM_GPUS" -gt 1 ]]; then
+    STRATEGY_ARG=(--trainer.strategy=ddp --trainer.devices="$NUM_GPUS")
+else
+    STRATEGY_ARG=(--trainer.devices=1)
+fi
+
+echo ">>> UCPE predict on PanShot:  run=$WANDB_RUN_ID  ckpt=$CKPT_PATH  videos=$VIDEO_SUBDIR"
+echo ">>>   gpus=$NUM_GPUS (CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES)  per_rank=${PER_RANK:-all}  effective_n=$EFFECTIVE_N"
+
+# DDP teardown can SIGPROF (exit 155) on this cluster after all videos are
+# already saved. Tolerate that — but re-raise any other failure.
+set +e
+python src/main.py predict \
+    --data=PanShotDataModule \
+    --data.video_subdir="$VIDEO_SUBDIR" \
+    --ckpt_path=null \
+    "${STRATEGY_ARG[@]}" \
+    "${LIMIT_ARG[@]}"
+rc=$?
+set -e
+if [[ $rc -ne 0 && $rc -ne 155 ]]; then
+    echo "ERROR: predict exited $rc" >&2
+    exit $rc
+fi
+[[ $rc -eq 155 ]] && echo ">>> (predict exited 155 / SIGPROF after saving — tolerated)"
+
+OUT_ROOT="logs/${WANDB_RUN_ID}/predict"
+GEN_DIR="$OUT_ROOT/t2v"
+REF_DIR="$OUT_ROOT/reference"
+CAP_DIR="$OUT_ROOT/caption"
+CMP_DIR="$OUT_ROOT/comparison"
+mkdir -p "$CMP_DIR"
+
+FONT="/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
+
+echo
+echo ">>> Stitching side-by-side comparisons -> $CMP_DIR"
+
+shopt -s nullglob
+count=0
+skipped=0
+for gen in "$GEN_DIR"/*.mp4; do
+    vid="$(basename "$gen" .mp4)"
+    ref="$REF_DIR/$vid.mp4"
+    out="$CMP_DIR/$vid.mp4"
+
+    if [[ ! -f "$ref" ]]; then
+        echo "  [skip] no reference for $vid"
+        skipped=$((skipped+1))
+        continue
+    fi
+    if [[ -f "$out" ]]; then
+        echo "  [skip] exists $vid"
+        skipped=$((skipped+1))
+        continue
+    fi
+
+    # Reference is at videos_704 native (1280x704); generated is at model size
+    # (832x480). Scale both to height=480 (preserve aspect via -2) before hstack.
+    ffmpeg -y -hide_banner -loglevel error \
+        -i "$ref" -i "$gen" \
+        -filter_complex "\
+[0:v]scale=-2:480,setsar=1,drawtext=fontfile=$FONT:text='Ground Truth':x=12:y=12:fontsize=28:fontcolor=white:box=1:boxborderw=6:boxcolor=black@0.55[gt];\
+[1:v]scale=-2:480,setsar=1,drawtext=fontfile=$FONT:text='UCPE Generated':x=12:y=12:fontsize=28:fontcolor=white:box=1:boxborderw=6:boxcolor=black@0.55[ge];\
+[gt][ge]hstack=inputs=2[v]" \
+        -map "[v]" -c:v libx264 -pix_fmt yuv420p -crf 18 -preset fast \
+        "$out"
+
+    count=$((count+1))
+    echo "  [ok]   $vid"
+done
+
+echo
+echo ">>> Wrote $count new comparisons (skipped $skipped)"
+echo "     gen:  $GEN_DIR"
+echo "     ref:  $REF_DIR"
+echo "     cap:  $CAP_DIR  (caption text per video_id)"
+echo "     out:  $CMP_DIR"

UCPE/scripts/demo.sh

@@ -0,0 +1,11 @@
+source scripts/set_Wan2.1-T2V-1.3B.sh
+export WANDB_MODE=disabled
+
+export WANDB_RUN_ID=6wodf04s
+export PL_PREDICT__DATA="DemoDataModule"
+export PL_PREDICT__MODEL__CKPT_PATH="logs/6wodf04s/checkpoints/pytorch_model.bin"
+# export PL_PREDICT__MODEL__NUM_PREDICT=5  # Number of predictions to generate per input
+
+python -m src.main predict --data.input_file="demo/lens.json"
+python -m src.main predict --data.input_file="demo/pose.json"
+python -m src.main predict --data.input_file="demo/teaser.json"

UCPE/scripts/evaluate.sh

@@ -0,0 +1,42 @@
+export EVAL_DATA_ROOT="data/UCPE"
+export EVAL_NUM_FRAMES=81
+# export EVAL_TEST_STEPS='["overall"]'
+# export EVAL_LIMIT_EVAL_VIDEOS=10
+export HF_HUB_OFFLINE=1
+
+
+# python src/evaluate.py --evaluate_gt True
+# WANDB_RUN_ID=lg1mxf9u python src/evaluate.py  # recammaster_norm
+# WANDB_RUN_ID=3yf7psvi python src/evaluate.py  # plucker_norm
+# WANDB_RUN_ID=khnmur4b python src/evaluate.py  # recammaster_noyaw
+# WANDB_RUN_ID=9hjx47bc python src/evaluate.py  # plucker_noyaw
+WANDB_RUN_ID=6wodf04s python src/evaluate.py  # relray_absmap_comp8
+# WANDB_RUN_ID=nv4al3mj python src/evaluate.py  # relray_absmap_comp4
+# WANDB_RUN_ID=lkxh4srz python src/evaluate.py  # relray_absmap_comp12
+# WANDB_RUN_ID=r0hmwcag python src/evaluate.py  # relray_absmap_comp2
+# WANDB_RUN_ID=p03o7rqy python src/evaluate.py  # relray_absmap_comp8_before
+# WANDB_RUN_ID=82awngqn python src/evaluate.py  # relray_absmap_comp8_after
+# WANDB_RUN_ID=coo9rjaq python src/evaluate.py  # relray
+# WANDB_RUN_ID=wekc4yx6 python src/evaluate.py  # prope_absmap
+# WANDB_RUN_ID=z0cfx65s python src/evaluate.py  # gta_absmap
+
+export EVAL_DATA="Re10kDataset"
+export EVAL_DATA_ROOT="data/RealEstate10k"
+export EVAL_POSE_FRAMES=16
+export EVAL_LIMIT_EVAL_VIDEOS=100
+# python src/evaluate.py --frame_stride=2 --test_res_path=/mnt/pfs/users/zhangchen/panshot/ac3d/out/5B/test/10000 # ac3d
+# python src/evaluate.py --frame_stride=1 --test_res_path=/mnt/pfs/users/zhangchen/panshot/CameraCtrl/out/re10k # cameractrl
+export EVAL_FRAME_STRIDE=4
+# WANDB_RUN_ID=lg1mxf9u python src/evaluate.py  # recammaster_norm
+# WANDB_RUN_ID=3yf7psvi python src/evaluate.py  # plucker_norm
+# WANDB_RUN_ID=khnmur4b python src/evaluate.py  # recammaster_noyaw
+# WANDB_RUN_ID=9hjx47bc python src/evaluate.py  # plucker_noyaw
+WANDB_RUN_ID=6wodf04s python src/evaluate.py  # relray_absmap_comp8
+# WANDB_RUN_ID=nv4al3mj python src/evaluate.py  # relray_absmap_comp4
+# WANDB_RUN_ID=lkxh4srz python src/evaluate.py  # relray_absmap_comp12
+# WANDB_RUN_ID=r0hmwcag python src/evaluate.py  # relray_absmap_comp2
+# WANDB_RUN_ID=p03o7rqy python src/evaluate.py  # relray_absmap_comp8_before
+# WANDB_RUN_ID=82awngqn python src/evaluate.py  # relray_absmap_comp8_after
+# WANDB_RUN_ID=coo9rjaq python src/evaluate.py  # relray
+# WANDB_RUN_ID=wekc4yx6 python src/evaluate.py  # prope_absmap
+# WANDB_RUN_ID=z0cfx65s python src/evaluate.py  # gta_absmap

UCPE/scripts/inference.py

@@ -0,0 +1,45 @@
+import torch
+from PIL import Image
+from diffsynth import save_video, VideoData, load_state_dict
+from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig
+
+
+prompt = "The camera smoothly arcs counterclockwise, offering a serene view of a secluded island surrounded by crystal-clear waters. The lush greenery and rugged cliffs create a picturesque contrast against the deep blue sea, while distant boats dot the horizon under a clear sky."
+
+# pipe = WanVideoPipeline.from_pretrained(
+#     torch_dtype=torch.bfloat16,
+#     device="cuda",
+#     model_configs=[
+#         ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="diffusion_pytorch_model*.safetensors", offload_device="cpu"),
+#         ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", offload_device="cpu"),
+#         ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="Wan2.1_VAE.pth", offload_device="cpu"),
+#     ],
+# )
+# pipe.enable_vram_management()
+
+# video = pipe(
+#     prompt=prompt,
+#     negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
+#     seed=1, tiled=True
+# )
+# save_video(video, "outputs/video_Wan2.1-T2V-1.3B.mp4", fps=16, quality=5)
+
+pipe = WanVideoPipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="diffusion_pytorch_model*.safetensors", offload_device="cpu"),
+        ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", offload_device="cpu"),
+        ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="Wan2.1_VAE.pth", offload_device="cpu"),
+    ],
+)
+state_dict = load_state_dict("models/train/Wan2.1-T2V-1.3B_full-1gpu/epoch-2.safetensors")
+pipe.dit.load_state_dict(state_dict)
+pipe.enable_vram_management()
+
+video = pipe(
+    prompt=prompt,
+    negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
+    seed=1, tiled=True
+)
+save_video(video, "outputs/video_Wan2.1-T2V-1.3B_full.mp4", fps=16, quality=5)

UCPE/scripts/predict.sh

@@ -0,0 +1,22 @@
+source scripts/set_Wan2.1-T2V-1.3B.sh
+export WANDB_MODE=disabled
+# export PL_PREDICT__TRAINER__LIMIT_PREDICT_BATCHES=20
+
+# WANDB_RUN_ID=46ooy8no python src/main.py predict \
+#     --model.camera_condition="ucpe"
+
+# WANDB_RUN_ID=7ur7pldr python src/main.py predict \
+#     --model.camera_condition="prope"
+
+# WANDB_RUN_ID=txloxo2j python src/main.py predict \
+#     --model.camera_condition="plucker"
+
+# WANDB_RUN_ID=gdlseut5 python src/main.py predict \
+#     --model.camera_condition="gta"
+
+# WANDB_RUN_ID=b3rv5pk8 python src/main.py predict \
+#     --model.camera_condition="recammaster"
+
+# WANDB_RUN_ID=shgomfd7 python src/main.py predict \
+#     --model.camera_condition="ucpe" \
+#     --model.attn_compress=4

UCPE/scripts/predict_one_sample.py

@@ -0,0 +1,192 @@
+"""Run UCPE's WanVideoPipeline on ONE PanShot sample and write the output mp4.
+
+Standalone (no Lightning trainer); reuses UCPE's `PanShotTrainModule.__init__` for
+the pipeline + camera-patch setup, then loads the DeepSpeed checkpoint manually,
+fetches a single sample by index from PanShotDataset, and calls `pipe(...)` the
+same way `PanShotTrainModule.forward` does in src/main.py.
+
+Used by ../cf_ucpe/scripts/compare_inference.py via subprocess.
+"""
+import argparse
+import os
+import sys
+from pathlib import Path
+
+import torch
+
+# Make UCPE imports available regardless of cwd at invocation time.
+HERE = Path(__file__).resolve().parent
+UCPE_ROOT = HERE.parent
+sys.path.insert(0, str(UCPE_ROOT))
+
+from diffsynth import save_video                                   # noqa: E402
+from src.main import PanShotTrainModule                                 # noqa: E402
+from src.dataset import PanShotDataset                             # noqa: E402
+
+NEGATIVE_PROMPT = (
+    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，"
+    "最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，"
+    "畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+)
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--sample_idx", type=int, default=None,
+                   help="Index into the test split (after filtering). "
+                        "Mutually exclusive with --video_id.")
+    p.add_argument("--video_id", default=None,
+                   help="Pick the exact PanShot video by id (recommended for "
+                        "cross-codebase reproducibility).")
+    p.add_argument("--ckpt_path", required=True,
+                   help="Path to UCPE checkpoint. Either a .ckpt file or a "
+                        "DeepSpeed folder (last.ckpt/checkpoint/mp_rank_00_model_states.pt).")
+    p.add_argument("--output_path", required=True, help="Where to write the mp4.")
+    p.add_argument("--data_root", default=str(UCPE_ROOT / "data" / "UCPE"),
+                   help="UCPE data root (parent of PanShot/).")
+    p.add_argument("--video_subdir", default="videos_704")
+    p.add_argument("--model_id", default=str(UCPE_ROOT / "Wan2.2-TI2V-5B"),
+                   help="Local Wan2.2-TI2V-5B model dir, or HF repo id.")
+    p.add_argument("--height", type=int, default=704)
+    p.add_argument("--width", type=int, default=1280)
+    p.add_argument("--num_frames", type=int, default=81)
+    p.add_argument("--num_inference_steps", type=int, default=50)
+    p.add_argument("--camera_condition", default="relray_absmap")
+    p.add_argument("--attn_compress", type=int, default=8)
+    p.add_argument("--adaptation_method", default="parallel",
+                   choices=["before", "after", "parallel"])
+    p.add_argument("--split", default="test")
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--fps", type=int, default=16)
+    return p.parse_args()
+
+
+def resolve_ckpt(path_str):
+    """Accept either a .pt file or a DeepSpeed folder (returns the .pt inside)."""
+    p = Path(path_str)
+    if p.is_dir():
+        cand = p / "checkpoint" / "mp_rank_00_model_states.pt"
+        if not cand.exists():
+            sys.exit(f"DeepSpeed folder {p} missing checkpoint/mp_rank_00_model_states.pt")
+        return str(cand)
+    return str(p)
+
+
+def main():
+    args = parse_args()
+    torch.set_grad_enabled(False)
+
+    # -------- 1. Build UCPE model (pipe + camera patch) --------
+    # PanShotTrainModule handles: from_pretrained, patch_dit, enable_grad. ckpt_path=None
+    # because we load it manually below to support DeepSpeed format.
+    print(f"[predict_one_sample] building model (model_id={args.model_id})")
+    model = PanShotTrainModule(
+        model_id=args.model_id,
+        ckpt_path=None,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.num_inference_steps,
+        camera_condition=args.camera_condition,
+        attn_compress=args.attn_compress,
+        adaptation_method=args.adaptation_method,
+    )
+    model = model.to("cuda")
+    model.pipe.device = torch.device("cuda")
+    model.eval()
+
+    # -------- 2. Load checkpoint (DeepSpeed or Lightning .ckpt) --------
+    ckpt_file = resolve_ckpt(args.ckpt_path)
+    print(f"[predict_one_sample] loading ckpt: {ckpt_file}")
+    sd = torch.load(ckpt_file, map_location="cpu", weights_only=False)
+    if "module" in sd:                  # DeepSpeed
+        sd = sd["module"]
+    elif "state_dict" in sd:            # Lightning
+        sd = sd["state_dict"]
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    print(f"[predict_one_sample] loaded; missing={len(missing)} unexpected={len(unexpected)}")
+    if unexpected:
+        print(f"[predict_one_sample] (first unexpected) {unexpected[0]}")
+
+    # patch_dit adds cam_self_attn modules in default fp32; pipe.from_pretrained
+    # already loaded the rest of the DiT in bf16. Mixed dtypes blow up at
+    # bf16-input × fp32-weight matmuls. Cast the whole DiT to bf16 after load
+    # so cam modules align with the rest.
+    model.pipe.dit = model.pipe.dit.to(torch.bfloat16)
+
+    # -------- 3. Build a small DataModule-equivalent args object --------
+    # PanShotDataset uses both attribute access (args.data_root) AND
+    # membership tests (`"model_id" in args`), so we need a Namespace-with-
+    # __contains__. omegaconf.DictConfig satisfies both.
+    from omegaconf import DictConfig
+    hp = DictConfig({
+        "data_root": args.data_root,
+        "video_subdir": args.video_subdir,
+        "zero_first_yaw": True,
+    })
+
+    if (args.video_id is None) == (args.sample_idx is None):
+        sys.exit("specify exactly one of --video_id or --sample_idx")
+    dataset_kwargs = dict(load_keys=["video", "pose", "input_image"])
+    if args.video_id is not None:
+        dataset_kwargs["video_ids"] = [args.video_id]
+    dataset = PanShotDataset(hp, split=args.split, **dataset_kwargs)
+    if args.video_id is not None:
+        if len(dataset) == 0 or dataset.metas[0]["video_id"] != args.video_id:
+            sys.exit(f"video_id {args.video_id!r} not found in {args.split} split")
+        idx = 0
+    else:
+        if not (0 <= args.sample_idx < len(dataset)):
+            sys.exit(f"sample_idx {args.sample_idx} out of range [0, {len(dataset)})")
+        idx = args.sample_idx
+    sample = dataset[idx]
+    video_id = sample["video_id"]
+    print(f"[predict_one_sample] picked video_id={video_id}")
+
+    # -------- 4. Build batch (single-element 'collated') --------
+    # The pipe was built with torch_dtype=bf16, so float tensors must be bf16
+    # to match the VAE / DiT / camera-encoder Linear weights. (Lightning's
+    # precision machinery handles this in normal training; we have to do it
+    # manually here.)
+    def _to_batch(v, cast_float=True):
+        import numpy as np
+        if isinstance(v, np.ndarray):
+            v = torch.from_numpy(v)
+        if isinstance(v, torch.Tensor):
+            t = v.unsqueeze(0).to("cuda")
+            if cast_float and t.is_floating_point():
+                t = t.to(dtype=torch.bfloat16)
+            return t
+        return [v]
+
+    batch = {
+        "caption": [sample["caption"]],
+        "input_image": _to_batch(sample["input_image"]),
+        "pose": _to_batch(sample["pose"]),
+        "xi": _to_batch(torch.tensor(float(sample["xi"]))),
+        "x_fov": _to_batch(torch.tensor(float(sample["x_fov"]))),
+    }
+
+    # -------- 5. Run inference (mirrors PanShotTrainModule.forward) --------
+    print(f"[predict_one_sample] running pipe ({args.num_inference_steps} steps)")
+    video = model.pipe(
+        prompt=batch["caption"][0],
+        input_image=batch.get("input_image", None),
+        camera_control_panshot={k: batch[k] for k in ["pose", "xi", "x_fov"]},
+        negative_prompt=NEGATIVE_PROMPT,
+        num_inference_steps=args.num_inference_steps,
+        tiled=False,
+        seed=args.seed,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+    )
+
+    out_path = Path(args.output_path)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    save_video(video, str(out_path), fps=args.fps, quality=8)
+    print(f"[predict_one_sample] wrote: {out_path}")
+
+
+if __name__ == "__main__":
+    main()

UCPE/scripts/set_Wan2.1-T2V-1.3B.sh

@@ -0,0 +1,12 @@
+for STAGE in FIT VALIDATE TEST PREDICT; do
+  export PL_${STAGE}__DATA="PanShotDataModule"
+  export PL_${STAGE}__CKPT_PATH="last"
+  export PL_${STAGE}__DATA__DATA_ROOT="data/UCPE"
+  export PL_${STAGE}__MODEL__FPS=16
+  export PL_${STAGE}__MODEL__HEIGHT=480
+  export PL_${STAGE}__MODEL__WIDTH=832
+  export PL_${STAGE}__MODEL__NUM_FRAMES=81
+  export PL_${STAGE}__MODEL__MODEL_ID="Wan-AI/Wan2.1-T2V-1.3B"
+done
+
+export HF_HUB_OFFLINE=1

UCPE/scripts/set_Wan2.2-TI2V-5B.sh

@@ -0,0 +1,13 @@
+for STAGE in FIT VALIDATE TEST PREDICT; do
+  export PL_${STAGE}__DATA="PanShotDataModule"
+  export PL_${STAGE}__CKPT_PATH="last"
+  export PL_${STAGE}__DATA__DATA_ROOT="data/UCPE"
+  export PL_${STAGE}__DATA__VIDEO_SUBDIR="videos_704"
+  export PL_${STAGE}__MODEL__FPS=16
+  export PL_${STAGE}__MODEL__HEIGHT=704
+  export PL_${STAGE}__MODEL__WIDTH=1280
+  export PL_${STAGE}__MODEL__NUM_FRAMES=81
+  export PL_${STAGE}__MODEL__MODEL_ID="Wan2.2-TI2V-5B"
+done
+
+export HF_HUB_OFFLINE=1

UCPE/scripts/train.sh

@@ -0,0 +1,66 @@
+source scripts/set_Wan2.1-T2V-1.3B.sh
+
+# export WANDB_NAME="plucker_noyaw"
+# python src/main.py fit --model.camera_condition="plucker" --model.learning_rate=1e-5
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="plucker"
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="plucker" --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="recammaster_noyaw"
+# python src/main.py fit --model.camera_condition="recammaster"
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="recammaster"
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="recammaster" --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="plucker_norm"
+# python src/main.py fit --model.camera_condition="plucker" --model.learning_rate=1e-5 --data.zero_first_yaw=False
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="plucker" --data.zero_first_yaw=False
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="plucker" --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="recammaster_norm"
+# python src/main.py fit --model.camera_condition="recammaster" --data.zero_first_yaw=False
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="recammaster" --data.zero_first_yaw=False
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="recammaster" --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+export WANDB_NAME="relray_absmap_comp8"
+python src/main.py fit --model.camera_condition="relray_absmap" --model.attn_compress=8
+WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.attn_compress=8
+WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.attn_compress=8 --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="relray_absmap_comp4"
+# python src/main.py fit --model.camera_condition="relray_absmap" --model.attn_compress=4
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.attn_compress=4
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.attn_compress=4 --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="relray_absmap_comp12"
+# python src/main.py fit --model.camera_condition="relray_absmap" --model.attn_compress=12
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.attn_compress=12
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.attn_compress=12 --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="relray_absmap_comp2"
+# python src/main.py fit --model.camera_condition="relray_absmap" --model.attn_compress=2
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.attn_compress=2
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.attn_compress=2 --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="relray_absmap_comp8_before"
+# python src/main.py fit --model.camera_condition="relray_absmap" --model.adaptation_method="before" --model.attn_compress=8
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.adaptation_method="before" --model.attn_compress=8
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.adaptation_method="before" --model.attn_compress=8 --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="relray_absmap_comp8_after"
+# python src/main.py fit --model.camera_condition="relray_absmap" --model.adaptation_method="after" --model.attn_compress=8
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.adaptation_method="after" --model.attn_compress=8
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray_absmap" --model.adaptation_method="after" --model.attn_compress=8 --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="relray"
+# python src/main.py fit --model.camera_condition="relray" --model.attn_compress=8
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray" --model.attn_compress=8
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="relray" --model.attn_compress=8 --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="prope_absmap"
+# python src/main.py fit --model.camera_condition="prope_absmap" --model.attn_compress=8
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="prope_absmap" --model.attn_compress=8
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="prope_absmap" --model.attn_compress=8 --data=Re10kDataModule --trainer.limit_predict_batches=13
+
+# export WANDB_NAME="gta_absmap"
+# python src/main.py fit --model.camera_condition="gta_absmap" --model.attn_compress=8
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="gta_absmap" --model.attn_compress=8
+# WANDB_MODE=offline python src/main.py predict --model.camera_condition="gta_absmap" --model.attn_compress=8 --data=Re10kDataModule --trainer.limit_predict_batches=13

UCPE/scripts/upload_704.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+export PATH=~/miniconda3/envs/UCPE/bin:$PATH
+export HF_TOKEN=$(cat /mnt/bn/foundation-ads3/weijie.lyu/UCPE/hf_token.txt)
+export HF_HUB_OFFLINE=0
+cd /mnt/bn/foundation-ads3/weijie.lyu/UCPE
+
+# Stage dir with hardlinks mirroring the repo layout (train/, test/).
+# upload_large_folder does not follow symlinks, so we use hardlinks instead.
+# Must live on the same filesystem as the source data.
+STAGE_DIR="upload_stage_704"
+mkdir -p "$STAGE_DIR/train" "$STAGE_DIR/test"
+
+while true; do
+    echo "$(date): Refreshing hardlinks..."
+    # cp -al: archive + hardlinks; -n: no-clobber (skip existing)
+    cp -aln data/UCPE/PanShot/videos_704-train/. "$STAGE_DIR/train/" 2>/dev/null
+    cp -aln data/UCPE/PanShot/videos_704-test/. "$STAGE_DIR/test/" 2>/dev/null
+
+    TRAIN_COUNT=$(ls "$STAGE_DIR/train/" 2>/dev/null | wc -l)
+    TEST_COUNT=$(ls "$STAGE_DIR/test/" 2>/dev/null | wc -l)
+    echo "$(date): Starting upload... (train: $TRAIN_COUNT, test: $TEST_COUNT)"
+
+    python -c "
+from huggingface_hub import HfApi
+api = HfApi()
+api.upload_large_folder(
+    folder_path='$STAGE_DIR',
+    repo_id='wlyu/ucpe_videos_704',
+    repo_type='dataset',
+)
+print('Upload complete')
+"
+
+    echo "$(date): Upload done. Sleeping 30 minutes..."
+    sleep 1800
+done

UCPE/setup.py

@@ -0,0 +1,30 @@
+import os
+from setuptools import setup, find_packages
+import pkg_resources
+
+# Path to the requirements file
+requirements_path = os.path.join(os.path.dirname(__file__), "requirements.txt")
+
+# Read the requirements from the requirements file
+if os.path.exists(requirements_path):
+    with open(requirements_path, 'r') as f:
+        install_requires = [str(r) for r in pkg_resources.parse_requirements(f)]
+else:
+    install_requires = []
+
+setup(
+    name="diffsynth",
+    version="1.1.8",
+    description="Enjoy the magic of Diffusion models!",
+    author="Artiprocher",
+    packages=find_packages(),
+    install_requires=install_requires,
+    include_package_data=True,
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: Apache Software License",
+        "Operating System :: OS Independent",
+    ],
+    package_data={"diffsynth": ["tokenizer_configs/**/**/*.*"]},
+    python_requires='>=3.6',
+)

UCPE/src/cache.py

@@ -0,0 +1,108 @@
+import lightning as pl
+from lightning.pytorch.cli import LightningCLI
+from torch.utils.data import DataLoader, Dataset
+from pathlib import Path
+import torch
+import lightning as pl
+from diffsynth.pipelines.wan_video_panshot import WanVideoPipeline, ModelConfig
+from types import SimpleNamespace
+from src.dataset import PanShotDataset
+
+
+
+class PanShotDataModule(pl.LightningDataModule):
+    def __init__(
+        self,
+        data_root: Path = Path("data/UCPE"),
+        batch_size: int = 1,
+        num_workers: int = 4,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+    def setup(self, stage):
+        self.hparams.model_id = self.trainer.model.hparams.model_id
+        load_keys = ["video"]
+        self.dataset = PanShotDataset(self.hparams, split="train", load_keys=load_keys, skip_cached=True)
+
+    def test_dataloader(self):
+        return DataLoader(self.dataset, batch_size=self.hparams.batch_size, shuffle=False, num_workers=self.hparams.num_workers)
+
+
+class PanShotCacheModule(pl.LightningModule):
+    def __init__(
+        self,
+        model_id: str = "Wan-AI/Wan2.1-T2V-1.3B",
+    ):
+        super().__init__()
+        model_configs=[
+            ModelConfig(model_id=model_id, origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth"),
+            ModelConfig(
+                model_id=model_id,
+                origin_file_pattern="Wan2.1_VAE.pth"
+            ),
+        ]
+        self.pipe = WanVideoPipeline.from_pretrained(
+            torch_dtype=torch.bfloat16,
+            device="cpu",
+            model_configs=model_configs,
+        )
+        self.pipe.dit = SimpleNamespace(
+            require_vae_embedding=True,
+            require_clip_embedding=True,
+            fuse_vae_embedding_in_latents=False,
+        )
+        self.pipe.scheduler.set_timesteps(1000, training=True)
+        self.save_hyperparameters()
+
+    def test_step(self, batch, batch_idx):
+        text, video, video_id = batch["caption"][0], batch["video"], batch["video_id"][0]
+        self.pipe.device = self.device
+        pth_path = self.trainer.datamodule.dataset.cache_folder / f"{video_id}.pth"
+        if pth_path.exists():
+            return
+        pth_path.parent.mkdir(parents=True, exist_ok=True)
+        _, _, num_frames, height, width = video.shape
+        inputs_posi = {"prompt": text}
+        inputs_nega = {}
+        inputs_shared = {
+            "input_video": video,
+            "input_image": batch.get("input_image", None),
+            "height": height,
+            "width": width,
+            "num_frames": num_frames,
+            "cfg_scale": 1,
+            "tiled": False,
+            "rand_device": None,
+            "use_gradient_checkpointing": False,
+            "use_gradient_checkpointing_offload": False,
+            "cfg_merge": False,
+        }
+        for unit in self.pipe.units:
+            inputs_shared, inputs_posi, inputs_nega = self.pipe.unit_runner(unit, self.pipe, inputs_shared, inputs_posi, inputs_nega)
+        inputs = {**inputs_shared, **inputs_posi}
+        data = {k: inputs[k][0] for k in ["input_latents", "context", "first_frame_latents"] if k in inputs}
+        torch.save(data, pth_path)
+
+
+def main():
+    cli = LightningCLI(
+        model_class=PanShotCacheModule,
+        datamodule_class=PanShotDataModule,
+        seed_everything_default=42,
+        run=False,
+        trainer_defaults={
+            "precision": "bf16-true",
+            "logger": False,
+        },
+        save_config_callback=None,
+    )
+    trainer = cli.trainer
+    model = cli.model
+    datamodule = cli.datamodule
+
+    trainer.test(model, datamodule=datamodule)
+
+
+if __name__ == "__main__":
+    main()

UCPE/src/camera_control.py

@@ -0,0 +1,678 @@
+import torch
+from torch import nn
+from thirdparty.prope.torch import PropeDotProductAttention
+from diffsynth.models.wan_video_dit import flash_attention
+from einops import rearrange, repeat, einsum
+import torch.nn.functional as F
+from typing import Tuple
+import numpy as np
+import os
+
+
+def patch_dit(pipe, method, height, width, attn_compress=1, adaptation_method="parallel"):
+    keywords = []
+    if method.startswith("recam"):
+        if method == "recammaster":
+            emb_dim = 14
+        elif method == "recam_plucker":
+            emb_dim = 6
+        else:
+            raise ValueError(f"Unknown method: {method}")
+
+        dim = pipe.dit.blocks[0].self_attn.q.weight.shape[0]
+        for block in pipe.dit.blocks:
+            block.cam_encoder = nn.Linear(emb_dim, dim)
+            block.projector = nn.Linear(dim, dim)
+            block.cam_encoder.weight.data.zero_()
+            block.cam_encoder.bias.data.zero_()
+            block.projector.weight = nn.Parameter(torch.eye(dim))
+            block.projector.bias = nn.Parameter(torch.zeros(dim))
+        keywords.extend(["cam_encoder", "projector", "self_attn"])
+
+    if method == "plucker":
+        from diffsynth.models.wan_video_camera_controller import SimpleAdapter
+        pipe.dit.control_adapter = SimpleAdapter(
+            24,
+            pipe.dit.dim,
+            kernel_size=[2, 2],
+            stride=[2, 2],
+            downscale_factor=pipe.vae.upsampling_factor,
+        )
+        pipe.dit.control_adapter.conv.weight.data.zero_()
+        pipe.dit.control_adapter.conv.bias.data.zero_()
+        for block in pipe.dit.control_adapter.residual_blocks:
+            block.conv2.weight.data.zero_()
+            block.conv2.bias.data.zero_()
+        keywords = "*"
+    elif any(k in method for k in ("gta", "prope", "relray")):
+        patch_factor = pipe.vae.upsampling_factor * 2
+        patches_x = width // patch_factor
+        patches_y = height // patch_factor
+
+        if "abs" in method:
+            if "absc2w" in method or "absray" in method:
+                emb_dim = 12
+            elif "absmap" in method:
+                emb_dim = 3
+            else:
+                raise ValueError(f"Unknown method: {method}")
+        else:
+            emb_dim = None
+
+        for block in pipe.dit.blocks:
+            block.cam_self_attn = UcpeSelfAttention(
+                pipe.dit.dim,
+                pipe.dit.dim // attn_compress,
+                block.num_heads // attn_compress,
+                patches_x=patches_x,
+                patches_y=patches_y,
+                image_width=width,
+                image_height=height,
+                emb_dim=emb_dim,
+                adaptation_method=adaptation_method,
+            )
+        keywords.append("cam_self_attn")
+
+    pipe.dit.camera_condition = method
+    return keywords
+
+
+def enable_grad(pipe, keywords):
+    pipe.eval()
+    pipe.requires_grad_(False)
+    if keywords == "*":
+        pipe.dit.train()
+        pipe.dit.requires_grad_(True)
+    else:
+        for name, module in pipe.dit.named_modules():
+            if any(keyword in name for keyword in keywords):
+                print(f"Trainable: {name}")
+                module.train()
+                module.requires_grad_(True)
+
+    trainable_params = 0
+    seen_params = set()
+    for name, module in pipe.dit.named_modules():
+        for param in module.parameters():
+            if param.requires_grad and param not in seen_params:
+                trainable_params += param.numel()
+                seen_params.add(param)
+    print(f"Total number of trainable parameters: {trainable_params}")
+
+
+def compute_fx_from_fov_xi(
+    x_fov: torch.Tensor | float,
+    xi: torch.Tensor | float,
+    width: int,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """
+    根据水平视场角 (x_fov) 和 UCM 参数 (xi) 计算相机焦距 fx。
+
+    Args:
+        x_fov: float 或 [B] Tensor，水平视场角（单位：度）
+        xi: float 或 [B] Tensor，UCM 镜面参数
+        width: 图像宽度（像素）
+        device: torch.device
+        dtype: torch.dtype
+
+    Returns:
+        fx: [B] Tensor，焦距（像素单位）
+    """
+    # --- 转为 Tensor ---
+    def to_tensor_1d(x):
+        if torch.is_tensor(x):
+            return x.to(device=device, dtype=dtype)
+        return torch.tensor([x], dtype=dtype, device=device)
+
+    x_fov = to_tensor_1d(x_fov)
+    xi = to_tensor_1d(xi)
+
+    # --- 自动广播 ---
+    B = max(x_fov.shape[0], xi.shape[0])
+    x_fov = x_fov.view(-1).expand(B)
+    xi = xi.view(-1).expand(B)
+
+    # --- 计算 fx ---
+    theta = torch.deg2rad(0.5 * x_fov)
+    eps = torch.finfo(dtype).eps
+    denom = torch.sin(theta).clamp_min(eps)
+    fx = (width * 0.5) * (torch.cos(theta) + xi) / denom
+    return fx
+
+
+def compute_fov_from_fx_xi(
+    fx: torch.Tensor | float,
+    xi: torch.Tensor | float,
+    width: int,
+    device="cpu",
+    dtype=torch.float32,
+):
+    """
+    根据 UCM 模型参数 fx, xi 计算水平 FOV（度）
+
+    Args:
+        fx: float 或 [B] Tensor, 焦距
+        xi: float 或 [B] Tensor, UCM xi 参数
+        width: 图像宽度
+    Returns:
+        x_fov: [B], 单位 degree
+    """
+    def to_tensor_1d(x):
+        if torch.is_tensor(x):
+            return x.to(device=device, dtype=dtype)
+        return torch.tensor([x], dtype=dtype, device=device)
+
+    fx  = to_tensor_1d(fx).view(-1)
+    xi  = to_tensor_1d(xi).view(-1)
+    B = max(fx.shape[0], xi.shape[0])
+    fx  = fx.expand(B)
+    xi  = xi.expand(B)
+
+    # A = 2 fx / W
+    A = 2.0 * fx / width
+
+    # phi = arctan(1/A)
+    phi = torch.atan(1.0 / A)
+
+    # sin(theta - phi) = xi / sqrt(A^2 + 1)
+    denom = torch.sqrt(A * A + 1.0)
+    ratio = (xi / denom).clamp(-1.0, 1.0)
+    theta = torch.asin(ratio) + phi
+
+    # x_fov = 2 * theta (rad → deg)
+    x_fov = torch.rad2deg(2.0 * theta)
+    return x_fov
+
+
+def ucm_unproject_grid_fov(
+    x_fov: float | torch.Tensor,
+    xi: float | torch.Tensor,
+    height: int,
+    width: int,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """
+    计算每个样本的相机方向向量 (UCM model, 用视场角定义)。
+    支持 float 或 [B] Tensor 的混合输入。
+    - 若全为 float → 返回 [H, W, 3]
+    - 若任意为 [B] → 返回 [B, H, W, 3]
+    """
+    # --- 判断是否 batched ---
+    is_batched = any(torch.is_tensor(p) and p.ndim == 1 for p in [x_fov, xi])
+
+    # --- 计算 fx, fy ---
+    fx = compute_fx_from_fov_xi(x_fov, xi, width, device, dtype)
+    fy = fx
+
+    # --- 调用 ucm_unproject_grid ---
+    from equilib.equi2pers.torch import ucm_unproject_grid
+    d_cam = ucm_unproject_grid(
+        height=height,
+        width=width,
+        fx=fx,
+        fy=fy,
+        cx=width / 2,
+        cy=height / 2,
+        xi=xi if torch.is_tensor(xi) else torch.tensor([xi], dtype=dtype, device=device),
+        dtype=dtype,
+        device=device,
+        y_down=True,
+    )
+
+    # --- 输出 shape 控制 ---
+    if not is_batched:
+        d_cam = d_cam[0]  # [H, W, 3]
+
+    return d_cam
+
+
+def project_ucm_points_fov(X, Y, Z, x_fov, xi, height, width):
+    """
+    Project 3D points in camera frame to UCM image plane using fov-based intrinsics.
+
+    Args:
+        X, Y, Z: torch.Tensor [..., 3D coordinates in camera frame]
+        x_fov: float or [B] —— horizontal field of view in degrees
+        xi: float or [B] —— UCM mirror parameter
+        height, width: target image dimensions
+
+    Returns:
+        du, dv: projected pixel coordinates [..., 2]
+    """
+    fx = compute_fx_from_fov_xi(x_fov, xi, width, X.device, X.dtype)
+    fy = fx
+    cx = width / 2
+    cy = height / 2
+
+    return project_ucm_points(X, Y, Z, fx, fy, cx, cy, xi)
+
+
+def project_ucm_points(X, Y, Z, fx, fy, cx, cy, xi):
+    """
+    Project 3D points in camera frame to UCM image plane.
+
+    Args:
+        X, Y, Z: torch.Tensor [..., 3D coordinates in camera frame]
+        fx, fy, cx, cy: intrinsics (scalars or tensors)
+        xi: UCM mirror parameter
+
+    Returns:
+        du, dv: projected pixel coordinates [..., 2]
+    """
+    r = torch.sqrt(X * X + Y * Y + Z * Z)
+    alpha = Z + xi * r
+    du = fx * (X / alpha) + cx
+    dv = fy * (Y / alpha) + cy
+    return du, dv
+
+
+def ray_condition_ucm(
+    x_fov,      # float or [B] —— same fov as used in equi2pers
+    xi,        # float or [B] —— same xi as used in equi2pers
+    pose,       # [B, V, 4, 4]
+    height, width,      # target height, width
+    device,
+):
+    """
+    ✅ UCM-based Plücker embedding, output format: [B, V, H, W, 6]
+    🔁 Internally uses your ucm_unproject_grid() for consistent ray geometry.
+    
+    Only required params:
+        fov_x  (degree)
+        xi
+        c2w    (camera-to-world pose, same as your exported pose)
+        H, W   (spatial resolution)
+        device
+    """
+
+    d_cam = ucm_unproject_grid_fov(
+        x_fov, xi, height, width, device, dtype=pose.dtype
+    )
+    d_cam = repeat(d_cam, "b ... -> b v ...", v=pose.shape[1])  # [B, V, H, W, 3]
+    mask = d_cam.isnan().any(-1)
+
+    # --- 4. Transform rays into world coordinates using c2w ---
+    R = pose[..., :3, :3]      # [B, V, 3, 3]
+    t = pose[..., :3, 3]       # [B, V, 3]
+
+    d_world = torch.einsum("bvij,bvhwj->bvhwi", R, d_cam)  # [B,V,H,W,3]
+    rays_o = t[..., None, None, :].expand_as(d_world)  # [B,V,H,W,3]
+
+    # --- 5. Plücker coordinates: m = o × d ---
+    m = torch.cross(rays_o, d_world, dim=-1)  # [B,V,H,W,3]
+
+    # --- 6. Final concat: [m, d] → [B,V,H,W,6]
+    plucker = torch.cat([m, d_world], dim=-1)
+    plucker[mask] = 0.
+    return plucker
+
+
+def d_cam_to_angles(d_cam: torch.Tensor) -> torch.Tensor:
+    """
+    将方向向量 [x, y, z] 转换为 [azimuth, elevation]。
+    坐标系：z前，x右，y下（符合 UCM 投影输出）
+    
+    输入: d_cam: [B, H, W, 3]
+    输出: angles: [B, H, W, 2] — azimuth, elevation （单位: 弧度）
+    """
+    d_unit = F.normalize(d_cam, dim=-1)  # [B, H, W, 3]
+
+    x = d_unit[..., 0]  # right
+    y = d_unit[..., 1]  # down
+    z = d_unit[..., 2]  # forward
+
+    # yaw / azimuth: angle in xz-plane
+    azimuth = torch.atan2(x, z)  # ∈ [-π, π]
+
+    # pitch / elevation: angle above xz-plane
+    elevation = -torch.asin(y)   # y 向下 → elevation = -asin(y)
+
+    return torch.stack([azimuth, elevation], dim=-1)  # [B, H, W, 2]
+
+
+def world_to_ray_mats(
+    d_cam: torch.Tensor,  # [B, H, W, 3]
+    c2w: torch.Tensor,    # [B, T, 4, 4]
+) -> torch.Tensor:
+    """
+    构造每条 ray 的世界到 ray 局部坐标系的变换矩阵 world2ray。
+    坐标系定义：
+        - z: ray direction
+        - x: cam_y × ray_dir
+        - y: z × x
+    返回:
+        raymats: [B, T, H, W, 4, 4]
+    """
+    B, H, W, _ = d_cam.shape
+    T = c2w.shape[1]
+    device = d_cam.device
+    dtype = d_cam.dtype
+
+    # --- Expand ray dirs across frames ---
+    # [B,H,W,3] -> [B,T,H,W,3]
+    d_cam = repeat(d_cam, 'b h w c -> b t h w c', t=T)
+
+    # extract camera R,t
+    R_cam = c2w[..., :3, :3]       # [B,T,3,3]
+    t_cam = c2w[..., :3, 3]        # [B,T,3]
+    
+    # --- d_world: rotate ray directions into world ---
+    d_world = einsum(R_cam, d_cam, 'b t i j, b t h w j -> b t h w i')
+
+    # camera y-axis from each view
+    cam_y = R_cam[..., :, 1]       # [B,T,3]
+    cam_y = repeat(cam_y, 'b t c -> b t h w c', h=H, w=W)
+
+    # === Construct orthonormal ray-local axes ===
+    z_ray = F.normalize(d_world, dim=-1, eps=1e-6)
+    x_ray = torch.cross(cam_y, z_ray, dim=-1)
+    x_ray = F.normalize(x_ray, dim=-1, eps=1e-6)
+    y_ray = torch.cross(z_ray, x_ray, dim=-1)
+    y_ray = F.normalize(y_ray, dim=-1, eps=1e-6)
+    
+    # local->world rotation
+    R_l2w = torch.stack([x_ray, y_ray, z_ray], dim=-1)  # [B,T,H,W,3,3]
+
+    # world->local rotation (transpose)
+    R_w2l = rearrange(R_l2w, 'b t h w i j -> b t h w j i')  # ✅
+
+    # broadcast camera center
+    t_world = repeat(t_cam, 'b t c -> b t h w c', h=H, w=W)
+
+    # world->local translation
+    t_w2l = -einsum(R_w2l, t_world, 'b t h w i j, b t h w j -> b t h w i')
+
+    # assemble transform matrix
+    raymats = torch.zeros(B, T, H, W, 4, 4, device=device, dtype=dtype)
+    raymats[..., :3, :3] = R_w2l
+    raymats[..., :3, 3] = t_w2l
+    raymats[..., 3, 3] = 1.0
+
+    # NaN handling
+    mask = torch.isnan(d_world).any(-1)
+    raymats[mask] = torch.eye(4, device=device, dtype=dtype)
+
+    return raymats
+
+
+def rope_precompute_coeffs(
+    positions: torch.Tensor,  # [B, H, W]
+    freq_base: float,
+    freq_scale: float,
+    feat_dim: int,
+    dtype: torch.dtype = torch.float32,
+) -> Tuple[torch.Tensor, torch.Tensor]:  # [B, 1, H*W, D], [B, 1, H*W, D]
+    """
+    批量计算每个样本对应的 RoPE 系数（cos, sin），用于 patch ray angle embedding。
+    输入:
+        positions: [B, H, W] —— 每个 patch 的 azimuth 或 elevation（单位弧度）
+    输出:
+        cos: [B, 1, H*W, feat_dim//2]
+        sin: [B, 1, H*W, feat_dim//2]
+    """
+    # 对 NaN 角度 patch，输出 cos=1, sin=0，即不做旋转，等价于保留原始 token 表示
+    mask = positions.isnan()
+    positions = positions.clone()
+    positions[mask] = 0.0
+
+    B, H, W = positions.shape
+    positions_flat = positions.view(B, H * W)  # [B, HW]
+    num_freqs = feat_dim // 2
+
+    freqs = freq_scale * (
+        freq_base ** (
+            -torch.arange(num_freqs, device=positions.device)[None, :]
+            / num_freqs
+        )  # [1, D]
+    )  # [1, D]
+
+    # Expand for batch & positions
+    angles = positions_flat[..., None] * freqs[None, :, :]  # [B, HW, D]
+    angles = angles.view(B, 1, H * W, num_freqs)
+
+    return torch.cos(angles).to(dtype), torch.sin(angles).to(dtype)
+
+
+def compute_up_lat_map(
+    R: torch.Tensor,
+    x_fov: torch.Tensor,
+    xi: torch.Tensor,
+    height: int,
+    width: int,
+    device: torch.device = torch.device("cpu"),
+    delta: float = 0.1,
+):
+    """
+    计算 up_map 和 lat_map。
+    
+    Args:
+        R: [B, T, 3, 3] 相机 c2w 旋转矩阵
+        x_fov: [B] 或 [B,T] 水平视场角（度）
+        xi:   [B] 或 [B,T] UCM 参数
+        height: int，图像/patch 高度
+        width:  int，图像/patch 宽度
+        device: torch.device
+        delta: float，小旋转角度（弧度）
+    Returns:
+        up_map: [B, T, H, W, 2] 单位向量 map
+        lat_map: [B, T, H, W, 1] 纬度 map
+    """
+    B, T, _, _ = R.shape
+    dtype = R.dtype
+    R = R.float()
+
+    # Step1：生成每像素射线方向（相机坐标系）
+    d_cam = ucm_unproject_grid_fov(
+        x_fov=x_fov,
+        xi=xi,
+        height=height,
+        width=width,
+        device=device,
+        dtype=torch.float32,
+    )  # [B, H, W, 3]
+    if d_cam.ndim == 3:
+        d_cam = d_cam.unsqueeze(0)  # [B, H, W, 3]
+    mask = d_cam.isnan().any(dim=-1, keepdim=True)  # [B, H, W, 1]
+
+    # Step2：从相机系旋转到世界���
+    d_cam_exp = repeat(d_cam, "B H W C -> B T H W C", T=T)  # [B, T, H, W, 3]
+    d_world = torch.einsum('btij,bthwj->bthwi', R, d_cam_exp)
+    d_world = d_world / torch.clamp_min(d_world.norm(dim=-1, keepdim=True), 1e-8)
+
+    # Step3：计算纬度 map
+    Xw, Yw, Zw = d_world[..., 0], d_world[..., 1], d_world[..., 2]
+    lat_map = torch.atan2(-Yw, torch.sqrt(Xw**2 + Zw**2)).unsqueeze(-1)  # [B, T, H, W, 1]
+
+    # Step4：计算 up_map
+    v = d_world  # 已归一化
+    up_world = torch.tensor([0, -1, 0], device=device, dtype=torch.float32)  # 世界上方方向（+Y 向下设定）
+    k = torch.cross(v, up_world.unsqueeze(0).unsqueeze(0).unsqueeze(0).expand_as(v), dim=-1)
+    k = k / torch.clamp_min(k.norm(dim=-1, keepdim=True), 1e-8)
+
+    delta = torch.tensor(delta, device=device, dtype=torch.float32)
+    cos_eps = torch.cos(delta)
+    sin_eps = torch.sin(delta)
+    # Rodrigues 公式旋转 v → v_rot
+    v_rot = v * cos_eps + torch.cross(k, v, dim=-1) * sin_eps + k * (k * (v * 1).sum(dim=-1, keepdim=True)) * (1 - cos_eps)
+
+    dirs_cam = torch.einsum('btij,bthwj->bthwi', R.transpose(-1, -2), v_rot)
+    Xs, Ys, Zs = dirs_cam[..., 0], dirs_cam[..., 1], dirs_cam[..., 2]
+
+    du, dv = project_ucm_points_fov(
+        Xs, Ys, Zs,
+        x_fov=x_fov.float(),
+        xi=xi.float(),
+        height=height,
+        width=width,
+    )
+    from equilib.torch_utils import create_grid
+    grid = create_grid(
+        height=height,
+        width=width,
+        batch=B,
+        dtype=torch.float32,
+        device=device,
+    )  # [B, H, W, 3]
+    grid_x = grid[..., 0].unsqueeze(1)  # [B,1,H,W]
+    grid_y = grid[..., 1].unsqueeze(1)
+
+    up_map = torch.stack((du - grid_x, dv - grid_y), dim=-1)  # [B, T, H, W, 2]
+    up_map = up_map / torch.clamp_min(up_map.norm(dim=-1, keepdim=True), 1e-8)
+
+    up_map = up_map.to(dtype=dtype)
+    lat_map = lat_map.to(dtype=dtype)
+
+    # 扩 mask 到同 shape
+    mask_exp2 = mask.unsqueeze(1).expand(B, T, height, width, 1)
+    up_map = up_map.masked_fill(mask_exp2, 0.0)
+    lat_map = lat_map.masked_fill(mask_exp2, 0.0)
+
+    return up_map, lat_map
+
+
+def visualize_up_lat_map(up_map: torch.Tensor, lat_map: torch.Tensor, save_path: str = None):
+    """
+    可视化 world-anchored 的 up_map 与 lat_map（GeoCalib-style overlay）。
+    仅依赖指定输入，其余设置在函数内固定。
+
+    Args:
+        up_map: [H, W, 2] 张量
+        lat_map: [H, W, 1] 张量
+        save_path: 保存文件路径
+    """
+    import matplotlib.pyplot as plt
+    from geocalib import viz2d
+
+    # --- 数据预处理 ---
+    up_vis = up_map.detach().float().cpu()       # [H, W, 2]
+    lat_vis = lat_map[..., 0].detach().float().cpu()  # [H, W]
+
+    # --- 绘图 ---
+    fig, ax = plt.subplots(figsize=(6, 4), dpi=200)
+    viz2d.plot_latitudes([lat_vis], axes=[ax])  # 绘制纬度热力图
+    viz2d.plot_vector_fields([up_vis.permute(2, 0, 1)], subsample=10, axes=[ax])  # 绘制up向量场
+
+    ax.set_axis_off()
+
+    # --- 保存结果 ---
+    if save_path is not None:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        fig.canvas.draw()
+        fig.savefig(save_path, dpi=200, bbox_inches="tight")
+        plt.close(fig)
+    else:
+        return fig
+
+
+class UcpeSelfAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        attn_dim: int,
+        num_heads: int,
+        patches_x: int = 8,
+        patches_y: int = 8,
+        image_width: int = 128,
+        image_height: int = 128,
+        freq_base: float = 100.0,
+        freq_scale: float = 1.0,
+        precompute_coeffs: bool = True,
+        emb_dim: int | None = None,
+        adaptation_method: str = "parallel",
+    ):
+        super().__init__()
+        assert dim % num_heads == 0
+        self.dim = dim
+        self.attn_dim = attn_dim
+        self.num_heads = num_heads
+        self.head_dim = attn_dim // num_heads
+        self.patches_x = patches_x
+        self.patches_y = patches_y
+        self.image_width = image_width
+        self.image_height = image_height
+        self.freq_base = freq_base
+        self.freq_scale = freq_scale
+        self.adaptation_method = adaptation_method
+
+        self.q_proj = nn.Linear(dim, attn_dim)
+        self.k_proj = nn.Linear(dim, attn_dim)
+        self.v_proj = nn.Linear(dim, attn_dim)
+        self.out_proj = nn.Linear(attn_dim, dim)
+        if emb_dim is not None:
+            self.cam_encoder = nn.Linear(emb_dim, dim)
+
+        # 初始化为零以增强 residual 训练稳定性
+        nn.init.zeros_(self.out_proj.weight)
+        nn.init.zeros_(self.out_proj.bias)
+
+        # 初始化 PRoPE attention 模块（带 precomputed coeffs）
+        self.prope_attn = PropeDotProductAttention(
+            head_dim=self.head_dim,
+            patches_x=patches_x,
+            patches_y=patches_y,
+            image_width=image_width,
+            image_height=image_height,
+            freq_base=freq_base,
+            freq_scale=freq_scale,
+            precompute_coeffs=precompute_coeffs,
+        )
+
+    def forward(self, x: torch.Tensor, control_camera_dit_input: dict):
+        """
+        Args:
+            x: (B, T, D) — input tokens
+            control_camera_dit_input: dict with keys:
+                - viewmats: (B, N, 4, 4)
+                - K: (B, N, 3, 3)
+        """
+        B, T, D = x.shape
+        N = control_camera_dit_input["viewmats"].shape[1]  # number of cameras
+        H, W = self.patches_y, self.patches_x
+        assert T == N * H * W or T == N, f"Expected token shape ({N}×{H}×{W} or {N}), got {T}"
+
+        if hasattr(self, "cam_encoder") and "cam_emb" in control_camera_dit_input:
+            cam_emb = control_camera_dit_input["cam_emb"]
+            y = self.cam_encoder(cam_emb)
+            if y.shape[1] != T:
+                hw = T // cam_emb.shape[1]
+                y = repeat(y, "b f d -> b (f hw) d", hw=hw)
+            x = x + y
+
+        # Project Q, K, V
+        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)  # [B, H, T, D_head]
+        k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Precompute camera-specific functions (only once per batch)
+        self.prope_attn._precompute_and_cache_apply_fns(
+            viewmats=control_camera_dit_input["viewmats"],
+            Ks=control_camera_dit_input.get("K", None),
+            coeffs_x=control_camera_dit_input.get("coeffs_x", None),
+            coeffs_y=control_camera_dit_input.get("coeffs_y", None),
+        )
+
+        # Apply RoPE-style positional encoding
+        q = self.prope_attn._apply_to_q(q)     # [B, H, T, D_head]
+        k = self.prope_attn._apply_to_kv(k)
+        v = self.prope_attn._apply_to_kv(v)
+
+        # Rearrange to [B, T, D] for flash_attention input
+        q = rearrange(q, "b h t d -> b t (h d)")
+        k = rearrange(k, "b h t d -> b t (h d)")
+        v = rearrange(v, "b h t d -> b t (h d)")
+
+        # Fast attention (Flash/Sage/SDPA fallback)
+        out = flash_attention(q, k, v, num_heads=self.num_heads)
+
+        # reshape back
+        out = rearrange(out, "b t (h d) -> b h t d", h=self.num_heads)
+
+        # Apply inverse transform for PRoPE
+        out = self.prope_attn._apply_to_o(out)
+
+        # Final projection
+        out = out.transpose(1, 2).reshape(B, T, -1)
+        return self.out_proj(out)

UCPE/src/dataset.py

@@ -0,0 +1,432 @@
+import torch
+from torch.utils.data import Dataset
+from einops import rearrange, repeat
+from pathlib import Path
+import jsonlines
+import json
+import numpy as np
+import math
+from tqdm.auto import tqdm
+from src import camera_control as ucpe
+
+
+class PanShotDataset(Dataset):
+    def __init__(self, args, split, load_keys=["video"], video_ids=None, skip_cached=False, result_root=None):
+        self.args = args
+        self.data_root = Path(args.data_root)
+        self.load_keys = load_keys
+        self.split = split  # "train" or "test"
+        
+        near_depth_file = self.data_root / "PanFlow" / f"near_plane_depth-{split}.jsonl"
+        near_depths = {}
+        with jsonlines.open(near_depth_file) as reader:
+            for obj in reader:
+                video_id = obj["video_id"]
+                clip_id = obj["clip_id"]
+                near_depths[f"{video_id}-{clip_id}"] = obj["near_depth"]
+        print(f"Loaded {len(near_depths)} near depth entries.")
+
+        meta_path = self.data_root / "PanShot" / f"meta-{split}"
+        meta_files = list(meta_path.glob("*.json"))
+        metas = {}
+        for meta_file in meta_files:
+            with open(meta_file, "r") as f:
+                matches = json.load(f)
+            for match in matches:
+                for video in match["videos"]:
+                    meta = {
+                        "pose_id": video["pose"],
+                        "x_fov": float(video["x_fov"]),
+                        "xi": video["xi"],
+                        "near_depth": near_depths[meta_file.stem],
+                    }
+                    # estimate y_fov by 16:9 aspect ratio
+                    fx = ucpe.compute_fx_from_fov_xi(meta["x_fov"], meta["xi"], 16)
+                    y_fov = ucpe.compute_fov_from_fx_xi(fx, meta["xi"], 9)
+                    meta["y_fov"] = float(y_fov)
+
+                    metas[video["video"]] = meta
+                    
+        print(f"Loaded {len(metas)} video metas.")
+
+        self.metas = []
+        caption_file = self.data_root / "PanShot" / f"captioned-{split}.jsonl"
+        with jsonlines.open(caption_file) as reader:
+            for obj in reader:
+                if obj["video"] not in metas:
+                    continue
+                meta = metas[obj["video"]]
+                meta["caption"] = obj["caption"]
+                meta["video_id"] = obj["video"]
+                self.metas.append(meta)
+        print(f"Loaded {len(self.metas)} captioned videos.")
+
+        if "model_id" in args:
+            self.cache_prefix = f"cache-{args.model_id.split('/')[-1]}"
+            self.cache_folder = self.data_root / "PanShot" / f"{self.cache_prefix}-{split}"
+            cache_names = set(c.stem for c in self.cache_folder.glob("*.pth"))
+            print(f"Found {len(cache_names)} cached videos.")
+            if skip_cached:
+                self.metas = [m for m in self.metas if m["video_id"] not in cache_names]
+                print(f"Skipped cached, {len(self.metas)} videos remaining.")
+            elif "cache" in self.load_keys:
+                self.metas = [m for m in self.metas if m["video_id"] in cache_names]
+                print(f"Only use cached, {len(self.metas)} videos remaining.")
+
+        if video_ids is not None:
+            self.metas = [meta for meta in self.metas if meta["video_id"] in video_ids]
+            print(f"Filtered by video_ids, {len(self.metas)} videos remaining.")
+
+        self.result_root = None
+        if result_root is not None:
+            self.result_root = Path(result_root)
+            video_ids = set(v.stem for v in self.result_root.glob("*.mp4"))
+            self.metas = [m for m in self.metas if m["video_id"] in video_ids]
+            print(f"Filtered by result_root, {len(self.metas)} videos remaining.")
+
+    def __len__(self):
+        return max(len(self.metas), 1)
+
+    def __getitem__(self, idx):
+        data = self.metas[idx].copy()
+
+        video_id = data["video_id"]
+        video_path = self.data_root / "PanShot" / f"videos-{self.split}" / f"{video_id}.mp4"
+        data["video_path"] = str(video_path)
+        data["result_path"] = data["video_path"] if self.result_root is None else str(self.result_root / f"{video_id}.mp4")
+
+        if "image_path" in self.load_keys:
+            image_path = self.data_root / "PanShot" / f"images-{self.split}" / f"{video_id}.png"
+            data["image_path"] = str(image_path)
+            if not image_path.exists():
+                from decord import VideoReader, cpu
+                from PIL import Image
+
+                vr = VideoReader(str(video_path), ctx=cpu(0), num_threads=1)
+                first_frame = vr[0].asnumpy()
+                first_frame = Image.fromarray(first_frame)
+                image_path.parent.mkdir(parents=True, exist_ok=True)
+                first_frame.save(image_path)
+
+        for key, path in [(k, data[f"{k}_path"]) for k in ["video", "result"] if k in self.load_keys]:
+            try:
+                from decord import VideoReader, cpu
+                vr = VideoReader(str(path), ctx=cpu(0), num_threads=1)
+                video = vr.get_batch(range(len(vr))).asnumpy()
+            except Exception as e:
+                alt_idx = (idx + 32) % len(self)
+                print(f"Error reading video {path}: {e}")
+                print(f"Use video {alt_idx} instead.")
+                return PanShotDataset.__getitem__(self, alt_idx)
+            video = video.astype(np.float32)
+            video = video / 255.0 * 2 - 1  # to [-1, 1]
+            video = rearrange(video, "T H W C -> C T H W")
+            data[key] = video
+            data["fps"] = float(vr.get_avg_fps())
+
+        if "input_image" in self.load_keys:
+            data["input_image"] = video[:, 0]
+        
+        if "cache" in self.load_keys:
+            cache_path = self.cache_folder / f"{data['video_id']}.pth"
+            data |= torch.load(cache_path, map_location="cpu")
+
+        if "pose" in self.load_keys:
+            pose_file = self.data_root / "PanShot" / f"pose-{self.split}" / data["pose_id"]
+            pose_file = pose_file.with_suffix(".npy")
+            pose = np.load(pose_file)
+            pose[..., 3] /= data["near_depth"]
+
+            if getattr(self.args, "zero_first_yaw", True):
+                # Rotate all cameras around world-y
+                # to move forward-z of the first camera to y-z plane
+                forward = pose[0, :, 2]  # (3,) the z-axis of the first camera in world
+                forward_xy = np.array([forward[0], 0, forward[2]])  # project to x-z plane
+                forward_xy /= np.linalg.norm(forward_xy) + 1e-8
+
+                # compute rotation angle theta = atan2(x, z)
+                theta = np.arctan2(forward_xy[0], forward_xy[2])
+
+                # rotation matrix around world-y by -theta
+                c, s = np.cos(-theta), np.sin(-theta)
+                R_y = np.array([[c, 0, s],
+                                [0, 1, 0],
+                                [-s, 0, c]], dtype=pose.dtype)
+
+                # apply rotation to all camera extrinsics
+                pose[..., :3] = (R_y[None] @ pose[..., :3])
+                pose[..., 3] = (R_y[None] @ pose[..., 3:4]).squeeze(-1)
+            else:
+                last_row = repeat(np.array([0,0,0,1], dtype=pose.dtype), "n -> t 1 n", t=pose.shape[0])
+                c2w = np.concatenate([pose, last_row], axis=-2)  # (T, 4, 4)
+                w2c0= np.linalg.inv(c2w[0])  # (4, 4)
+                c2w = w2c0[None] @ c2w  # (T, 4, 4)
+                pose = c2w[:, :3]  # (T, 3, 4)
+
+            data["pose"] = pose
+                
+        return data
+
+
+class Re10kDataset(Dataset):
+    def __init__(self, args, split, load_keys=["pose"], video_ids=None, result_root=None):
+        self.args = args
+        self.data_root = Path(args.data_root)
+        self.load_keys = load_keys
+        self.split = split  # "train" or "test"
+        self.normalize_traj = getattr(args, "normalize_traj", None)
+
+        self.pose_path = self.data_root / "pose_files" / split
+
+        caption_file = self.data_root / "captions" / f"{split}.json"
+        with open(caption_file, "r") as f:
+            captions = json.load(f)
+        print(f"Loaded {len(captions)} captions.")
+
+        self.metas = []
+        for video_name, caption in captions.items():
+            video_id = Path(video_name).stem
+            pose_file = self.pose_path / f"{video_id}.txt"
+            if not pose_file.exists():
+                continue
+            self.metas.append({
+                "video_id": video_id,
+                "caption": caption[0],
+            })
+        print(f"Total {len(self.metas)} videos with poses.")
+
+        self.result_root = None
+        filter_file = Path(args.data_root) / "filter_files" / f"filter_{split}_{args.num_frames}.txt"
+        if not filter_file.exists():
+            self.filter_frames(filter_file)
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            torch.distributed.barrier()
+        with open(filter_file, "r") as f:
+            filtered_ids = set(ln.strip() for ln in f.readlines() if ln.strip())
+        self.metas = [m for m in self.metas if m["video_id"] in filtered_ids]
+        print(f"After filtering, {len(self.metas)} videos remaining.")
+
+        if video_ids is not None:
+            self.metas = [meta for meta in self.metas if meta["video_id"] in video_ids]
+            print(f"Filtered by video_ids, {len(self.metas)} videos remaining.")
+
+        if result_root is not None:
+            self.result_root = Path(result_root)
+            video_ids = set(v.stem for v in self.result_root.glob("*.mp4"))
+            self.metas = [m for m in self.metas if m["video_id"] in video_ids]
+            print(f"Filtered by result_root, {len(self.metas)} videos remaining.")
+
+    def __len__(self):
+        return len(self.metas)
+
+    def filter_frames(self, filter_file):
+        if torch.distributed.is_available() and torch.distributed.is_initialized() and torch.distributed.get_rank() != 0:
+            return
+        Path(filter_file).parent.mkdir(parents=True, exist_ok=True)
+        with open(filter_file, "w") as f:
+            for data in tqdm(self, desc=f"Filtering {self.split} set"):
+                video_id = data["video_id"]
+                if len(data["pose"]) >= self.args.num_frames:
+                    f.write(f"{video_id}\n")
+        print(f"Filter file saved to {filter_file}.")
+
+    def __getitem__(self, idx):
+        data = self.metas[idx].copy()
+        pose_file = self.pose_path / f"{data['video_id']}.txt"
+
+        if self.result_root is not None:
+            video_id = data["video_id"]
+            data["result_path"] = str(self.result_root / f"{video_id}.mp4")
+
+        if "result" in self.load_keys:
+            path = data["result_path"]
+            try:
+                from decord import VideoReader, cpu
+                vr = VideoReader(str(path), ctx=cpu(0), num_threads=1)
+                video = vr.get_batch(range(len(vr))).asnumpy()
+            except Exception as e:
+                alt_idx = (idx + 32) % len(self)
+                print(f"Error reading video {path}: {e}")
+                print(f"Use video {alt_idx} instead.")
+                return PanShotDataset.__getitem__(self, alt_idx)
+            video = video.astype(np.float32)
+            video = video / 255.0 * 2 - 1  # to [-1, 1]
+            video = rearrange(video, "T H W C -> C T H W")
+            data["result"] = video
+            data["fps"] = float(vr.get_avg_fps())
+
+        if "pose" in self.load_keys:
+            num_frames = self.args.num_frames
+        else:
+            num_frames = 1
+        with open(pose_file, "r") as f:
+            # lines = [ln.strip() for ln in f.readlines() if ln.strip() and not ln.startswith("http")]
+            lines = []
+            for line in f:
+                line = line.strip()
+                if line and not line.startswith("http"):
+                    lines.append(line)
+                    if len(lines) >= num_frames:
+                        break
+
+        if "pose" in self.load_keys:
+            poses = []
+            for line in lines:
+                parts = line.split()
+                pose = np.array(list(map(float, parts[7:]))).reshape(3, 4)
+                poses.append(pose)
+            poses = np.stack(poses, axis=0)  # [T, 3, 4]
+            last_row = repeat(np.array([0,0,0,1], dtype=poses.dtype), "n -> t 1 n", t=poses.shape[0])
+            w2c = np.concatenate([poses, last_row], axis=-2)  # (T, 4, 4)
+            c2w = np.linalg.inv(w2c)  # (T, 4, 4)
+            w2c0= np.linalg.inv(c2w[0])  # (4, 4)
+            c2w = w2c0[None] @ c2w  # (T, 4, 4)
+            poses = c2w[:, :3]  # (T, 3, 4)
+            if self.normalize_traj is not None:
+                poses[..., 3] *= self.normalize_traj
+            data["pose"] = poses
+
+        fx = float(lines[0].split()[1])
+        fy = float(lines[0].split()[2])
+        x_fov = float(2 * math.atan(0.5 / fx) * 180 / math.pi)
+        y_fov = float(2 * math.atan(0.5 / fy) * 180 / math.pi)
+        overwrite_xfov = getattr(self.args, "overwrite_xfov", None)
+        data["x_fov"] = x_fov if overwrite_xfov is None else overwrite_xfov
+        data["y_fov"] = y_fov
+        data["xi"] = 0.0  # pinhole camera
+
+        return data
+
+
+class DemoDataset(Dataset):
+    def __init__(self, args, split=None, load_keys=["pose"], video_ids=None, result_root=None):
+        self.args = args
+        self.panshot_data_root = Path(args.panshot_data_root)
+        self.load_keys = load_keys
+        with open(args.input_file, "r") as f:
+            self.metas = json.load(f)
+        self.normalize_traj = getattr(args, "re10k_normalize_traj", None)
+
+        near_depth_file = self.panshot_data_root / "PanFlow" / "near_plane_depth-test.jsonl"
+        near_depths = {}
+        with jsonlines.open(near_depth_file) as reader:
+            for obj in reader:
+                video_id = obj["video_id"]
+                clip_id = obj["clip_id"]
+                near_depths[f"{video_id}-{clip_id}"] = obj["near_depth"]
+        print(f"Loaded {len(near_depths)} near depth entries.")
+        for meta in self.metas:
+            pose_file = Path(meta["pose_path"])
+            if pose_file.suffix == ".npy":
+                clip_name = pose_file.stem.rsplit("-", 2)[0]
+                meta["near_depth"] = near_depths[clip_name]
+
+        for idx, data in enumerate(self.metas):
+            pose_file = Path(data["pose_path"])
+            prefix = f"{idx}-{pose_file.stem}-fov{int(data['x_fov'])}-xi{data['xi']:.2f}-"
+            data["video_id"] = prefix + data["caption"][:50].replace(" ", "_")
+
+        if video_ids is not None:
+            self.metas = [meta for meta in self.metas if meta["video_id"] in video_ids]
+            print(f"Filtered by video_ids, {len(self.metas)} videos remaining.")
+
+        self.result_root = None
+        if result_root is not None:
+            self.result_root = Path(result_root)
+            new_metas = []
+            metas = {m["video_id"]: m for m in self.metas}
+            results = list(self.result_root.glob("*.mp4"))
+            results.sort()
+            for v in results:
+                video_id = v.stem.rsplit("-", 1)[0]
+                if video_id in metas:
+                    meta = metas[video_id].copy()
+                    meta["result_path"] = str(v)
+                    new_metas.append(meta)
+            self.metas = new_metas
+            print(f"Filtered by result_root, {len(self.metas)} videos remaining.")
+
+    def __len__(self):
+        return len(self.metas)
+
+    def __getitem__(self, idx):
+        data = self.metas[idx].copy()
+
+        if "result" in self.load_keys:
+            path = data["result_path"]
+            try:
+                from decord import VideoReader, cpu
+                vr = VideoReader(str(path), ctx=cpu(0), num_threads=1)
+                video = vr.get_batch(range(len(vr))).asnumpy()
+            except Exception as e:
+                alt_idx = (idx + 32) % len(self)
+                print(f"Error reading video {path}: {e}")
+                print(f"Use video {alt_idx} instead.")
+                return PanShotDataset.__getitem__(self, alt_idx)
+            video = video.astype(np.float32)
+            video = video / 255.0 * 2 - 1  # to [-1, 1]
+            video = rearrange(video, "T H W C -> C T H W")
+            data["result"] = video
+            data["fps"] = float(vr.get_avg_fps())
+
+        if "pose" in self.load_keys:
+            pose_file = Path(data["pose_path"])
+            if pose_file.suffix == ".txt":
+                with open(pose_file, "r") as f:
+                    lines = []
+                    for line in f:
+                        line = line.strip()
+                        if line and not line.startswith("http"):
+                            lines.append(line)
+                            if len(lines) >= self.args.num_frames:
+                                break
+                poses = []
+                for line in lines:
+                    parts = line.split()
+                    pose = np.array(list(map(float, parts[7:]))).reshape(3, 4)
+                    poses.append(pose)
+                poses = np.stack(poses, axis=0)  # [T, 3, 4]
+                last_row = repeat(np.array([0,0,0,1], dtype=poses.dtype), "n -> t 1 n", t=poses.shape[0])
+                w2c = np.concatenate([poses, last_row], axis=-2)  # (T, 4, 4)
+                c2w = np.linalg.inv(w2c)  # (T, 4, 4)
+                w2c0= np.linalg.inv(c2w[0])  # (4, 4)
+                c2w = w2c0[None] @ c2w  # (T, 4, 4)
+                poses = c2w[:, :3]  # (T, 3, 4)
+                if self.normalize_traj is not None:
+                    poses[..., 3] *= self.normalize_traj
+                data["pose"] = poses
+            elif pose_file.suffix == ".npy":
+                pose = np.load(pose_file)
+                pose[..., 3] /= data["near_depth"]
+                    
+                if getattr(self.args, "zero_first_yaw", True):
+                    # Rotate all cameras around world-y
+                    # to move forward-z of the first camera to y-z plane
+                    forward = pose[0, :, 2]  # (3,) the z-axis of the first camera in world
+                    forward_xy = np.array([forward[0], 0, forward[2]])  # project to x-z plane
+                    forward_xy /= np.linalg.norm(forward_xy) + 1e-8
+
+                    # compute rotation angle theta = atan2(x, z)
+                    theta = np.arctan2(forward_xy[0], forward_xy[2])
+
+                    # rotation matrix around world-y by -theta
+                    c, s = np.cos(-theta), np.sin(-theta)
+                    R_y = np.array([[c, 0, s],
+                                    [0, 1, 0],
+                                    [-s, 0, c]], dtype=pose.dtype)
+
+                    # apply rotation to all camera extrinsics
+                    pose[..., :3] = (R_y[None] @ pose[..., :3])
+                    pose[..., 3] = (R_y[None] @ pose[..., 3:4]).squeeze(-1)
+                else:
+                    last_row = repeat(np.array([0,0,0,1], dtype=pose.dtype), "n -> t 1 n", t=pose.shape[0])
+                    c2w = np.concatenate([pose, last_row], axis=-2)  # (T, 4, 4)
+                    w2c0= np.linalg.inv(c2w[0])  # (4, 4)
+                    c2w = w2c0[None] @ c2w  # (T, 4, 4)
+                    pose = c2w[:, :3]  # (T, 3, 4)
+
+                data["pose"] = pose
+            else:
+                raise NotImplementedError(f"Unsupported pose file: {pose_file}")
+
+        return data

UCPE/src/evaluate.py

@@ -0,0 +1,870 @@
+import os
+import sys
+sys.path.append(os.getcwd())
+import numpy as np
+import torch
+from torch import Tensor
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pathlib import Path
+from datetime import datetime
+from typing import Any, List, Literal, Tuple, Optional, Dict
+from src.dataset import PanShotDataset, Re10kDataset
+from collections import defaultdict
+from torch.utils.data import DataLoader, Subset
+import json
+from tqdm.auto import tqdm
+import subprocess
+from scipy.spatial.transform import Rotation as R
+
+
+class Args(BaseSettings):
+    data: str = "PanShotDataset"
+    num_frames: int = 81  # for Re10kDataset
+    test_steps: List[str] = ["qalign", "video", "vipe", "pose", "overall"]
+    conda_envs: Dict = {"qalign": "qalign"}
+    data_root: Path = Path("data/UCPE")
+    num_workers: int = 2
+    test_device: Literal["cuda", "cpu"] = "cuda"
+    test_res_path: Optional[Path] = None
+    evaluate_gt: bool = False
+    test_name: str = datetime.now().strftime('%Y-%m-%d-%H:%M:%S')
+    limit_eval_videos: Optional[int] = None
+    save_last: bool = True
+    load_last: bool = True
+    jitter_filter_percent: float = 0.8
+
+    # qalign
+    qalign_fps: float = 4.0
+
+    # video
+    test_chunk_size: int = 8
+
+    # pose
+    valid_pose_percent: float = 0.5
+    frame_stride: Optional[int] = None
+    pose_frames: Optional[int] = None
+
+    model_config = SettingsConfigDict(
+        env_prefix="EVAL_",
+        cli_parse_args=True,
+        cli_ignore_unknown_args=False,
+    )
+
+
+def get_path(args):
+    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")}
+    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}")}
+    else:
+        run_id = os.environ.get('WANDB_RUN_ID', None)
+        assert run_id is not None, "WANDB_RUN_ID environment variable must be set."
+        paths = {}
+        split = "predict" if args.data == "PanShotDataset" else Path(args.data_root).name
+        predict_dir = Path("logs") / run_id / split
+        for task in ["t2v", "i2v"]:
+            task_path = predict_dir / task
+            if task_path.exists():
+                paths[task] = (task_path, predict_dir / f"evaluate_{task}")
+    print(f"Evaluation paths: {paths}")
+    return paths
+
+
+def collate_fn(samples):
+    data = samples[0]
+    return data
+
+
+def filter_jitter(args):
+    if args.jitter_filter_percent >= 1.0:
+        return
+
+    # Fix jittering rotation issue
+    # Filter out videos with rapid rotations
+    max_rotation_file = args.data_root / "PanShot" / "max_rotation-test.json"
+    if not max_rotation_file.exists() and args.jitter_filter_percent < 1.:
+        max_rotations = {}
+        dataset = PanShotDataset(args, "test", load_keys=["pose"])
+        for data in dataset:
+            R_all = data["pose"][:, :3, :3]  # (T, 3, 3)
+            # 计算相邻帧之间的相对旋转
+            rel_rot = np.einsum("tij,tjk->tik", np.linalg.inv(R_all[:-1]), R_all[1:])
+            # 将相对旋转矩阵转换为旋转角度（度数）
+            rel_angles = R.from_matrix(rel_rot).magnitude() * 180.0 / np.pi
+            max_rotations[data['video_id']] = float(np.max(rel_angles))
+        max_rotations = dict(sorted(max_rotations.items(), key=lambda x: x[1]))
+        with open(max_rotation_file, "w") as f:
+            json.dump(max_rotations, f, indent=4)
+    else:
+        with open(max_rotation_file, "r") as f:
+            max_rotations = json.load(f)
+
+    num_videos = int(len(max_rotations) * args.jitter_filter_percent)
+    valid_video_ids = set(list(max_rotations.keys())[:num_videos])
+    print(f"Filtered {len(max_rotations) - num_videos} videos with high jittering rotations.")
+
+    return valid_video_ids
+
+
+def prepare_dataloader(args, load_keys, result_root=None, video_ids=None):
+    dataset_class = globals().get(args.data, None)
+
+    if dataset_class is PanShotDataset:
+        valid_video_ids = filter_jitter(args)
+        if valid_video_ids is not None:
+            if video_ids is not None:
+                video_ids = set(video_ids) & valid_video_ids
+            else:
+                video_ids = valid_video_ids
+
+    dataset = dataset_class(args, "test", load_keys=load_keys, result_root=result_root, video_ids=video_ids)
+    if args.limit_eval_videos and args.limit_eval_videos < len(dataset):
+        print(f"Limiting evaluation to {args.limit_eval_videos} videos.")
+        sample_ids = np.linspace(0, len(dataset) - 1, args.limit_eval_videos).astype(int).tolist()
+        dataset = Subset(dataset, sample_ids)
+    dataloader = DataLoader(
+        dataset,
+        collate_fn=collate_fn,
+        batch_size=1,
+        num_workers=args.num_workers,
+        shuffle=False,
+    )
+    return dataloader
+
+
+def link_last(output_path):
+    last_path = output_path.parent / "last.json"
+    if last_path.exists():
+        os.remove(last_path)
+    os.symlink(output_path.name, last_path)
+    print(f"Saved last evaluation results to {last_path}")
+
+
+def save_evaluation(args, test_dir, eval_results, subfolder):
+    for key, values in eval_results.items():
+        if isinstance(values, list):
+            results = [v["video_results"] for v in values]
+            results = float(np.mean(results))
+            eval_results[key] = [results, values]
+        else:
+            eval_results[key] = [values]
+
+    output_folder = test_dir / subfolder
+    output_folder.mkdir(parents=True, exist_ok=True)
+    output_path = output_folder / f"{args.test_name}_eval_results.json"
+    with open(output_path, "w") as f:
+        json.dump(eval_results, f, indent=4)
+    print(f"Evaluation results saved to {output_path}")
+
+    if args.save_last:
+        link_last(output_path)
+
+
+@torch.inference_mode()
+def qalign(args):
+    from q_align import QAlignVideoScorer, QAlignAestheticScorer, QAlignScorer
+    from PIL import Image
+    from einops import rearrange, repeat
+
+    print("Running QAlign evaluation...")
+    tasks = get_path(args)
+
+    video_scorer = QAlignVideoScorer()
+    scorers = {
+        "image_aesthetic": QAlignAestheticScorer(
+            tokenizer=video_scorer.tokenizer,
+            model=video_scorer.model,
+            image_processor=video_scorer.image_processor
+        ),
+        "image_quality": QAlignScorer(
+            tokenizer=video_scorer.tokenizer,
+            model=video_scorer.model,
+            image_processor=video_scorer.image_processor
+        ),
+        "video_quality": video_scorer,
+    }
+
+    for task, (test_res_path, test_dir) in tasks.items():
+        print(f"Evaluating task: {task}")
+        dataloader = prepare_dataloader(args, ["result"], test_res_path)
+
+        eval_results = defaultdict(list)
+        for data in tqdm(dataloader, desc="Evaluating videos"):
+            frames = data["result"]
+            frames = rearrange(frames, "C T H W -> T H W C")
+            frame_count = len(frames) * args.qalign_fps / data["fps"]
+            frame_count = min(int(frame_count), len(frames))
+            frame_count = max(frame_count, 1)
+            frame_indices = np.linspace(0, len(frames) - 1, num=frame_count)
+            frame_indices = np.round(frame_indices).astype(int)
+            frames = frames[frame_indices]
+            frames = frames / 2. + 0.5
+            frames = (frames * 255.0).astype(np.uint8)
+            frames = [Image.fromarray(frame) for frame in frames]
+            video = [video_scorer.expand2square(frame, tuple(int(x*255) for x in video_scorer.image_processor.image_mean)) for frame in frames]
+            video_tensors = video_scorer.image_processor.preprocess(video, return_tensors="pt")["pixel_values"].half()
+
+            video_tensors = video_tensors.to(video_scorer.model.device)
+            for key, scorer in scorers.items():
+                images = video_tensors if "image" in key else [video_tensors]
+                output_logits = scorer.model(
+                    scorer.input_ids.repeat(len(images), 1),
+                    images=images
+                )["logits"][:, -1, scorer.preferential_ids_]
+
+                values = torch.softmax(output_logits, -1) @ scorer.weight_tensor
+                score = values.mean().cpu().item()
+                eval_results[key].append({
+                    "video_id": data["video_id"],
+                    "video_results": score,
+                })
+
+        save_evaluation(args, test_dir, eval_results, "qalign")
+
+
+@torch.inference_mode()
+def video(args):
+    import src.camera_control as ucpe
+    from unik3d.models import UniK3D
+    from unik3d.utils.evaluation_depth import rho
+    from torchmetrics import MeanMetric, Metric
+    from einops import rearrange, repeat
+    from torchmetrics.image import (
+        LearnedPerceptualImagePatchSimilarity,
+        PeakSignalNoiseRatio,
+        StructuralSimilarityIndexMeasure,
+    )
+    from torchmetrics.image.fid import FrechetInceptionDistance, _compute_fid
+    from torchmetrics.image.inception import InceptionScore
+    from torchmetrics.multimodal import CLIPScore
+    from thirdparty.fvd.fvd import (
+        load_i3d_pretrained,
+        get_fvd_logits,
+    )
+    from geocalib import GeoCalib
+    sys.path.append("thirdparty/GeoCalib")
+    from siclib.models.utils.metrics import (
+        gravity_error,
+        latitude_error,
+        pitch_error,
+        roll_error,
+        up_error,
+        vfov_error,
+    )
+
+    class GeoCalibError(Metric):
+        higher_is_better: bool = False
+
+        def __init__(
+            self,
+            chunk_size: int | None = None,
+            skip_frames: int = 4,
+        ):
+            super().__init__()
+            self.gc = GeoCalib(weights="distorted")
+            self.errors = torch.nn.ModuleDict({k: MeanMetric() for k in [
+                "pitch", "roll", "gravity", "vfov", "k1", "k2", "latitude", "up"
+            ]})
+            self.chunk_size = chunk_size
+            self.skip_frames = skip_frames
+
+        def update(
+            self,
+            pred: torch.Tensor,  # [B, 3, H, W]
+            gt: torch.Tensor,    # [B, 3, H, W]
+        ):
+            if self.skip_frames > 1:
+                pred = pred[::self.skip_frames]
+                gt = gt[::self.skip_frames]
+            chunk_size = len(pred) if self.chunk_size is None else self.chunk_size
+            for pred_chunk, gt_chunk in zip(
+                pred.split(chunk_size, dim=0),
+                gt.split(chunk_size, dim=0),
+            ):
+                pred_result = self.gc.calibrate(pred_chunk, camera_model="radial", shared_intrinsics=True)
+                gt_result = self.gc.calibrate(gt_chunk, camera_model="radial", shared_intrinsics=True)
+
+                pred_gravity, gt_gravity = pred_result["gravity"], gt_result["gravity"]
+                self.errors["pitch"].update(pitch_error(pred_gravity, gt_gravity))
+                self.errors["roll"].update(roll_error(pred_gravity, gt_gravity))
+                self.errors["gravity"].update(gravity_error(pred_gravity, gt_gravity))
+
+                pred_cam, gt_cam = pred_result["camera"], gt_result["camera"]
+                self.errors["vfov"].update(vfov_error(pred_cam, gt_cam))
+                self.errors["k1"].update(torch.abs(pred_cam.k1 - gt_cam.k1))
+                self.errors["k2"].update(torch.abs(pred_cam.k2 - gt_cam.k2))
+
+                self.errors["latitude"].update(latitude_error(
+                    pred_result["latitude_field"],
+                    gt_result["latitude_field"],
+                ).mean(axis=(1, 2)))
+                self.errors["up"].update(up_error(
+                    pred_result["up_field"],
+                    gt_result["up_field"],
+                ).mean(axis=(1, 2)))
+
+        def compute(self):
+            return {f"{k}_err": v.compute() for k, v in self.errors.items()}
+
+    class UcmCameraRayAngularErrorRho(Metric):
+        higher_is_better = True
+
+        def __init__(
+            self,
+            model_id: str = "lpiccinelli/unik3d-vitl",
+            chunk_size: int = 16,
+            resolution_level: int = 0,
+        ):
+            super().__init__()
+            self.model = UniK3D.from_pretrained(model_id)
+            self.rho = torch.nn.ModuleDict({k: MeanMetric() for k in [
+                "gt", "pred"
+            ]})
+            self.model.resolution_level = resolution_level
+            self.chunk_size = chunk_size
+
+        def update(
+            self,
+            pred: torch.Tensor,  # [B, 3, H, W]
+            gt: torch.Tensor,    # [B, 3, H, W]
+            x_fov: float,
+            xi: float,
+        ):
+            _, _, height, width = pred.shape
+            d_cam = ucpe.ucm_unproject_grid_fov(
+                x_fov=x_fov,
+                xi=xi,
+                height=height,
+                width=width,
+                device=pred.device,
+            )
+            for pred_chunk, gt_chunk in zip(
+                pred.split(self.chunk_size, dim=0),
+                gt.split(self.chunk_size, dim=0),
+            ):
+                pred_result = self.model.infer(pred_chunk)
+                rays = pred_result["rays"]
+                rays = rearrange(rays, "B C H W -> B H W C")
+                d_cams = repeat(d_cam, "... -> B ...", B=rays.shape[0])
+                rho_errors = rho(d_cams, rays)  # [B]
+                self.rho["gt"].update(rho_errors)
+
+                gt_result = self.model.infer(gt_chunk)
+                gt_rays = gt_result["rays"]
+                gt_rays = rearrange(gt_rays, "B C H W -> B H W C")
+                rho_errors = rho(gt_rays, rays)  # [B]
+                self.rho["pred"].update(rho_errors)
+
+        def compute(self):
+            return {f"rho_{k}": v.compute() for k, v in self.rho.items()}
+
+    class FrechetVideoDistance(Metric):
+        higher_is_better: bool = False
+        full_state_update: bool = False
+
+        def __init__(
+            self,
+            crop_center: bool = True,
+            batch_size: int = 10,
+        ):
+            super().__init__()
+            self.crop_center = crop_center
+            self.batch_size = batch_size
+            self.i3d = load_i3d_pretrained()
+
+            num_features = 400
+            mx_num_feats = (num_features, num_features)
+            self.add_state("real_features_sum", torch.zeros(num_features).double(), dist_reduce_fx="sum")
+            self.add_state("real_features_cov_sum", torch.zeros(mx_num_feats).double(), dist_reduce_fx="sum")
+            self.add_state("real_features_num_samples", torch.tensor(0).long(), dist_reduce_fx="sum")
+
+            self.add_state("fake_features_sum", torch.zeros(num_features).double(), dist_reduce_fx="sum")
+            self.add_state("fake_features_cov_sum", torch.zeros(mx_num_feats).double(), dist_reduce_fx="sum")
+            self.add_state("fake_features_num_samples", torch.tensor(0).long(), dist_reduce_fx="sum")
+
+        def update(self, videos: Tensor, real: bool) -> None:
+            features = get_fvd_logits(videos, self.i3d, self.device, bs=self.batch_size, crop_center=self.crop_center)
+            self.orig_dtype = features.dtype
+            features = features.double()
+
+            if features.dim() == 1:
+                features = features.unsqueeze(0)
+            if real:
+                self.real_features_sum += features.sum(dim=0)
+                self.real_features_cov_sum += features.t().mm(features)
+                self.real_features_num_samples += videos.shape[0]
+            else:
+                self.fake_features_sum += features.sum(dim=0)
+                self.fake_features_cov_sum += features.t().mm(features)
+                self.fake_features_num_samples += videos.shape[0]
+
+        def compute(self) -> Tensor:
+            if self.real_features_num_samples < 2 or self.fake_features_num_samples < 2:
+                raise RuntimeError("More than one sample is required for both the real and fake distributed to compute FID")
+            mean_real = (self.real_features_sum / self.real_features_num_samples).unsqueeze(0)
+            mean_fake = (self.fake_features_sum / self.fake_features_num_samples).unsqueeze(0)
+
+            cov_real_num = self.real_features_cov_sum - self.real_features_num_samples * mean_real.t().mm(mean_real)
+            cov_real = cov_real_num / (self.real_features_num_samples - 1)
+            cov_fake_num = self.fake_features_cov_sum - self.fake_features_num_samples * mean_fake.t().mm(mean_fake)
+            cov_fake = cov_fake_num / (self.fake_features_num_samples - 1)
+            return _compute_fid(mean_real.squeeze(0), cov_real, mean_fake.squeeze(0), cov_fake).to(self.orig_dtype)
+
+    print("Running video evaluation...")
+    tasks = get_path(args)
+
+    for task, (test_res_path, test_dir) in tasks.items():
+        print(f"Evaluating task: {task}")
+        dataloader = prepare_dataloader(args, ["video", "result"], test_res_path)
+
+        image_metrics = {
+            "geocalib": GeoCalibError(),
+            "rho": UcmCameraRayAngularErrorRho(),
+            "lpips": LearnedPerceptualImagePatchSimilarity(
+                net_type="vgg",
+                normalize=True,
+            ),
+            "psnr": PeakSignalNoiseRatio(
+                data_range=1.,
+                dim=(1, 2, 3)
+            ),
+            "ssim": StructuralSimilarityIndexMeasure(
+                data_range=1.,
+            ),
+            "cs_text": CLIPScore(
+                model_name_or_path="zer0int/LongCLIP-L-Diffusers",
+            ),
+            "cs_image": CLIPScore(),
+        }
+        image_metrics = {k: v.to(args.test_device) for k, v in image_metrics.items()}
+        data_metrics = {
+            "fvd_center": FrechetVideoDistance(),
+            "fvd": FrechetVideoDistance(
+                crop_center=False,
+            ),
+            "fid": FrechetInceptionDistance(
+                normalize=True,
+            ),
+            "is": InceptionScore(
+                normalize=True,
+            )
+        }
+        data_metrics = {k: v.to(args.test_device) for k, v in data_metrics.items()}
+
+        eval_results = defaultdict(list)
+        for data in tqdm(dataloader, desc="Evaluating videos"):
+            if "video" in data:
+                gt_video = torch.from_numpy(data["video"]).to(args.test_device)  # [C, T, H, W]
+                gt_video = rearrange(gt_video, "C T H W -> T C H W")  # [T, C, H, W]
+                gt_video = gt_video / 2. + 0.5  # to [0, 1]
+
+            video = torch.from_numpy(data["result"]).to(args.test_device)  # [C, T, H, W]
+            video = rearrange(video, "C T H W -> T C H W")  # [T, C, H, W]
+            video = video / 2. + 0.5  # to [0, 1]
+
+            for metric_name, metric in image_metrics.items():
+                if metric_name == "geocalib":
+                    if "video" not in data:
+                        continue
+                    metric.update(
+                        pred=video,
+                        gt=gt_video,
+                    )
+                elif metric_name == "rho":
+                    if "video" not in data:
+                        continue
+                    metric.update(
+                        pred=video,
+                        gt=gt_video,
+                        x_fov=data["x_fov"],
+                        xi=data["xi"],
+                    )
+                elif metric_name == "cs_text":
+                    for pred in video.split(args.test_chunk_size, dim=0):
+                        pred = pred * 255.0
+                        metric.update(
+                            pred.to(torch.uint8),
+                            [data["caption"]] * len(pred),
+                        )
+                elif metric_name == "cs_image":
+                    for pred, gt in zip(
+                        video[:-1].split(args.test_chunk_size, dim=0),
+                        video[1:].split(args.test_chunk_size, dim=0),
+                    ):
+                        pred, gt = pred * 255.0, gt * 255.0
+                        metric.update(
+                            pred.to(torch.uint8),
+                            gt.to(torch.uint8),
+                        )
+                elif metric_name in ["lpips", "psnr", "ssim"]:
+                    if "video" not in data:
+                        continue
+                    for pred, gt in zip(
+                        video.split(args.test_chunk_size, dim=0),
+                        gt_video.split(args.test_chunk_size, dim=0),
+                    ):
+                        metric.update(
+                            pred.contiguous(),
+                            gt
+                        )
+                else:
+                    raise NotImplementedError(f"Image metric {metric_name} not implemented.")
+
+                if metric_name in ("geocalib", "rho"):
+                    results = metric.compute()
+                    for key, value in results.items():
+                        eval_results[key].append({
+                            "video_id": data["video_id"],
+                            "video_results": value.cpu().item(),
+                        })
+                else:
+                    eval_results[metric_name].append({
+                        "video_id": data["video_id"],
+                        "video_results": metric.compute().cpu().item(),
+                    })
+                metric.reset()
+
+            for metric_name, metric in data_metrics.items():
+                if metric_name == "is":
+                    for pred in video.split(args.test_chunk_size, dim=0):
+                        metric.update(pred)
+                if "video" not in data:
+                    continue
+                if metric_name == "fid":
+                    for pred, gt in zip(
+                        video.split(args.test_chunk_size, dim=0),
+                        gt_video.split(args.test_chunk_size, dim=0),
+                    ):
+                        metric.update(pred, real=False)
+                        metric.update(gt, real=True)
+                elif metric_name in ("fvd", "fvd_center"):
+                    metric.update(video.unsqueeze(0), real=False)
+                    metric.update(gt_video.unsqueeze(0), real=True)
+
+        for metric_name, metric in data_metrics.items():
+            if not metric.update_called:
+                continue
+            if metric_name in ("fid", "fvd", "fvd_center"):
+                eval_results[metric_name] = metric.compute().item()
+            elif metric_name == "is":
+                eval_results[metric_name], eval_results[f"{metric_name}_std"] = metric.compute()
+                eval_results[metric_name] = eval_results[metric_name].cpu().item()
+                eval_results[f"{metric_name}_std"] = eval_results[f"{metric_name}_std"].cpu().item()
+
+        save_evaluation(args, test_dir, eval_results, "video_metrics")
+
+
+def overall(args):
+    tasks = get_path(args)
+    eval_res_name = "last.json" if args.load_last else f"{args.test_name}_eval_results.json"
+    for task, (_, test_dir) in tasks.items():
+        overall_res = {}
+        for key in ["qalign", "video_metrics", "pose"]:
+            eval_res_path = test_dir / key / eval_res_name
+            if not eval_res_path.exists():
+                print(f"Evaluation results for {key} not found at {eval_res_path}. Skipping.")
+                continue
+            with open(eval_res_path, "r") as f:
+                eval_res = json.load(f)
+            for metric, values in eval_res.items():
+                overall_res[f"{key}/{metric}"] = values[0]
+        overall_res_path = test_dir / "overall" / f"{args.test_name}.json"
+        overall_res_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(overall_res_path, "w") as f:
+            json.dump(overall_res, f, indent=4)
+        print(f"Overall evaluation results saved to {overall_res_path}")
+        print(json.dumps(overall_res, indent=4))
+
+        if args.save_last:
+            link_last(overall_res_path)
+
+
+def vipe(args):
+    from einops import rearrange, repeat
+    import ffmpeg
+    import torch.nn.functional as F
+    import src.camera_control as ucpe
+
+    def rectify_ucm_to_pinhole(video, x_fov, xi, max_xfov=100.0):
+        """
+        UCM video → rectified pinhole video (undistortion)
+        Args:
+            video: torch.Tensor [T, C, H, W], dtype=float32, range [-1,1]
+            x_fov: float, horizontal field of view (deg) in UCM model
+            xi:    float, UCM mirror parameter
+            max_xfov: float, limit effective horizontal FOV (deg)
+        Returns:
+            rectified: numpy array [T, H, W, 3], uint8, rectified pinhole video
+        """
+
+        T, C, H, W = video.shape
+        device = video.device
+
+        # Normalize input to [0,1]
+        video = (video + 1.0) / 2.0
+
+        # ---------- 1) UCM camera intrinsics ----------
+        theta = torch.deg2rad(torch.tensor(x_fov / 2, device=device))
+        # Limit maximal horizontal FOV (helps reduce distortion & black edges)
+        max_theta = torch.deg2rad(torch.tensor(max_xfov / 2, device=device))
+        theta_x = torch.min(theta, max_theta)
+
+        # ---------- 2) Compute vertical FOV from UCM physical rays ----------
+        d_cam = ucpe.ucm_unproject_grid_fov(
+            x_fov=x_fov,
+            xi=xi,
+            height=H,
+            width=W,
+            device=device,
+        )  # [H,W,3]
+
+        mid_x = W // 2
+        verts = d_cam[:, mid_x, :]  # sample center column rays
+
+        # vertical angle wrt Z forward axis
+        theta_y_rc = torch.atan2(
+            torch.abs(verts[:, 1]),  # vertical component (Y)
+            verts[:, 2].clamp(min=1e-8)  # forward Z
+        )
+        theta_y_eff = torch.max(theta_y_rc) * 0.98  # avoid edge overflow
+
+        # ---------- 3) Target pinhole intrinsics ----------
+        fx_p = fy_p = torch.max(
+            (W * 0.5) / torch.tan(theta_x),
+            (H * 0.5) / torch.tan(theta_y_eff)
+        )
+        cx_p = (W - 1) * 0.5
+        cy_p = (H - 1) * 0.5
+
+        # ✅ pinhole grid coordinates
+        u = torch.linspace(0, W - 1, W, device=device)
+        v = torch.linspace(0, H - 1, H, device=device)
+        uu, vv = torch.meshgrid(u, v, indexing="xy")  # [W,H]
+
+        X = (uu - cx_p) / fx_p
+        Y = (vv - cy_p) / fy_p
+        Z = torch.ones_like(X)
+
+        # ---------- 5) Map pinhole rays → UCM pixels ----------
+        du, dv = ucpe.project_ucm_points_fov(X, Y, Z, x_fov, xi, H, W)
+
+        # grid normalize to [-1,1]
+        grid_x = 2.0 * (du / (W - 1)) - 1.0
+        grid_y = 2.0 * (dv / (H - 1)) - 1.0
+
+        grid = torch.stack([grid_x, grid_y], dim=-1)  # [H,W,2]
+        grid = grid.unsqueeze(0).expand(T, -1, -1, -1)  # [T,H,W,2]
+
+        # ---------- 6) Warp ----------
+        rectified = F.grid_sample(
+            video,
+            grid,
+            mode="bilinear",
+            align_corners=False,
+        ).clamp(0, 1)
+
+        rectified = (rectified * 255.0).byte()
+        rectified = rectified.permute(0, 2, 3, 1).contiguous()  # [T,H,W,3]
+        return rectified.cpu().numpy()
+
+    print("Running Vipe pose generation...")
+    tasks = get_path(args)
+
+    for task, (test_res_path, test_dir) in tasks.items():
+        print(f"Evaluating task: {task}")
+        dataloader = prepare_dataloader(args, ["result"], test_res_path)
+        rectify_res_path = test_res_path.parent / f"{test_res_path.name}_rectify"
+        rectify_res_path.mkdir(parents=True, exist_ok=True)
+
+        for data in tqdm(dataloader, desc="Exporting videos"):
+            video = torch.from_numpy(data["result"]).to(args.test_device)  # [C, T, H, W]
+            video = rearrange(video, "C T H W -> T C H W")  # [T, C, H, W]
+            _, _, height, width = video.shape
+            x_fov = data["x_fov"]
+            xi = data["xi"]
+
+            rectify_video = rectify_ucm_to_pinhole(video, x_fov, xi)
+
+            out_file = rectify_res_path / f"{data['video_id']}.mp4"
+            process = (
+                ffmpeg
+                .input('pipe:', format='rawvideo', pix_fmt='rgb24', s=f'{width}x{height}', framerate=data["fps"])
+                .output(str(out_file), pix_fmt='yuv420p', vcodec='libx264', r=data["fps"], crf=16, preset='slow')
+                .overwrite_output()
+                .run_async(pipe_stdin=True, quiet=True)
+            )
+            process.stdin.write(rectify_video.tobytes())
+            process.stdin.close()
+            process.wait()
+        
+        vipe_path = test_dir / "vipe"
+        cmd = [
+            "conda", "run", "-n", "vipe",
+            "--no-capture-output",
+            "python", "/mnt/pfs/users/zhangchen/panshot/UCPE/thirdparty/vipe/run.py",
+            "pipeline=default",
+            "streams=raw_mp4_stream",
+            f"streams.base_path={rectify_res_path}",
+            f"pipeline.output.path={vipe_path}",
+            "pipeline.output.save_artifacts=true",
+            "pipeline.post.depth_align_model=null",
+        ]
+        print(f"[CMD] {' '.join(cmd)}")
+        subprocess.run(cmd, check=True)
+
+
+def pose(args):
+    from torchmetrics import MeanMetric, Metric
+    from einops import rearrange, repeat
+
+    print("Running pose evaluation...")
+    tasks = get_path(args)
+
+    if not args.evaluate_gt and args.valid_pose_percent < 1.0:
+        pose_eval_path = args.data_root / "evaluate" / "pose" / "last.json"
+        if not pose_eval_path.exists():
+            print(f"GT pose evaluation results not found at {pose_eval_path}. Cannot limit to valid poses.")
+            valid_video_ids = None
+        else:
+            with open(pose_eval_path, "r") as f:
+                gt_eval_res = json.load(f)
+            cammc = {v["video_id"]: v["video_results"] for v in gt_eval_res["cammc"][1]}
+            sorted_videos = sorted(cammc.items(), key=lambda x: x[1])
+            sorted_videos = sorted_videos[:int(len(sorted_videos) * args.valid_pose_percent)]
+            valid_video_ids = set(v[0] for v in sorted_videos)
+    else:
+        valid_video_ids = None
+
+
+    def normalize_t(rt):
+        # normalize translation by max-norm within the same trajectory
+        t = rt[:, :3, 3]
+        scale = np.max(np.linalg.norm(t, axis=-1)) + 1e-9
+        rt[:, :3, 3] /= scale
+        return rt
+
+    def relative_pose(rt):
+        # C2W → relative to first frame
+        rel = np.zeros_like(rt)
+        rel[0] = np.eye(4)
+        inv0 = np.linalg.inv(rt[0])
+        rel[1:] = inv0 @ rt[1:]
+        return rel
+
+    def calc_rot_err(r1, r2):
+        # r1, r2: (T, 3, 3)
+        R = np.matmul(np.transpose(r1, (0,2,1)), r2)
+        trace = np.trace(R, axis1=-2, axis2=-1)
+        angle = np.arccos(np.clip((trace - 1) / 2, -1, 1))  # (T)
+        return np.sum(angle)
+
+    def calc_trans_err(t1, t2):
+        return np.sum(np.linalg.norm(t1 - t2, axis=-1))
+
+    def calc_cammc(rt1, rt2):
+        # flatten camera motion difference
+        diff = (rt2 - rt1).reshape(rt1.shape[0], -1)
+        return np.sum(np.linalg.norm(diff, axis=-1))
+
+    for task, (test_res_path, test_dir) in tasks.items():
+        print(f"Evaluating task: {task}")
+        vipe_path = test_dir / "vipe"
+        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:
+        #     vipe_video_ids = vipe_video_ids.intersection(valid_video_ids)
+        #     print(f"Evaluating {len(vipe_video_ids)} valid videos with high-quality GT poses.")
+        dataloader = prepare_dataloader(args, ["pose"], video_ids=vipe_video_ids)
+
+        eval_results = defaultdict(list)
+        for data in tqdm(dataloader, desc="Evaluating poses"):
+            gt_c2w = data["pose"]
+            last_row = repeat(np.array([0,0,0,1], dtype=gt_c2w.dtype), "n -> t 1 n", t=gt_c2w.shape[0])
+            gt_c2w = np.concatenate([gt_c2w, last_row], axis=-2)  # (T, 4, 4)
+
+            pred_c2w = np.load(vipe_pose_path / f"{data['video_id']}.npz")["data"]  # (T, 4, 4)
+
+            if args.frame_stride is not None:
+                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]
+
+            # Relative + translation normalized
+            gt_rel = normalize_t(relative_pose(gt_c2w.copy()))
+            pred_rel = normalize_t(relative_pose(pred_c2w.copy()))
+
+            # Metrics
+            rot_err = calc_rot_err(gt_rel[:, :3, :3], pred_rel[:, :3, :3])
+            trans_err = calc_trans_err(gt_rel[:, :3, 3], pred_rel[:, :3, 3])
+            cammc = calc_cammc(gt_rel[:, :3, :4], pred_rel[:, :3, :4])
+
+            results = {
+                "rot_err": rot_err,
+                "trans_err": trans_err,
+                "cammc": cammc,
+            }
+
+            vipe_gt_path = args.data_root / "evaluate" / "vipe" / "pose"
+            if not args.evaluate_gt and vipe_gt_path.exists() \
+                    and valid_video_ids is not None and data["video_id"] in valid_video_ids:
+                gt_c2w = np.load(vipe_gt_path / f"{data['video_id']}.npz")["data"]  # (T, 4, 4)
+                gt_rel = normalize_t(relative_pose(gt_c2w.copy()))
+
+                # Metrics
+                rot_err = calc_rot_err(gt_rel[:, :3, :3], pred_rel[:, :3, :3])
+                trans_err = calc_trans_err(gt_rel[:, :3, 3], pred_rel[:, :3, 3])
+                cammc = calc_cammc(gt_rel[:, :3, :4], pred_rel[:, :3, :4])
+                results.update({
+                    "rot_err_vipe": rot_err,
+                    "trans_err_vipe": trans_err,
+                    "cammc_vipe": cammc,
+                })
+
+            for key, value in results.items():
+                eval_results[key].append({
+                    "video_id": data["video_id"],
+                    "video_results": float(value),
+                })
+
+        save_evaluation(args, test_dir, eval_results, "pose")
+
+
+def main():
+    args = Args()
+
+    for step in args.test_steps:
+        if args.conda_envs and step in args.conda_envs:
+            conda_env = args.conda_envs[step]
+            print(f"[INFO] Running step '{step}' in conda env: {conda_env}")
+
+            # 当前脚本路径
+            script_path = Path(__file__).resolve()
+            script_path = script_path.relative_to(Path.cwd())
+
+            # 构造命令：使用 conda run 调用
+            cmd = [
+                "conda", "run", "-n", conda_env,
+                "--no-capture-output",
+                "python", str(script_path),
+                f"--test_steps=[{step}]",   # 只运行该 step
+                "--conda_envs={}",        # 避免递归调用 conda
+            ]
+
+            # 把其他命令行参数透传下去
+            # 注意：Args 使用了 pydantic-settings + tyro 等 CLI 解析工具，
+            # 你可以根据需要加上传入的 CLI 参数，这里简化为当前 sys.argv
+            extra_args = sys.argv[1:]
+            cmd.extend(extra_args)
+
+            print(f"[CMD] {' '.join(cmd)}")
+            subprocess.run(cmd, check=True)
+        else:
+            globals()[step](args)
+
+
+if __name__ == "__main__":
+    main()

UCPE/src/main.py

@@ -0,0 +1,387 @@
+import lightning as pl
+from lightning.pytorch.cli import LightningCLI
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
+from jsonargparse import lazy_instance
+from torch.utils.data import DataLoader
+from pathlib import Path
+import torch
+from diffsynth.pipelines.wan_video_panshot import WanVideoPipeline, ModelConfig
+import wandb
+import os
+from src.dataset import PanShotDataset, Re10kDataset, DemoDataset
+from diffsynth import save_video
+from src.camera_control import patch_dit, enable_grad
+from typing import Literal
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
+import numpy as np
+from tqdm.auto import tqdm
+from typing import Optional
+
+
+class PanShotDataModule(pl.LightningDataModule):
+    def __init__(
+        self,
+        data_root: Path = Path("data/UCPE"),
+        batch_size: int = 1,
+        num_workers: int = 4,
+        zero_first_yaw: bool = True,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+        self.test_load_keys = ["video", "pose"]
+
+    def setup(self, stage):
+        self.hparams.model_id = self.trainer.model.hparams.model_id
+
+    def train_dataloader(self):
+        dataset = PanShotDataset(self.hparams, split="train", load_keys=["cache", "pose"])
+        return DataLoader(dataset, batch_size=self.hparams.batch_size, shuffle=True, num_workers=self.hparams.num_workers)
+
+    def val_dataloader(self):
+        dataset = PanShotDataset(self.hparams, split="test", load_keys=self.test_load_keys)
+        return DataLoader(dataset, batch_size=1, shuffle=False, num_workers=self.hparams.num_workers)
+
+    def test_dataloader(self):
+        dataset = PanShotDataset(self.hparams, split="test", load_keys=self.test_load_keys)
+        return DataLoader(dataset, batch_size=1, shuffle=False, num_workers=self.hparams.num_workers)
+
+    def predict_dataloader(self):
+        dataset = PanShotDataset(self.hparams, split="test", load_keys=self.test_load_keys)
+        return DataLoader(dataset, batch_size=1, shuffle=False, num_workers=self.hparams.num_workers)
+
+
+class Re10kDataModule(pl.LightningDataModule):
+    def __init__(
+        self,
+        data_root: Path = Path("data/RealEstate10k"),
+        batch_size: int = 1,
+        num_workers: int = 4,
+        overwrite_xfov: float = 100.0,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+    @rank_zero_only
+    def normalize_traj(self, normalization_file):
+        tgt_data = PanShotDataModule()
+        tgt_data.test_load_keys = ["pose"]
+        tgt_dataloader = tgt_data.predict_dataloader()
+        src_dataloader = self.predict_dataloader()
+        traj_length_mean = []
+        for dataloader in (tgt_dataloader, src_dataloader):
+            traj_length_sum = 0.
+            traj_num = 0
+            for data in tqdm(dataloader, desc="Calculating trajectory length"):
+                pose = data["pose"]  # (B, T, 3, 4)
+                traj = pose[..., 3]  # (B, T, 3)
+                traj_length = torch.sum(torch.linalg.norm(traj[:, 1:] - traj[:, :-1], dim=-1), dim=-1)  # (B,)
+                traj_length_sum += traj_length.sum().item()
+                traj_num += traj.shape[0]
+            traj_length_mean.append(traj_length_sum / traj_num)
+        normalize_traj = traj_length_mean[0] / traj_length_mean[1]
+        print(f"Trajectory length normalization factor: {normalize_traj}")
+        Path(normalization_file).parent.mkdir(parents=True, exist_ok=True)
+        np.savetxt(normalization_file, np.array([normalize_traj], dtype=np.float32))
+
+    def setup(self, stage):
+        self.hparams.num_frames = self.trainer.model.hparams.num_frames        
+        normalization_file = Path(self.hparams.data_root) / "traj_normalization.txt"
+        if not normalization_file.exists():
+            self.normalize_traj(normalization_file)
+        self.trainer.strategy.barrier()
+        self.hparams.normalize_traj = float(np.loadtxt(normalization_file))
+
+    def predict_dataloader(self):
+        dataset = Re10kDataset(self.hparams, split="test")
+        return DataLoader(dataset, batch_size=self.hparams.batch_size, shuffle=False, num_workers=self.hparams.num_workers)
+
+
+class DemoDataModule(pl.LightningDataModule):
+    def __init__(
+        self,
+        panshot_data_root: Path = Path("data/UCPE"),
+        re10k_data_root: Path = Path("data/RealEstate10k"),
+        input_file: Path = Path("demo/teaser.json"),
+        batch_size: int = 1,
+        num_workers: int = 1,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+    def setup(self, stage):
+        self.hparams.num_frames = self.trainer.model.hparams.num_frames        
+        normalization_file = Path(self.hparams.re10k_data_root) / "traj_normalization.txt"
+        self.hparams.re10k_normalize_traj = float(np.loadtxt(normalization_file))
+
+    def predict_dataloader(self):
+        dataset = DemoDataset(self.hparams)
+        return DataLoader(dataset, batch_size=self.hparams.batch_size, shuffle=False, num_workers=self.hparams.num_workers)
+
+
+class PanShotTrainModule(pl.LightningModule):
+    def __init__(
+        self,
+        model_id: str = "Wan-AI/Wan2.1-T2V-1.3B",
+        learning_rate: float = 1e-4,
+        use_gradient_checkpointing: bool = True,
+        use_gradient_checkpointing_offload: bool = False,
+        ckpt_path: Path = None,
+        fps: int = 16,
+        height: int = 480,
+        width: int = 832,
+        num_frames: int = 81,
+        num_inference_steps: int = 50,
+        tiled: bool = False,
+        camera_condition: str = "relray_absmap",
+        adaptation_method: Literal[
+            "before",
+            "after",
+            "parallel",
+        ] = "parallel",
+        ti2v_input_image_prob: float = 0.5,
+        attn_compress: int = 8,
+        num_predict: Optional[int] = None,
+    ):
+        super().__init__()
+        file_patterns = [
+            "models_t5_umt5-xxl-enc-bf16.pth",
+            "diffusion_pytorch_model*.safetensors",
+            "Wan2.1_VAE.pth",
+        ]
+        self.pipe = WanVideoPipeline.from_pretrained(
+            torch_dtype=torch.bfloat16,
+            device="cpu",
+            model_configs=[
+                ModelConfig(model_id=model_id, origin_file_pattern=pattern, offload_device="cpu")
+                for pattern in file_patterns
+            ]
+        )
+
+        keywords = patch_dit(
+            self.pipe, camera_condition, height, width,
+            attn_compress=attn_compress, adaptation_method=adaptation_method
+        )
+        enable_grad(self.pipe, keywords)
+
+        self.strict_loading = False
+        if ckpt_path is not None:
+            print(f"Loading weights from {ckpt_path}")
+            state_dict = torch.load(ckpt_path, map_location="cpu")
+            if "state_dict" in state_dict:
+                state_dict = state_dict["state_dict"]
+            self.load_state_dict(state_dict, strict=False)
+
+        self.save_hyperparameters()
+
+    def setup(self, stage=None):
+        self.pipe.device = self.device
+
+    def predict_step(self, batch, batch_idx, dataloader_idx=0):
+        for i in range(self.hparams.num_predict or 1):
+            video = self(batch, seed=i)
+            is_i2v = "input_image" in batch and self.pipe.dit.fuse_vae_embedding_in_latents
+            video_folder = "i2v" if is_i2v else "t2v"
+            if isinstance(self.trainer.datamodule, PanShotDataModule):
+                split = "predict"
+            elif isinstance(self.trainer.datamodule, DemoDataModule):
+                split = "demo"
+            else:
+                split = Path(self.trainer.datamodule.hparams.data_root).name
+            self.save_output(
+                video,
+                batch,
+                split=split,
+                video_folder=video_folder,
+                quality=8,
+                suffix=f"-{i}" if self.hparams.num_predict else None
+            )
+            if is_i2v:
+                del batch["input_image"]
+                self.predict_step(batch, batch_idx, dataloader_idx)
+
+    def save_output(self, video, batch, split, video_folder, step=None, quality=5, suffix=None):
+        video_id = batch["video_id"][0]
+
+        video_prefix = os.path.join(self.logger.save_dir, split, video_folder, video_id)
+        if step is not None:
+            video_prefix = f"{video_prefix}-{step}"
+        if suffix is not None:
+            video_prefix = video_prefix + suffix
+        video_path = video_prefix + ".mp4"
+        os.makedirs(os.path.dirname(video_path), exist_ok=True)
+        save_video(video, video_path, fps=self.hparams.fps, quality=quality)
+
+        reference_path = os.path.join(self.logger.save_dir, split, "reference", f"{video_id}.mp4")
+        os.makedirs(os.path.dirname(reference_path), exist_ok=True)
+        if not os.path.exists(reference_path) and "video" in batch:
+            reference_video = self.pipe.vae_output_to_video(batch["video"])
+            save_video(reference_video, reference_path, fps=self.hparams.fps, quality=quality)
+
+        caption_path = os.path.join(self.logger.save_dir, split, "caption", f"{video_id}.txt")
+        os.makedirs(os.path.dirname(caption_path), exist_ok=True)
+        if not os.path.exists(caption_path):
+            with open(caption_path, "w") as f:
+                f.write(batch["caption"][0])
+
+        print(f"Saved video to {video_path}")
+
+        return video_path, reference_path
+
+    def forward(self, batch, seed=None):
+        video = self.pipe(
+            prompt=batch["caption"][0],
+            input_image=batch.get("input_image", None),
+            camera_control_panshot={k: batch[k] for k in ["pose", "xi", "x_fov"]},
+            negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，��止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
+            num_inference_steps=self.hparams.num_inference_steps,
+            tiled=self.hparams.tiled,
+            seed=seed,
+            height=self.hparams.height,
+            width=self.hparams.width,
+            num_frames=self.hparams.num_frames,
+        )
+        return video
+
+    def on_fit_start(self):
+        if self.trainer.is_global_zero and hasattr(self.logger, "experiment"):
+            self.logger.experiment.watch(self, log_graph=False, log_freq=1000)
+
+    def training_step(self, batch, batch_idx):
+        self.pipe.scheduler.set_timesteps(1000, training=True)
+
+        # Data
+        _, _, length, height, width = batch["input_latents"].shape
+        num_frames = (length - 1) * 4 + 1
+        height = height * self.pipe.vae.upsampling_factor
+        width = width * self.pipe.vae.upsampling_factor
+        inputs_posi = {}
+        inputs_nega = {}
+        inputs_shared = {
+            "camera_control_panshot": {k: batch[k] for k in ["pose", "xi", "x_fov"]},
+            "input_latents": batch["input_latents"],
+            "context": batch["context"],
+            "height": height,
+            "width": width,
+            "num_frames": num_frames,
+            "cfg_scale": 1,
+            "rand_device": self.device,
+            "use_gradient_checkpointing": self.hparams.use_gradient_checkpointing,
+            "use_gradient_checkpointing_offload": self.hparams.use_gradient_checkpointing_offload,
+            "cfg_merge": False,
+        }
+
+        if "first_frame_latents" in batch \
+            and self.pipe.dit.fuse_vae_embedding_in_latents \
+                and torch.rand(1).item() < self.hparams.ti2v_input_image_prob:
+            inputs_shared["first_frame_latents"] = batch["first_frame_latents"]
+
+        for unit in self.pipe.units:
+            inputs_shared, inputs_posi, inputs_nega = self.pipe.unit_runner(unit, self.pipe, inputs_shared, inputs_posi, inputs_nega)
+        inputs = {**inputs_shared, **inputs_posi}
+
+        # Compute loss
+        models = {name: getattr(self.pipe, name) for name in self.pipe.in_iteration_models}
+        loss = self.pipe.training_loss(**models, **inputs)
+
+        # Record log
+        self.log("train/loss", loss, prog_bar=True)
+        return loss
+
+    @torch.no_grad()
+    def validation_step(self, batch, batch_idx):
+        video = self(batch, seed=0)
+        is_i2v = "input_image" in batch and self.pipe.dit.fuse_vae_embedding_in_latents
+        video_folder = "i2v" if is_i2v else "t2v"
+        video_path, reference_path = self.save_output(
+            video, batch, split="validation", video_folder=video_folder, step=self.global_step)
+        log_dict = self.visualize(video_path, reference_path, batch)
+        log_dict = {f"val/{k}": v for k, v in log_dict.items()}
+        self.logger.experiment.log(log_dict)
+        if is_i2v:
+            del batch["input_image"]
+            self.validation_step(batch, batch_idx)
+
+    def visualize(self, video_path, reference_path, batch):
+        log_dict = {}
+        log_dict["video"] = wandb.Video(
+            video_path,
+            caption=batch["caption"][0],
+            format="mp4",
+        )
+        log_dict["reference"] = wandb.Video(
+            reference_path,
+            caption=batch["video_id"][0],
+            format="mp4",
+        )
+        return log_dict
+
+    def configure_optimizers(self):
+        trainable_modules = filter(lambda p: p.requires_grad, self.pipe.dit.parameters())
+        optimizer = torch.optim.AdamW(trainable_modules, lr=self.hparams.learning_rate)
+        return optimizer
+
+    def on_save_checkpoint(self, checkpoint):
+        for key in list(checkpoint["state_dict"].keys()):
+            if not key.startswith("pipe.dit."):
+                del checkpoint["state_dict"][key]
+
+
+class MyCLI(LightningCLI):
+    def add_arguments_to_parser(self, parser):
+        parser.add_lightning_class_args(ModelCheckpoint, "checkpoint")
+        parser.set_defaults({
+            # "data": "PanShotDataModule",
+            "checkpoint.dirpath": os.path.join(self.trainer_defaults["default_root_dir"], "checkpoints"),
+            "checkpoint.save_last": True,
+            # "checkpoint.every_n_train_steps": 10000,
+            # "checkpoint.every_n_epochs": 1,
+        })
+
+
+def main():
+    torch.set_float32_matmul_precision('high')
+
+    wandb_id = os.environ.get("WANDB_RUN_ID", wandb.util.generate_id())
+    exp_dir = os.path.join("logs", wandb_id)
+    wandb_logger = lazy_instance(
+        WandbLogger,
+        # entity="pidan1231239",
+        project="ucpe",
+        id=wandb_id,
+        save_dir=exp_dir,
+        resume="allow",
+    )
+
+    lr_monitor = LearningRateMonitor(logging_interval="step")
+
+    cli = MyCLI(
+        model_class=PanShotTrainModule,
+        # datamodule_class=PanShotDataModule,
+        save_config_kwargs={"overwrite": True},
+        parser_kwargs={"parser_mode": "omegaconf", "default_env": True},
+        seed_everything_default=int(os.environ.get("LOCAL_RANK", 0)),
+        trainer_defaults={
+            "accelerator": "gpu",
+            "devices": "auto",
+            "strategy": "deepspeed_stage_1",
+            "log_every_n_steps": 10,
+            "num_sanity_val_steps": 1,
+            "limit_train_batches": 1000,
+            "limit_val_batches": 3,
+            # "limit_predict_batches": 10,
+            "limit_test_batches": 10,
+            "benchmark": True,
+            "max_epochs": 10,
+            # "accumulate_grad_batches": 16,
+            "precision": "bf16-true",
+            "callbacks": [lr_monitor],
+            "logger": wandb_logger,
+            "default_root_dir": exp_dir,
+        },
+        
+    )
+
+
+if __name__ == "__main__":
+    main()

UCPE/thirdparty/GeoCalib/.gitattributes

@@ -0,0 +1 @@
+*.ipynb linguist-documentation

UCPE/tools/align_panflow.py

@@ -0,0 +1,599 @@
+from pathlib import Path
+import jsonlines
+from tqdm.auto import tqdm
+from collections import defaultdict
+import matplotlib.pyplot as plt
+import ffmpeg
+import json
+import csv
+import numpy as np
+import cv2
+from einops import rearrange, repeat
+from visualize_pose import vis_to_html
+from decord import VideoReader, cpu
+from copy import deepcopy
+
+
+panflow_root = Path("data/360-1M")
+panshot_root = Path("data/UCPE")
+debug_root = Path("debug/align_panflow")
+debug_root.mkdir(parents=True, exist_ok=True)
+
+filter_root = panshot_root / "PanFlow" / "filtered"
+filter_clips_thres = 0.5  # reject if >50% clips are filtered
+
+score_jsonl = panshot_root / "PanFlow" / "scores.jsonl"
+qalign_keys = ["image_aesthetic", "image_quality", "video_quality"]
+max_clips_per_video = 10
+max_clips_per_poi = 5
+
+split = "train"  # "train" or "test"
+camerabench_root = panshot_root / "CameraBench"
+output_root = panshot_root / "PanFlow" / f"align_to_camerabench-{split}"
+output_root.mkdir(parents=True, exist_ok=True)
+summary_root = output_root.parent / f"{output_root.name}-summary"
+summary_root.mkdir(parents=True, exist_ok=True)
+static_words = ["static", "fixed"]
+rotation_score_thres = 5.0  # degrees, reject if rotation_score < threshold
+rotating_clips_thres = 0.5  # reject if >50% clips have large rotation
+
+top_k = 30 if split == "train" else 10  # how many matches to keep per 360-1M clip
+target_fps = 16
+match_step = 5  # sweep step in frames
+
+visualize_gravity = False
+visualize_pose = False
+visualize_motion = False
+visualize_rotation = False
+visualize_fps = 1.
+
+cb_geocalib_file = camerabench_root / "geocalib.jsonl"
+with jsonlines.open(cb_geocalib_file, "r") as reader:
+    cb_geocalib = {obj["video"]: obj for obj in reader}
+
+# 读取 CameraBench 的位姿数据
+cb_meta_file = camerabench_root / f"processed_{split}.jsonl"
+with jsonlines.open(cb_meta_file, "r") as reader:
+    cb_meta_all = list(reader)
+cb_R_gravity = []
+cb_poses = []
+cb_meta = []
+for obj in tqdm(cb_meta_all, desc="Loading CameraBench poses"):
+    obj["video"] = Path(obj["path"]).stem
+    video_id = obj["video"]
+    camera_caption = obj["caption"].lower()
+    if any(word in camera_caption for word in static_words):
+        tqdm.write(f"Skipping static camera {video_id} ({camera_caption})")
+        continue
+
+    pose_file = camerabench_root / "vipe" / "pose" / f"{video_id}.npz"
+    if not pose_file.exists():
+        tqdm.write(f"Pose file not found: {pose_file}, skipping.")
+        continue
+
+    cb_meta.append(obj)
+    pose = np.load(pose_file)["data"]  # (T, 4, 4)
+    cb_poses.append(pose)
+    cb_R_gravity.append(cb_geocalib[video_id]["R"])
+
+print(f"Loaded {len(cb_poses)} / {len(cb_meta_all)} CameraBench poses.")
+cb_poses = np.array(cb_poses)  # (N, T, 4, 4)
+cb_R_gravity = np.array(cb_R_gravity)        # (N, 3, 3)
+
+cb_w2c0 = np.linalg.inv(cb_poses[:, 0])        # (N, 4, 4) batched inv
+cb_poses_origin = cb_w2c0[:, None] @ cb_poses       # (N, T, 4, 4)  第一帧在原点
+
+# rotate cb_poses based on cb_R to world align with gravity direction
+cb_T_gravity = repeat(np.eye(4), 'h w -> n h w', n=cb_R_gravity.shape[0])  # (N,4,4)
+cb_T_gravity[:, :3, :3] = cb_R_gravity
+cb_poses_gravity = cb_T_gravity[:, None, :, :] @ cb_poses_origin
+
+# save cb_poses_gravity
+cb_pose_root = camerabench_root / "pose"
+cb_pose_root.mkdir(parents=True, exist_ok=True)
+for i, obj in enumerate(cb_meta):
+    cb_pose_dir = cb_pose_root / f"{obj['video']}.npy"
+    np.save(cb_pose_dir, cb_poses_gravity[i])
+
+cb_pos = cb_poses_gravity[:, :, :3, 3]               # (N, T, 3)
+target_frames = cb_pos.shape[1]
+
+if visualize_gravity:
+    for i, obj in enumerate(cb_meta[:3]):
+        video_id = obj["video"]
+        cb_pose_dir = debug_root / "gravity" / video_id
+        cb_pose_dir.mkdir(parents=True, exist_ok=True)
+        origin_pose_file = cb_pose_dir / "origin.npy"
+        np.save(origin_pose_file, cb_poses_gravity[i])
+        gravity_pose_file = cb_pose_dir / "gravity_aligned.npy"
+        np.save(gravity_pose_file, cb_poses_origin[i])
+        combined_file = cb_pose_dir / "comparison.npy"
+        combined_pose = np.concatenate([cb_poses_gravity[i], cb_poses_origin[i]], axis=0)
+        np.save(combined_file, combined_pose)
+        # vis_to_html(cb_pose_dir, [origin_pose_file, gravity_pose_file])
+        vis_to_html(cb_pose_dir, [combined_file])
+
+# 读取 PanFlow 的质量分数
+with jsonlines.open(score_jsonl, "r") as reader:
+    qalign_scores = {
+        (obj["video_id"], obj["clip_id"]): {k: obj[k] for k in qalign_keys}
+        for obj in reader
+    }
+
+
+def get_traj_align(A, B, allow_scale=True, eps=1e-12):
+    """
+    计算将轨迹 B 对齐到轨迹 A 的相似变换 (R, s)，
+    但 R 被约束为仅绕世界 y 轴的旋转（偏航）。
+    假设所有轨迹的第一帧已在原点。
+    A: (N1, T, 3)
+    B: (N2, T, 3)
+    Returns:
+        R: (N1, N2, 3, 3)  使得 A ≈ s * R * B
+        s: (N1, N2)
+    """
+    N1, T, _ = A.shape
+    N2 = B.shape[0]
+
+    # 扩维到配对形状 (N1, N2, T, 3)
+    A_rel = A[:, None, :, :]          # (N1, N2, T, 3)
+    B_rel = B[None, :, :, :]          # (N1, N2, T, 3)
+
+    # 仅取 x,z 分量：[..., 0] 为 x，[..., 2] 为 z
+    Ax = A_rel[..., 0]                # (N1, N2, T)
+    Az = A_rel[..., 2]
+    Bx = B_rel[..., 0]
+    Bz = B_rel[..., 2]
+
+    # 计算 H_xz 的四个元素（按时间平均）
+    h11 = np.einsum("nmt,nmt->nm", Ax, Bx) / T
+    h12 = np.einsum("nmt,nmt->nm", Ax, Bz) / T
+    h21 = np.einsum("nmt,nmt->nm", Az, Bx) / T
+    h22 = np.einsum("nmt,nmt->nm", Az, Bz) / T
+
+    # 最优偏航角 theta（只绕 y 轴）
+    theta = np.arctan2(h12 - h21, h11 + h22)    # (N1, N2)
+
+    c = np.cos(theta)
+    s_th = np.sin(theta)
+
+    # 组装 3x3 的绕 y 轴旋转矩阵
+    R = np.zeros((N1, N2, 3, 3), dtype=A.dtype)
+    R[..., 0, 0] = c
+    R[..., 0, 2] = s_th
+    R[..., 1, 1] = 1.0
+    R[..., 2, 0] = -s_th
+    R[..., 2, 2] = c
+
+    # 尺度：只用 xz 平面的能量（与 yaw-only 一致）
+    if allow_scale:
+        var_B_xz = (np.sum(Bx**2 + Bz**2, axis=2) / T) + eps  # (N1, N2) 通过 broadcast
+        # tr(R2D*H_xz) = c*(h11+h22) + s*(h21 - h12)
+        numer = c * (h11 + h22) + s_th * (h21 - h12)
+        s = numer / var_B_xz
+    else:
+        s = np.ones((N1, N2), dtype=A.dtype)
+
+    return R, s
+
+
+def apply_traj_align(B, R, s):
+    """
+    应用 (R, s) 到轨迹 B，使其对齐到 A。
+    B: (N2, T, 3)
+    R: (N1, N2, 3, 3)
+    s: (N1, N2)
+    return: (N1, N2, T, 3)
+    """
+    # (N1,N2,3,3) @ (N2,T,3) -> (N1,N2,T,3)
+    rotated = np.einsum("nmij,mtj->nmti", R, B)
+    return s[..., None, None] * rotated
+
+
+def apply_pose_align(c2w, R, s):
+    """
+    应用 (R, s) 到 c2w 外参，左乘 R，平移乘 s：
+    R'_cw = R R_cw,  t'_cw = s R t_cw
+    c2w: (N2,T,4,4) 或 (T,4,4)
+    R: (N1,N2,3,3), s: (N1,N2)
+    return: (N1,N2,T,4,4)
+    """
+    if c2w.ndim == 3:  # (T,4,4) -> (1,T,4,4)
+        c2w = c2w[None, ...]
+    N2, T, _, _ = c2w.shape
+
+    Rc2w = c2w[..., :3, :3]     # (N2,T,3,3)
+    tc2w = c2w[..., :3, 3]      # (N2,T,3)
+
+    Rc2w_new = np.einsum("nmij,mtjk->nmtik", R, Rc2w)          # (N1,N2,T,3,3)
+    tc2w_new = s[..., None, None] * np.einsum("nmij,mtj->nmti", R, tc2w)  # (N1,N2,T,3)
+
+    c2w_aligned = np.zeros((R.shape[0], N2, T, 4, 4), dtype=c2w.dtype)
+    c2w_aligned[..., :3, :3] = Rc2w_new
+    c2w_aligned[..., :3, 3]  = tc2w_new
+    c2w_aligned[..., 3, 3]   = 1.0
+    return c2w_aligned
+
+
+def compute_rmse(A, B_aligned):
+    """
+    计算对齐后轨迹之间的 RMSE (Root Mean Square Error)。
+
+    Args:
+        A: (N1, T, 3)           # 目标轨迹
+        B_aligned: (N1, N2, T, 3)  # 已对齐到 A 的轨迹
+
+    Returns:
+        rmse: (N1, N2)          # 每对 (A_i, B_j) 的误差
+    """
+    if A.ndim == 2:  # 单条轨迹
+        A = A[None, ...]
+
+    # 误差 (N1,N2,T)
+    diff = A[:, None, :, :] - B_aligned
+    sqerr = np.sum(diff**2, axis=-1)  # (N1,N2,T)
+    mse = np.mean(sqerr, axis=-1)     # (N1,N2)
+    rmse = np.sqrt(mse)
+
+    return rmse
+
+
+def traj_length(traj):
+    """
+    traj: (..., T, 3)
+    return: (...,), 每条轨迹路径长度
+    """
+    diffs = traj[..., 1:] - traj[..., :-1]       # (..., T-1, 3)
+    lengths = np.linalg.norm(diffs, axis=-1).sum(axis=-1)  # (...)
+    return lengths
+
+
+def normalize_traj(traj, eps=1e-8):
+    """
+    traj: (..., T, 3)
+    return: (..., T, 3), 每条轨迹路径长度归一化
+    """
+    return traj / (traj_length(traj)[..., None, None] + eps)
+
+
+def visualize_clip(video_file, frames, out_clip_file):
+    vr = VideoReader(str(video_file), ctx=cpu(0), num_threads=1)
+    start_frame, end_frame = frames
+    sample_frames = np.arange(start_frame, end_frame, fps / visualize_fps)
+    sample_frames = np.round(sample_frames).astype(int)
+    clip_data = vr.get_batch(sample_frames).asnumpy()
+
+    out_clip_file.parent.mkdir(parents=True, exist_ok=True)
+    process = (
+        ffmpeg
+        .input("pipe:", format="rawvideo", pix_fmt="rgb24", s=f"{clip_data.shape[2]}x{clip_data.shape[1]}", framerate=visualize_fps)
+        .output(str(out_clip_file), pix_fmt="yuv420p", vcodec="libx264", r=visualize_fps, crf=23, preset="medium")
+        .overwrite_output()
+        .run_async(pipe_stdin=True, quiet=True)
+    )
+    process.stdin.write(clip_data.tobytes())
+    process.stdin.close()
+    process.wait()
+
+
+def max_rot_from_anchor(rot_seq, degrees=True, robust_percentile=None):
+    """
+    计算一个相机姿态序列中相对于首帧的最大旋转角。
+
+    Args:
+        rot_seq: (T,3,3) 单个clip的旋转矩阵序列
+        degrees: True=返回角度, False=弧度
+        robust_percentile: 如果指定, 用分位数而不是max
+
+    Returns:
+        float: 最大(或分位数)旋转角
+    """
+    T = rot_seq.shape[0]
+    R0 = rot_seq[0]
+    # 相对旋转: R0^T @ Rt
+    R_rel = R0.T @ rot_seq              # (3,3)@(T,3,3) -> (T,3,3)
+    # np.matmul 自动广播: (3,3)@(T,3,3)不可直接，需要einsum
+    R_rel = np.einsum('ij,tjk->tik', R0.T, rot_seq)
+
+    trace = np.trace(R_rel, axis1=-2, axis2=-1)
+    cos_theta = np.clip((trace - 1.0) / 2.0, -1.0, 1.0)
+    theta = np.arccos(cos_theta)        # 弧度, shape (T,)
+    theta = theta[1:]                   # 去掉首帧
+
+    if robust_percentile is None:
+        val = np.max(theta)
+    else:
+        val = np.percentile(theta, robust_percentile)
+    if degrees:
+        val = np.degrees(val)
+    return val
+
+
+# 读取 360-1M 的位姿数据
+pf_meta_root = panflow_root / "meta"
+pf_meta_files = list((pf_meta_root).glob("*.json"))
+pf_meta_files.sort()
+print(f"Found {len(pf_meta_files)} PanFlow meta files.")
+# pf_meta_files = pf_meta_files[:100]
+
+filter_summary = defaultdict(lambda: 0)
+camera_summary = defaultdict(lambda: 0)
+for meta_file in tqdm(pf_meta_files, desc="Matching 360-1M poses"):
+    # tqdm.write(f"Processing {meta_file}")
+
+    with open(meta_file, "r") as f:
+        pf_meta = json.load(f)
+    if "slam_clips" not in pf_meta:
+        tqdm.write(f"No slam_clips in {meta_file}, skipping.")
+        continue
+
+    filter_file = filter_root / meta_file.name
+    if not filter_file.exists():
+        tqdm.write(f"Filter file not found: {filter_file}, skipping.")
+        continue
+    with open(filter_file, "r") as f:
+        filter_meta = json.load(f)
+    if not filter_meta:
+        tqdm.write(f"Empty filter file: {filter_file}, skipping.")
+        continue
+    reject_clips = [any(clip["filter"].values()) for clip in filter_meta]
+    reject_ratio = np.mean(reject_clips)
+    if reject_ratio > filter_clips_thres:
+        tqdm.write(f"Reject video {meta_file} due to too many ({reject_ratio:.2%}) filtered clips.")
+        filter_summary["filter"] += len(pf_meta["slam_clips"]["clips"])
+        continue
+
+    video_id = meta_file.stem
+    video_file = panflow_root / "videos" / video_id
+    video_file = video_file.with_suffix(".mp4")
+    cap = cv2.VideoCapture(str(video_file))
+    if not cap.isOpened():
+        tqdm.write(f"Failed to open video: {video_file}, skipping.")
+        continue
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    if "motion_score" not in pf_meta:
+        tqdm.write(f"No motion_score in {meta_file}, skipping.")
+        continue
+
+    if "watermark_score" not in pf_meta:
+        tqdm.write(f"No watermark_score in {meta_file}, skipping.")
+        continue
+
+    clips = deepcopy(pf_meta["slam_clips"]["clips"])
+    missing_qalign = 0
+    for clip, slam_pose, motion_score, watermark_score in zip(
+        clips,
+        pf_meta["slam_pose"]["clips"],
+        pf_meta["motion_score"]["clips"],
+        pf_meta["watermark_score"]["clips"],
+    ):
+        clip["scores"] = {}
+        clip["scores"]["motion_score"] = motion_score["score"]
+        clip["scores"]["watermark_score"] = watermark_score["score"]
+        clip["info"] = slam_pose["info"]
+        if (video_id, clip["clip_id"]) in qalign_scores:
+            scores = qalign_scores[(video_id, clip["clip_id"])]
+            clip["scores"] |= scores
+            clip["scores"]["avg_qalign"] = float(np.mean(list(scores.values())))
+        else:
+            clip["scores"]["avg_qalign"] = -1
+            missing_qalign += 1
+    if missing_qalign > 0:
+        tqdm.write(f"Warning: {missing_qalign} / {len(clips)} clips missing q-align scores in {video_id}.")
+    clips.sort(key=lambda x: x["scores"]["avg_qalign"], reverse=True)
+    
+    ps_meta = {
+        "fps": fps,
+        "clips": [],
+    }
+    poi_counter = defaultdict(lambda: 0)
+    rotation_counter = 0
+    for i_clip, clip in enumerate(tqdm(clips, desc="Processing clips", leave=False)):
+        if len(ps_meta["clips"]) >= max_clips_per_video:
+            tqdm.write(f"Reached max clips per video ({max_clips_per_video}), stopping.")
+            filter_summary["max_clips_per_video"] += len(clips) - i_clip
+            break
+
+        clip_name = clip["clip_name"]
+        clip_id = clip["clip_id"]
+
+        clip_filter = filter_meta[clip_id - 1]
+        assert clip_filter["clip_id"] == clip_id
+        if any(clip_filter["filter"].values()):
+            reasons = [k for k, v in clip_filter["filter"].items() if v]
+            # tqdm.write(f"Clip {video_id}/{clip_name} filtered due to {reasons}, skipping.")
+            filter_summary["filter"] += 1
+            continue
+
+        poi_category = clip_filter["poi_category"]
+        if all(poi_counter[c] >= max_clips_per_poi for c in poi_category):
+            # tqdm.write(f"Clip {video_id}/{clip_name} skipped due to max clips per POI {poi_category}.")
+            filter_summary["max_clips_per_poi"] += 1
+            continue
+
+        slam_info = clip["info"]
+        motion_score = clip["scores"]["motion_score"]
+        frames = clip["frames"]
+        clip_dict = {
+            "video_id": video_id,
+            "clip_id": clip_id,
+            "clip_name": clip_name,
+            "frames": frames,
+            "scores": clip["scores"].copy(),
+            "poi_category": poi_category,
+            "slam_info": slam_info,
+        }
+
+        num_frames = frames[-1] - frames[0] + 1
+        num_frames_sampled = int(round(num_frames / fps * target_fps))
+        if num_frames_sampled < target_frames:
+            # tqdm.write(f"Clip {video_id}/{clip_name} too short ({num_frames_sampled} < {target_frames}), skipping.")
+            filter_summary["too_short"] += 1
+            continue
+
+        if slam_info == "Small camera motion":
+            ps_meta["clips"].append(clip_dict)
+            if visualize_motion:
+                out_clip_file = debug_root / "static_clips" / f"{motion_score:.4f}-{video_id}-{clip_name}.mp4"
+                visualize_clip(video_file, frames, out_clip_file)
+            filter_summary["small_camera_motion"] += 1
+            continue
+
+        if visualize_motion:
+            continue  # 只保留静止片段
+
+        if slam_info != "Success":
+            # tqdm.write(f"Clip {video_id}/{clip_name} SLAM not successful ({slam_info}), skipping.")
+            filter_summary["slam_fail"] += 1
+            continue
+
+        pose_file = panflow_root / "slam_pose" / video_id / clip_name
+        pose_file = pose_file.with_suffix(".npy")
+        if not pose_file.exists():
+            tqdm.write(f"Pose file not found: {pose_file}, skipping.")
+            continue
+        pf_pose = np.load(pose_file)  # (T, 3, 4)
+
+        rot_seg = pf_pose[:, :3, :3]  # (num_segs,T,3,3)
+        rotation_score = max_rot_from_anchor(rot_seg, degrees=True, robust_percentile=95)  # (num_segs,)
+        rotation_score = float(rotation_score)
+        if visualize_rotation:
+            out_clip_file = debug_root / "rotation_score" / f"{rotation_score:.4f}-{video_id}-{clip_name}.mp4"
+            visualize_clip(video_file, frames, out_clip_file)
+        if rotation_score > rotation_score_thres:
+            # tqdm.write(f"Clip {video_id}/{clip_name} rejected due to large rotation ({rotation_score:.2f}° > {rotation_score_thres}°), skipping.")
+            rotation_counter += 1
+            filter_summary["large_rotation"] += 1
+            continue
+        clip_dict["scores"]["rotation_score"] = rotation_score
+
+        sample_frames = np.linspace(0, num_frames - 1, num_frames_sampled)
+        sample_frames = np.round(sample_frames).astype(int)
+
+        # 1. 提取所有片段索引
+        max_start = num_frames_sampled - target_frames
+        num_segs = max_start // match_step + 1
+        starts = np.arange(0, max_start+1, match_step)
+        num_segs = len(starts)
+        idx = starts[:, None] + np.arange(target_frames)[None, :]
+        idx = sample_frames[idx]
+
+        # 2. 取出 c2w
+        c2w = pf_pose[idx]  # (num_segs, T, 3, 4) 或 (num_segs, T, 4, 4)
+        if c2w.shape[-2:] == (3, 4):
+            # 补成 4x4
+            last_row = repeat(np.array([0,0,0,1], dtype=c2w.dtype), "n -> s t 1 n", s=c2w.shape[0], t=c2w.shape[1])
+            c2w = np.concatenate([c2w, last_row], axis=-2)  # (num_segs,T,4,4)
+
+        # 3. 每段归一化到第一帧
+        w2c0 = np.linalg.inv(c2w[:, 0])  # (num_segs,4,4)
+        c2w = w2c0[:, None] @ c2w  # (num_segs,T,4,4)
+
+        # 4. 提取相机中心轨迹 (num_segs, T, 3)
+        pos_seg = c2w[:, :, :3, 3]
+
+        # 5. 批量对齐 + RMSE
+        # pos_seg -> (num_segs,T,3)，扩展成 (num_segs,1,T,3)，与 cb_pos (N_cb,T,3) 对齐
+        R, s = get_traj_align(cb_pos, pos_seg)          # R:(N_cb,num_segs,3,3), s:(N_cb,num_segs)
+        pos_aligned = apply_traj_align(pos_seg, R, s)   # (N_cb,num_segs,T,3)
+        rmse = compute_rmse(normalize_traj(cb_pos), normalize_traj(pos_aligned))        # (N_cb,num_segs)
+
+        best_idx = np.argmin(rmse, axis=1)     # (N_cb,)
+        best_seg = idx[best_idx]         # (N_cb, T)
+        best_rmse = rmse[np.arange(len(cb_pos)), best_idx]  # (N_cb,)
+
+        topk_cb = np.argsort(best_rmse)[:top_k]
+
+        if visualize_pose:
+            pose_aligned = apply_pose_align(c2w, R, s)  # (N_cb,num_segs,T,4,4)
+
+        clip_dict["matches"] = []
+        for i in topk_cb:
+            j = best_idx[i]
+            cb_name = cb_meta[i]["video"]
+            if cb_meta[i].get("camera_labels", False):
+                for label in cb_meta[i]["camera_labels"]:
+                    camera_summary[label] += 1
+            else:
+                camera_summary[cb_name] += 1
+            clip_dict["matches"].append({
+                "video": cb_name,
+                "frames": (int(best_seg[i, 0]), int(best_seg[i, -1])),
+                "rmse": float(rmse[i, j]),
+                "R": R[i, j].tolist(),
+                "s": float(s[i, j]),
+            })
+
+            if visualize_pose:
+                cb_pose_dir = debug_root / video_id / clip_name / cb_name
+                cb_pose_dir.mkdir(parents=True, exist_ok=True)
+                cb_pose_file = cb_pose_dir / "target.npy"
+                np.save(cb_pose_file, cb_poses_gravity[i])
+                pf_pose_file = cb_pose_dir / "aligned.npy"
+                np.save(pf_pose_file, pose_aligned[i, j])
+                combined_file = cb_pose_dir / "comparison.npy"
+                combined_pose = np.concatenate([cb_poses_gravity[i], pose_aligned[i, j]], axis=0)
+                np.save(combined_file, combined_pose)
+                # vis_to_html(cb_pose_dir, [cb_pose_file, pf_pose_file])
+                vis_to_html(cb_pose_dir, [combined_file])
+
+        for c in poi_category:
+            poi_counter[c] += 1
+        ps_meta["clips"].append(clip_dict)
+
+    if not ps_meta["clips"]:
+        tqdm.write(f"No valid clips in {meta_file}, skipping.")
+        continue
+
+    if rotation_counter / (len(clips) + rotation_counter) > rotating_clips_thres:
+        tqdm.write(f"Reject video {meta_file} due to too many ({rotation_counter}/{len(clips)}) high-rotation clips.")
+        filter_summary["large_rotation"] += len(ps_meta["clips"])
+        continue
+
+    filter_summary["success"] += len(ps_meta["clips"])
+    out_file = output_root / f"{video_id}.json"
+    with open(out_file, "w") as f:
+        json.dump(ps_meta, f, indent=4)
+
+for name, summary in [
+    ("filter_summary", filter_summary),
+    ("camera_summary", camera_summary),
+]:
+    summary_file = summary_root / f"{name}.json"
+    with open(summary_file, "w") as f:
+        json.dump(summary, f, indent=4)
+    print(f"Wrote summary to {summary_file}")
+
+# 可视化统计结果
+fig, ax = plt.subplots(figsize=(10, 6))
+items = sorted(camera_summary.items(), key=lambda x: x[1], reverse=True)
+values = list(camera_summary.values())
+ax.hist(values, bins='auto', color="skyblue", edgecolor="black")
+
+ax.set_xlabel("Match Count")
+ax.set_ylabel("Number of Videos")
+ax.set_title(f"CameraBench Match Distribution")
+
+plt.tight_layout()
+summary_file = summary_root / "camera_summary.png"
+fig.savefig(summary_file, dpi=300, bbox_inches="tight")
+plt.close(fig)
+print(f"Saved camera summary histogram to {summary_file}")
+
+# 按数量从大到小排序
+items = sorted(filter_summary.items(), key=lambda x: x[1], reverse=True)
+labels, counts = zip(*items)
+
+fig, ax = plt.subplots(figsize=(10, 6))
+# 建议用水平条形图，长标签也容易展示
+ax.barh(labels, counts, color="salmon")
+ax.invert_yaxis()  # 让数量最多的排在最上
+ax.set_xlabel("Number of filtered clips")
+ax.set_title("Filter Summary")
+
+plt.tight_layout()
+summary_file = summary_root / "filter_summary.png"
+fig.savefig(summary_file, dpi=300, bbox_inches="tight")
+plt.close(fig)
+print(f"Saved filter summary plot to {summary_file}")

UCPE/tools/caption_camerabench.py

@@ -0,0 +1,195 @@
+from vllm import LLM, SamplingParams
+
+from pathlib import Path
+import jsonlines
+from tqdm.auto import tqdm
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import os
+import json
+import re
+
+from transformers import AutoProcessor
+from qwen_vl_utils import process_vision_info
+
+
+# -----------------------------
+# basic configuration
+# -----------------------------
+data_root = Path("data/UCPE/CameraBench")
+split = "train"  # "train" or "test"
+output_jsonl = data_root / f"captioned_{split}.jsonl"
+meta_file = data_root / f"processed_{split}.jsonl"
+
+model_id = "Qwen/Qwen2.5-VL-7B-Instruct"   # try 7B first; switch to 32B if resources allow
+# model_id = "chancharikm/qwen2.5-vl-7b-cam-motion-preview"
+nframes = 32                                # hint for frame sampling inside qwen_vl_utils
+fps_hint = None                             # None or a small integer like 1/2/4 (optional)
+batch_size = 8                              # how many videos per vLLM.generate batch
+max_workers = min(8, os.cpu_count() or 4)   # 线程数按机器调整
+inflight_limit = batch_size * 2             # 同时在制的样本上限
+print(f"Using max_workers={max_workers}, inflight_limit={inflight_limit}")
+
+
+max_new_tokens = 512
+temperature = 0.2
+top_p = 0.9
+repetition_penalty = 1.05
+gpu_memory_utilization = 0.9
+tensor_parallel_size = 1
+limit_mm_per_prompt = {"video": 1}
+
+
+def build_llm_input(video_path: Path, prompt_text: str, processor: AutoProcessor):
+    # compose messages (system + user text + video item)
+    video_item = {"type": "video", "video": str(video_path), "nframes": nframes}
+    if fps_hint is not None:
+        video_item["fps"] = fps_hint
+
+    messages = [
+        {"role": "system", "content": "You are a helpful video captioning assistant."},
+        {"role": "user", "content": [
+            {"type": "text", "text": prompt_text},
+            video_item
+        ]}
+    ]
+
+    # extract frames / prepare tensors for the model (CPU/I/O-heavy)
+    image_inputs, video_inputs = process_vision_info(messages)
+
+    # text template → prompt string
+    prompt = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    mm_data = {}
+    if video_inputs is not None:
+        mm_data["video"] = video_inputs
+
+    return {"prompt": prompt, "multi_modal_data": mm_data}
+
+
+def prepare_one(obj, processor: AutoProcessor):
+    """单样本：基于 labels/seed_caption 构造提示 → 抽帧预处理 → 组装 vLLM 输入"""
+    vpath = data_root / obj["path"]
+    prompt_text = "Please describe this video in detail."
+    llm_in = build_llm_input(vpath, prompt_text, processor)
+    return obj, llm_in
+
+
+# -----------------------------
+# load metadata
+# -----------------------------
+with jsonlines.open(meta_file, "r") as reader:
+    metadata = list(reader)
+for obj in tqdm(metadata, desc="checking files"):
+    assert (data_root / obj["path"]).exists(), f"File not found: {obj['path']}"
+
+
+# -----------------------------
+# init vLLM + processor
+# -----------------------------
+llm = LLM(
+    model=model_id,
+    tensor_parallel_size=tensor_parallel_size,
+    gpu_memory_utilization=gpu_memory_utilization,
+    # enforce_eager=True,
+    limit_mm_per_prompt=limit_mm_per_prompt,
+)
+processor = AutoProcessor.from_pretrained(model_id)
+sampling_params = SamplingParams(
+    max_tokens=max_new_tokens,
+    temperature=temperature,
+    top_p=top_p,
+    repetition_penalty=repetition_penalty,
+)
+
+
+# -----------------------------
+# pipeline：边准备边推理边写出（动态提交 + 行缓冲）
+# -----------------------------
+output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+prepared_buffer = []  # 缓存已准备好的 (obj, llm_in)
+
+
+def infer_and_flush(buffer, writer):
+    """对 buffer 中的若干样本推理，并写出结果"""
+    if not buffer:
+        return
+    batch_objs = [it[0] for it in buffer]
+    batch_inputs = [it[1] for it in buffer]
+    gens = llm.generate(batch_inputs, sampling_params)
+
+    for ob, g in zip(batch_objs, gens):
+        text_out = g.outputs[0].text.strip()
+
+        writer.write({
+            "path": ob["path"],
+            "labels": ob.get("labels", []),   # test 集本身有 labels
+            "caption": text_out
+        })
+
+try:
+    # 用行缓冲打开文件，便于“边写边可见”
+    f = open(output_jsonl, "w", buffering=1, encoding="utf-8")
+    writer = jsonlines.Writer(f)
+
+    with ThreadPoolExecutor(max_workers=max_workers) as ex:
+        pbar = tqdm(total=len(metadata), desc="preparing & inferring")
+
+        # 动态 pending 集合
+        pending = set()
+        i_submit = 0
+
+        # 先填满 in-flight
+        while i_submit < len(metadata) and len(pending) < inflight_limit:
+            fut = ex.submit(prepare_one, metadata[i_submit], processor)
+            pending.add(fut)
+            i_submit += 1
+
+        # 循环直到所有任务完成
+        while pending:
+            # 只等待当前 pending 集合���的任务
+            for fut in as_completed(list(pending), timeout=None):
+                pending.remove(fut)
+                obj, llm_in = fut.result()
+                prepared_buffer.append((obj, llm_in))
+                pbar.update(1)
+
+                # 满一批就立刻推理并清空对应部分
+                if len(prepared_buffer) >= batch_size:
+                    infer_and_flush(prepared_buffer[:batch_size], writer)
+                    prepared_buffer = prepared_buffer[batch_size:]
+
+                # 补交新任务，保持 in-flight 数量
+                while i_submit < len(metadata) and len(pending) < inflight_limit:
+                    fut_new = ex.submit(prepare_one, metadata[i_submit], processor)
+                    pending.add(fut_new)
+                    i_submit += 1
+
+                # 跳出到 while pending，重新评估 pending 集合（已更新）
+                break
+
+        # 把“尾巴”按 batch 循环清空，确保不丢最后一个或多个 batch
+        while prepared_buffer:
+            chunk = prepared_buffer[:batch_size]
+            infer_and_flush(chunk, writer)
+            prepared_buffer = prepared_buffer[len(chunk):]
+
+        pbar.close()
+finally:
+    # 关闭 writer / 文件句柄
+    try:
+        writer.close()
+    except Exception:
+        pass
+    try:
+        f.close()
+    except Exception:
+        pass
+    # 优雅关闭 vLLM 引擎
+    try:
+        llm.shutdown()
+    except Exception:
+        pass
+
+print(f"done. captions saved to: {output_jsonl}")

UCPE/tools/caption_panshot.py

@@ -0,0 +1,210 @@
+from vllm import LLM, SamplingParams
+
+from pathlib import Path
+import jsonlines
+from tqdm.auto import tqdm
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import os
+import json
+import re
+
+from transformers import AutoProcessor
+from qwen_vl_utils import process_vision_info
+
+
+# -----------------------------
+# basic configuration
+# -----------------------------
+split = "train"  # "train" or "test"
+data_root = Path("data/UCPE/PanShot")
+video_dir = data_root / f"videos-{split}"
+output_jsonl = data_root / f"captioned-{split}.jsonl"
+
+model_id = "Qwen/Qwen2.5-VL-7B-Instruct"   # try 7B first; switch to 32B if resources allow
+# model_id = "chancharikm/qwen2.5-vl-7b-cam-motion-preview"
+nframes = 32                                # hint for frame sampling inside qwen_vl_utils
+fps_hint = None                             # None or a small integer like 1/2/4 (optional)
+batch_size = 8                              # how many videos per vLLM.generate batch
+max_workers = min(8, os.cpu_count() or 4)   # 线程数按机器调整
+inflight_limit = batch_size * 2             # 同时在制的样本上限
+print(f"Using max_workers={max_workers}, inflight_limit={inflight_limit}")
+
+
+max_new_tokens = 512
+temperature = 0.2
+top_p = 0.9
+repetition_penalty = 1.05
+gpu_memory_utilization = 0.9
+tensor_parallel_size = 1
+limit_mm_per_prompt = {"video": 1}
+
+
+def build_llm_input(video_path: Path, prompt_text: str, processor: AutoProcessor):
+    # compose messages (system + user text + video item)
+    video_item = {"type": "video", "video": str(video_path), "nframes": nframes}
+    if fps_hint is not None:
+        video_item["fps"] = fps_hint
+
+    messages = [
+        {"role": "system", "content": "You are a helpful video captioning assistant."},
+        {"role": "user", "content": [
+            {"type": "text", "text": prompt_text},
+            video_item
+        ]}
+    ]
+
+    # extract frames / prepare tensors for the model (CPU/I/O-heavy)
+    image_inputs, video_inputs = process_vision_info(messages)
+
+    # text template → prompt string
+    prompt = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    mm_data = {}
+    if video_inputs is not None:
+        mm_data["video"] = video_inputs
+
+    return {"prompt": prompt, "multi_modal_data": mm_data}
+
+
+def prepare_one(video, processor: AutoProcessor):
+    """单样本：基于 labels/seed_caption 构造提示 → 抽帧预处理 → 组装 vLLM 输入"""
+    vpath = video_dir / f"{video}.mp4"
+    prompt_text = "Please describe this video in detail."
+    try:
+        llm_in = build_llm_input(vpath, prompt_text, processor)
+    except Exception as e:
+        print(f"Error processing {video}: {e}")
+        return None
+    return video, llm_in
+
+
+# -----------------------------
+# load metadata
+# -----------------------------
+videos = [v.stem for v in video_dir.glob("*.mp4")]
+print(f"Found {len(videos)} videos in {video_dir}")
+processed = set()
+if output_jsonl.exists():
+    print(f"Resuming from {output_jsonl}")
+    with open(output_jsonl, "r", encoding="utf-8") as f_in:
+        for line in f_in:
+            try:
+                rec = json.loads(line)
+                processed.add(rec["video"])
+            except Exception:
+                continue
+    print(f"Found {len(processed)} processed videos to skip.")
+videos = [v for v in videos if v not in processed]
+print(f"Total videos to process: {len(videos)}")
+
+# -----------------------------
+# init vLLM + processor
+# -----------------------------
+llm = LLM(
+    model=model_id,
+    tensor_parallel_size=tensor_parallel_size,
+    gpu_memory_utilization=gpu_memory_utilization,
+    # enforce_eager=True,
+    limit_mm_per_prompt=limit_mm_per_prompt,
+)
+processor = AutoProcessor.from_pretrained(model_id)
+sampling_params = SamplingParams(
+    max_tokens=max_new_tokens,
+    temperature=temperature,
+    top_p=top_p,
+    repetition_penalty=repetition_penalty,
+)
+
+
+# -----------------------------
+# pipeline：边准备边推理边写出（动态提交 + 行缓冲）
+# -----------------------------
+output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+prepared_buffer = []  # 缓存已准备好的 (obj, llm_in)
+
+
+def infer_and_flush(buffer, writer):
+    """对 buffer 中的若干样本推理，并写出结果"""
+    if not buffer:
+        return
+    batch_videos = [it[0] for it in buffer]
+    batch_inputs = [it[1] for it in buffer]
+    gens = llm.generate(batch_inputs, sampling_params)
+
+    for video, g in zip(batch_videos, gens):
+        text_out = g.outputs[0].text.strip()
+
+        writer.write({
+            "video": video,
+            "caption": text_out
+        })
+
+try:
+    # 用行缓冲打开文件，便于“边写边可见”
+    f = open(output_jsonl, "a", buffering=1, encoding="utf-8")
+    writer = jsonlines.Writer(f)
+
+    with ThreadPoolExecutor(max_workers=max_workers) as ex:
+        pbar = tqdm(total=len(videos), desc="preparing & inferring")
+
+        # 动态 pending 集合
+        pending = set()
+        i_submit = 0
+
+        # 先填满 in-flight
+        while i_submit < len(videos) and len(pending) < inflight_limit:
+            fut = ex.submit(prepare_one, videos[i_submit], processor)
+            pending.add(fut)
+            i_submit += 1
+
+        # 循环直到所有任务完成
+        while pending:
+            # 只等待当前 pending 集合中的任务
+            for fut in as_completed(list(pending), timeout=None):
+                pending.remove(fut)
+                result = fut.result()
+                pbar.update(1)
+                if result is None:
+                    continue
+                prepared_buffer.append(result)
+
+                # 满一批就立刻推理并清空对应部分
+                if len(prepared_buffer) >= batch_size:
+                    infer_and_flush(prepared_buffer[:batch_size], writer)
+                    prepared_buffer = prepared_buffer[batch_size:]
+
+                # 补交新任务，保持 in-flight 数量
+                while i_submit < len(videos) and len(pending) < inflight_limit:
+                    fut_new = ex.submit(prepare_one, videos[i_submit], processor)
+                    pending.add(fut_new)
+                    i_submit += 1
+
+                # 跳出到 while pending，重新评估 pending 集合（已更新）
+                break
+
+        # 把“尾巴”按 batch 循环清空，确保不丢最后一个或多个 batch
+        while prepared_buffer:
+            chunk = prepared_buffer[:batch_size]
+            infer_and_flush(chunk, writer)
+            prepared_buffer = prepared_buffer[len(chunk):]
+
+        pbar.close()
+finally:
+    # 关闭 writer / 文件句柄
+    try:
+        writer.close()
+    except Exception:
+        pass
+    try:
+        f.close()
+    except Exception:
+        pass
+    # 优雅关闭 vLLM 引擎
+    try:
+        llm.shutdown()
+    except Exception:
+        pass
+
+print(f"done. captions saved to: {output_jsonl}")

UCPE/tools/dataset_statistics.py

@@ -0,0 +1,304 @@
+import json
+import pandas as pd
+from pathlib import Path
+from typing import Optional
+import tyro
+from pydantic import BaseModel
+import numpy as np
+from tqdm.auto import tqdm
+from torch.utils.data import DataLoader, Subset
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+from src.dataset import PanShotDataset, Re10kDataset, DemoDataset
+from einops import rearrange
+
+
+# ======================================================
+#   seaborn global theme (官方模板)
+# ======================================================
+sns.set_theme(
+    context="paper",
+    style="whitegrid",
+    palette="deep",
+    font_scale=0.9,
+)
+
+
+# ======================================================
+#   Args
+# ======================================================
+class Args(BaseModel):
+    data: str = "PanShotDataset"
+    num_frames: int = 81
+    data_root: Path = Path("data/UCPE")
+    num_workers: int = 4
+    zero_first_yaw: bool = True
+    output_dir: Path = Path("outputs/suppl")
+    num_samples: Optional[int] = None
+    split: str = "train"
+    color: str = "C0"   # NEW: color
+
+
+# ======================================================
+#   DataLoader
+# ======================================================
+def collate_fn(samples):
+    return samples[0]
+
+
+def prepare_dataloader(args):
+    dataset_class = globals().get(args.data, None)
+    dataset = dataset_class(
+        args, args.split,
+        load_keys=["pose", "xi", "y_fov"]
+    )
+
+    if args.num_samples is not None:
+        dataset = Subset(dataset, list(range(args.num_samples)))
+
+    return DataLoader(
+        dataset,
+        collate_fn=collate_fn,
+        batch_size=1,
+        num_workers=args.num_workers,
+        shuffle=False,
+    )
+
+
+# ======================================================
+#   Camera Euler: x-right, y-down, z-forward
+# ======================================================
+def rotmat_to_euler_cam(R):
+    fx, fy, fz = R[:, 2]   # camera forward in world
+    yaw = np.arctan2(fx, fz)
+    pitch = np.arctan2(-fy, np.sqrt(fx**2 + fz**2))
+    roll = np.arctan2(R[1, 0], R[0, 0])
+    return np.degrees([yaw, pitch, roll])
+
+
+def rot_angle(R0, Ri):
+    R = R0.T @ Ri
+    cos_theta = (np.trace(R) - 1) / 2
+    cos_theta = np.clip(cos_theta, -1, 1)
+    return np.degrees(np.arccos(cos_theta))
+
+
+# ======================================================
+#   MAIN
+# ======================================================
+def main():
+    args = tyro.cli(Args)
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    dataloader = prepare_dataloader(args)
+
+    # seaborn palette → choose refined colors
+    COLOR_MAP = {
+        "C0": sns.color_palette("deep")[0],
+        "C1": sns.color_palette("deep")[1],
+        "C2": sns.color_palette("deep")[2],
+        "C3": sns.color_palette("deep")[3],
+    }
+    COLOR = COLOR_MAP.get(args.color, args.color)
+
+    # ======================================================
+    # containers
+    # ======================================================
+    final_rel_yaw = []   # displacement-based azimuth
+    init_pitch = []
+    init_roll = []
+    rotation_magnitude = []
+    max_rel_yaw = []     # NEW: per-clip max camera relative yaw (signed)
+    xis = []
+    fovs = []
+
+    # ======================================================
+    # iterate dataset
+    # ======================================================
+    for data in tqdm(dataloader, desc=f"Processing {args.data}"):
+
+        poses = data["pose"]  # (N, 3, 4)
+        xi = float(data["xi"])
+        fov = float(data["y_fov"])
+
+        xis.append(xi)
+        fovs.append(fov)
+
+        R_all = poses[:, :, :3]
+        R0 = R_all[0]
+
+        # Initial orientation
+        yaw0, pitch0, roll0 = rotmat_to_euler_cam(R0)
+        init_pitch.append(pitch0)
+        init_roll.append(roll0)
+
+        # ------- Displacement-based azimuth (position) -------
+        p0 = poses[0, :, 3]
+        pN = poses[-1, :, 3]
+        v = pN - p0
+        yaw_pos = np.arctan2(v[0], v[2])
+        final_rel_yaw.append(np.degrees(yaw_pos))
+
+        # ------- Rotation magnitude wrt frame 0 (unsigned) -------
+        rotation_magnitude.append(max(
+            rot_angle(R0, Ri) for Ri in R_all
+        ))
+
+        # ------- NEW: max camera relative yaw (signed) -------
+        rel_yaws = []
+        for Ri in R_all:
+            yaw_i, _, _ = rotmat_to_euler_cam(Ri)
+            rel_yaws.append(yaw_i)
+        rel_yaws = np.array(rel_yaws)
+
+        # pick the frame with largest |relative yaw|, keep sign
+        idx_max = np.argmax(np.abs(rel_yaws))
+        max_rel_yaw.append(rel_yaws[idx_max])
+
+    # convert to numpy
+    final_rel_yaw = np.array(final_rel_yaw)
+    init_pitch = np.array(init_pitch)
+    init_roll = np.array(init_roll)
+    rotation_magnitude = np.array(rotation_magnitude)
+    max_rel_yaw = np.array(max_rel_yaw)
+    xis = np.array(xis)
+    fovs = np.array(fovs)
+
+    # remove top 1% roll outliers
+    roll_thr = np.percentile(init_roll, 99)
+    init_roll = init_roll[init_roll <= roll_thr]
+
+    # remove bottom 1% roll outliers
+    roll_thr = np.percentile(init_roll, 1)
+    init_roll = init_roll[init_roll >= roll_thr]
+
+    # ======================================================
+    # save helper
+    # ======================================================
+    def save_pdf(fig, name):
+        fig.savefig(args.output_dir / f"{name}.pdf",
+                    dpi=300,
+                    bbox_inches="tight",
+                    format="pdf")
+        plt.close(fig)
+
+    # ======================================================
+    # 1. Rose plot: displacement azimuth (position)
+    # ======================================================
+    fig = plt.figure(figsize=(2.0, 2.0))
+    ax = plt.subplot(111, polar=True)
+
+    rad = np.radians(final_rel_yaw)
+
+    ax.grid(True, linewidth=0.4, alpha=0.5)
+    ax.set_facecolor("white")
+
+    ax.hist(rad, bins=36, alpha=0.55, color=COLOR,
+            edgecolor=".3", linewidth=0.4)
+
+    ax.set_theta_zero_location("N")
+    ax.set_theta_direction(-1)
+
+    ax.set_thetagrids(
+        angles=np.arange(0, 360, 45),
+        labels=[f"{(a if a <= 180 else a - 360)}°" for a in np.arange(0, 360, 45)]
+    )
+
+    # ax.set_title("Azimuth", pad=8)
+    save_pdf(fig, "rose_direction")
+
+    # ======================================================
+    # NEW: 1b. Rose plot: max camera relative yaw (signed)
+    # ======================================================
+    fig = plt.figure(figsize=(2.0, 2.0))
+    ax = plt.subplot(111, polar=True)
+
+    rad_max_yaw = np.radians(max_rel_yaw)
+
+    ax.grid(True, linewidth=0.4, alpha=0.5)
+    ax.set_facecolor("white")
+
+    ax.hist(rad_max_yaw, bins=36, alpha=0.55, color=COLOR,
+            edgecolor=".3", linewidth=0.4)
+
+    ax.set_theta_zero_location("N")
+    ax.set_theta_direction(-1)
+
+    ax.set_thetagrids(
+        angles=np.arange(0, 360, 45),
+        labels=[f"{(a if a <= 180 else a - 360)}°" for a in np.arange(0, 360, 45)]
+    )
+
+    # ax.set_title("Max Relative Yaw", pad=8)
+    save_pdf(fig, "rose_rotation")  # 文件名仍叫 rose_rotation
+
+    # ======================================================
+    # 2. Pitch histogram (seaborn)
+    # ======================================================
+    fig = plt.figure(figsize=(2.2, 2.0))
+    sns.histplot(init_pitch, bins=40, kde=False, stat="density",
+                 color=COLOR, edgecolor=".3", linewidth=0.4, alpha=0.6)
+    plt.xlabel("Pitch (°)")
+    plt.ylabel("Density")
+    # plt.title("Initial Pitch")
+    save_pdf(fig, "pitch_hist")
+
+
+    # ======================================================
+    # 3. Roll histogram
+    # ======================================================
+    fig = plt.figure(figsize=(2.2, 2.0))
+    sns.histplot(init_roll, bins=40, kde=False, stat="density",
+                 color=COLOR, edgecolor=".3", linewidth=0.4, alpha=0.6)
+    plt.xlabel("Roll (°)")
+    plt.ylabel("Density")
+    # plt.title("Initial Roll")
+    save_pdf(fig, "roll_hist")
+
+
+    # ======================================================
+    # 4. Rotation magnitude
+    # ======================================================
+    fig = plt.figure(figsize=(2.2, 2.0))
+    sns.histplot(rotation_magnitude, bins=40, kde=False, stat="density",
+                 color=COLOR, edgecolor=".3", linewidth=0.4, alpha=0.6)
+    plt.xlabel("Maximum Rotation (°)")
+    plt.ylabel("Density")
+    # plt.title("Rotation Magnitude")
+    save_pdf(fig, "rotation_magnitude_hist")
+
+
+    # ======================================================
+    # 5. ξ histogram
+    # ======================================================
+    fig = plt.figure(figsize=(2.2, 2.0))
+    sns.histplot(xis, bins=30, kde=False, stat="density",
+                 color=COLOR, edgecolor=".3", linewidth=0.4, alpha=0.6)
+    plt.xlabel("ξ")
+    plt.ylabel("Density")
+    # plt.title("ξ Distribution")
+    save_pdf(fig, "xi_hist")
+
+
+    # ======================================================
+    # 6. FoV histogram
+    # ======================================================
+    # remove top 1% FoV outliers
+    fov_thr = np.percentile(fovs, 99)
+    fovs = fovs[fovs <= fov_thr]
+    # remove bottom 1% FoV outliers
+    fov_thr = np.percentile(fovs, 1)
+    fovs = fovs[fovs >= fov_thr]
+
+    fig = plt.figure(figsize=(2.2, 2.0))
+    sns.histplot(fovs, bins=30, kde=False, stat="density",
+                 color=COLOR, edgecolor=".3", linewidth=0.4, alpha=0.6)
+    plt.xlabel("FoV (°)")
+    plt.ylabel("Density")
+    # plt.title("FoV Distribution")
+    save_pdf(fig, "fov_hist")
+
+
+# ======================================================
+if __name__ == "__main__":
+    main()

UCPE/tools/download_panflow.py

@@ -0,0 +1,57 @@
+from pathlib import Path
+from tqdm.auto import tqdm
+import json
+import yt_dlp
+
+
+panflow_root = Path("data/360-1M")
+panshot_root = Path("data/UCPE")
+pf_clip_root = panshot_root / "PanFlow" / "match_clips"
+pf_video_root = panshot_root / "PanFlow" / "videos"
+pf_video_root.mkdir(parents=True, exist_ok=True)
+
+clip_metas = list(pf_clip_root.glob("*.json"))
+clip_metas.sort()
+print(f"Found {len(clip_metas)} PanFlow clip files.")
+
+videos = set()
+for clip_meta in clip_metas:
+    with open(clip_meta, "r") as f:
+        meta = json.load(f)
+    videos.add(meta["video_id"])
+print(f"Found {len(videos)} unique PanFlow videos.")
+
+downloaded_videos = pf_video_root.glob("*.mp4")
+downloaded_videos = set([p.stem for p in downloaded_videos])
+print(f"Found {len(downloaded_videos)} already downloaded videos.")
+
+videos = videos - downloaded_videos
+print(f"{len(videos)} videos to download.")
+
+
+def download_video(video_id, output_path):
+    video_url = f"https://www.youtube.com/watch?v={video_id}"
+    ydl_opts = {
+        "outtmpl": str(output_path),
+        "format": "bestvideo[height<=2500][height>1500]",
+        "quiet": False,
+        "no_warnings": True,
+        "simulate": False,
+        "cookiefile": "~/.config/cookies.txt",
+        "print": [
+            "before_dl:Format: %(format_id)s | Res: %(resolution)s | FPS: %(fps)s",
+            "after_move:Size: %(filesize:.2fMB)s"
+        ],
+        "user_agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
+    }
+
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        ydl.download([video_url])
+
+for video_id in tqdm(videos, desc="Downloading videos"):
+    video_path = Path(pf_video_root) / f"{video_id}.mp4"
+    try:
+        download_video(video_id, video_path)
+    except Exception as e:
+        tqdm.write(f"Failed to download {video_id}: {e}")
+        continue

UCPE/tools/export_camerabench.py

@@ -0,0 +1,64 @@
+import csv
+import jsonlines
+from pathlib import Path
+from tqdm.auto import tqdm
+import jsonlines
+
+# -----------------------------
+# paths
+# -----------------------------
+data_root = Path("data/UCPE/CameraBench")
+filtered_jsonl = data_root / "filtered.jsonl"
+split = "train"  # "train" or "test"
+input_jsonl = data_root / f"captioned_{split}.jsonl"
+camera_jsonl = data_root / f"processed_{split}.jsonl"
+output_csv = data_root / f"metadata_{split}.csv"
+output_jsonl = data_root / f"metadata_{split}.jsonl"
+
+with jsonlines.open(filtered_jsonl, "r") as reader:
+    filtered_videos = {obj["path"] for obj in reader if any(obj["filter"].values())}
+
+# 保证输出目录存在
+output_csv.parent.mkdir(parents=True, exist_ok=True)
+
+total = 0
+written = 0
+skipped = 0
+filtered = 0
+
+with jsonlines.open(input_jsonl, "r") as reader, \
+     open(output_csv, "w", newline="", encoding="utf-8") as fout, \
+        jsonlines.open(output_jsonl, "w") as jsonl_writer, \
+            jsonlines.open(camera_jsonl, "r") as camera_reader:
+    camera_metas = {obj["path"]: obj for obj in camera_reader}
+    writer = csv.writer(fout)
+    # 表头
+    writer.writerow(["video", "prompt"])
+
+    for obj in tqdm(reader, desc="Converting JSONL → CSV"):
+        total += 1
+        path = obj.get("path", None)
+        caption = obj.get("caption", None)
+
+        if path is None or caption is None:
+            skipped += 1
+            continue
+
+        if path in filtered_videos:
+            filtered += 1
+            continue
+
+        # 规范化 caption 的空白字符，避免 CSV 里出现杂乱换行
+        caption_norm = " ".join(str(caption).split())
+        jsonl_writer.write({
+            "video": Path(path).stem,
+            "prompt": caption_norm,
+            "camera_caption": camera_metas[path]["caption"],
+            "camera_labels": camera_metas[path].get("labels", []),
+        })
+        writer.writerow([path, caption_norm])
+        written += 1
+
+print(f"Done. Total lines: {total}, written: {written}, skipped (missing fields): {skipped}, filtered: {filtered}")
+print(f"CSV saved to: {output_csv}")
+print(f"JSONL saved to: {output_jsonl}")

UCPE/tools/export_camerabench_for_rerender.py

@@ -0,0 +1,140 @@
+"""
+Export CameraBench videos needed for PanShot re-rendering.
+
+This script replaces the train-split part of process_camerabench.py
+by reading the list of needed cb_videos directly from PanShot meta files,
+bypassing the need for cam_motion/captionset.json.
+
+Usage:
+    conda activate UCPE
+    python tools/export_camerabench_for_rerender.py
+"""
+
+from pathlib import Path
+import json
+import jsonlines
+import ffmpeg
+import cv2
+from tqdm.auto import tqdm
+
+
+# ===== Configuration =====
+panshot_root = Path("/tmp/data/UCPE/PanShot")
+camerabench_data_root = Path("data/CameraBench")
+output_root = Path("data/UCPE/CameraBench")
+
+target_height = 720
+target_width = 1280
+target_frames = 81
+target_fps = 16
+
+output_root.mkdir(parents=True, exist_ok=True)
+
+
+# ===== Collect needed cb_videos from PanShot meta files =====
+cb_videos_needed = set()
+for split in ["train", "test"]:
+    meta_root = panshot_root / f"meta-{split}"
+    for f in meta_root.glob("*.json"):
+        with open(f) as fp:
+            for m in json.load(fp):
+                cb_videos_needed.add(m["cb_video"])
+print(f"Total unique cb_videos needed: {len(cb_videos_needed)}")
+
+
+# ===== Find video files =====
+# Videos come from two sources:
+#   - data/CameraBench/videos/ (from videos.zip, train split)
+#   - data/CameraBench/data/videos/ (from Videos4CameraBnech, test split)
+video_search_dirs = [
+    camerabench_data_root / "videos",
+    camerabench_data_root / "data" / "videos",
+]
+
+cb_video_paths = {}
+for cb_video in cb_videos_needed:
+    for search_dir in video_search_dirs:
+        video_path = search_dir / f"{cb_video}.mp4"
+        if video_path.exists():
+            cb_video_paths[cb_video] = video_path
+            break
+
+found = len(cb_video_paths)
+missing = len(cb_videos_needed) - found
+print(f"Found {found} video files, missing {missing}")
+if missing > 0:
+    not_found = cb_videos_needed - set(cb_video_paths.keys())
+    print(f"Missing videos: {list(not_found)[:10]}...")
+
+
+# ===== Export videos at target resolution =====
+exported = 0
+skipped = 0
+filtered = 0
+processed_meta = []
+
+for cb_video, video_path in tqdm(sorted(cb_video_paths.items()), desc="Exporting videos"):
+    output_video = output_root / "videos" / f"{cb_video}.mp4"
+
+    # Check video properties
+    cap = cv2.VideoCapture(str(video_path))
+    if not cap.isOpened():
+        tqdm.write(f"  Failed to open {video_path}, skipping")
+        skipped += 1
+        continue
+    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    cap.release()
+
+    num_frames_sampled = int(num_frames / fps * target_fps) if fps > 0 else 0
+
+    # Filter
+    if num_frames_sampled < target_frames:
+        tqdm.write(f"  {cb_video}: filtered by frames {num_frames_sampled} < {target_frames}")
+        filtered += 1
+        continue
+    if h < target_height or w < target_width:
+        tqdm.write(f"  {cb_video}: filtered by resolution {w}x{h}")
+        filtered += 1
+        continue
+
+    # Export
+    if not output_video.exists():
+        output_video.parent.mkdir(parents=True, exist_ok=True)
+        in_stream = ffmpeg.input(str(video_path))
+        v = (
+            in_stream.video
+            .filter('fps', fps=target_fps)
+            .filter('scale',
+                    f'if(gt(a,{target_width}/{target_height}),-2,{target_width})',
+                    f'if(gt(a,{target_width}/{target_height}),{target_height},-2)')
+            .filter('crop',
+                    target_width, target_height,
+                    f'(in_w-{target_width})/2', f'(in_h-{target_height})/2')
+        )
+        out = ffmpeg.output(
+            v,
+            str(output_video),
+            vcodec='libx264',
+            pix_fmt='yuv420p',
+            r=target_fps,
+            vframes=target_frames,
+        )
+        ffmpeg.run(out, overwrite_output=True, quiet=True)
+
+    processed_meta.append({
+        "path": f"videos/{cb_video}.mp4",
+        "video": cb_video,
+    })
+    exported += 1
+
+print(f"Exported: {exported}, skipped: {skipped}, filtered: {filtered}")
+
+# Save processed metadata for both splits (combined since vipe processes all at once)
+for split in ["train", "test"]:
+    jsonl_file = output_root / f"processed_{split}.jsonl"
+    with jsonlines.open(jsonl_file, "w") as writer:
+        writer.write_all(processed_meta)
+    print(f"Saved {jsonl_file}")

UCPE/tools/export_figure.py

@@ -0,0 +1,192 @@
+import json
+import pandas as pd
+from pathlib import Path
+from typing import Optional
+import tyro
+from pydantic import BaseModel
+import numpy as np
+from PIL import Image
+import shutil
+from tqdm.auto import tqdm
+from torch.utils.data import DataLoader
+from src.dataset import PanShotDataset, Re10kDataset, DemoDataset
+from einops import rearrange
+import src.camera_control as ucpe
+import torch
+import imageio
+from matplotlib.backends.backend_agg import FigureCanvasAgg
+import matplotlib.pyplot as plt
+
+
+class Args(BaseModel):
+    methods: list[str]
+    data: str = "DemoDataset"
+    num_frames: int = 81
+    data_root: Path = Path("data/UCPE")
+    panshot_data_root: Path = Path("data/UCPE")
+    re10k_data_root: Path = Path("data/RealEstate10k")
+    input_file: Path = Path("demo/teaser.json")
+    num_workers: int = 2
+    zero_first_yaw: bool = True
+    output_dir: Path = Path("outputs/figures")
+    sample_frames: Optional[int] = 4
+    padding: int = 25
+    quality: int = 95
+    video_ids: Optional[list[str]] = None
+    fps: int = 16
+    animate_latup: bool = False
+
+
+def collate_fn(samples):
+    data = samples[0]
+    return data
+
+
+def prepare_dataloader(args, result_root=None, video_ids=None):
+    dataset_class = globals().get(args.data, None)
+    dataset = dataset_class(args, "test", load_keys=["pose", "result"], result_root=result_root, video_ids=video_ids)
+    dataloader = DataLoader(
+        dataset,
+        collate_fn=collate_fn,
+        batch_size=1,
+        num_workers=args.num_workers,
+        shuffle=False,
+    )
+    return dataloader
+
+
+def main():
+    args = tyro.cli(Args)
+
+    dataloaders = {}
+    for method, result in zip(args.methods[::2], args.methods[1::2]):
+        dataloader = prepare_dataloader(args, result, args.video_ids)
+        dataloaders[method] = dataloader
+
+    if args.video_ids is None:
+        args.video_ids = None
+        for dataloader in dataloaders.values():
+            video_ids = set()
+            for meta in dataloader.dataset.metas:
+                video_ids.add(meta["video_id"])
+            args.video_ids = video_ids if args.video_ids is None else video_ids & args.video_ids
+        args.video_ids = list(args.video_ids)
+        dataloaders = {}
+        for method, result in zip(args.methods[::2], args.methods[1::2]):
+            dataloader = prepare_dataloader(args, result, args.video_ids)
+            dataloaders[method] = dataloader
+
+    print(f"Found {len(args.video_ids)} videos to process.")
+
+    for datas in tqdm(zip(*dataloaders.values()), desc=f"Exporting figures", total=len(dataloader)):
+        data = datas[0]
+        video_id = data["video_id"]
+        result_id = Path(data["result_path"]).stem
+
+        # Save prompts
+        prompt_dir = args.output_dir / "prompts"
+        prompt_path = prompt_dir / f"{video_id}.txt"
+        if not prompt_path.exists():
+            prompt_dir.mkdir(parents=True, exist_ok=True)
+            with open(prompt_path, "w") as f:
+                f.write(data["caption"])
+
+        # Save Lat-up map visualization
+        lat_up_dir = args.output_dir / "lat_up_map"
+        lat_up_path = lat_up_dir / f"{video_id}.png"
+        if not lat_up_path.exists():
+            rot = torch.from_numpy(data["pose"][..., :3, :3]).float()  # [T, 3, 3]
+            rot = rot[:1].unsqueeze(0)  # [B=1, 1, 3, 3]
+            up_map, lat_map = ucpe.compute_up_lat_map(
+                R=rot,  # [B, T, 3, 3]
+                x_fov=torch.tensor(data["x_fov"]).float().unsqueeze(0),  # [B=1, 1]
+                xi=torch.tensor(data["xi"]).float().unsqueeze(0),  # [B=1, 1]
+                height=30,
+                width=52,
+            )
+
+            lat_up_dir.mkdir(parents=True, exist_ok=True)
+            ucpe.visualize_up_lat_map(
+                up_map[0, 0],
+                lat_map[0, 0],
+                str(lat_up_path),
+            )
+            tqdm.write(f"Saved lat-up map to {lat_up_path}")
+
+        lat_up_dir = args.output_dir / "lat_up_video"
+        lat_up_path = lat_up_dir / f"{video_id}.mp4"
+        if args.animate_latup and not lat_up_path.exists():
+            rot = torch.from_numpy(data["pose"][..., :3, :3]).float()  # [T, 3, 3]
+            rot = rot.unsqueeze(0)  # [B=1, T, 3, 3]
+            up_map, lat_map = ucpe.compute_up_lat_map(
+                R=rot,  # [B, T, 3, 3]
+                x_fov=torch.tensor(data["x_fov"]).float().unsqueeze(0),  # [B=1, 1]
+                xi=torch.tensor(data["xi"]).float().unsqueeze(0),  # [B=1, 1]
+                height=30,
+                width=52,
+            )
+
+            lat_up_dir.mkdir(parents=True, exist_ok=True)
+            writer = imageio.get_writer(lat_up_path, fps=args.fps)
+            for up, lat in zip(up_map[0], lat_map[0]):
+                fig = ucpe.visualize_up_lat_map(up, lat)
+                canvas = FigureCanvasAgg(fig)
+                canvas.draw()
+                buf = canvas.buffer_rgba()
+                img = np.asarray(buf, dtype=np.uint8)
+                writer.append_data(img[:, :, :3])
+                plt.close(fig)
+            writer.close()
+            tqdm.write(f"Saved lat-up video to {lat_up_path}")
+
+        grid_image_path = args.output_dir / "grid" / f"{result_id}.jpg"
+        if not grid_image_path.exists():
+            frame_methods = []
+            H, W = datas[-1]["result"].shape[2:4]
+            for method, data in zip(dataloaders.keys(), datas):
+                frames = data["result"]
+                frames = (frames + 1.0) / 2.0 * 255.0
+                frames = frames.astype(np.uint8)
+                frames = rearrange(frames, "C T H W -> T H W C")  # (T, H, W, 3)
+                total_frames = len(frames)
+                if args.sample_frames < total_frames:
+                    frame_indices = np.linspace(0, total_frames - 1, args.sample_frames, dtype=int)
+                else:
+                    frame_indices = np.arange(total_frames)
+                frames = frames[frame_indices]  # (sample_frames, H, W, 3)
+
+                # Save frames as images
+                output_frames_dir = args.output_dir / method / result_id
+                output_frames_dir.mkdir(parents=True, exist_ok=True)
+                scaled = []
+                for i, frame in enumerate(frames):
+                    frame_path = output_frames_dir / f"{i}.jpg"
+                    frame = Image.fromarray(frame)
+                    frame.save(frame_path, quality=args.quality)
+                    frame = frame.resize((W, H), Image.LANCZOS)
+                    scaled.append(frame)
+
+                frame_methods.append(scaled)
+
+            # Save frames as a grid image
+            grid_image = Image.new(
+                'RGB',
+                (
+                    W * len(frame_indices) + args.padding * (len(frame_indices) - 1),
+                    H * len(frame_methods) + args.padding * (len(frame_methods) - 1),
+                ),
+                (255, 255, 255)
+            )
+            for j, frames in enumerate(frame_methods):
+                for i, frame in enumerate(frames):    
+                    grid_image.paste(frame, (
+                        i * (W + args.padding),
+                        j * (H + args.padding),
+                    ))
+            grid_image_path.parent.mkdir(parents=True, exist_ok=True)
+            grid_image.save(grid_image_path, quality=args.quality)
+            tqdm.write(f"Saved grid image to {grid_image_path}")
+
+
+if __name__ == "__main__":
+    main()

UCPE/tools/export_table.py

@@ -0,0 +1,261 @@
+import json
+import pandas as pd
+from pathlib import Path
+from typing import Optional
+import tyro
+from pydantic import BaseModel
+
+
+class Args(BaseModel):
+    methods: list[str]
+    metrics: Optional[list[str]] = None
+    pad_cols: int = 0
+
+
+metric_infos = {
+    "video_metrics/rho_pred": {
+        "name": "${\\rho}_{\\text{A}}$",
+        "group": "Camera Lens Control",
+        "higher_is_better": True,
+        "scale": 100,
+        "decimal_places": 2,
+    },
+    "video_metrics/rho_gt": {
+        "name": "${\\rho}_{\\text{A-gt}}$",
+        "group": "Camera Lens Control",
+        "higher_is_better": True,
+        "scale": 100,
+        "decimal_places": 2,
+    },
+    "video_metrics/vfov_err": {
+        "name": "FoV (°)",
+        "group": "Camera Lens Control",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "video_metrics/k1_err": {
+        "name": "${k}_{1}$",
+        "group": "Camera Lens Control",
+        "higher_is_better": False,
+        "decimal_places": 3,
+    },
+    "video_metrics/k2_err": {
+        "name": "${k}_{2}$",
+        "group": "Camera Lens Control",
+        "higher_is_better": False,
+        "decimal_places": 3,
+    },
+    "video_metrics/pitch_err": {
+        "name": "Pitch (°)",
+        "group": "Absolute Orientation",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "video_metrics/roll_err": {
+        "name": "Roll (°)",
+        "group": "Absolute Orientation",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "video_metrics/gravity_err": {
+        "name": "Gravity (°)",
+        "group": "Absolute Orientation",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "video_metrics/latitude_err": {
+        "name": "Latitude (°)",
+        "group": "Absolute Orientation",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "video_metrics/up_err": {
+        "name": "Up (°)",
+        "group": "Absolute Orientation",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "video_metrics/lpips": {
+        "name": "LPIPS",
+        "group": "Relative Camera Pose Control",
+        "higher_is_better": False,
+        "decimal_places": 3,
+    },
+    "video_metrics/psnr": {
+        "name": "PSNR",
+        "group": "Relative Camera Pose Control",
+        "higher_is_better": True,
+        "decimal_places": 2,
+    },
+    "video_metrics/ssim": {
+        "name": "SSIM",
+        "group": "Relative Camera Pose Control",
+        "higher_is_better": True,
+        "decimal_places": 3,
+    },
+    "pose/rot_err": {
+        "name": "RotErr (°)",
+        "group": "Relative Camera Pose Control",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "pose/trans_err": {
+        "name": "TransErr",
+        "group": "Relative Camera Pose Control",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "pose/cammc": {
+        "name": "CamMC",
+        "group": "Relative Camera Pose Control",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "pose/rot_err_vipe": {
+        "name": "RotErr - Vipe (°)",
+        "group": "Relative Camera Pose Control",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "pose/trans_err_vipe": {
+        "name": "TransErr - Vipe",
+        "group": "Relative Camera Pose Control",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "pose/cammc_vipe": {
+        "name": "CamMC - Vipe",
+        "group": "Relative Camera Pose Control",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "video_metrics/fvd_center": {
+        "name": "FVD-center",
+        "group": "Video Generation Quality",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "video_metrics/fvd": {
+        "name": "FVD",
+        "group": "Video Generation Quality",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "video_metrics/fid": {
+        "name": "FID",
+        "group": "Video Generation Quality",
+        "higher_is_better": False,
+        "decimal_places": 2,
+    },
+    "video_metrics/cs_text": {
+        "name": "CLIP",
+        "group": "Video Generation Quality",
+        "higher_is_better": True,
+        "decimal_places": 2,
+    },
+    "video_metrics/cs_image": {
+        "name": "CLIP-image",
+        "group": "Video Generation Quality",
+        "higher_is_better": True,
+        "decimal_places": 2,
+    },
+    "video_metrics/is": {
+        "name": "IS",
+        "group": "Video Generation Quality",
+        "higher_is_better": True,
+        "decimal_places": 2,
+        "std_dev": "video_metrics/is_std"
+    },
+    "qalign/image_quality": {
+        "name": "Image Quality",
+        "group": "Video Generation Quality",
+        "higher_is_better": True,
+        "decimal_places": 4,
+    },
+    "qalign/image_aesthetic": {
+        "name": "Image Aesthetic",
+        "group": "Video Generation Quality",
+        "higher_is_better": True,
+        "decimal_places": 4,
+    },
+    "qalign/video_quality": {
+        "name": "Video Quality",
+        "group": "Video Generation Quality",
+        "higher_is_better": True,
+        "decimal_places": 4,
+    },
+}
+
+color_cells = [
+    "firstcell",
+    "secondcell",
+    "thirdcell",
+]
+
+
+def main():
+    args = tyro.cli(Args)
+
+    rows = []
+    metrics = metric_infos.keys() if args.metrics is None else [
+        metric for metric in metric_infos.keys() if metric in args.metrics
+    ]
+    for method, result in zip(args.methods[::2], args.methods[1::2]):
+        if result == Path(""):
+            rows.append({
+                "Method": method,
+                **{metric_infos[metric]["name"]: None for metric in metrics}
+            })
+            continue
+
+        with open(result, "r") as f:
+            data = json.load(f)
+
+        row = {"Method": method}
+        for metric in metrics:
+            metric_name = metric_infos[metric]["name"]
+            row[metric_name] = data[metric] * metric_infos[metric].get("scale", 1.)
+        rows.append(row)
+    df = pd.DataFrame(rows)
+    df = df.round({metric_infos[metric]["name"]: metric_infos[metric]["decimal_places"] for metric in metrics})
+    print(df)
+
+    for metric_info in metric_infos.values():
+        col = metric_info["name"]
+        if col not in df.columns:
+            continue
+
+        ranks = df[col].dropna().rank(
+            method="dense",
+            ascending=not metric_info["higher_is_better"],
+        ).astype(int) - 1
+
+        df[col] = df.apply(
+            lambda row: (f"\\{color_cells[ranks[row.name]]} " if pd.notna(row[col]) and ranks[row.name] < len(color_cells) else "") +
+                (f"{row[col]:.{metric_info['decimal_places']}f}" if pd.notna(row[col]) else ""),
+            axis=1
+        )
+
+    for i in range(args.pad_cols):
+        df.insert(loc=i, column=i, value="")
+
+    tuples = [("", "")] * args.pad_cols + [("", "Method")]
+    for metric in metrics:
+        info = metric_infos[metric]
+        group = info["group"]
+        name_with_arrow = info["name"] + ("$\\uparrow$" if info["higher_is_better"] else "$\\downarrow$")
+        tuples.append((group, name_with_arrow))
+    df.columns = pd.MultiIndex.from_tuples(tuples, names=["Group", "Metric"])
+
+    latex_table = df.to_latex(
+        index=False,
+        multicolumn=True,
+        multicolumn_format="c",
+        multirow=True,
+        column_format="r" * args.pad_cols + "l" + "c" * (len(df.columns) - 1 - args.pad_cols),
+    )
+    print(latex_table)
+
+
+if __name__ == "__main__":
+    main()

UCPE/tools/extract_camerabench_poses.py

@@ -0,0 +1,67 @@
+"""
+Extract CameraBench gravity-aligned poses from vipe output + geocalib.
+This is the first part of align_panflow.py, extracted to avoid needing PanFlow data.
+"""
+
+from pathlib import Path
+import numpy as np
+import jsonlines
+from tqdm.auto import tqdm
+from einops import repeat
+
+camerabench_root = Path("data/UCPE/CameraBench")
+
+# Load geocalib data
+cb_geocalib_file = camerabench_root / "geocalib.jsonl"
+with jsonlines.open(cb_geocalib_file, "r") as reader:
+    cb_geocalib = {obj["video"]: obj for obj in reader}
+
+# Load CameraBench processed metadata (both splits have same content)
+cb_meta_file = camerabench_root / "processed_train.jsonl"
+with jsonlines.open(cb_meta_file, "r") as reader:
+    cb_meta_all = list(reader)
+
+cb_R_gravity = []
+cb_poses = []
+cb_meta = []
+static_words = ["static", "still", "stationary", "fixed", "no motion", "no camera motion"]
+
+for obj in tqdm(cb_meta_all, desc="Loading CameraBench poses"):
+    obj["video"] = Path(obj["path"]).stem
+    video_id = obj["video"]
+
+    pose_file = camerabench_root / "vipe" / "pose" / f"{video_id}.npz"
+    if not pose_file.exists():
+        tqdm.write(f"Pose file not found: {pose_file}, skipping.")
+        continue
+
+    if video_id not in cb_geocalib:
+        tqdm.write(f"Geocalib not found for {video_id}, skipping.")
+        continue
+
+    cb_meta.append(obj)
+    pose = np.load(pose_file)["data"]  # (T, 4, 4)
+    cb_poses.append(pose)
+    cb_R_gravity.append(cb_geocalib[video_id]["R"])
+
+print(f"Loaded {len(cb_poses)} / {len(cb_meta_all)} CameraBench poses.")
+cb_poses = np.array(cb_poses)  # (N, T, 4, 4)
+cb_R_gravity = np.array(cb_R_gravity)  # (N, 3, 3)
+
+# Normalize: first frame at origin
+cb_w2c0 = np.linalg.inv(cb_poses[:, 0])  # (N, 4, 4)
+cb_poses_origin = cb_w2c0[:, None] @ cb_poses  # (N, T, 4, 4)
+
+# Apply gravity alignment
+cb_T_gravity = repeat(np.eye(4), 'h w -> n h w', n=cb_R_gravity.shape[0])
+cb_T_gravity[:, :3, :3] = cb_R_gravity
+cb_poses_gravity = cb_T_gravity[:, None, :, :] @ cb_poses_origin
+
+# Save poses
+cb_pose_root = camerabench_root / "pose"
+cb_pose_root.mkdir(parents=True, exist_ok=True)
+for i, obj in enumerate(cb_meta):
+    out_file = cb_pose_root / f"{obj['video']}.npy"
+    np.save(out_file, cb_poses_gravity[i])
+
+print(f"Saved {len(cb_meta)} pose files to {cb_pose_root}")

UCPE/tools/filter_camerabench.py

@@ -0,0 +1,252 @@
+from vllm import LLM, SamplingParams
+
+from pathlib import Path
+import jsonlines
+from tqdm.auto import tqdm
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import os
+import json
+import re
+
+from transformers import AutoProcessor
+from qwen_vl_utils import process_vision_info
+
+
+# -----------------------------
+# basic configuration
+# -----------------------------
+data_root = Path("data/UCPE/CameraBench")
+videos = sorted((data_root / "videos").glob("*.mp4"))  # 直接遍历视频文件
+output_jsonl = data_root / f"filtered.jsonl"
+
+model_id = "Qwen/Qwen2.5-VL-7B-Instruct"   # try 7B first; switch to 32B if resources allow
+# model_id = "chancharikm/qwen2.5-vl-7b-cam-motion-preview"
+nframes = 32                                # hint for frame sampling inside qwen_vl_utils
+fps_hint = None                             # None or a small integer like 1/2/4 (optional)
+batch_size = 8                              # how many videos per vLLM.generate batch
+max_workers = min(8, os.cpu_count() or 4)   # 线程数按机器调整
+inflight_limit = batch_size * 2  # 同时在制的样本上限
+print(f"Using max_workers={max_workers}, inflight_limit={inflight_limit}")
+
+max_new_tokens = 512
+temperature = 0.2
+top_p = 0.9
+repetition_penalty = 1.05
+gpu_memory_utilization = 0.9
+tensor_parallel_size = 1
+limit_mm_per_prompt = {"video": 1}
+
+prompt_text = """
+You are a video filtering assistant.
+Your task is to analyze the given panoramic video and decide whether it meets certain quality and format requirements.
+Check the following conditions and output results strictly in JSON format, with boolean values (true/false) for each label:
+
+- non_fullscreen_or_black_borders: true if the video is not full-screen, has black borders, or is vertical instead of wide.
+- has_subtitles_or_watermarks: true if subtitles, captions, or watermarks are visible.
+- is_cartoon_or_flat_style: true if the video is animated, cartoon-like, 2D flat style, or non-photorealistic.
+
+Important:
+– Output strictly in JSON format.
+– Do not add extra text or explanation.
+– Each key must exist, with true/false values.
+
+Example output:
+{
+  "non_fullscreen_or_black_borders": false,
+  "has_subtitles_or_watermarks": true,
+  "is_cartoon_or_flat_style": false
+}
+"""
+
+
+# ===================================================================
+# 从这里开始：重写为“VLM 布尔标签 + JSONL 写出”的完整流程
+# ===================================================================
+
+def build_llm_input(video_path: Path, processor: AutoProcessor):
+    """构造单条输入，包含视频与文本 prompt。"""
+    video_item = {"type": "video", "video": str(video_path), "nframes": nframes}
+    if fps_hint is not None:
+        video_item["fps"] = fps_hint
+
+    messages = [
+        {"role": "system", "content": "You are a helpful video filtering assistant."},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt_text},
+                video_item
+            ]
+        }
+    ]
+
+    # 预处理多模态输入（抽帧/张量装配）
+    _, video_inputs = process_vision_info(messages)
+
+    # 模板化为可生成的字符串
+    prompt = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    mm_data = {}
+    if video_inputs is not None:
+        mm_data["video"] = video_inputs
+
+    return {"prompt": prompt, "multi_modal_data": mm_data}
+
+
+def prepare_one(vpath: Path, processor: AutoProcessor):
+    """单样本：基于固定过滤 prompt → 组装 vLLM 输入。"""
+    llm_in = build_llm_input(vpath, processor)
+    obj = {"path": vpath.relative_to(data_root)}
+    return obj, llm_in
+
+
+# -----------------------------
+# init vLLM + processor
+# -----------------------------
+llm = LLM(
+    model=model_id,
+    tensor_parallel_size=tensor_parallel_size,
+    gpu_memory_utilization=gpu_memory_utilization,
+    # enforce_eager=True,
+    limit_mm_per_prompt=limit_mm_per_prompt,
+)
+processor = AutoProcessor.from_pretrained(model_id)
+sampling_params = SamplingParams(
+    max_tokens=max_new_tokens,
+    temperature=temperature,
+    top_p=top_p,
+    repetition_penalty=repetition_penalty,
+)
+
+
+# -----------------------------
+# utils
+# -----------------------------
+def clean_json_output(text: str) -> str:
+    """清理 Qwen 输出中的 ```json / ``` 包裹，并尝试只保留第一段 JSON。"""
+    s = text.strip()
+
+    # 去除 Markdown 代码围栏
+    s = re.sub(r"```(?:json)?", "", s, flags=re.IGNORECASE).strip()
+    s = s.replace("```", "").strip()
+
+    # 尝试截取最外层花括号的 JSON 片段
+    # 找到第一个 '{' 与最后一个 '}' 之间的内容
+    first = s.find("{")
+    last = s.rfind("}")
+    if first != -1 and last != -1 and last > first:
+        s = s[first:last + 1].strip()
+    return s
+
+
+# -----------------------------
+# 推理与写出
+# -----------------------------
+output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+prepared_buffer = []  # 缓存已准备好的 (obj, llm_in)
+
+
+def infer_and_flush(buffer, writer):
+    """对 buffer 中的若干样本推理，并写出结果（VLM 布尔标签 JSON）。"""
+    if not buffer:
+        return
+    batch_objs = [it[0] for it in buffer]
+    batch_inputs = [it[1] for it in buffer]
+
+    gens = llm.generate(batch_inputs, sampling_params)
+
+    for ob, g in zip(batch_objs, gens):
+        # vLLM 生成的第一个候选
+        text_out = (g.outputs[0].text if g.outputs and g.outputs[0].text is not None else "").strip()
+
+        try:
+            cleaned = clean_json_output(text_out)
+            parsed = json.loads(cleaned)
+
+            # 写出：path + 模型布尔标签
+            writer.write({
+                "path": str(ob["path"]),
+                "filter": parsed
+            })
+
+        except Exception as e:
+            # 输出无法解析成 JSON：打印日志并跳过
+            print(f"[skip] JSON parse failed for {ob['path']}: {e}")
+            print(f"Raw output: {text_out}")
+            continue
+
+
+# -----------------------------
+# pipeline：边准备边推理边写出（动态提交 + 行缓冲）
+# -----------------------------
+# 基础检查
+for v in videos:
+    assert v.exists(), f"File not found: {v}"
+
+try:
+    # 用行缓冲打开文件，便于“边写边可见”
+    f = open(output_jsonl, "w", buffering=1, encoding="utf-8")
+    writer = jsonlines.Writer(f)
+
+    with ThreadPoolExecutor(max_workers=max_workers) as ex:
+        pbar = tqdm(total=len(videos), desc="preparing & inferring")
+
+        # 动态 pending 集合
+        pending = set()
+        i_submit = 0
+
+        # 先填满 in-flight
+        while i_submit < len(videos) and len(pending) < inflight_limit:
+            fut = ex.submit(prepare_one, videos[i_submit], processor)
+            pending.add(fut)
+            i_submit += 1
+
+        # 循环直到所有任务完成
+        while pending:
+            # 只等待当前 pending 集合中的任务
+            for fut in as_completed(list(pending), timeout=None):
+                pending.remove(fut)
+                obj, llm_in = fut.result()
+                prepared_buffer.append((obj, llm_in))
+                pbar.update(1)
+
+                # 满一批就立刻推理并清空对应部分
+                if len(prepared_buffer) >= batch_size:
+                    infer_and_flush(prepared_buffer[:batch_size], writer)
+                    prepared_buffer = prepared_buffer[batch_size:]
+
+                # 补交新任务，保持 in-flight 数量
+                while i_submit < len(videos) and len(pending) < inflight_limit:
+                    fut_new = ex.submit(prepare_one, videos[i_submit], processor)
+                    pending.add(fut_new)
+                    i_submit += 1
+
+                # 跳出到 while pending，重新评估 pending 集合（已更新）
+                break
+
+        # 把“尾巴”按 batch 循环清空，确保不丢最后一个或多个 batch
+        while prepared_buffer:
+            chunk = prepared_buffer[:batch_size]
+            infer_and_flush(chunk, writer)
+            prepared_buffer = prepared_buffer[len(chunk):]
+
+        pbar.close()
+finally:
+    # 关闭 writer / 文件句柄
+    try:
+        writer.close()
+    except Exception:
+        pass
+    try:
+        f.close()
+    except Exception:
+        pass
+    # 优雅关闭 vLLM 引擎
+    try:
+        llm.shutdown()
+    except Exception:
+        pass
+
+print(f"done. VLM labels saved to: {output_jsonl}")

UCPE/tools/filter_panflow.py

@@ -0,0 +1,254 @@
+from vllm import LLM, SamplingParams
+
+from pathlib import Path
+from tqdm.auto import tqdm
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import os
+import json
+import re
+from decord import VideoReader, cpu
+import numpy as np
+import torch
+from PIL import Image
+ngpus = torch.cuda.device_count()
+
+from transformers import AutoProcessor
+from qwen_vl_utils import process_vision_info
+
+
+# -----------------------------
+# basic configuration
+# -----------------------------
+panflow_root = Path("data/360-1M")
+panshot_root = Path("data/UCPE")
+output_root = panshot_root / "PanFlow" / "filtered"
+output_root.mkdir(parents=True, exist_ok=True)
+
+model_id = "Qwen/Qwen2.5-VL-7B-Instruct"   # try 7B first; switch to 32B if resources allow
+batch_size = 16 * ngpus                             # how many videos per vLLM.generate batch
+max_workers = min(4 * ngpus, os.cpu_count() or 4)
+inflight_limit = max_workers * 2
+print(f"Using max_workers={max_workers}, inflight_limit={inflight_limit}")
+
+max_new_tokens = 512
+temperature = 0.2
+top_p = 0.9
+repetition_penalty = 1.05
+gpu_memory_utilization = 0.9
+tensor_parallel_size = ngpus
+limit_mm_per_prompt = {"image": 1}
+
+meta_root = panflow_root / "meta"
+meta_files = list((meta_root).glob("*.json"))
+meta_files.sort()
+print(f"Found {len(meta_files)} PanFlow meta files.")
+
+existing_out = list(output_root.glob("*.json"))
+existing_ids = {f.stem for f in existing_out}
+meta_files = [f for f in meta_files if f.stem not in existing_ids]
+print(f"{len(meta_files)} files to process after skipping existing.")
+
+prompt_text = """
+You are a video understanding assistant specialized in analyzing panoramic ERP-format videos.  
+Given one frame of a panoramic video, your tasks are:  
+
+1. **Filtering**: Identify if the video should be filtered out.  
+Output boolean flags for the following conditions (true if the issue exists, false otherwise):  
+
+- non_ERP_format: The video is **not in ERP (Equirectangular Projection)** panoramic format. For example, if the video looks like a flat perspective, fisheye, cube-map, or any projection other than ERP, set this to true.  
+
+- has_subtitle_or_watermark: The video contains **text overlays, subtitles, logos, or watermarks**. Look carefully for visible text at the bottom, center, or corners of the video. If such elements are present and not part of the real scene, set this to true.  
+
+- edge_missing: The top or bottom edges of the ERP panorama are **cut off, blacked out, cropped, or covered by logos/watermarks**, so the full 360° vertical coverage is missing or obstructed. If you cannot clearly see the poles (sky/ground) or if the edges are hidden by overlays, set this to true.  
+
+- has_overlay: The frame contains **artificial overlays**, such as embedded UI elements, pop-up graphics, stickers, video-in-video inserts, menus, or other synthetic elements that are not part of the natural scene. If you see signs of AR/VR interface, streaming UI, or added images, set this to true.  
+
+- low_quality: The video is of **poor visual quality**, such as being blurry, noisy, heavily pixelated, very low resolution, or distorted in a way that prevents recognizing the scene. If the content is hard to interpret due to quality issues, set this to true.  
+
+- unnatural_content: The video contains **cartoons, animations, CGI, synthetic 3D renderings, or game engine graphics** rather than real-world panoramic footage. If the content is not realistic, set this to true.  
+
+2. **POI Categorization**: From the provided list of categories, select **one or more most relevant** labels that best describe the scene.  
+   Only use the given categories, do not invent new ones.  
+
+---
+
+**poi_category list (choose only from below):**
+
+Restaurant, Coffee-Shop, Bars-and-Pubs, Residential-area, Hotels-Motels, Vaccation-Rentals, Hospitals-Clinics, Pharmacies, Dentists, School-Universities, Library, Supermarkets, Shopping-Malls, Clothing-Stores, Shoe-Stores, Bookstores, Flowerstore, Furniture-Stores, Electorical-Store, Pet-Store, Toy-Shop, Airports, Train-Stations, Bus-Stops, Gas-Station, Car-Rental-Agencies, Theaters, Concert-Halls, Sports-Stadiums, Parks-and-Recreation-Areas, Museums, Art-Galleries, Zoos-Aquariums, Botanical-Gardens, Landmarks, Cultural-Centers, Post-Offices, Police-Stations, Courthouses, CityHalls, Banks-ATMs, Events-Conferences-halls, Beaches, Hiking-Trails, Campgrounds, Lakes, Mountains, Forest-Mountains, Farms, Street-View, Square, Business-Centers, Tech-Companies, Co-working-Spaces, Gyms-and-Fitness-Centers, Sports-Clubs, Swimming-Pools, Tennis-Courts, Auto-Repair-Shops, Car-Washes, Parking-Lots, Churches, Mosques, Temples, Graveyards.
+
+---
+
+**Output strictly in JSON format** as follows:
+
+```json
+{
+  "filter": {
+    "non_ERP_format": false,
+    "has_subtitle_or_watermark": false,
+    "edge_missing": false,
+    "has_overlay": false,
+    "low_quality": false,
+    "unnatural_content": false
+  },
+  "poi_category": ["Mountains"]
+}
+```
+"""
+
+# ================================================================
+# utils
+# ================================================================
+
+def clean_json_output(text: str) -> str:
+    s = text.strip()
+    s = re.sub(r"```(?:json)?", "", s, flags=re.IGNORECASE).strip()
+    s = s.replace("```", "").strip()
+    first = s.find("{")
+    last = s.rfind("}")
+    if first != -1 and last != -1 and last > first:
+        s = s[first:last + 1].strip()
+    return s
+
+
+def process_one_video(meta_file, processor):
+    """子进程：为一个视频的 clips 抽中间帧，生成 LLM 输入"""
+    pf_meta = json.load(open(meta_file))
+    video_id = pf_meta.get("video_id", Path(meta_file).stem)
+    video_path = panflow_root / "videos" / f"{video_id}.mp4"
+
+    if "slam_clips" not in pf_meta or "clips" not in pf_meta["slam_clips"]:
+        print(f"[skip] No slam_clips in meta: {meta_file}")
+        return None
+
+    if not video_path.exists():
+        print(f"[skip] Video file not found: {video_path}")
+        return None
+
+    try:
+        vr = VideoReader(str(video_path), ctx=cpu(0), num_threads=1)
+    except Exception as e:
+        print(f"[skip] Failed to open video {video_path}: {e}")
+        return None
+
+    clip_infos, llm_inputs = [], []
+
+    for clip in pf_meta["slam_clips"]["clips"]:
+        start, end = clip["frames"][0], clip["frames"][-1]
+        mid = (start + end) // 2
+
+        frame = vr[mid].asnumpy().astype(np.uint8)
+        frame_pil = Image.fromarray(frame)
+
+        image_item = {"type": "image", "image": frame_pil}
+        messages = [
+            {"role": "system", "content": "You are a helpful video filtering assistant."},
+            {"role": "user", "content": [{"type": "text", "text": prompt_text}, image_item]},
+        ]
+
+        image_inputs, _ = process_vision_info(messages)
+        prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+        mm_data = {}
+        if image_inputs is not None:
+            mm_data["image"] = image_inputs
+
+        llm_inputs.append({"prompt": prompt, "multi_modal_data": mm_data})
+        clip_infos.append(clip)
+
+    return video_id, clip_infos, llm_inputs
+
+
+# ================================================================
+# init vLLM
+# ================================================================
+llm = LLM(
+    model=model_id,
+    tensor_parallel_size=tensor_parallel_size,
+    gpu_memory_utilization=gpu_memory_utilization,
+    # enforce_eager=True,
+    limit_mm_per_prompt=limit_mm_per_prompt,
+)
+processor = AutoProcessor.from_pretrained(model_id)
+sampling_params = SamplingParams(
+    max_tokens=max_new_tokens,
+    temperature=temperature,
+    top_p=top_p,
+    repetition_penalty=repetition_penalty,
+)
+
+
+# ================================================================
+# flush results
+# ================================================================
+def flush_one_video(video_id, clip_infos, llm_inputs):
+    """对单个视频的所有 clips 批量推理并写 JSON。"""
+    video_results = []
+    total_clips = len(llm_inputs)
+
+    # 二级进度条：按 clip 数量更新
+    with tqdm(total=total_clips, desc=f"Video {video_id}", leave=False) as pbar:
+        for i in range(0, total_clips, batch_size):
+            chunk = llm_inputs[i:i + batch_size]
+            gens = llm.generate(chunk, sampling_params)
+            for clip, g in zip(clip_infos[i:i + batch_size], gens):
+                text_out = g.outputs[0].text.strip() if g.outputs else ""
+                try:
+                    cleaned = clean_json_output(text_out)
+                    parsed = json.loads(cleaned)
+                    video_results.append({
+                        "clip_id": clip["clip_id"],
+                        "clip_name": clip["clip_name"],
+                        "filter": parsed.get("filter", {}),
+                        "poi_category": parsed.get("poi_category", [])
+                    })
+                except Exception as e:
+                    print(f"[skip] JSON parse failed for {clip['clip_name']}: {e}")
+                    print(f"Raw output: {text_out}")
+                finally:
+                    pbar.update(1)  # 每处理一个 clip 就更新进度
+
+    out_file = output_root / f"{video_id}.json"
+    with open(out_file, "w", encoding="utf-8") as f:
+        json.dump(video_results, f, indent=4)
+
+
+# ================================================================
+# 主流程
+# ================================================================
+
+with ThreadPoolExecutor(max_workers=max_workers) as ex:
+    pbar = tqdm(total=len(meta_files), desc="processing videos")
+    pending = set()
+    i_submit = 0
+
+    # 初始填充 inflight
+    while i_submit < len(meta_files) and len(pending) < inflight_limit:
+        fut = ex.submit(process_one_video, meta_files[i_submit], processor)
+        pending.add(fut)
+        i_submit += 1
+
+    while pending:
+        for fut in as_completed(list(pending), timeout=None):
+            pending.remove(fut)
+            result = fut.result()
+            if result is None:
+                continue
+            video_id, clip_infos, llm_inputs = result
+            flush_one_video(video_id, clip_infos, llm_inputs)
+
+            pbar.update(1)
+
+            # 补交新任务
+            while i_submit < len(meta_files) and len(pending) < inflight_limit:
+                fut_new = ex.submit(process_one_video, meta_files[i_submit], processor)
+                pending.add(fut_new)
+                i_submit += 1
+            break
+    pbar.close()
+
+try:
+    llm.shutdown()
+except Exception:
+    pass
+
+print(f"done. Results saved to {output_root}")

UCPE/tools/geocalib_camerabench.py

@@ -0,0 +1,46 @@
+import jsonlines
+from pathlib import Path
+from tqdm.auto import tqdm
+import jsonlines
+from geocalib import GeoCalib
+from decord import VideoReader, cpu
+import torch
+from einops import rearrange
+
+
+data_root = Path("data/UCPE/CameraBench")
+videos = sorted((data_root / "videos").glob("*.mp4"))  # 直接遍历视频文件
+output_jsonl = data_root / "geocalib.jsonl"
+
+processed = set()
+if output_jsonl.exists():
+    print(f"Resuming from {output_jsonl}")
+    with jsonlines.open(output_jsonl, "r") as reader:
+        for obj in reader:
+            processed.add(obj["video"])
+    print(f"Found {len(processed)} processed videos to skip.")
+videos = [v for v in videos if v.stem not in processed]
+print(f"Total videos to process: {len(videos)}")
+
+gc = GeoCalib(weights="pinhole").cuda()
+
+with jsonlines.open(output_jsonl, "a") as writer:
+    for video in tqdm(videos, desc="Calibrating videos"):
+        vr = VideoReader(str(video), ctx=cpu(0), num_threads=1)
+        frames = vr.get_batch(range(0, len(vr))).asnumpy()
+        frames = torch.from_numpy(frames)
+        frames = rearrange(frames, "n h w c -> n c h w")
+        frames = frames.float() / 255.0
+        frames = frames.cuda()
+
+        result = gc.calibrate(frames, shared_intrinsics=True)
+        gravity = result["gravity"][0]
+        R = gravity.R.cpu().numpy().tolist()
+        roll, pitch = gravity.rp.cpu().numpy().tolist()
+
+        writer.write({
+            "video": video.stem,
+            "R": R,
+            "roll": roll,
+            "pitch": pitch,
+        })

UCPE/tools/match_panflow.py

@@ -0,0 +1,376 @@
+from pathlib import Path
+import jsonlines
+from tqdm.auto import tqdm
+from collections import defaultdict
+import matplotlib.pyplot as plt
+import ffmpeg
+import json
+import csv
+import numpy as np
+import cv2
+from einops import rearrange, repeat
+from visualize_pose import vis_to_html
+from decord import VideoReader, cpu
+import os
+import shutil
+from scipy.spatial.transform import Rotation as R
+import seaborn as sns
+import pandas as pd
+
+
+split = "train"  # "train" or "test"
+panflow_root = Path("data/360-1M")
+panshot_root = Path("data/UCPE")
+camerabench_root = panshot_root / "CameraBench"
+debug_root = Path("debug/match_panflow")
+match_cb_root = panshot_root / "PanFlow" / f"align_to_camerabench-{split}"
+match_pf_root = panshot_root / "PanFlow" / f"match_to_camerabench-{split}"
+match_pf_root.mkdir(parents=True, exist_ok=True)
+pf_clip_root = panshot_root / "PanFlow" / f"match_clips-{split}"
+pf_clip_root.mkdir(parents=True, exist_ok=True)
+summary_root = match_pf_root.parent / f"{match_pf_root.name}-summary"
+summary_root.mkdir(parents=True, exist_ok=True)
+
+max_pf_per_cb = 100 if split == "train" else 1
+max_cb_per_pf = 5 if split == "train" else 3
+matches_per_poi = 1000 if split == "train" else 1
+rotation_score_thres = 5.0
+avg_qalign_thres = 0.7
+watermark_score_thres = 0.3
+motion_score_bins = 10
+max_match_per_bin = 2000 if split == "train" else 10
+visualize_video = False
+visualize_scores = ["motion_score", "watermark_score", "avg_qalign", "rotation_score"]
+skip_undownloaded = True
+filter_cb_per_pf = True
+filter_cb_percentile = 0.8
+
+if skip_undownloaded:
+    pf_video_root = panshot_root / "PanFlow" / "videos"
+    downloaded_videos = set([p.stem for p in pf_video_root.glob("*.mp4")])
+    print(f"Found {len(downloaded_videos)} downloaded PanFlow videos.")
+
+match_meta_files = list(match_cb_root.glob("*.json"))
+match_meta_files.sort()
+print(f"Found {len(match_meta_files)} match meta files.")
+
+if split == "test":
+    pf_clip_train = panshot_root / "PanFlow" / "match_clips-train"
+    pf_clip_train = pf_clip_train.glob("*.json")
+    pf_clip_train = set([p.stem for p in pf_clip_train])
+    print(f"Found {len(pf_clip_train)} training PanFlow clips.")
+
+    print(f"Loading CameraBench poses for filtering...")
+    cb_meta_file = camerabench_root / f"processed_{split}.jsonl"
+    with jsonlines.open(cb_meta_file, "r") as reader:
+        cb_meta_all = list(reader)
+
+    cb_max_rotations = {}
+    for obj in tqdm(cb_meta_all, desc="Loading CameraBench poses"):
+        video_id = Path(obj["path"]).stem
+        pose_file = camerabench_root / "vipe" / "pose" / f"{video_id}.npz"
+        if not pose_file.exists():
+            tqdm.write(f"Pose file not found: {pose_file}, skipping.")
+            continue
+
+        pose = np.load(pose_file)["data"]  # (T, 4, 4)
+        R_all = pose[:, :3, :3]  # (T, 3, 3)
+        # 计算相邻帧之间的相对旋转
+        rel_rot = np.einsum("tij,tjk->tik", np.linalg.inv(R_all[:-1]), R_all[1:])
+        # 将相对旋转矩阵转换为旋转角度（度数）
+        rel_angles = R.from_matrix(rel_rot).magnitude() * 180.0 / np.pi
+        # 取最大旋转角
+        cb_max_rotations[video_id] = float(np.max(rel_angles))
+
+    # sort by max rotation
+    cb_max_rotations = sorted(cb_max_rotations.items(), key=lambda x: x[1])
+    num_cb_to_keep = int(len(cb_max_rotations) * filter_cb_percentile)
+    cb_videos_to_keep = set([v[0] for v in cb_max_rotations[:num_cb_to_keep]])
+    print(f"Keeping {len(cb_videos_to_keep)} CameraBench videos for testing.")
+
+
+def plot_match_histogram(summary, summary_name, summary_file):
+    fig, ax = plt.subplots(figsize=(10, 6))
+
+    items = sorted(summary.items(), key=lambda x: -x[1])
+    labels, counts = zip(*items)
+    ax.barh(labels, counts, color="salmon")
+
+    ax.set_xlabel("Match Count")
+    ax.set_ylabel("Number of Videos")
+    ax.set_title(f"{summary_name} Distribution")
+
+    plt.tight_layout()
+    fig.savefig(summary_file, dpi=300, bbox_inches="tight")
+    plt.close(fig)
+    print(f"Saved camera summary histogram to {summary_file}")
+
+
+def plot_match_histogram_category(summary, summary_name, summary_file, max_categories=30):
+    """
+    seaborn 版本的类别分布绘图：
+      - 横轴为类别
+      - 纵轴为 density（相对频率）
+      - 超过 max_categories 的类别合并为 "others"
+      - 输出 PDF（矢量图）
+    """
+
+    # ---------------------------
+    # Step 1: 排序并保留 top-K 类别
+    # ---------------------------
+    items = sorted(summary.items(), key=lambda x: -x[1])
+
+    if len(items) > max_categories:
+        top_items = items[:max_categories]
+        others_total = sum([c for _, c in items[max_categories:]])
+        top_items.append(("others", others_total))
+    else:
+        top_items = items
+
+    labels, counts = zip(*top_items)
+
+    # ---------------------------
+    # Step 2: density = count / total
+    # ---------------------------
+    total = sum(counts)
+    density = [c / total for c in counts]
+
+    # ---------------------------
+    # Step 3: seaborn 绘图
+    # ---------------------------
+    df = pd.DataFrame({
+        "category": labels,
+        "density": density,
+        "count": counts,
+    })
+
+    fig, ax = plt.subplots(figsize=(12, 6))
+    sns.barplot(
+        data=df,
+        x="category",
+        y="density",
+        ax=ax,
+        palette="viridis"
+    )
+
+    ax.set_xlabel("Category", fontsize=12)
+    ax.set_ylabel("Proportion", fontsize=12)
+    # ax.set_title(f"{summary_name} (Proportion)", fontsize=14)
+    plt.xticks(rotation=45, ha="right")
+
+    plt.tight_layout()
+
+    # ---------------------------
+    # Step 4: 保存为 PDF（矢量图）
+    # ---------------------------
+    summary_file = Path(summary_file)
+    summary_file = summary_file.with_suffix(".pdf")
+    fig.savefig(summary_file, dpi=300, bbox_inches="tight", format="pdf")
+
+    plt.close(fig)
+    print(f"[Saved PDF] {summary_file}")
+
+
+def plot_score_histogram(motion_scores, summary_file):
+    fig, ax = plt.subplots(figsize=(10, 6))
+    ax.hist(motion_scores, bins="auto", color="skyblue", edgecolor="black")
+    ax.set_xlabel("Score")
+    ax.set_ylabel("Number of Clips")
+    ax.set_title("PanFlow Clip Score Distribution")
+
+    plt.tight_layout()
+    fig.savefig(summary_file, dpi=300, bbox_inches="tight")
+    plt.close(fig)
+    print(f"Saved score histogram to {summary_file}")
+
+
+# match_meta_files = match_meta_files[:100]
+poi_category_count = defaultdict(int)
+metas = []
+for meta_file in tqdm(match_meta_files, desc=f"Loading match metadata"):
+    with open(meta_file, "r") as f:
+        meta = json.load(f)
+    meta["clips"] = [clip for clip in meta["clips"] if "matches" in clip]
+    for pf_clip in meta["clips"]:
+        for c in pf_clip["poi_category"]:
+            poi_category_count[c] += 1
+    metas.append(meta)
+        
+summary_file = summary_root / "poi_category_count.pdf"
+plot_match_histogram_category(poi_category_count, "poi_category_count", summary_file)
+num_clips = sum(poi_category_count.values())
+print(f"Found {len(poi_category_count)} POI categories, covering {num_clips} clips.")
+
+match_cands = []
+pf_clips = {}
+for meta in tqdm(metas, desc=f"Preprocessing match metadata"):
+    for pf_clip in meta["clips"]:
+        pf_key = f"{pf_clip['video_id']}-{pf_clip['clip_id']}"
+        if split == "test" and pf_key in pf_clip_train:
+            continue
+
+        pf_clips[pf_key] = {k: pf_clip[k] for k in [
+            "scores", "video_id", "clip_id", "clip_name", "poi_category", "frames"
+        ]}
+        pf_clips[pf_key]["fps"] = meta["fps"]
+        pf_clips[pf_key]["time_range"] = (
+            pf_clip["frames"][0] / meta["fps"],
+            pf_clip["frames"][-1] / meta["fps"],
+        )
+
+        scores = pf_clip["scores"]
+        for cb_clip in pf_clip["matches"]:
+            if split == "test" and cb_clip["video"] not in cb_videos_to_keep:
+                continue
+
+            if scores["watermark_score"] > watermark_score_thres:
+                continue
+            if scores["avg_qalign"] < avg_qalign_thres:
+                continue
+            
+            rotation_score = (rotation_score_thres - scores["rotation_score"]) / rotation_score_thres
+            total_score = \
+                0.3 * (1 - scores["watermark_score"]) \
+                + 0.5 * scores["avg_qalign"] \
+                + 0.2 * rotation_score
+
+            match_info = {
+                "total_score": total_score,
+                "video_id": pf_clip["video_id"],
+                "clip_id": pf_clip["clip_id"],
+                "frames": cb_clip["frames"],
+                "rmse": cb_clip["rmse"],
+                "R": cb_clip["R"],
+                "s": cb_clip["s"],
+            }
+            match_info["time_range"] = (
+                cb_clip["frames"][0] / meta["fps"],
+                cb_clip["frames"][-1] / meta["fps"],
+            )
+            match_cands.append((cb_clip["video"], match_info))
+
+print(f"Found {len(match_cands)} match candidates.")
+
+score_summary = defaultdict(list)
+for pf_clip in pf_clips.values():
+    for key in visualize_scores:
+        score_summary[key].append(pf_clip["scores"][key])
+for key, score in score_summary.items():
+    summary_file = summary_root / f"{key}_before.png"
+    plot_score_histogram(score, summary_file)
+
+motion_scores = score_summary["motion_score"]
+motion_bin_edges = np.quantile(motion_scores, np.linspace(0, 1, motion_score_bins + 1))
+motion_bin_edges = np.unique(motion_bin_edges)
+print(f"Motion score histogram bins: {motion_bin_edges}")
+
+cb_matches = defaultdict(list)
+panflow_summary = defaultdict(int)
+poi_category_summary = defaultdict(int)
+motion_bins = defaultdict(int)
+for cb_clip, match_info in tqdm(match_cands, desc=f"Analysing match metadata"):
+    pf_key = f"{match_info['video_id']}-{match_info['clip_id']}"
+    pf_clip = pf_clips[pf_key]
+    motion_score = pf_clip["scores"]["motion_score"]
+    bin_idx = np.searchsorted(motion_bin_edges, motion_score, side='right') - 1
+    bin_idx = min(bin_idx, len(motion_bin_edges) - 2)
+
+    if skip_undownloaded and match_info["video_id"] not in downloaded_videos:
+        continue
+    if motion_bins[bin_idx] >= max_match_per_bin:
+        continue
+    if pf_key in panflow_summary and panflow_summary[pf_key] >= max_cb_per_pf:
+        continue
+    if cb_clip in cb_matches and len(cb_matches[cb_clip]) >= max_pf_per_cb:
+        continue
+    if all(poi_category_summary[c] > matches_per_poi for c in pf_clip["poi_category"]):
+        continue
+
+    for c in pf_clip["poi_category"]:
+        poi_category_summary[c] += 1
+    cb_matches[cb_clip].append(match_info)
+    panflow_summary[pf_key] += 1
+    motion_bins[bin_idx] += 1
+
+camerabench_summary = {k: len(v) for k, v in cb_matches.items()}
+for summary_name, summary in [
+    ("camerabench_summary", camerabench_summary),
+    ("panflow_summary", panflow_summary),
+]:
+    summary_file = summary_root / f"{summary_name}.png"
+    plot_match_histogram(summary, summary_name, summary_file)
+
+plot_match_histogram_category(
+    poi_category_summary,
+    "poi_category_summary",
+    summary_root / "poi_category_summary.pdf"
+)
+
+print(f"{len(pf_clips)} PanFlow clips in total.")
+print(f"{len(panflow_summary)} PanFlow clips have >= 1 CameraBench match.")
+print(f"Selected {len(cb_matches)} CameraBench videos.")
+if split == "train" and filter_cb_per_pf:
+    pf_keys = {k: v for k, v in panflow_summary.items() if v >= max_cb_per_pf}
+    print(f"{len(pf_keys)} PanFlow clips have >= {max_cb_per_pf} CameraBench matches.")
+else:
+    pf_keys = panflow_summary
+pf_clips = {k: v for k, v in pf_clips.items() if k in pf_keys}
+panflow_videos = set([c["video_id"] for c in pf_clips.values()])
+print(f"Selected {len(panflow_videos)} PanFlow videos.")
+print(f"Selected {sum([len(v) for v in cb_matches.values()])} total matches.")
+
+score_summary = defaultdict(list)
+for pf_clip in pf_clips.values():
+    for key in visualize_scores:
+        score_summary[key].append(pf_clip["scores"][key])
+for key, score in score_summary.items():
+    summary_file = summary_root / f"{key}_after.png"
+    plot_score_histogram(score, summary_file)
+
+for pf_key, pf_clip in tqdm(pf_clips.items(), desc="Saving selected PanFlow clips"):
+    out_file = pf_clip_root / f"{pf_key}.json"
+    with open(out_file, "w") as f:
+        json.dump(pf_clip, f, indent=4)
+
+for cb_video, match_list in tqdm(cb_matches.items(), desc="Saving match results"):
+    match_list = [m for m in match_list if f"{m['video_id']}-{m['clip_id']}" in pf_keys]
+    out_file = match_pf_root / f"{cb_video}.json"
+    with open(out_file, "w") as f:
+        json.dump(match_list, f, indent=4)
+    # tqdm.write(f"Wrote {len(match_list)} matches to {out_file}")
+    # tqdm.write(f"Best score {match_list[0]['total_score']:.4f}, worst score {match_list[-1]['total_score']:.4f}")
+
+    if visualize_video:
+        cb_video_dir = debug_root / cb_video
+        cb_video_dir.mkdir(parents=True, exist_ok=True)
+        src_video_file = panshot_root / "CameraBench" / "videos" / f"{cb_video}.mp4"
+        tgt_video_file = cb_video_dir / "CameraBench.mp4"
+        shutil.copy2(src_video_file, tgt_video_file)
+
+        cap = cv2.VideoCapture(str(src_video_file))
+        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        cap.release()
+
+        for match in match_list:
+            pf_video_file = panflow_root / "videos" / f"{match['video_id']}.mp4"
+            vr = VideoReader(str(pf_video_file), ctx=cpu(0), num_threads=1)
+            pf_clip = pf_clips[f"{match['video_id']}/{match['clip_id']}"]
+            clip_start = pf_clip["frames"][0]
+            start_frame, end_frame = match["frames"]
+            start_frame = start_frame + clip_start
+            end_frame = end_frame + clip_start
+            sample_frames = np.linspace(start_frame, end_frame, num=frame_count)
+            sample_frames = np.round(sample_frames).astype(int)
+            clip = vr.get_batch(sample_frames).asnumpy()
+            out_clip_file = cb_video_dir / f"{match['video_id']}-{start_frame}-{end_frame}.mp4"
+
+            process = (
+                ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", s=f"{clip.shape[2]}x{clip.shape[1]}", r=fps)
+                .output(str(out_clip_file), pix_fmt="yuv420p", vcodec="libx264", r=fps, crf=23)
+                .overwrite_output()
+                .run_async(pipe_stdin=True, quiet=True)
+            )
+            process.stdin.write(clip.tobytes())
+            process.stdin.close()
+            process.wait()

UCPE/tools/normalize_panflow.py

@@ -0,0 +1,200 @@
+from pathlib import Path
+import jsonlines
+from tqdm.auto import tqdm
+from collections import defaultdict
+import matplotlib.pyplot as plt
+import ffmpeg
+import json
+import csv
+import numpy as np
+import cv2
+from einops import einsum, rearrange, repeat
+from visualize_pose import vis_to_html
+import decord
+from decord import VideoReader, cpu
+decord.bridge.set_bridge("torch")
+import os
+import shutil
+import torch
+from thirdparty.PanFlow.utils.erp_utils import transformation_to_flow
+from thirdparty.PanoFlowAPI.apis.PanoRaft import PanoRAFTAPI
+
+
+split = "train"  # "train" or "test"
+panflow_root = Path("data/360-1M")
+panshot_root = Path("data/UCPE")
+debug_root = Path("debug/match_panflow")
+pf_clip_root = panshot_root / "PanFlow" / f"match_clips-{split}"
+output_jsonl = panshot_root / "PanFlow" / f"near_plane_depth-{split}.jsonl"
+output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+summary_root = output_jsonl.parent / f"{output_jsonl.stem}-summary"
+summary_root.mkdir(parents=True, exist_ok=True)
+
+flow_height = 512
+flow_width = 1024
+epipole_thres = 30
+upper_edge_mask = 0.35
+lower_edge_mask = 0.25
+sample_fps = 2
+frame_near_quantile = 25
+video_near_quantile = 50
+batch_size = 24
+
+visualize_disp = False
+if visualize_disp:
+    disp_root = summary_root / "disp"
+    disp_root.mkdir(parents=True, exist_ok=True)
+
+device = torch.device("cuda")
+flow_estimater_ckpt = "models/PanoFlow/PanoFlow(RAFT)-wo-CFE.pth"
+
+flow_estimater = PanoRAFTAPI(
+    device=device, model_path=flow_estimater_ckpt
+)
+
+def epipole_flow(pose):
+    pose = pose[1:].inverse() @ pose[:-1]
+    return transformation_to_flow(pose, (flow_height, flow_width))
+
+
+def flow_depth(flow, pose, eps=1e-6):
+    epi_flow = epipole_flow(pose)
+    dot = einsum(flow, epi_flow, "n c h w, n c h w -> n h w")
+    epi_norm = epi_flow.norm(dim=1)
+    depth = epi_norm ** 2 / dot.clamp_min(eps)
+
+    flow_norm = flow.norm(dim=1)
+    cos = dot / (epi_norm * flow_norm).clamp_min(eps)
+    cos = cos.clamp(-1, 1)
+    degree = torch.rad2deg(torch.acos(cos))
+    invalid = degree > epipole_thres
+
+    depth[invalid] = float("nan")
+    depth[dot < eps] = float("nan")
+    height = depth.shape[-2]
+    depth[:, :int(upper_edge_mask * height)] = float("nan")
+    depth[:, int((1 - lower_edge_mask) * height):] = float("nan")
+    return depth
+
+
+def near_plane_depth(depth):
+    depth = rearrange(depth, "n h w -> n (h w)")
+    near_depth = torch.nanquantile(depth, frame_near_quantile / 100, dim=-1)
+    near_depth = torch.nanquantile(near_depth, video_near_quantile / 100)
+    return near_depth.item()
+
+
+clip_metas = list(pf_clip_root.glob("*.json"))
+clip_metas.sort()
+print(f"Found {len(clip_metas)} PanFlow clip files.")
+# clip_metas = clip_metas[:10]  # DEBUG
+
+near_depths = []
+try:
+    f = open(output_jsonl, "w", buffering=1, encoding="utf-8")
+    writer = jsonlines.Writer(f)
+
+    for clip_meta in tqdm(clip_metas, desc=f"Processing matched clips"):
+        with open(clip_meta, "r") as f:
+            meta = json.load(f)
+
+        video_id = meta["video_id"]
+        video_file = panflow_root / "videos" / f"{video_id}.mp4"
+        vr = VideoReader(str(video_file), width=flow_width, height=flow_height, ctx=cpu(0), num_threads=1)
+
+        clip_id = meta["clip_id"]
+        clip_name = meta["clip_name"]
+        frames = meta["frames"]
+        num_frames = frames[-1] - frames[0] + 1
+        fps = meta["fps"]
+        num_frames_sampled = int(round(num_frames / fps * sample_fps))
+        if num_frames_sampled < 2:
+            tqdm.write(f"Too few frames ({num_frames}) in clip {clip_name}, skipping.")
+            continue
+        sample_frames = np.linspace(frames[0], frames[-1], num_frames_sampled)
+        sample_frames = np.round(sample_frames).astype(int)
+
+        pose_file = panflow_root / "slam_pose" / video_id / clip_name
+        pose_file = pose_file.with_suffix(".npy")
+        if not pose_file.exists():
+            tqdm.write(f"Pose file not found: {pose_file}, skipping.")
+            continue
+        c2w = np.load(pose_file)  # (T, 3, 4)
+        c2w = c2w[sample_frames - frames[0]]  # (N, 3, 4)
+        c2w_4x4 = np.eye(4, dtype=np.float32)
+        c2w = np.hstack((c2w, repeat(c2w_4x4[-1], "n -> f 1 n", f=c2w.shape[0])))
+        pf_pose = torch.from_numpy(c2w)
+
+        video = vr.get_batch(sample_frames)
+        flow_in = video.float()
+        flow_in = rearrange(flow_in, "n h w c -> n c h w")
+
+        flow_in = flow_in.to(device)
+        flows = flow_estimater.chunk_estimate_flow_cfe(flow_in, chunk_size=batch_size)
+        flows = rearrange(flows, "n h w c -> n c h w")
+        pf_pose = pf_pose.to(device)
+        depth = flow_depth(flows, pf_pose)
+        near_depth = near_plane_depth(depth)
+
+        if visualize_disp:
+            # 保存 disparity
+            disp_file = disp_root / f"{video_id}-{clip_id}-disp.png"
+            rgb_file  = disp_root / f"{video_id}-{clip_id}-rgb.png"
+
+            # ---------- disparity ----------
+            depth_vis = depth[0].detach().cpu().numpy()  # [H,W]
+            disp_vis = np.zeros_like(depth_vis, dtype=np.float32)
+            valid = np.isfinite(depth_vis) & (depth_vis > 1e-6)
+            disp_vis[valid] = 1.0 / depth_vis[valid]
+
+            if valid.any():
+                dmin = np.nanpercentile(disp_vis[valid], 1)
+                dmax = np.nanpercentile(disp_vis[valid], 99)
+                dmin = max(dmin, 0)
+                norm = np.clip((disp_vis - dmin) / (dmax - dmin + 1e-6), 0, 1)
+                disp_color = (plt.cm.magma(norm)[..., :3] * 255).astype(np.uint8)
+            else:
+                disp_color = np.zeros((*disp_vis.shape, 3), dtype=np.uint8)
+
+            cv2.imwrite(str(disp_file), cv2.cvtColor(disp_color, cv2.COLOR_RGB2BGR))
+
+            # ---------- 对应的 RGB ----------
+            # video: decord 返回 [N,H,W,C], 取第一帧
+            if video.shape[0] > 0:
+                rgb = video[0].cpu().numpy()  # [H,W,C], float32 0~255?
+                if rgb.dtype != np.uint8:
+                    rgb = np.clip(rgb, 0, 255).astype(np.uint8)
+                cv2.imwrite(str(rgb_file), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))
+
+        near_depths.append(near_depth)
+        result = {
+            "video_id": video_id,
+            "clip_name": clip_name,
+            "clip_id": clip_id,
+            "near_depth": near_depth,
+        }
+        writer.write(result)
+finally:
+    # 关闭 writer / 文件句柄
+    try:
+        writer.close()
+    except Exception:
+        pass
+    try:
+        f.close()
+    except Exception:
+        pass
+
+near_depths = [d for d in near_depths if np.isfinite(d)]
+summary_file = summary_root / "near_plane_depth.png"
+
+fig, ax = plt.subplots(figsize=(10, 6))
+ax.hist(near_depths, bins="auto", color="skyblue", edgecolor="black")
+ax.set_xlabel("Near Plane Depth")
+ax.set_ylabel("Number of Clips")
+ax.set_title("Near Plane Depth Distribution")
+
+plt.tight_layout()
+fig.savefig(summary_file, dpi=300, bbox_inches="tight")
+plt.close(fig)
+print(f"Saved near plane depth histogram to {summary_file}")

UCPE/tools/pre_normalize_panflow.py

@@ -0,0 +1,191 @@
+from pathlib import Path
+import jsonlines
+from tqdm.auto import tqdm
+from collections import defaultdict
+import matplotlib.pyplot as plt
+import ffmpeg
+import json
+import csv
+import numpy as np
+import cv2
+from einops import einsum, rearrange, repeat
+from visualize_pose import vis_to_html
+import decord
+from decord import VideoReader, cpu
+decord.bridge.set_bridge("torch")
+import os
+import shutil
+import torch
+from thirdparty.PanFlow.utils.erp_utils import transformation_to_flow
+from thirdparty.PanoFlowAPI.apis.PanoRaft import PanoRAFTAPI
+
+
+panflow_root = Path("data/360-1M")
+panshot_root = Path("data/UCPE")
+debug_root = Path("debug/match_panflow")
+match_cb_root = panshot_root / "PanFlow" / "align_to_camerabench"
+output_jsonl = panshot_root / "PanFlow" / "near_plane_depth.jsonl"
+output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+summary_root = output_jsonl.parent / f"{output_jsonl.stem}-summary"
+summary_root.mkdir(parents=True, exist_ok=True)
+
+flow_height = 512
+flow_width = 1024
+epipole_thres = 30
+upper_edge_mask = 0.35
+lower_edge_mask = 0.2
+sample_fps = 2
+frame_near_quantile = 5
+video_near_quantile = 10
+batch_size = 24
+
+visualize_disp = True
+if visualize_disp:
+    disp_root = summary_root / "disp"
+    disp_root.mkdir(parents=True, exist_ok=True)
+
+device = torch.device("cuda")
+flow_estimater_ckpt = "models/PanoFlow/PanoFlow(RAFT)-wo-CFE.pth"
+
+flow_estimater = PanoRAFTAPI(
+    device=device, model_path=flow_estimater_ckpt
+)
+
+def epipole_flow(pose):
+    pose = pose[1:].inverse() @ pose[:-1]
+    return transformation_to_flow(pose, (flow_height, flow_width))
+
+
+def flow_depth(flow, pose, eps=1e-6):
+    epi_flow = epipole_flow(pose)
+    dot = einsum(flow, epi_flow, "n c h w, n c h w -> n h w")
+    epi_norm = epi_flow.norm(dim=1)
+    depth = epi_norm ** 2 / dot.clamp_min(eps)
+
+    flow_norm = flow.norm(dim=1)
+    cos = dot / (epi_norm * flow_norm).clamp_min(eps)
+    cos = cos.clamp(-1, 1)
+    degree = torch.rad2deg(torch.acos(cos))
+    invalid = degree > epipole_thres
+
+    depth[invalid] = float("nan")
+    depth[dot < eps] = float("nan")
+    height = depth.shape[-2]
+    depth[:, :int(upper_edge_mask * height)] = float("nan")
+    depth[:, int((1 - lower_edge_mask) * height):] = float("nan")
+    return depth
+
+
+def near_plane_depth(depth):
+    depth = rearrange(depth, "n h w -> n (h w)")
+    near_depth = torch.nanquantile(depth, frame_near_quantile / 100, dim=-1)
+    near_depth = torch.nanquantile(near_depth, video_near_quantile / 100)
+    return near_depth.item()
+
+
+match_meta_files = list(match_cb_root.glob("*.json"))
+match_meta_files.sort()
+match_meta_files = match_meta_files[:10]  # DEBUG
+print(f"Found {len(match_meta_files)} match meta files.")
+
+near_depths = []
+with jsonlines.open(output_jsonl, "w") as writer:
+    for meta_file in tqdm(match_meta_files, desc=f"Processing matched clips"):
+        with open(meta_file, "r") as f:
+            meta = json.load(f)
+
+        video_id = meta_file.stem
+        video_file = panflow_root / "videos" / f"{video_id}.mp4"
+        vr = VideoReader(str(video_file), width=flow_width, height=flow_height, ctx=cpu(0), num_threads=1)
+
+        for pf_clip in meta["clips"]:
+            if "matches" not in pf_clip:
+                tqdm.write(f"No matches in clip {pf_clip['clip_name']}, skipping.")
+                continue
+
+            clip_id = pf_clip["clip_id"]
+            clip_name = pf_clip["clip_name"]
+            frames = pf_clip["frames"]
+            num_frames = frames[-1] - frames[0] + 1
+            fps = meta["fps"]
+            num_frames_sampled = int(round(num_frames / fps * sample_fps))
+            if num_frames_sampled < 2:
+                tqdm.write(f"Too few frames ({num_frames}) in clip {clip_name}, skipping.")
+                continue
+            sample_frames = np.linspace(frames[0], frames[-1], num_frames_sampled)
+            sample_frames = np.round(sample_frames).astype(int)
+
+            pose_file = panflow_root / "slam_pose" / video_id / clip_name
+            pose_file = pose_file.with_suffix(".npy")
+            if not pose_file.exists():
+                tqdm.write(f"Pose file not found: {pose_file}, skipping.")
+                continue
+            c2w = np.load(pose_file)  # (T, 3, 4)
+            c2w = c2w[sample_frames - frames[0]]  # (N, 3, 4)
+            c2w_4x4 = np.eye(4, dtype=np.float32)
+            c2w = np.hstack((c2w, repeat(c2w_4x4[-1], "n -> f 1 n", f=c2w.shape[0])))
+            pf_pose = torch.from_numpy(c2w)
+
+            video = vr.get_batch(sample_frames)
+            flow_in = video.float()
+            flow_in = rearrange(flow_in, "n h w c -> n c h w")
+
+            flow_in = flow_in.to(device)
+            flows = flow_estimater.chunk_estimate_flow_cfe(flow_in, chunk_size=batch_size)
+            flows = rearrange(flows, "n h w c -> n c h w")
+            pf_pose = pf_pose.to(device)
+            depth = flow_depth(flows, pf_pose)
+            near_depth = near_plane_depth(depth)
+
+            if visualize_disp:
+                # 保存 disparity
+                disp_file = disp_root / f"{video_id}-{clip_id}-disp.png"
+                rgb_file  = disp_root / f"{video_id}-{clip_id}-rgb.png"
+
+                # ---------- disparity ----------
+                depth_vis = depth[0].detach().cpu().numpy()  # [H,W]
+                disp_vis = np.zeros_like(depth_vis, dtype=np.float32)
+                valid = np.isfinite(depth_vis) & (depth_vis > 1e-6)
+                disp_vis[valid] = 1.0 / depth_vis[valid]
+
+                if valid.any():
+                    dmin = np.nanpercentile(disp_vis[valid], 1)
+                    dmax = np.nanpercentile(disp_vis[valid], 99)
+                    dmin = max(dmin, 0)
+                    norm = np.clip((disp_vis - dmin) / (dmax - dmin + 1e-6), 0, 1)
+                    disp_color = (plt.cm.magma(norm)[..., :3] * 255).astype(np.uint8)
+                else:
+                    disp_color = np.zeros((*disp_vis.shape, 3), dtype=np.uint8)
+
+                cv2.imwrite(str(disp_file), cv2.cvtColor(disp_color, cv2.COLOR_RGB2BGR))
+
+                # ---------- 对应的 RGB ----------
+                # video: decord 返回 [N,H,W,C], 取第一帧
+                if video.shape[0] > 0:
+                    rgb = video[0].cpu().numpy()  # [H,W,C], float32 0~255?
+                    if rgb.dtype != np.uint8:
+                        rgb = np.clip(rgb, 0, 255).astype(np.uint8)
+                    cv2.imwrite(str(rgb_file), cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))
+
+            near_depths.append(near_depth)
+            result = {
+                "video_id": video_id,
+                "clip_name": clip_name,
+                "clip_id": clip_id,
+                "near_depth": near_depth,
+            }
+            writer.write(result)
+
+near_depths = [d for d in near_depths if np.isfinite(d)]
+summary_file = summary_root / "near_plane_depth.png"
+
+fig, ax = plt.subplots(figsize=(10, 6))
+ax.hist(near_depths, bins="auto", color="skyblue", edgecolor="black")
+ax.set_xlabel("Near Plane Depth")
+ax.set_ylabel("Number of Clips")
+ax.set_title("Near Plane Depth Distribution")
+
+plt.tight_layout()
+fig.savefig(summary_file, dpi=300, bbox_inches="tight")
+plt.close(fig)
+print(f"Saved near plane depth histogram to {summary_file}")

UCPE/tools/process_camerabench.py

@@ -0,0 +1,195 @@
+from pathlib import Path
+import jsonlines
+from tqdm.auto import tqdm
+from collections import defaultdict
+import matplotlib.pyplot as plt
+import ffmpeg
+import json
+import cv2
+
+
+# 数据与路径设置
+data_root = Path("data/CameraBench")
+output_root = Path("data/UCPE/CameraBench")
+target_height = 720
+target_width = 1280
+target_frames = 81
+target_fps = 16
+banned_labels = ["zoom-in", "zoom-out"]
+banned_keywords = ["zoom"]
+split = "train"  # "train" or "test"
+dryrun = False
+
+# 读取数据
+if split == "test":
+    meta_file = data_root / "test.jsonl"
+    with jsonlines.open(meta_file, "r") as reader:
+        metadata = list(reader)
+else:
+    meta_file = data_root / "cam_motion" / "captionset.json"
+    with open(meta_file, "r") as f:
+        captionset = json.load(f)
+    metadata = []
+    videos = set()
+    for obj in captionset:
+        video = obj["videos"][0]
+        if video in videos:
+            continue
+        videos.add(video)
+        metadata.append({
+            "path": video,
+            "caption": obj["messages"][1]["content"]
+        })
+
+# 处理与过滤视频
+labels = defaultdict(int)
+frames = []
+heights = []
+fps_list = []
+filtered_meta = []
+for obj in tqdm(metadata, desc="Reading videos"):
+    video_file = obj["path"]
+    if split == "test":
+        for label in obj["labels"]:
+            labels[label] += 1
+    video_path = data_root / video_file
+
+    # 获取视频信息
+    # meta = ffmpeg.probe(str(video_path))  # 使用 ffprobe 获取全 metadata
+    # vstream = next(s for s in meta["streams"] if s["codec_type"] == "video")
+    # w = int(vstream["width"])
+    # h = int(vstream["height"])
+    # fps_str = vstream.get("avg_frame_rate", "0/1")
+    # num, den = map(int, fps_str.split('/'))
+    # fps = num / den if den != 0 else None
+    # num_frames = int(vstream.get("nb_frames"))
+
+    cap = cv2.VideoCapture(str(video_path))
+    if not cap.isOpened():
+        tqdm.write(f"  - failed to open {video_path}, skipping")
+        continue
+    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    cap.release()
+
+    # 打印视频信息
+    num_frames_sampled = int(num_frames / fps * target_fps)
+    frames.append(num_frames)
+    heights.append(h)
+    fps_list.append(fps)
+    tqdm.write(f"{video_path} — {num_frames} frames, fps ~ {fps:.2f}, sampled to {num_frames_sampled}, height {h}")
+
+    # 过滤条件
+    if num_frames_sampled < target_frames:
+        tqdm.write(f"  - filtered by frames: {num_frames_sampled} < {target_frames}")
+        continue
+    if h < target_height:
+        tqdm.write(f"  - filtered by height: {h} < {target_height}")
+        continue
+    if w < target_width:
+        tqdm.write(f"  - filtered by width: {w} < {target_width}")
+        continue
+    if split == "test":
+        if any(label in banned_labels for label in obj["labels"]):
+            tqdm.write(f"  - filtered by label: {obj['labels']}")
+            continue
+    else:
+        caption = obj["caption"]
+        if any(kw in caption.lower() for kw in banned_keywords):
+            tqdm.write(f"  - filtered by keyword in caption: {caption}")
+            continue
+    filtered_meta.append(obj)
+
+    # 导出视频
+    output_video = (output_root / video_file).with_suffix(".mp4")
+    tqdm.write(f"  - exporting to {output_video}")
+    if dryrun:
+        continue
+    output_video.parent.mkdir(parents=True, exist_ok=True)
+
+    # 说明：
+    # 1) fps 过滤器把时间轴重采样为 target_fps（不会改变播放速度）
+    # 2) scale 先把画面按长宽比“铺满”16:9 画幅（a=iw/ih）：
+    #    - 如果原视频更宽(gt(a,16/9))：固定高=720，宽按比例（-2 表示自动，且保证可被2整除）
+    #    - 否则固定宽=1280，高按比例
+    # 3) 再做中心裁剪到 1280x720
+    # 4) vframes 只写前 target_frames 帧
+    in_stream = ffmpeg.input(str(video_path))
+    v = (
+        in_stream.video
+        .filter('fps', fps=target_fps)  # 重新采样到目标帧率
+        .filter('scale',
+                f'if(gt(a,{target_width}/{target_height}),-2,{target_width})',
+                f'if(gt(a,{target_width}/{target_height}),{target_height},-2)')
+        .filter('crop',
+                target_width, target_height,
+                f'(in_w-{target_width})/2', f'(in_h-{target_height})/2')
+    )
+
+    # 如无需音频可去掉 audio；若要保留音频，建议也做 asetpts=PTS-STARTPTS
+    out = ffmpeg.output(
+        v,
+        str(output_video),
+        vcodec='libx264',
+        pix_fmt='yuv420p',
+        r=target_fps,             # 容器帧率元数据
+        vframes=target_frames     # 只导出前 N 帧
+    )
+    ffmpeg.run(out, overwrite_output=True, quiet=True)
+
+print(f"Total videos: {len(frames)}, after filtering: {len(filtered_meta)}")
+
+# 导出 jsonl
+if not dryrun:
+    jsonl_file = output_root / f"processed_{split}.jsonl"
+    with jsonlines.open(jsonl_file, "w") as writer:
+        writer.write_all(filtered_meta)
+
+# 创建输出目录
+output_dir = Path(f"debug/summarize_camerabench_{split}")
+output_dir.mkdir(parents=True, exist_ok=True)
+
+# 1. labels 柱状图
+if labels:
+    fig1, ax1 = plt.subplots(figsize=(8, 6))
+    label_names = list(labels.keys())
+    label_counts = [labels[k] for k in label_names]
+    ax1.bar(label_names, label_counts, color='skyblue')
+    ax1.set_xticklabels(label_names, rotation=45, ha='right')
+    ax1.set_ylabel('Count')
+    ax1.set_title('Label Frequencies')
+    plt.tight_layout()
+    fig1.savefig(output_dir / "label_frequencies.png", dpi=300, bbox_inches='tight')
+    plt.close(fig1)
+
+# 2. frames 直方图
+fig2, ax2 = plt.subplots(figsize=(8, 6))
+ax2.hist(frames, bins=30, color='orange', edgecolor='black')
+ax2.set_xlabel('Number of Frames')
+ax2.set_ylabel('Frequency')
+ax2.set_title('Distribution of Frames')
+plt.tight_layout()
+fig2.savefig(output_dir / "frames_distribution.png", dpi=300, bbox_inches='tight')
+plt.close(fig2)
+
+# 3. heights 直方图
+fig3, ax3 = plt.subplots(figsize=(8, 6))
+ax3.hist(heights, bins=30, color='green', edgecolor='black')
+ax3.set_xlabel('Height (pixels)')
+ax3.set_ylabel('Frequency')
+ax3.set_title('Distribution of Frame Heights')
+plt.tight_layout()
+fig3.savefig(output_dir / "heights_distribution.png", dpi=300, bbox_inches='tight')
+plt.close(fig3)
+
+# 4. fps 直方图
+fig4, ax4 = plt.subplots(figsize=(8, 6))
+ax4.hist(fps_list, bins=30, color='purple', edgecolor='black')
+ax4.set_xlabel('FPS (frames per second)')
+ax4.set_ylabel('Frequency')
+ax4.set_title('FPS Distribution')
+plt.tight_layout()
+fig4.savefig(output_dir / "fps_distribution.png", dpi=300, bbox_inches='tight')
+plt.close(fig4)

UCPE/tools/process_panshot.py

@@ -0,0 +1,470 @@
+from pathlib import Path
+import jsonlines
+from tqdm.auto import tqdm
+from collections import defaultdict
+import matplotlib.pyplot as plt
+import ffmpeg
+import json
+import csv
+import numpy as np
+import cv2
+from einops import rearrange, repeat
+from visualize_pose import vis_to_html
+import decord
+from decord import VideoReader, cpu
+from copy import deepcopy
+from equilib import equi2pers
+from thirdparty.PanFlow.utils.erp_utils import equilib_rotation
+from scipy.spatial.transform import Rotation as R, Slerp
+import shutil
+import torch
+import random
+from concurrent.futures import ProcessPoolExecutor, as_completed
+import yt_dlp
+from threading import Lock
+import gc
+
+
+split = "train"  # "train" or "test"
+panflow_root = Path("data/360-1M")
+pf_pose_root = panflow_root / "slam_pose"
+ucpe_root = Path("data/UCPE")
+camerabench_root = ucpe_root / "CameraBench"
+match_pf_root = ucpe_root / "PanFlow" / f"match_to_camerabench-{split}"
+pf_clip_root = ucpe_root / "PanFlow" / f"match_clips-{split}"
+pf_video_root = ucpe_root / "PanFlow" / "videos"
+ps_video_root = ucpe_root / "PanShot" / f"videos-{split}"
+ps_video_root.mkdir(parents=True, exist_ok=True)
+ps_pose_root = ucpe_root / "PanShot" / f"pose-{split}"
+ps_pose_root.mkdir(parents=True, exist_ok=True)
+ps_meta_root = ucpe_root / "PanShot" / f"meta-{split}"
+ps_meta_root.mkdir(parents=True, exist_ok=True)
+debug_root = Path("debug/process_panshot")
+
+random_seed = 42
+num_workers = 8
+target_fps = 16
+target_height = 480
+target_width = 832
+
+x_fovs = [
+    [90, 110],    # 典型 pinhole / 准广角
+    [110, 140],   # 广角（轻微畸变）
+    [140, 180],   # 常见鱼眼（GoPro / 全景相机）
+    [160, 200],   # 极鱼眼 / 安防全景
+]
+
+xis = [
+    [0.0, 0.0],   # pinhole
+    [0.5, 0.95],   # 广角，畸变很小
+    [1.05, 2.0],   # 常见鱼眼
+    [1.5, 2.3],   # 极鱼眼
+]
+rot_augs = {
+    "no_rot_aug": {
+        "num": 1,
+    },
+    "yaw_aug": {
+        "num": 1,
+        "fixed_yaw": 180,
+    },
+    "yaw_pitch_aug": {
+        "num": 1,
+        "fixed_yaw": 180,
+        "fixed_pitch": 80,
+    },
+    "linear_aug": {
+        "num": 1,
+        "fixed_yaw": 90,
+        "fixed_pitch": 40,
+        "fixed_roll": 30,
+        "linear_yaw": 90,
+        "linear_pitch": 40,
+        "linear_roll": 30,
+    },
+}
+
+visualize_ref = False
+debug_aug = False
+debug_aug_lower = False
+debug = False
+pose_only = False
+overwrite = False
+offline = False
+
+clip_metas = list(pf_clip_root.glob("*.json"))
+clip_metas.sort()
+print(f"Found {len(clip_metas)} PanFlow clip files.")
+video_metas = defaultdict(dict)
+seed = random_seed
+for clip_meta in clip_metas:
+    with open(clip_meta, "r") as f:
+        meta = json.load(f)
+    video_metas[meta["video_id"]][meta["clip_id"]] = meta
+    meta["matches"] = []
+    meta["seed"] = seed
+    seed += 1
+print(f"Found {len(video_metas)} unique PanFlow videos.")
+
+match_metas = list(match_pf_root.glob("*.json"))
+match_metas.sort()
+cb_poses = {}
+num_matches = 0
+for match_meta in tqdm(match_metas, desc="Loading match metas"):
+    with open(match_meta, "r") as f:
+        cb_matches = json.load(f)
+    cb_video = match_meta.stem
+
+    if cb_video not in cb_poses:
+        pose_file = camerabench_root / "pose" / f"{cb_video}.npy"
+        pose= np.load(pose_file)  # (T, 4, 4)
+        cb_poses[cb_video] = pose
+
+    for cb_match in cb_matches:
+        matches = video_metas[cb_match['video_id']][cb_match['clip_id']]["matches"]
+        matches.append({
+            "match_id": len(matches),
+            "cb_video": cb_video,
+            "frames": cb_match["frames"],
+            "R": cb_match["R"],
+        })
+        num_matches += 1
+print(f"Found {num_matches} matches to {len(cb_poses)} CameraBench videos.")
+
+num_rot_augs = sum(aug["num"] for aug in rot_augs.values()) + 1
+estimated_clips = num_rot_augs * num_matches
+print(f"Estimated total {estimated_clips} PanShot clips to be generated.")
+
+
+def apply_rotation_align(rotation, R):
+    """
+    将旋转对齐矩阵 R 应用于相机轨迹的旋转矩阵序列。
+    左乘 R: R'_cw = R * R_cw
+
+    Args:
+        rotation: (T, 3, 3)   # 待对齐的旋转序列
+        R: (3, 3)             # 对齐用的旋转矩阵
+
+    Returns:
+        rotation_aligned: (T, 3, 3)
+    """
+    # einsum: i j, t j k -> t i k
+    return np.einsum("ij,tjk->tik", R, rotation)
+
+
+def make_mixed_rotation(rotation: np.ndarray,
+                        yaw_end: float = 0.0,
+                        pitch_end: float = 0.0,
+                        roll_end: float = 0.0,
+                        sample_frames: np.ndarray = None) -> np.ndarray:
+    """
+    混合增强: yaw 在世界坐标系, pitch/roll 在相机坐标系
+    rotation: (T,3,3)   采样后的原始 R_cw
+    yaw_end/pitch_end/roll_end: 增强角度(度)
+    sample_frames: (T,) 原始帧索引; 如果提供, 将按首尾 idx 做连续插值后再采样
+    """
+
+    # === 确定插值时间轴 ===
+    if sample_frames is not None:
+        # 原始轨迹长度
+        start_idx = int(sample_frames[0])
+        end_idx = int(sample_frames[-1])
+        full_T = end_idx - start_idx + 1
+        full_times = np.linspace(0, 1, full_T)
+    else:
+        full_T = rotation.shape[0]
+        full_times = np.linspace(0, 1, full_T)
+
+    # ===== 1) 世界坐标系 yaw 插值 =====
+    R_yaw_start = R.identity()
+    R_yaw_end = R.from_euler('y', yaw_end, degrees=True)
+    slerp_yaw = Slerp([0, 1], R.from_matrix([R_yaw_start.as_matrix(), R_yaw_end.as_matrix()]))
+    R_yaw_full = slerp_yaw(full_times).as_matrix()  # (full_T,3,3)
+
+    # ===== 2) 相机坐标系 pitch/roll 插值 =====
+    R_pr_start = R.identity()
+    R_pr_end = R.from_euler('xz', [pitch_end, roll_end], degrees=True)
+    slerp_pr = Slerp([0, 1], R.from_matrix([R_pr_start.as_matrix(), R_pr_end.as_matrix()]))
+    R_pr_full = slerp_pr(full_times).as_matrix()    # (full_T,3,3)
+
+    # ===== 3) 对齐采样帧 =====
+    if sample_frames is not None:
+        # 通过偏移索引选出采样帧对应的旋转
+        idxs = (sample_frames - sample_frames[0]).astype(int)
+        R_yaw = R_yaw_full[idxs]
+        R_pr = R_pr_full[idxs]
+    else:
+        R_yaw = R_yaw_full
+        R_pr = R_pr_full
+
+    # ===== 4) 合成 =====
+    # yaw 世界系 → 左乘, pitch/roll 相机系 → 右乘
+    R_out = np.einsum('tij,tjk,tkl->til', R_yaw, rotation, R_pr)
+    return R_out
+
+
+def download_video(video_id, output_path):
+    video_url = f"https://www.youtube.com/watch?v={video_id}"
+    ydl_opts = {
+        "outtmpl": str(output_path),
+        "format": "bestvideo[height<=2500][height>1500]",
+        "quiet": False,
+        "no_warnings": True,
+        "simulate": False,
+        "cookiefile": "~/.config/cookies.txt",
+        "print": [
+            "before_dl:Format: %(format_id)s | Res: %(resolution)s | FPS: %(fps)s",
+            "after_move:Size: %(filesize:.2fMB)s"
+        ],
+        "user_agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
+        "retries": 3,
+    }
+
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        ydl.download([video_url])
+
+
+download_lock = Lock()
+
+def process_one_video(video_id, video_meta):
+    video_path = Path(pf_video_root) / f"{video_id}.mp4"
+    if not video_path.exists() and not pose_only:
+        if offline:
+            tqdm.write(f"Offline mode, skipping download of {video_id}.")
+            return False
+        try:
+            with download_lock:
+                download_video(video_id, video_path)
+        except Exception as e:
+            tqdm.write(f"Failed to download {video_id}: {e}")
+            return False
+
+    if video_path.exists():
+        for clip_meta in video_meta.values():
+            if not process_one_clip(clip_meta):
+                return False
+    return True
+
+
+def process_one_clip(clip_meta):
+    video_id = clip_meta["video_id"]
+    clip_id = clip_meta["clip_id"]
+    pf_key = f"{video_id}-{clip_id}"
+    meta_file = ps_meta_root / f"{pf_key}.json"
+    if not overwrite and meta_file.exists():
+        tqdm.write(f"Meta file {meta_file} exists, skipping.")
+        return
+
+    seed = clip_meta["seed"]
+    np.random.seed(seed)
+    random.seed(seed)
+
+    video_file = pf_video_root / f"{video_id}.mp4"
+    try:
+        vr = VideoReader(str(video_file), ctx=cpu(0), num_threads=1)
+    except Exception as e:
+        tqdm.write(f"Failed to read video {video_file}: {e}")
+        return False
+    total_frames = len(vr)
+    del vr
+
+    matches = deepcopy(clip_meta["matches"])
+    for match in matches:
+        match_id = match["match_id"]
+        cb_video = match["cb_video"]
+        cb_pose = cb_poses[cb_video]
+        clip_start = clip_meta["frames"][0]
+        start_frame, end_frame = match["frames"]
+        start_frame += clip_start
+        end_frame += clip_start
+
+        if visualize_ref:
+            vis_file = debug_root / "cb_videos" / f"{pf_key}-{match_id}.mp4"
+            vis_file.parent.mkdir(parents=True, exist_ok=True)
+            if not vis_file.exists():
+                src_file = camerabench_root / "videos" / f"{cb_video}.mp4"
+                shutil.copy2(src_file, vis_file)
+
+        if end_frame >= total_frames:
+            tqdm.write(f"End frame {end_frame} > total frames {total_frames} in {video_file}, skipping.")
+            continue
+
+        num_frames_sampled = len(cb_pose)
+        sample_frames = np.linspace(start_frame, end_frame, num_frames_sampled)
+        sample_frames = np.round(sample_frames).astype(int)
+
+        if not pose_only:
+            vr = VideoReader(str(video_file), ctx=cpu(0), num_threads=1)
+            decord.bridge.set_bridge("torch")
+            erp_frames = vr.get_batch(sample_frames)  # (T, H, W, 3)
+            del vr
+            gc.collect()
+            erp_frames = rearrange(erp_frames, "t h w c -> t c h w")
+
+            if visualize_ref:
+                out_path = debug_root / "pf_videos" / f"{pf_key}-{match_id}.mp4"
+                out_path.parent.mkdir(parents=True, exist_ok=True)
+
+                # ---- 1. 转换��� float 并降采样 ----
+                erp_down = torch.nn.functional.interpolate(
+                    erp_frames.float(),  # [1, T, C, H, W]
+                    size=(target_height, target_width),
+                    mode='bilinear',
+                    align_corners=False
+                ).clamp(0, 255).to(torch.uint8)
+
+                # ---- 2. 改为 [T, H, W, C] ----
+                frames_np = rearrange(erp_down, "t c h w -> t h w c").contiguous().numpy()
+
+                # ---- 3. 写入视频 ----
+                process = (
+                    ffmpeg
+                    .input(
+                        'pipe:', 
+                        format='rawvideo',
+                        pix_fmt='rgb24',
+                        s=f'{target_width}x{target_height}',
+                        framerate=target_fps
+                    )
+                    .output(str(out_path), vcodec='libx264', pix_fmt='yuv420p', crf=18, preset='slow')
+                    .overwrite_output()
+                    .run_async(pipe_stdin=True, quiet=True)
+                )
+                process.stdin.write(frames_np.tobytes())
+                process.stdin.close()
+                process.wait()
+                print(f"✅ saved to {out_path}")
+
+        clip_name = f"Clip-{clip_id:03d}"
+        pose_file = pf_pose_root / video_id / f"{clip_name}.npy"
+        pf_pose = np.load(pose_file)  # (N, 3, 4)
+        c2w = pf_pose[sample_frames - clip_start]  # (T, 3, 4)
+        last_row = repeat(np.array([0,0,0,1], dtype=c2w.dtype), "n -> t 1 n", t=c2w.shape[0])
+        c2w = np.concatenate([c2w, last_row], axis=-2)  # (T, 4, 4)
+        w2c0 = np.linalg.inv(c2w[0])  # (4, 4)
+        c2w = w2c0[None] @ c2w  # (T, 4, 4)
+        pf_pose = c2w[:, :3, :]  # (T, 3, 4)
+
+        rotation = cb_pose[:, :3, :3]  # (T, 3, 3)
+        R_align = np.array(match["R"]).T
+        rotation = apply_rotation_align(rotation, R_align)
+
+        videos = []
+        i_rot = 0
+        for aug_key, aug_setting in rot_augs.items():
+            for i_aug in range(aug_setting["num"]):
+                aug_meta = deepcopy(aug_setting)
+                del aug_meta["num"]
+
+                rotation_aug = rotation.copy()
+                if "fixed_yaw" in aug_setting:
+                    yaw_aug = int(np.random.randint(-aug_setting["fixed_yaw"], aug_setting["fixed_yaw"]))
+                    yaw_rad = float(np.deg2rad(yaw_aug))
+                    aug_meta["yaw"] = yaw_aug
+                    R_yaw_world = R.from_euler("y", yaw_rad, degrees=False).as_matrix()
+                    rotation_aug = np.einsum("ij,tjk->tik", R_yaw_world, rotation_aug)
+                if "fixed_pitch" in aug_setting:
+                    pitch_aug = int(np.random.randint(-aug_setting["fixed_pitch"], aug_setting["fixed_pitch"]))
+                    pitch_rad = float(np.deg2rad(pitch_aug))
+                    aug_meta["pitch"] = pitch_aug
+                    R_pitch_cam = R.from_euler("x", pitch_rad, degrees=False).as_matrix()
+                    rotation_aug = np.einsum("tij,jk->tik", rotation_aug, R_pitch_cam)
+                if "fixed_roll" in aug_setting:
+                    roll_aug = int(np.random.randint(-aug_setting["fixed_roll"], aug_setting["fixed_roll"]))
+                    roll_rad = float(np.deg2rad(roll_aug))
+                    aug_meta["roll"] = roll_aug
+                    R_roll_cam = R.from_euler("z", roll_rad, degrees=False).as_matrix()
+                    rotation_aug = np.einsum("tij,jk->tik", rotation_aug, R_roll_cam)
+                if any(k in aug_setting for k in ("linear_yaw", "linear_pitch", "linear_roll")):
+                    yaw_range = aug_setting.get("linear_yaw", 0)
+                    pitch_rang = aug_setting.get("linear_pitch", 0)
+                    roll_rang = aug_setting.get("linear_roll", 0)
+                    yaw_end = int(np.random.randint(-yaw_range, yaw_range)) if yaw_range > 0 else 0
+                    pitch_end = int(np.random.randint(-pitch_rang, pitch_rang)) if pitch_rang > 0 else 0
+                    roll_end = int(np.random.randint(-roll_rang, roll_rang)) if roll_rang > 0 else 0
+                    rotation_aug = make_mixed_rotation(rotation_aug, yaw_end, pitch_end, roll_end, sample_frames)
+                    aug_meta.update({"linear_yaw": yaw_end, "linear_pitch": pitch_end, "linear_roll": roll_end})
+                # rotation_aug = repeat(np.eye(3), "i j -> t i j", t=rotation_aug.shape[0])  # Debug
+
+                # generate camera pose
+                ps_pose_file = Path(f"{pf_key}-{match_id}-{aug_key}_{i_aug}.npy")
+                ps_pose = pf_pose.copy()
+                ps_pose[..., :3] = pf_pose[..., :3] @ rotation_aug
+                np.save(ps_pose_root / ps_pose_file, ps_pose)
+
+                idx_lens = i_rot % len(x_fovs) if debug_aug else random.randint(0, len(x_fovs) - 1)
+                x_fov_range = x_fovs[idx_lens]
+                xi_range = xis[idx_lens]
+                if debug_aug:
+                    if debug_aug_lower:
+                        x_fov = x_fov_range[0]
+                        xi = xi_range[0]
+                    else:
+                        x_fov = x_fov_range[1]
+                        xi = xi_range[1]
+                else:
+                    x_fov = int(np.round(np.random.uniform(*x_fov_range)))
+                    xi = float(np.round(np.random.uniform(*xi_range), 2))
+
+                if not pose_only:
+                    rotation_aug = equilib_rotation(rotation_aug)
+                    pers_frames = equi2pers(erp_frames, rotation_aug, target_height, target_width, x_fov, xi=xi)
+                    pers_frames = rearrange(pers_frames, "t c h w -> t h w c")
+                    pers_frames = pers_frames.numpy()
+
+                    out_file = ps_video_root / f"{pf_key}-{match_id}-{aug_key}_{i_aug}-fov{x_fov}-xi{xi:.2f}.mp4"
+                    process = (
+                        ffmpeg
+                        .input('pipe:', format='rawvideo', pix_fmt='rgb24', s=f'{target_width}x{target_height}', framerate=target_fps)
+                        .output(str(out_file), pix_fmt='yuv420p', vcodec='libx264', r=target_fps, crf=16, preset='slow')
+                        .overwrite_output()
+                        .run_async(pipe_stdin=True, quiet=True)
+                    )
+                    process.stdin.write(pers_frames.tobytes())
+                    del pers_frames
+                    process.stdin.close()
+                    process.wait()
+
+                    videos.append({
+                        "video": str(out_file.stem),
+                        "pose": str(ps_pose_file.stem),
+                        "x_fov": x_fov,
+                        "xi": xi,
+                        "rot_aug": aug_meta,
+                    })
+
+                gc.collect()
+                torch.cuda.empty_cache()
+
+                i_rot += 1
+
+        match["videos"] = videos
+
+    if not pose_only:
+        with open(meta_file, "w") as f:
+            json.dump(matches, f, indent=4)
+
+    return True
+
+success = []
+if debug:
+    # === 单线程调试模式 ===
+    for video_id, video_meta in tqdm(
+        video_metas.items(),
+        desc="Processing clips (debug single-thread)"
+    ):
+        success.append(process_one_video(video_id, video_meta))
+else:
+    # === 正常并行模式 ===
+    with ProcessPoolExecutor(max_workers=num_workers) as ex:
+        futures = [
+            ex.submit(process_one_video, video_id, video_meta)
+            for video_id, video_meta in video_metas.items()
+        ]
+        for fut in tqdm(as_completed(futures), total=len(futures), desc="Processing clips"):
+            try:
+                success.append(fut.result())
+            except Exception as e:
+                print("Error:", e)
+
+print(f"All done. {sum(success)}/{len(success)} videos processed successfully.")

UCPE/tools/rerender_panshot.py

@@ -0,0 +1,328 @@
+"""
+Re-render PanShot videos at 704x1280 using existing meta/pose files
+and CameraBench poses. Downloads YouTube source videos on-the-fly.
+
+Usage:
+    conda activate UCPE
+    python tools/rerender_panshot.py
+
+Prerequisites:
+    - PanShot meta files at {panshot_root}/meta-{split}/
+    - CameraBench poses at data/UCPE/CameraBench/pose/
+    - YouTube cookies at ~/.config/cookies.txt
+"""
+
+from pathlib import Path
+import json
+import numpy as np
+import torch
+import ffmpeg
+import gc
+from einops import rearrange
+from tqdm.auto import tqdm
+from collections import defaultdict
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from threading import Lock
+from scipy.spatial.transform import Rotation as R, Slerp
+import yt_dlp
+import decord
+from decord import VideoReader, cpu
+from equilib import equi2pers
+from thirdparty.PanFlow.utils.erp_utils import equilib_rotation
+
+
+# ===== Configuration =====
+split = "test"  # "train" or "test"
+panshot_root = Path("/tmp/data/UCPE/PanShot")
+camerabench_root = Path("data/UCPE/CameraBench")
+video_cache_root = Path("youtube_cache")
+output_video_root = panshot_root / f"videos_704-{split}"
+persistent_video_root = Path(f"data/UCPE/PanShot/videos_704-{split}")  # persistent copy
+meta_root = panshot_root / f"meta-{split}"
+
+target_fps = 16
+target_height = 704
+target_width = 1280
+num_workers = 16  # parallel workers
+overwrite = False
+
+video_cache_root.mkdir(parents=True, exist_ok=True)
+output_video_root.mkdir(parents=True, exist_ok=True)
+persistent_video_root.mkdir(parents=True, exist_ok=True)
+
+
+# ===== Helper functions =====
+
+def download_video(video_id, output_path):
+    video_url = f"https://www.youtube.com/watch?v={video_id}"
+    ydl_opts = {
+        "outtmpl": str(output_path),
+        "format": "bestvideo[height>=1080]",
+        "quiet": False,
+        "no_warnings": True,
+        "simulate": False,
+        "cookiefile": str(Path("~/.config/cookies.txt").expanduser()),
+        "print": [
+            "before_dl:Format: %(format_id)s | Res: %(resolution)s | FPS: %(fps)s",
+        ],
+        "user_agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
+        "retries": 3,
+    }
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        ydl.download([video_url])
+
+
+def apply_rotation_align(rotation, R_mat):
+    return np.einsum("ij,tjk->tik", R_mat, rotation)
+
+
+def make_mixed_rotation(rotation, yaw_end, pitch_end, roll_end, sample_frames):
+    if sample_frames is not None:
+        start_idx = int(sample_frames[0])
+        end_idx = int(sample_frames[-1])
+        full_T = end_idx - start_idx + 1
+        full_times = np.linspace(0, 1, full_T)
+    else:
+        full_T = rotation.shape[0]
+        full_times = np.linspace(0, 1, full_T)
+
+    R_yaw_start = R.identity()
+    R_yaw_end_r = R.from_euler('y', yaw_end, degrees=True)
+    slerp_yaw = Slerp([0, 1], R.from_matrix([R_yaw_start.as_matrix(), R_yaw_end_r.as_matrix()]))
+    R_yaw_full = slerp_yaw(full_times).as_matrix()
+
+    R_pr_start = R.identity()
+    R_pr_end_r = R.from_euler('xz', [pitch_end, roll_end], degrees=True)
+    slerp_pr = Slerp([0, 1], R.from_matrix([R_pr_start.as_matrix(), R_pr_end_r.as_matrix()]))
+    R_pr_full = slerp_pr(full_times).as_matrix()
+
+    if sample_frames is not None:
+        idxs = (sample_frames - sample_frames[0]).astype(int)
+        R_yaw = R_yaw_full[idxs]
+        R_pr = R_pr_full[idxs]
+    else:
+        R_yaw = R_yaw_full
+        R_pr = R_pr_full
+
+    R_out = np.einsum('tij,tjk,tkl->til', R_yaw, rotation, R_pr)
+    return R_out
+
+
+def reconstruct_rotation_aug(base_rotation, rot_aug, sample_frames):
+    """Reconstruct rotation_aug from base rotation + stored augmentation params."""
+    rotation_aug = base_rotation.copy()
+
+    if "yaw" in rot_aug:
+        yaw_rad = float(np.deg2rad(rot_aug["yaw"]))
+        R_yaw_world = R.from_euler("y", yaw_rad, degrees=False).as_matrix()
+        rotation_aug = np.einsum("ij,tjk->tik", R_yaw_world, rotation_aug)
+
+    if "pitch" in rot_aug:
+        pitch_rad = float(np.deg2rad(rot_aug["pitch"]))
+        R_pitch_cam = R.from_euler("x", pitch_rad, degrees=False).as_matrix()
+        rotation_aug = np.einsum("tij,jk->tik", rotation_aug, R_pitch_cam)
+
+    if "roll" in rot_aug:
+        roll_rad = float(np.deg2rad(rot_aug["roll"]))
+        R_roll_cam = R.from_euler("z", roll_rad, degrees=False).as_matrix()
+        rotation_aug = np.einsum("tij,jk->tik", rotation_aug, R_roll_cam)
+
+    if any(k in rot_aug for k in ("linear_yaw", "linear_pitch", "linear_roll")):
+        yaw_end = rot_aug.get("linear_yaw", 0)
+        pitch_end = rot_aug.get("linear_pitch", 0)
+        roll_end = rot_aug.get("linear_roll", 0)
+        rotation_aug = make_mixed_rotation(rotation_aug, yaw_end, pitch_end, roll_end, sample_frames)
+
+    return rotation_aug
+
+
+download_lock = Lock()
+
+
+def process_one_video(video_id, clips_info):
+    """Download YouTube video and re-render all its clips."""
+    video_path = video_cache_root / f"{video_id}.mp4"
+
+    if not video_path.exists():
+        try:
+            with download_lock:
+                if not video_path.exists():  # double check after lock
+                    download_video(video_id, video_path)
+        except Exception as e:
+            tqdm.write(f"Failed to download {video_id}: {e}")
+            return 0
+
+    if not video_path.exists():
+        return 0
+
+    rendered = 0
+    for clip_info in clips_info:
+        try:
+            rendered += process_one_clip(video_id, video_path, clip_info)
+        except Exception as e:
+            tqdm.write(f"Error processing {clip_info['meta_stem']}: {e}")
+            import traceback
+            traceback.print_exc()
+
+    return rendered
+
+
+def process_one_clip(video_id, video_path, clip_info):
+    """Re-render all video variants for one clip at the target resolution."""
+    meta_stem = clip_info["meta_stem"]
+    matches = clip_info["matches"]
+    rendered = 0
+
+    # Read source ERP video metadata
+    try:
+        vr = VideoReader(str(video_path), ctx=cpu(0), num_threads=1)
+    except Exception as e:
+        tqdm.write(f"Failed to read {video_path}: {e}")
+        return 0
+    total_frames = len(vr)
+    del vr
+
+    for match in matches:
+        cb_video = match["cb_video"]
+        R_align = np.array(match["R"]).T
+        start_frame, end_frame = match["frames"]
+
+        if end_frame >= total_frames:
+            tqdm.write(f"End frame {end_frame} >= total {total_frames} for {video_id}, skipping match.")
+            continue
+
+        # Load CameraBench pose to get num_frames and base rotation
+        cb_pose_file = camerabench_root / "pose" / f"{cb_video}.npy"
+        if not cb_pose_file.exists():
+            tqdm.write(f"CameraBench pose not found: {cb_pose_file}, skipping.")
+            continue
+        cb_pose = np.load(cb_pose_file)  # (T, 4, 4)
+
+        num_frames = len(cb_pose)
+        sample_frames = np.linspace(start_frame, end_frame, num_frames)
+        sample_frames = np.round(sample_frames).astype(int)
+
+        # Base rotation: CameraBench rotation aligned
+        base_rotation = cb_pose[:, :3, :3]  # (T, 3, 3)
+        base_rotation = apply_rotation_align(base_rotation, R_align)
+
+        # Read ERP frames (shared across all variants of this match)
+        erp_frames = None
+
+        for video_info in match.get("videos", []):
+            video_name = video_info["video"]
+            out_file = output_video_root / f"{video_name}.mp4"
+
+            persistent_file = persistent_video_root / f"{video_name}.mp4"
+            if not overwrite and (out_file.exists() or persistent_file.exists()):
+                rendered += 1
+                continue
+
+            x_fov = video_info["x_fov"]
+            xi = video_info["xi"]
+            rot_aug = video_info.get("rot_aug", {})
+
+            # Reconstruct rotation_aug
+            rotation_aug = reconstruct_rotation_aug(base_rotation, rot_aug, sample_frames)
+
+            # Lazy load ERP frames
+            if erp_frames is None:
+                vr = VideoReader(str(video_path), ctx=cpu(0), num_threads=1)
+                decord.bridge.set_bridge("torch")
+                erp_frames = vr.get_batch(sample_frames)  # (T, H, W, 3)
+                del vr
+                gc.collect()
+                erp_frames = rearrange(erp_frames, "t h w c -> t c h w")
+
+            # Project ERP to perspective
+            rot_for_equi2pers = equilib_rotation(rotation_aug)
+            pers_frames = equi2pers(erp_frames, rot_for_equi2pers, target_height, target_width, x_fov, xi=xi)
+            pers_frames = rearrange(pers_frames, "t c h w -> t h w c")
+            pers_frames = pers_frames.numpy()
+
+            # Write video
+            out_file.parent.mkdir(parents=True, exist_ok=True)
+            process = (
+                ffmpeg
+                .input('pipe:', format='rawvideo', pix_fmt='rgb24',
+                       s=f'{target_width}x{target_height}', framerate=target_fps)
+                .output(str(out_file), pix_fmt='yuv420p', vcodec='libx264',
+                        r=target_fps, crf=16, preset='fast')
+                .overwrite_output()
+                .run_async(pipe_stdin=True, quiet=True)
+            )
+            process.stdin.write(pers_frames.tobytes())
+            del pers_frames
+            process.stdin.close()
+            process.wait()
+
+            # Save persistent copy
+            import shutil
+            persistent_file = persistent_video_root / f"{video_name}.mp4"
+            if not persistent_file.exists():
+                shutil.copy2(str(out_file), str(persistent_file))
+
+            rendered += 1
+
+            gc.collect()
+            torch.cuda.empty_cache()
+
+        if erp_frames is not None:
+            del erp_frames
+            gc.collect()
+
+    return rendered
+
+
+# ===== Main =====
+
+print(f"Loading meta files from {meta_root}...")
+meta_files = sorted(meta_root.glob("*.json"))
+print(f"Found {len(meta_files)} meta files.")
+
+# Group by YouTube video ID (first 11 chars of meta filename)
+video_clips = defaultdict(list)
+total_variants = 0
+for meta_file in meta_files:
+    with open(meta_file) as f:
+        matches = json.load(f)
+    stem = meta_file.stem
+    youtube_id = stem[:11]
+    for match in matches:
+        total_variants += len(match.get("videos", []))
+    video_clips[youtube_id].append({
+        "meta_stem": stem,
+        "matches": matches,
+    })
+
+print(f"Found {len(video_clips)} unique YouTube videos, {total_variants} total variants to render.")
+
+# Check how many already exist
+existing = sum(1 for f in output_video_root.glob("*.mp4"))
+print(f"Already rendered: {existing}")
+
+# Process
+if num_workers <= 1:
+    for video_id, clips in tqdm(sorted(video_clips.items()), desc="Processing"):
+        process_one_video(video_id, clips)
+else:
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        futures = {
+            executor.submit(process_one_video, video_id, clips): video_id
+            for video_id, clips in sorted(video_clips.items())
+        }
+        pbar = tqdm(total=len(futures), desc="Processing videos")
+        total_rendered = 0
+        for future in as_completed(futures):
+            video_id = futures[future]
+            try:
+                n = future.result()
+                total_rendered += n
+            except Exception as e:
+                tqdm.write(f"Error on {video_id}: {e}")
+            pbar.update(1)
+            pbar.set_postfix(rendered=total_rendered)
+        pbar.close()
+
+final_count = sum(1 for f in output_video_root.glob("*.mp4"))
+print(f"Done. Total rendered videos: {final_count}")

UCPE/tools/score_panflow.py

@@ -0,0 +1,244 @@
+from q_align import QAlignVideoScorer, QAlignAestheticScorer, QAlignScorer
+from pathlib import Path
+from tqdm.auto import tqdm
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import os
+import json
+from decord import VideoReader, cpu
+import jsonlines
+from PIL import Image
+import torch
+import numpy as np
+
+# -----------------------------
+# basic configuration
+# -----------------------------
+panflow_root = Path("data/360-1M")
+panshot_root = Path("data/UCPE")
+output_root = panshot_root / "PanFlow"
+output_root.mkdir(parents=True, exist_ok=True)
+output_jsonl = output_root / "scores.jsonl"
+batch_size = 16
+max_workers = min(64, os.cpu_count() or 4)
+inflight_limit = max_workers * 2
+max_frames = 10
+print(f"Using max_workers={max_workers}, inflight_limit={inflight_limit}")
+
+meta_root = panflow_root / "meta"
+meta_files = list((meta_root).glob("*.json"))
+meta_files.sort()
+print(f"Found {len(meta_files)} PanFlow meta files.")
+
+def load_one_meta(meta_file):
+    try:
+        with open(meta_file, "r") as f:
+            meta = json.load(f)
+        video_id = meta.get("video_id", Path(meta_file).stem)
+        video_path = panflow_root / "videos" / f"{video_id}.mp4"
+
+        if "slam_clips" not in meta or "clips" not in meta["slam_clips"]:
+            return []  # skip
+        if not video_path.exists():
+            return []  # skip
+
+        clips_local = []
+        for clip in meta["slam_clips"]["clips"]:
+            clips_local.append({
+                "clip_id": clip["clip_id"],
+                "frames": clip["frames"],
+                "video_id": video_id,
+                "video_path": str(video_path),
+            })
+        return clips_local
+    except Exception as e:
+        tqdm.write(f"Error reading {meta_file}: {e}")
+        return []
+
+
+# ================================================================
+# 已处理 clip 检查
+# ================================================================
+processed = set()
+if output_jsonl.exists():
+    print(f"Resuming from {output_jsonl}")
+    with open(output_jsonl, "r", encoding="utf-8") as f_in:
+        for line in f_in:
+            try:
+                rec = json.loads(line)
+                processed.add((rec["video_id"], rec["clip_id"]))
+            except Exception:
+                continue
+    print(f"Found {len(processed)} processed clips to skip.")
+
+
+# -----------------------------
+# 并行加载 meta 文件
+# -----------------------------
+clips = []
+with ThreadPoolExecutor(max_workers=8) as ex:
+    futures = [ex.submit(load_one_meta, mf) for mf in meta_files]
+    for fut in tqdm(as_completed(futures), total=len(futures), desc="Loading PanFlow metadata"):
+        clips.extend(fut.result())
+
+# 过滤掉已处理的 clips
+clips = [c for c in clips if (c["video_id"], c["clip_id"]) not in processed]
+print(f"Total {len(clips)} remaining clips from {len(meta_files)} videos.")
+
+
+# -----------------------------
+# scorer 初始化
+# -----------------------------
+video_scorer = QAlignVideoScorer()
+scorers = {
+    "image_aesthetic": QAlignAestheticScorer(
+        tokenizer=video_scorer.tokenizer,
+        model=video_scorer.model,
+        image_processor=video_scorer.image_processor
+    ),
+    "image_quality": QAlignScorer(
+        tokenizer=video_scorer.tokenizer,
+        model=video_scorer.model,
+        image_processor=video_scorer.image_processor
+    ),
+    "video_quality": video_scorer,
+}
+
+# ================================================================
+# util functions
+# ================================================================
+def process_one_clip(clip):
+    """
+    从 clip 中读取指定起止帧范围内的帧（1fps），返回 {clip, frames(list of PIL)}.
+    """
+    video_path = clip["video_path"]
+    frame_range = clip.get("frames", None)
+
+    try:
+        vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
+    except Exception as e:
+        print(f"[process_one_clip] cannot open {video_path}: {e}")
+        return None
+
+    fps = vr.get_avg_fps()
+    start, end = frame_range[0], frame_range[-1]
+    start = max(0, start)
+    end = min(len(vr) - 1, end)
+
+    # 按 1fps 抽帧
+    frame_count = min((end - start + 1) / fps, max_frames)
+    frame_indices = np.linspace(start, end, num=max(int(frame_count), 1))
+    frame_indices = np.round(frame_indices).astype(int)
+
+    try:
+        frames_np = vr.get_batch(frame_indices).asnumpy()
+    except Exception as e:
+        print(f"[process_one_clip] error decoding {video_path}: {e}")
+        return None
+
+    frames = [Image.fromarray(frames_np[i]) for i in range(frames_np.shape[0])]
+    video = [video_scorer.expand2square(frame, tuple(int(x*255) for x in video_scorer.image_processor.image_mean)) for frame in frames]
+    video_tensors = video_scorer.image_processor.preprocess(video, return_tensors="pt")["pixel_values"].half()
+    return {"clip": clip, "frames": video_tensors}
+
+
+def infer_and_flush(results, writer):
+    """
+    批量推理并写出结果。
+    results: list of {clip, frames}
+    """
+    if not results:
+        return
+
+    video_batch = [r["frames"] for r in results]
+    scores = {}
+    with torch.inference_mode():
+        video_tensors = [vid.to(video_scorer.model.device) for vid in video_batch]
+        video_frames = [len(vid) for vid in video_tensors]
+        for key, scorer in scorers.items():
+            # image 分支拼接所有帧，video 分支传列表
+            images = torch.cat(video_tensors) if "image" in key else video_tensors
+            output_logits = scorer.model(
+                scorer.input_ids.repeat(len(images), 1),
+                images=images
+            )["logits"][:, -1, scorer.preferential_ids_]
+
+            values = torch.softmax(output_logits, -1) @ scorer.weight_tensor
+
+            if "image" in key:
+                # 按每个 clip 的帧数切开并取均值
+                values_split = torch.split(values, video_frames)
+                values = torch.stack([v.mean(0) for v in values_split])
+
+            scores[key] = values  # shape: [n_clips]
+
+    # 遍历每个 clip，把三个分数写出
+    n_clips = len(results)
+    for i in range(n_clips):
+        clip_info = results[i]["clip"]
+        out_obj = {
+            "video_id": clip_info.get("video_id"),
+            "clip_id": clip_info.get("clip_id"),
+        }
+        for key, value in scores.items():
+            out_obj[key] = float(value[i].item())
+        writer.write(out_obj)
+
+
+# ================================================================
+# 主流程
+# ================================================================
+output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+prepared_buffer = []
+
+try:
+    f = open(output_jsonl, "a", buffering=1, encoding="utf-8")
+    writer = jsonlines.Writer(f)
+
+    with ThreadPoolExecutor(max_workers=max_workers) as ex:
+        pbar = tqdm(total=len(clips), desc="processing clips")
+
+        pending = set()
+        i_submit = 0
+
+        # 先提交一些任务
+        while i_submit < len(clips) and len(pending) < inflight_limit:
+            fut = ex.submit(process_one_clip, clips[i_submit])
+            pending.add(fut)
+            i_submit += 1
+
+        while pending:
+            for fut in as_completed(list(pending), timeout=None):
+                pending.remove(fut)
+                result = fut.result()
+                if result is not None:
+                    prepared_buffer.append(result)
+
+                # 满一批就推理一次
+                if len(prepared_buffer) >= batch_size:
+                    infer_and_flush(prepared_buffer[:batch_size], writer)
+                    prepared_buffer = prepared_buffer[batch_size:]
+                    pbar.update(batch_size)
+
+                # 提交新的任务
+                while i_submit < len(clips) and len(pending) < inflight_limit:
+                    fut_new = ex.submit(process_one_clip, clips[i_submit])
+                    pending.add(fut_new)
+                    i_submit += 1
+                break
+
+        # flush 剩余 buffer
+        while prepared_buffer:
+            chunk = prepared_buffer[:batch_size]
+            infer_and_flush(chunk, writer)
+            prepared_buffer = prepared_buffer[len(chunk):]
+
+        pbar.close()
+finally:
+    try:
+        writer.close()
+    except Exception:
+        pass
+    try:
+        f.close()
+    except Exception:
+        pass

UCPE/tools/visualize_pose.py

@@ -0,0 +1,1231 @@
+#!/usr/bin/env python3
+"""
+Camera Pose Visualization Module
+
+This module provides comprehensive tools for visualizing camera poses and trajectories
+in 3D space using Plotly. It supports both static and animated visualizations with
+automatic camera view optimization.
+"""
+
+import argparse
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+import plotly.graph_objs as go
+import plotly.io as pio
+from tqdm import tqdm
+import einops
+import torch
+from einops import repeat
+
+# Use non-interactive backend for matplotlib to avoid display issues
+matplotlib.use("agg")
+
+
+class Pose:
+    """
+    A class of operations on camera poses (numpy arrays with shape [...,3,4]).
+    Each [3,4] camera pose takes the form of [R|t].
+    """
+
+    def __call__(self, R=None, t=None):
+        """
+        Construct a camera pose from the given rotation matrix R and/or translation vector t.
+        """
+        assert R is not None or t is not None
+        if R is None:
+            if not isinstance(t, np.ndarray):
+                t = np.array(t)
+            R = np.eye(3).repeat(*t.shape[:-1], 1, 1)
+        elif t is None:
+            if not isinstance(R, np.ndarray):
+                R = np.array(R)
+            t = np.zeros(R.shape[:-1])
+        else:
+            if not isinstance(R, np.ndarray):
+                R = np.array(R)
+            if not isinstance(t, np.ndarray):
+                t = np.array(t)
+        assert R.shape[:-1] == t.shape and R.shape[-2:] == (3, 3)
+        R = R.astype(np.float32)
+        t = t.astype(np.float32)
+        pose = np.concatenate([R, t[..., None]], axis=-1)  # [...,3,4]
+        assert pose.shape[-2:] == (3, 4)
+        return pose
+
+    def invert(self, pose, use_inverse=False):
+        """
+        Invert a camera pose.
+        """
+        R, t = pose[..., :3], pose[..., 3:]
+        R_inv = np.linalg.inv(R) if use_inverse else R.transpose(0, 2, 1)
+        t_inv = (-R_inv @ t)[..., 0]
+        pose_inv = self(R=R_inv, t=t_inv)
+        return pose_inv
+
+    def compose(self, pose_list):
+        """
+        Compose a sequence of poses together.
+        pose_new(x) = poseN o ... o pose2 o pose1(x)
+        """
+        pose_new = pose_list[0]
+        for pose in pose_list[1:]:
+            pose_new = self.compose_pair(pose_new, pose)
+        return pose_new
+
+    def compose_pair(self, pose_a, pose_b):
+        """
+        Compose two poses together.
+        """
+        R_a, t_a = pose_a[..., :3], pose_a[..., 3:]
+        R_b, t_b = pose_b[..., :3], pose_b[..., 3:]
+        R_new = R_b @ R_a
+        t_new = (R_b @ t_a + t_b)[..., 0]
+        pose_new = self(R=R_new, t=t_new)
+        return pose_new
+
+    def scale_center(self, pose, scale):
+        """
+        Scale the camera center from the origin.
+        0 = R@c+t --> c = -R^T@t (camera center in world coordinates)
+        0 = R@(sc)+t' --> t' = -R@(sc) = -R@(-R^T@st) = st
+        """
+        R, t = pose[..., :3], pose[..., 3:]
+        pose_new = np.concatenate([R, t * scale], axis=-1)
+        return pose_new
+
+
+def to_hom(X):
+    """
+    Convert points to homogeneous coordinates by appending ones.
+    """
+    X_hom = np.concatenate([X, np.ones_like(X[..., :1])], axis=-1)
+    return X_hom
+
+
+def cam2world(X, pose):
+    """
+    Transform points from camera coordinates to world coordinates.
+    """
+    X_hom = to_hom(X)
+    pose_inv = Pose().invert(pose)
+    return X_hom @ pose_inv.transpose(0, 2, 1)
+
+
+def get_camera_mesh(pose, depth=1):
+    """
+    Create a 3D mesh representation of camera frustums for visualization.
+    """
+    # Define camera frustum geometry: 4 corners of image plane + camera center
+    vertices = (
+        np.array(
+            [[-0.5, -0.5, 1], [0.5, -0.5, 1], [0.5, 0.5, 1], [-0.5, 0.5, 1], [0, 0, 0]]
+        )
+        * depth
+    )  # Shape: [5, 3] - 4 image plane corners + camera center
+
+    # Define triangular faces for the camera frustum mesh
+    faces = np.array(
+        [[0, 1, 2], [0, 2, 3], [0, 1, 4], [1, 2, 4], [2, 3, 4], [3, 0, 4]]
+    )  # Shape: [6, 3] - 6 triangular faces forming the pyramid
+
+    # Transform vertices from camera space to world space
+    vertices = cam2world(vertices[None], pose)  # Shape: [N, 5, 3]
+
+    # Create wireframe lines connecting: corners -> center -> next corner
+    wireframe = vertices[:, [0, 1, 2, 3, 0, 4, 1, 2, 4, 3]]  # Shape: [N, 10, 3]
+
+    return vertices, faces, wireframe
+
+
+# def merge_xyz_indicators_plotly(xyz):
+#     """Merge xyz coordinate indicators for plotly visualization."""
+#     xyz = xyz[:, [[-1, 0], [-1, 1], [-1, 2]]]  # [N,3,2,3]
+#     xyz_0, xyz_1 = unbind_np(xyz, axis=2)  # [N,3,3]
+#     xyz_dummy = xyz_0 * np.nan
+#     xyz_merged = np.stack([xyz_0, xyz_1, xyz_dummy], axis=2)  # [N,3,3,3]
+#     xyz_merged = xyz_merged.reshape(-1, 3)
+#     return xyz_merged
+
+
+# def get_xyz_indicators(pose, length=0.1):
+#     """Get xyz coordinate axis indicators for a camera pose."""
+#     xyz = np.eye(4, 3)[None] * length
+#     xyz = cam2world(xyz, pose)
+#     return xyz
+
+
+def merge_wireframes_plotly(wireframe):
+    """
+    Merge camera wireframes for efficient Plotly visualization.
+    """
+    wf_dummy = wireframe[:, :1] * np.nan  # Create NaN separators
+    wireframe_merged = np.concatenate([wireframe, wf_dummy], axis=1).reshape(-1, 3)
+    return wireframe_merged
+
+
+def merge_meshes(vertices, faces):
+    """
+    Merge multiple camera meshes into a single mesh for efficient rendering.
+    """
+    mesh_N, vertex_N = vertices.shape[:2]
+    # Adjust face indices for each mesh by adding vertex offset
+    faces_merged = np.concatenate([faces + i * vertex_N for i in range(mesh_N)], axis=0)
+    # Flatten all vertices into single array
+    vertices_merged = vertices.reshape(-1, vertices.shape[-1])
+    return vertices_merged, faces_merged
+
+
+def unbind_np(array, axis=0):
+    """
+    Split numpy array along specified axis into a list of arrays.
+    """
+    if axis == 0:
+        return [array[i, :] for i in range(array.shape[0])]
+    elif axis == 1 or (len(array.shape) == 2 and axis == -1):
+        return [array[:, j] for j in range(array.shape[1])]
+    elif axis == 2 or (len(array.shape) == 3 and axis == -1):
+        return [array[:, :, j] for j in range(array.shape[2])]
+    else:
+        raise ValueError("Invalid axis. Use 0 for rows, 1 for columns, or 2 for depth.")
+
+
+def plotly_visualize_pose(
+    poses, vis_depth=0.5, xyz_length=0.5, center_size=2, xyz_width=5, mesh_opacity=0.05
+):
+    """
+    Create comprehensive Plotly visualization traces for camera poses.
+    """
+    N = len(poses)
+
+    # Calculate camera centers in world coordinates
+    centers_cam = np.zeros([N, 1, 3])  # Camera centers in camera space (origin)
+    centers_world = cam2world(centers_cam, poses)  # Transform to world space
+    centers_world = centers_world[:, 0]  # Remove extra dimension [N, 3]
+
+    # Generate camera frustum geometry
+    vertices, faces, wireframe = get_camera_mesh(poses, depth=vis_depth)
+
+    # Merge all camera meshes into single arrays for efficient rendering
+    vertices_merged, faces_merged = merge_meshes(vertices, faces)
+    wireframe_merged = merge_wireframes_plotly(wireframe)
+
+    # Extract x, y, z coordinates for Plotly
+    wireframe_x, wireframe_y, wireframe_z = unbind_np(wireframe_merged, axis=-1)
+    centers_x, centers_y, centers_z = unbind_np(centers_world, axis=-1)
+    vertices_x, vertices_y, vertices_z = unbind_np(vertices_merged, axis=-1)
+
+    # Set up rainbow color mapping for trajectory progression
+    color_map = plt.get_cmap("gist_rainbow")  # red -> yellow -> green -> blue -> purple
+    center_color = []
+    faces_merged_color = []
+    wireframe_color = []
+
+    # Determine quarter positions for emphasis (start, 1/3, 2/3, end)
+    quarter_indices = set([0])  # Always include start
+    if N >= 3:
+        quarter_indices.add(N // 3)
+        quarter_indices.add(2 * N // 3)
+    quarter_indices.add(N - 1)  # Always include end
+
+    # Apply colors with emphasis on key trajectory points
+    for i in range(N):
+        # Emphasize quarter positions with higher opacity and brightness
+        is_quarter = i in quarter_indices
+        alpha = 6.0 if is_quarter else 0.4  # Higher opacity for key points
+
+        # Generate color from rainbow colormap
+        r, g, b, _ = color_map(i / (N - 1))
+        rgb = np.array([r, g, b]) * (1.2 if is_quarter else 0.8)  # Brighten key points
+        rgba = np.concatenate([rgb, [alpha]])
+
+        # Apply colors to all visualization elements
+        wireframe_color += [rgba] * 11  # 11 line segments per camera wireframe
+        center_color += [rgba]
+        faces_merged_color += [rgba] * 6  # 6 triangular faces per camera frustum
+
+    # Create Plotly trace objects
+    plotly_traces = [
+        # Camera wireframe outlines
+        go.Scatter3d(
+            x=wireframe_x,
+            y=wireframe_y,
+            z=wireframe_z,
+            mode="lines",
+            line=dict(color=wireframe_color, width=1),
+            name="Camera Wireframes",
+        ),
+        # Camera center points
+        go.Scatter3d(
+            x=centers_x,
+            y=centers_y,
+            z=centers_z,
+            mode="markers",
+            marker=dict(color=center_color, size=center_size, opacity=1),
+            name="Camera Centers",
+        ),
+        # Camera frustum mesh faces
+        go.Mesh3d(
+            x=vertices_x,
+            y=vertices_y,
+            z=vertices_z,
+            i=[f[0] for f in faces_merged],
+            j=[f[1] for f in faces_merged],
+            k=[f[2] for f in faces_merged],
+            facecolor=faces_merged_color,
+            opacity=mesh_opacity,
+            name="Camera Frustums",
+        ),
+    ]
+    return plotly_traces
+
+
+def compute_optimal_camera_view(poses):
+    """
+    Compute optimal camera view parameters to ensure the entire trajectory is visible
+    and aesthetically pleasing.
+    """
+    # Calculate all camera positions in world coordinates
+    centers_cam = np.zeros([len(poses), 1, 3])
+    centers_world = cam2world(centers_cam, poses)[:, 0]
+
+    # Compute bounding box of the trajectory
+    min_coords = np.min(centers_world, axis=0)
+    max_coords = np.max(centers_world, axis=0)
+    ranges = max_coords - min_coords
+
+    # Calculate trajectory center point
+    trajectory_center = (min_coords + max_coords) / 2
+
+    # Calculate maximum range for adaptive scaling
+    max_range = np.max(ranges)
+
+    # Set minimum range to avoid division by zero for very small trajectories
+    if max_range < 1e-6:
+        max_range = 1.0
+        ranges = np.ones(3)
+
+    # Calculate principal direction of trajectory using PCA (Principal Component Analysis)
+    if len(centers_world) > 1:
+        # Center the points by subtracting the mean
+        centered_points = centers_world - trajectory_center
+
+        # Compute covariance matrix for PCA
+        cov_matrix = np.cov(centered_points.T)
+
+        # Calculate eigenvalues and eigenvectors
+        eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)
+
+        # Sort by eigenvalues in descending order
+        idx = np.argsort(eigenvalues)[::-1]
+        eigenvalues = eigenvalues[idx]
+        eigenvectors = eigenvectors[:, idx]
+
+        # Main direction is the first eigenvector (highest variance)
+        main_direction = eigenvectors[:, 0]
+
+        # Ensure main direction points towards trajectory's positive direction
+        start_to_end = centers_world[-1] - centers_world[0]
+        if np.dot(main_direction, start_to_end) < 0:
+            main_direction = -main_direction
+
+    else:
+        # Default direction for single pose or insufficient data
+        main_direction = np.array([1, 0, 0])
+
+    # Calculate optimal camera distance
+    # Based on trajectory range and field of view, using smaller factor for better screen filling
+    fov_factor = (
+        0.8  # Reduced field of view factor to make trajectory occupy more screen space
+    )
+    base_distance = max_range * fov_factor
+
+    # Consider trajectory aspect ratio and adjust distance accordingly
+    aspect_ratios = ranges / max_range
+    distance_scale = 1.0 + 0.1 * np.std(
+        aspect_ratios
+    )  # Reduced distance adjustment magnitude
+    camera_distance = base_distance * distance_scale
+
+    # Calculate optimal camera position
+    # Method 1: Diagonal viewing angle based on main direction
+    up_vector = np.array([0, 0, 1])  # World up direction (Z-axis)
+
+    # Adjust strategy if main direction is nearly vertical
+    if abs(np.dot(main_direction, up_vector)) > 0.9:
+        # Main direction is nearly vertical, use side view
+        view_direction = np.cross(main_direction, np.array([1, 0, 0]))
+        if np.linalg.norm(view_direction) < 0.1:
+            view_direction = np.cross(main_direction, np.array([0, 1, 0]))
+        view_direction = view_direction / np.linalg.norm(view_direction)
+    else:
+        # Calculate diagonal view direction perpendicular to main direction
+        # Combine horizontal component of main direction with tilt angle
+        horizontal_component = (
+            main_direction - np.dot(main_direction, up_vector) * up_vector
+        )
+        horizontal_component = horizontal_component / (
+            np.linalg.norm(horizontal_component) + 1e-8
+        )
+
+        # Add some tilt angles for better 3D perspective
+        elevation_angle = np.pi / 6  # 30 degrees elevation angle
+        azimuth_offset = np.pi / 4  # 45 degrees azimuth offset
+
+        # Create tilted view direction for optimal 3D perspective
+        view_direction = (
+            horizontal_component * np.cos(azimuth_offset) * np.cos(elevation_angle)
+            + np.cross(horizontal_component, up_vector)
+            * np.sin(azimuth_offset)
+            * np.cos(elevation_angle)
+            + up_vector * np.sin(elevation_angle)
+        )
+
+    # Calculate camera eye position
+    camera_eye = trajectory_center + view_direction * camera_distance
+
+    # Fine-tune camera position to ensure entire trajectory is within view
+    # Calculate vectors from camera position to all trajectory points
+    view_vectors = centers_world - camera_eye
+    view_distances = np.linalg.norm(view_vectors, axis=1)
+
+    # Adjust camera distance moderately if some points are too close
+    min_distance = camera_distance * 0.3  # Reduced minimum distance ratio
+    if np.min(view_distances) < min_distance:
+        distance_adjustment = min_distance / np.min(view_distances)
+        # Limit adjustment magnitude to avoid excessive scaling
+        distance_adjustment = min(
+            distance_adjustment, 1.2
+        )  # Further limit adjustment range
+        camera_eye = (
+            trajectory_center + view_direction * camera_distance * distance_adjustment
+        )
+
+    # Calculate adaptive parameters with appropriate proportions
+    auto_vis_depth = max_range * 0.08  # Moderately reduced camera frustum size
+    auto_center_size = max_range * 1.5  # Moderately reduced center point size
+
+    # Ensure parameters are within reasonable bounds
+    auto_vis_depth = max(0.01, min(auto_vis_depth, max_range * 0.2))
+    auto_center_size = max(0.1, min(auto_center_size, max_range * 2.0))
+
+    return {
+        "camera_eye": camera_eye,
+        "trajectory_center": trajectory_center,
+        "auto_vis_depth": auto_vis_depth,
+        "auto_center_size": auto_center_size,
+        "max_range": max_range,
+        "ranges": ranges,
+        "main_direction": main_direction,
+    }
+
+
+def compute_multiple_camera_views(poses):
+    """
+    Compute multiple optimized camera view angles, providing different viewing options.
+    """
+    base_params = compute_optimal_camera_view(poses)
+
+    trajectory_center = base_params["trajectory_center"]
+    max_range = base_params["max_range"]
+    main_direction = base_params["main_direction"]
+
+    # Calculate multiple view options
+    views = {}
+
+    # 1. Best automatic view (original optimal view)
+    views["optimal"] = base_params
+
+    # 2. Top-down bird's eye view
+    top_distance = max_range * 1.5  # Further reduced top-down view distance
+    views["top"] = {
+        **base_params,
+        "camera_eye": trajectory_center + np.array([0, 0, top_distance]),
+        "description": "Top-down view",
+    }
+
+    # 3. Side view perspective
+    side_distance = max_range * 1.3  # Further reduced side view distance
+    side_direction = np.cross(main_direction, np.array([0, 0, 1]))
+    if np.linalg.norm(side_direction) < 0.1:
+        side_direction = np.array([1, 0, 0])
+    else:
+        side_direction = side_direction / np.linalg.norm(side_direction)
+
+    views["side"] = {
+        **base_params,
+        "camera_eye": trajectory_center + side_direction * side_distance,
+        "description": "Side view",
+    }
+
+    # 4. Diagonal view (45-degree elevation)
+    diagonal_distance = max_range * 1.4  # Further reduced diagonal view distance
+    elevation = np.pi / 4  # 45 degrees elevation
+    azimuth = np.pi / 4  # 45 degrees azimuth angle
+
+    diagonal_direction = np.array(
+        [
+            np.cos(elevation) * np.cos(azimuth),
+            np.cos(elevation) * np.sin(azimuth),
+            np.sin(elevation),
+        ]
+    )
+
+    views["diagonal"] = {
+        **base_params,
+        "camera_eye": trajectory_center + diagonal_direction * diagonal_distance,
+        "description": "Diagonal view (45° elevation)",
+    }
+
+    # 5. Trajectory start-oriented view
+    if len(poses) > 1:
+        start_to_center = trajectory_center - base_params["camera_eye"]
+        start_distance = max_range * 1.2  # Further reduced start view distance
+        start_direction = start_to_center / (np.linalg.norm(start_to_center) + 1e-8)
+
+        views["trajectory_start"] = {
+            **base_params,
+            "camera_eye": trajectory_center + start_direction * start_distance,
+            "description": "View from trajectory start direction",
+        }
+
+    # 6. Compact view - ensure entire trajectory is fully visible
+    fit_distance = max_range * 0.6  # Very compact distance for close-up view
+    fit_direction = np.array([0.7, 0.7, 0.5])  # Stable viewing direction
+    fit_direction = fit_direction / np.linalg.norm(fit_direction)
+
+    views["fit_all"] = {
+        **base_params,
+        "camera_eye": trajectory_center + fit_direction * fit_distance,
+        "description": "Fit all trajectory in view",
+    }
+
+    return views
+
+
+def add_view_selector_to_html(html_str, views):
+    """
+    Add interactive view selector to HTML visualization.
+
+    This function injects JavaScript code into the HTML to provide an interactive
+    interface for switching between different camera views and enabling auto-rotation.
+
+    Args:
+        html_str: Original HTML string containing the Plotly visualization
+        views: Dictionary of view configurations
+
+    Returns:
+        str: Enhanced HTML string with view selector and controls
+    """
+
+    # Generate JavaScript code for view selector
+    view_selector_js = """
+    <div id="view-selector" style="position: fixed; top: 10px; left: 10px; background: rgba(255,255,255,0.9); padding: 15px; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); font-family: Arial, sans-serif; font-size: 12px; z-index: 1000; min-width: 120px;">
+        <button onclick="autoRotate()" style="background: #ffc107; color: black; border: none; padding: 8px 12px; border-radius: 4px; cursor: pointer; width: 100%;">Auto Rotate</button>
+    </div>
+    
+    <script>
+    // Pre-defined view configurations
+    const views = {"""
+
+    # Add view data to JavaScript
+    for view_name, view_data in views.items():
+        eye = view_data["camera_eye"]
+        center = view_data["trajectory_center"]
+        view_selector_js += f"""
+        {view_name}: {{
+            eye: {{x: {eye[0]:.6f}, y: {eye[1]:.6f}, z: {eye[2]:.6f}}},
+            center: {{x: {center[0]:.6f}, y: {center[1]:.6f}, z: {center[2]:.6f}}},
+            up: {{x: 0, y: 0, z: 1}}
+        }},"""
+
+    view_selector_js += """
+    };
+    
+    let rotationInterval = null;
+    
+    function autoRotate() {
+        if (rotationInterval) {
+            clearInterval(rotationInterval);
+            rotationInterval = null;
+            return;
+        }
+        
+        var plotlyDiv = document.querySelector('.plotly-graph-div');
+        if (!plotlyDiv) return;
+        
+        var currentView = views.fit_all;
+        var center = currentView.center;
+        var radius = Math.sqrt(
+            Math.pow(currentView.eye.x - center.x, 2) + 
+            Math.pow(currentView.eye.y - center.y, 2) + 
+            Math.pow(currentView.eye.z - center.z, 2)
+        );
+        
+        var angle = 0;
+        rotationInterval = setInterval(function() {
+            angle += 0.02; // Rotation speed
+            
+            var newEye = {
+                x: center.x + radius * Math.cos(angle) * 0.7,
+                y: center.y + radius * Math.sin(angle) * 0.7,
+                z: center.z + radius * 0.5
+            };
+            
+            var update = {
+                'scene.camera.eye': newEye
+            };
+            
+            Plotly.relayout(plotlyDiv, update);
+        }, 50);
+    }
+    
+    // Set default view after page loading is complete
+    document.addEventListener('DOMContentLoaded', function() {
+        setTimeout(function() {
+            // Use Fit All as default view, no button operation required
+            var plotlyDiv = document.querySelector('.plotly-graph-div');
+            if (plotlyDiv && views.fit_all) {
+                var update = {
+                    'scene.camera': views.fit_all
+                };
+                Plotly.relayout(plotlyDiv, update);
+            }
+        }, 1000);
+    });
+    </script>
+    """
+
+    # Add view selector to the beginning of HTML
+    return view_selector_js + html_str
+
+
+def write_html(poses, file, vis_depth=1, xyz_length=0.2, center_size=0.01, xyz_width=2):
+    """
+    Write camera pose visualization to HTML file with optimized camera view.
+    """
+    # Calculate basic optimal view parameters
+    base_view = compute_optimal_camera_view(poses)
+
+    # Extract trajectory information
+    trajectory_center = base_view["trajectory_center"]
+    max_range = base_view["max_range"]
+    ranges = base_view["ranges"]
+    auto_vis_depth = base_view["auto_vis_depth"]
+    auto_center_size = base_view["auto_center_size"]
+
+    # Calculate optimal view to see entire trajectory
+    # Use larger distance to ensure entire trajectory is visible with better angles
+    optimal_distance = (
+        max_range * 1.8 * 10
+    )  # Increase distance by 10x for better overall view
+
+    # Choose ideal angle that can see the full trajectory
+    # Use combination of 45-degree elevation and azimuth for good 3D perspective
+    elevation = np.pi / 4  # 45-degree elevation angle
+    azimuth = np.pi / 4  # 45-degree azimuth angle
+
+    # Calculate optimal viewing direction
+    optimal_direction = np.array(
+        [
+            np.cos(elevation) * np.cos(azimuth),
+            np.cos(elevation) * np.sin(azimuth),
+            np.sin(elevation),
+        ]
+    )
+
+    # Calculate optimal camera position
+    camera_eye = trajectory_center + optimal_direction * optimal_distance
+
+    # Verify view coverage - ensure all trajectory points are within reasonable distance
+    centers_cam = np.zeros([len(poses), 1, 3])
+    centers_world = cam2world(centers_cam, poses)[:, 0]
+
+    # Calculate distances from optimal camera position to all trajectory points
+    distances_to_points = np.linalg.norm(centers_world - camera_eye, axis=1)
+    max_distance_to_point = np.max(distances_to_points)
+    min_distance_to_point = np.min(distances_to_points)
+
+    # If distance variation is too large, the view might not be ideal, adjust accordingly
+    if max_distance_to_point / min_distance_to_point > 3.0:
+        # Recalculate more balanced distance
+        optimal_distance = max_range * 2.2 * 10  # Further increase distance (10x)
+        camera_eye = trajectory_center + optimal_direction * optimal_distance
+
+    # Create view dictionary with only optimal view for Auto Rotate
+    views = {
+        "fit_all": {
+            "camera_eye": camera_eye,
+            "trajectory_center": trajectory_center,
+            "auto_vis_depth": auto_vis_depth,
+            "auto_center_size": auto_center_size,
+            "max_range": max_range,
+            "ranges": ranges,
+            "description": "Optimal view to see entire trajectory",
+        }
+    }
+
+    print(f"Trajectory ranges: x={ranges[0]:.3f}, y={ranges[1]:.3f}, z={ranges[2]:.3f}")
+    print(f"Max range: {max_range:.3f}")
+    print(f"Auto vis_depth: {auto_vis_depth:.3f}, center_size: {auto_center_size:.3f}")
+    print(
+        f"Trajectory center: ({trajectory_center[0]:.3f}, {trajectory_center[1]:.3f}, {trajectory_center[2]:.3f})"
+    )
+    print(
+        f"Optimal camera position for full trajectory view: ({camera_eye[0]:.3f}, {camera_eye[1]:.3f}, {camera_eye[2]:.3f})"
+    )
+    print(f"Camera distance from trajectory center: {optimal_distance:.3f}")
+    print(
+        f"Distance range to trajectory points: {min_distance_to_point:.3f} - {max_distance_to_point:.3f}"
+    )
+
+    xyz_length = xyz_length / 3
+    xyz_width = xyz_width
+    vis_depth = auto_vis_depth  # Use automatically computed depth
+    center_size = auto_center_size  # Use automatically computed size
+
+    traces_poses = plotly_visualize_pose(
+        poses,
+        vis_depth=vis_depth,
+        xyz_length=xyz_length,
+        center_size=center_size,
+        xyz_width=xyz_width,
+        mesh_opacity=0.05,
+    )
+    traces_all2 = traces_poses
+    layout2 = go.Layout(
+        scene=dict(
+            xaxis=dict(visible=False),
+            yaxis=dict(visible=False),
+            zaxis=dict(visible=False),
+            dragmode="orbit",
+            aspectratio=dict(x=1, y=1, z=1),
+            aspectmode="data",
+            # Set initial camera view to fully see the trajectory with optimized positioning
+            camera=dict(
+                eye=dict(x=camera_eye[0], y=camera_eye[1], z=camera_eye[2]),
+                center=dict(
+                    x=trajectory_center[0],
+                    y=trajectory_center[1],
+                    z=trajectory_center[2],
+                ),
+                up=dict(x=0, y=0, z=1),
+            ),
+        ),
+        height=800,
+        width=1200,
+        showlegend=False,
+    )
+
+    fig2 = go.Figure(data=traces_all2, layout=layout2)
+    html_str2 = pio.to_html(fig2, full_html=False)
+
+    # Add real-time camera view display functionality
+    camera_info_html = """
+    <div id="camera-info" style="position: fixed; top: 10px; right: 10px; background: rgba(255,255,255,0.9); padding: 15px; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); font-family: monospace; font-size: 12px; z-index: 1000; min-width: 250px;">
+        <h4 style="margin: 0 0 10px 0; color: #333;">Camera Info</h4>
+        <div><strong>Eye:</strong></div>
+        <div>x: <span id="eye-x">2.000</span></div>
+        <div>y: <span id="eye-y">2.000</span></div>
+        <div>z: <span id="eye-z">1.000</span></div>
+        <br>
+        <div><strong>Center:</strong></div>
+        <div>x: <span id="center-x">0.000</span></div>
+        <div>y: <span id="center-y">0.000</span></div>
+        <div>z: <span id="center-z">0.000</span></div>
+        <br>
+        <div><strong>Up:</strong></div>
+        <div>x: <span id="up-x">0.000</span></div>
+        <div>y: <span id="up-y">0.000</span></div>
+        <div>z: <span id="up-z">1.000</span></div>
+        <br>
+        <button onclick="copyToClipboard()" style="background: #007bff; color: white; border: none; padding: 5px 10px; border-radius: 4px; cursor: pointer; width: 100%;">Copy to Clipboard</button>
+    </div>
+    
+    <script>
+    function updateCameraInfo() {
+        // Get Plotly chart
+        var plotlyDiv = document.querySelector('.plotly-graph-div');
+        if (!plotlyDiv) return;
+        
+        // Listen for camera change events
+        plotlyDiv.on('plotly_relayout', function(eventData) {
+            if (eventData['scene.camera']) {
+                var camera = eventData['scene.camera'];
+                updateCameraDisplay(camera);
+            }
+        });
+        
+        // Initial display
+        setTimeout(function() {
+            var gd = plotlyDiv;
+            if (gd.layout && gd.layout.scene && gd.layout.scene.camera) {
+                updateCameraDisplay(gd.layout.scene.camera);
+            }
+        }, 1000);
+    }
+    
+    function updateCameraDisplay(camera) {
+        if (camera.eye) {
+            document.getElementById('eye-x').textContent = camera.eye.x.toFixed(3);
+            document.getElementById('eye-y').textContent = camera.eye.y.toFixed(3);
+            document.getElementById('eye-z').textContent = camera.eye.z.toFixed(3);
+        }
+        if (camera.center) {
+            document.getElementById('center-x').textContent = camera.center.x.toFixed(3);
+            document.getElementById('center-y').textContent = camera.center.y.toFixed(3);
+            document.getElementById('center-z').textContent = camera.center.z.toFixed(3);
+        }
+        if (camera.up) {
+            document.getElementById('up-x').textContent = camera.up.x.toFixed(3);
+            document.getElementById('up-y').textContent = camera.up.y.toFixed(3);
+            document.getElementById('up-z').textContent = camera.up.z.toFixed(3);
+        }
+    }
+    
+    function copyToClipboard() {
+        var eyeX = document.getElementById('eye-x').textContent;
+        var eyeY = document.getElementById('eye-y').textContent;
+        var eyeZ = document.getElementById('eye-z').textContent;
+        var centerX = document.getElementById('center-x').textContent;
+        var centerY = document.getElementById('center-y').textContent;
+        var centerZ = document.getElementById('center-z').textContent;
+        var upX = document.getElementById('up-x').textContent;
+        var upY = document.getElementById('up-y').textContent;
+        var upZ = document.getElementById('up-z').textContent;
+        
+        var cameraConfig = `camera=dict(
+    eye=dict(x=${eyeX}, y=${eyeY}, z=${eyeZ}),
+    center=dict(x=${centerX}, y=${centerY}, z=${centerZ}),
+    up=dict(x=${upX}, y=${upY}, z=${upZ})
+)`;
+        
+        navigator.clipboard.writeText(cameraConfig).then(function() {
+            alert('Copy to clipboard successful!');
+        }).catch(function(err) {
+            console.error('Copy failed:', err);
+            // Fallback: Create a temporary textarea
+            var textArea = document.createElement('textarea');
+            textArea.value = cameraConfig;
+            document.body.appendChild(textArea);
+            textArea.select();
+            document.execCommand('copy');
+            document.body.removeChild(textArea);
+            alert('Copy to clipboard successful!');
+        });
+    }
+
+    // Initialize camera info display
+    document.addEventListener('DOMContentLoaded', function() {
+        updateCameraInfo();
+    });
+
+    // If the page has already loaded
+    if (document.readyState === 'complete') {
+        updateCameraInfo();
+    }
+    </script>
+    """
+
+    # Add view selector and camera info to HTML
+    enhanced_html = add_view_selector_to_html(camera_info_html + html_str2, views)
+
+    file.write(enhanced_html)
+
+    print(f"Enhanced visualized poses are saved to {file.name}")
+    # Removed redundant view options printing
+
+
+def plotly_visualize_pose_animated(
+    poses_full,
+    vis_depth=0.5,
+    xyz_length=0.5,
+    center_size=2,
+    xyz_width=5,
+    mesh_opacity=0.05,
+):
+    """
+    Create plotly visualization traces for camera poses, frame by frame for animation.
+    Now shows the full trajectory with future poses as completely transparent.
+    """
+    N_total = len(poses_full)
+    plotly_frames = []
+
+    # Pre-compute data for all poses to ensure consistent layout
+    centers_cam = np.zeros([N_total, 1, 3])
+    centers_world = cam2world(centers_cam, poses_full)
+    centers_world = centers_world[:, 0]
+    # Get the camera wireframes for all poses
+    vertices, faces, wireframe = get_camera_mesh(poses_full, depth=vis_depth)
+    vertices_merged, faces_merged = merge_meshes(vertices, faces)
+    wireframe_merged = merge_wireframes_plotly(wireframe)
+    # Break up (x,y,z) coordinates.
+    wireframe_x, wireframe_y, wireframe_z = unbind_np(wireframe_merged, axis=-1)
+    centers_x, centers_y, centers_z = unbind_np(centers_world, axis=-1)
+    vertices_x, vertices_y, vertices_z = unbind_np(vertices_merged, axis=-1)
+
+    # Initial frame showing all poses with appropriate transparency
+    initial_data = []
+
+    for i in tqdm(range(1, N_total + 1), desc="Generating animation frames"):
+        current_frame = i - 1  # Current frame index (0-based)
+
+        # Set the color map for the camera trajectory
+        color_map = plt.get_cmap("gist_rainbow")
+        center_color = []
+        faces_merged_color = []
+        wireframe_color = []
+
+        for k in range(N_total):  # Process all poses
+            # Set the camera pose colors (with a smooth gradient color map).
+            r, g, b, _ = color_map(k / (N_total - 1))
+            rgb = np.array([r, g, b]) * 0.8
+
+            # Set transparency based on current frame
+            if k < current_frame:  # Past poses - visible with reduced opacity
+                # Set transparency based on temporal distance, more distant = more transparent
+                time_distance = (current_frame - k) / max(current_frame, 1)
+                alpha = 0.15 + 0.25 * (1 - time_distance)  # Transparency range 0.15-0.4
+                wireframe_alpha = alpha
+                mesh_alpha = alpha * 0.4
+            elif k == current_frame:  # Current pose - fully visible
+                alpha = 0.8  # Fully opaque, dark display
+                wireframe_alpha = 0.8
+                mesh_alpha = 0.6
+            else:  # Future poses - completely transparent
+                alpha = 0.0  # Completely transparent
+                wireframe_alpha = 0.0
+                mesh_alpha = 0.0
+
+            # Set colors and transparency
+            wireframe_color += [np.concatenate([rgb, [wireframe_alpha]])] * 11
+            center_color += [np.concatenate([rgb, [alpha]])]
+            faces_merged_color += [np.concatenate([rgb, [mesh_alpha]])] * 6
+
+        frame_data = [
+            go.Scatter3d(
+                x=wireframe_x,
+                y=wireframe_y,
+                z=wireframe_z,
+                mode="lines",
+                line=dict(color=wireframe_color, width=1),
+            ),
+            go.Scatter3d(
+                x=centers_x,
+                y=centers_y,
+                z=centers_z,
+                mode="markers",
+                marker=dict(color=center_color, size=center_size),
+            ),
+            go.Mesh3d(
+                x=vertices_x,
+                y=vertices_y,
+                z=vertices_z,
+                i=[f[0] for f in faces_merged],
+                j=[f[1] for f in faces_merged],
+                k=[f[2] for f in faces_merged],
+                facecolor=faces_merged_color,
+                opacity=0.6,  # Set base opacity for mesh
+            ),
+        ]
+
+        if i == 1:  # Set initial data for the first frame
+            initial_data = frame_data
+
+        plotly_frames.append(go.Frame(data=frame_data, name=str(i)))
+
+    return initial_data, plotly_frames
+
+
+def write_html_animated(
+    poses, file, vis_depth=1, xyz_length=0.2, center_size=0.01, xyz_width=2
+):
+    """
+    Write camera pose visualization with animation to HTML file with optimized camera view.
+    """
+    # Calculate basic optimal view parameters
+    base_view = compute_optimal_camera_view(poses)
+
+    # Extract trajectory information
+    trajectory_center = base_view["trajectory_center"]
+    max_range = base_view["max_range"]
+    ranges = base_view["ranges"]
+    auto_vis_depth = base_view["auto_vis_depth"]
+    auto_center_size = base_view["auto_center_size"]
+
+    # Calculate optimal view to see entire trajectory
+    # Use larger distance to ensure entire trajectory is visible with better angles
+    optimal_distance = (
+        max_range * 1.8 * 10
+    )  # Increase distance by 10x for better overall view
+
+    # Choose ideal angle that can see the full trajectory
+    # Use combination of 45-degree elevation and azimuth for good 3D perspective
+    elevation = np.pi / 4  # 45-degree elevation angle
+    azimuth = np.pi / 4  # 45-degree azimuth angle
+
+    # Calculate optimal viewing direction
+    optimal_direction = np.array(
+        [
+            np.cos(elevation) * np.cos(azimuth),
+            np.cos(elevation) * np.sin(azimuth),
+            np.sin(elevation),
+        ]
+    )
+
+    # Calculate optimal camera position
+    camera_eye = trajectory_center + optimal_direction * optimal_distance
+
+    # Verify view coverage - ensure all trajectory points are within reasonable distance
+    centers_cam = np.zeros([len(poses), 1, 3])
+    centers_world = cam2world(centers_cam, poses)[:, 0]
+
+    # Calculate distances from optimal camera position to all trajectory points
+    distances_to_points = np.linalg.norm(centers_world - camera_eye, axis=1)
+    max_distance_to_point = np.max(distances_to_points)
+    min_distance_to_point = np.min(distances_to_points)
+
+    # If distance variation is too large, the view might not be ideal, adjust accordingly
+    if max_distance_to_point / min_distance_to_point > 3.0:
+        # Recalculate more balanced distance
+        optimal_distance = max_range * 2.2 * 10  # Further increase distance (10x)
+        camera_eye = trajectory_center + optimal_direction * optimal_distance
+
+    # Adjust parameters for animation
+    xyz_length = xyz_length / 3
+    xyz_width = xyz_width
+    vis_depth = auto_vis_depth  # Use automatically computed depth
+    center_size = auto_center_size  # Use automatically computed size
+
+    print(
+        f"Animation - Trajectory ranges: x={ranges[0]:.3f}, y={ranges[1]:.3f}, z={ranges[2]:.3f}"
+    )
+    print(f"Animation - Max range: {max_range:.3f}")
+    print(
+        f"Animation - Auto vis_depth: {auto_vis_depth:.3f}, center_size: {auto_center_size:.3f}"
+    )
+    print(
+        f"Animation - Trajectory center: ({trajectory_center[0]:.3f}, {trajectory_center[1]:.3f}, {trajectory_center[2]:.3f})"
+    )
+    print(
+        f"Animation - Optimal camera position for full trajectory view: ({camera_eye[0]:.3f}, {camera_eye[1]:.3f}, {camera_eye[2]:.3f})"
+    )
+    print(f"Animation - Camera distance from trajectory center: {optimal_distance:.3f}")
+    print(
+        f"Animation - Distance range to trajectory points: {min_distance_to_point:.3f} - {max_distance_to_point:.3f}"
+    )
+
+    initial_data, plotly_frames = plotly_visualize_pose_animated(
+        poses,
+        vis_depth=vis_depth,
+        xyz_length=xyz_length,
+        center_size=center_size,
+        xyz_width=xyz_width,
+        mesh_opacity=0.05,
+    )
+
+    layout = go.Layout(
+        scene=dict(
+            xaxis=dict(visible=False),
+            yaxis=dict(visible=False),
+            zaxis=dict(visible=False),
+            dragmode="orbit",
+            aspectratio=dict(x=1, y=1, z=1),
+            aspectmode="data",
+            # Use optimized camera view settings (same 10x distance as write_html)
+            camera=dict(
+                eye=dict(x=camera_eye[0], y=camera_eye[1], z=camera_eye[2]),
+                center=dict(
+                    x=trajectory_center[0],
+                    y=trajectory_center[1],
+                    z=trajectory_center[2],
+                ),
+                up=dict(x=0, y=0, z=1),
+            ),
+        ),
+        height=800,  # Increased height for better animation display
+        width=1200,  # Increased width for better animation display
+        showlegend=False,
+        updatemenus=[
+            dict(
+                type="buttons",
+                buttons=[
+                    dict(
+                        label="Play",
+                        method="animate",
+                        args=[
+                            None,
+                            {
+                                "frame": {"duration": 50, "redraw": True},
+                                "fromcurrent": True,
+                                "transition": {"duration": 0},
+                            },
+                        ],
+                    )
+                ],
+            )
+        ],
+    )
+
+    fig = go.Figure(data=initial_data, layout=layout, frames=plotly_frames)
+    html_str = pio.to_html(fig, full_html=False)
+    file.write(html_str)
+
+    print(f"Visualized poses are saved to {file}")
+
+
+def quaternion_to_matrix(quaternions, eps: float = 1e-8):
+    """
+    Convert 4-dimensional quaternions to 3x3 rotation matrices.
+
+    Reference:
+        https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py
+    """
+
+    # Order changed to match scipy format: (i, j, k, r)
+    i, j, k, r = torch.unbind(quaternions, dim=-1)
+    two_s = 2 / ((quaternions * quaternions).sum(dim=-1) + eps)
+
+    # Construct rotation matrix elements using quaternion algebra
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),  # R[0,0]
+            two_s * (i * j - k * r),  # R[0,1]
+            two_s * (i * k + j * r),  # R[0,2]
+            two_s * (i * j + k * r),  # R[1,0]
+            1 - two_s * (i * i + k * k),  # R[1,1]
+            two_s * (j * k - i * r),  # R[1,2]
+            two_s * (i * k - j * r),  # R[2,0]
+            two_s * (j * k + i * r),  # R[2,1]
+            1 - two_s * (i * i + j * j),  # R[2,2]
+        ),
+        -1,
+    )
+    return einops.rearrange(o, "... (i j) -> ... i j", i=3, j=3)
+
+
+def pose_from_quaternion(pose):
+    """
+    Convert quaternion-based pose representation to 4x4 transformation matrices.
+
+    Reference:
+        https://github.com/pointrix-project/Geomotion/blob/6ab0c364f1b44ab4ea190085dbf068f62b42727c/geomotion/model/cameras.py#L6
+    """
+    # Convert numpy array to torch tensor if needed
+    if type(pose) == np.ndarray:
+        pose = torch.tensor(pose)
+    # Add batch dimension if input is 1D
+    if len(pose.shape) == 1:
+        pose = pose[None]
+    # Extract translation and quaternion components
+    quat_t = pose[..., :3]  # Translation components [tx, ty, tz]
+    quat_r = pose[..., 3:]  # Quaternion components [qi, qj, qk, qr]
+
+    # Initialize world-to-camera transformation matrix
+    w2c_matrix = torch.zeros((*list(pose.shape)[:-1], 3, 4), device=pose.device)
+    w2c_matrix[..., :3, 3] = quat_t  # Set translation part
+    w2c_matrix[..., :3, :3] = quaternion_to_matrix(quat_r)  # Set rotation part
+    return w2c_matrix
+
+
+def viz_poses(i, pth, file, scale_factor, dynamic, vis_depth):
+    """
+    Visualize camera poses for a sequence and write to HTML file.
+    """
+    file.write(f"<span style='font-size: 18pt;'>{i} {pth}</span><br>")
+
+    # Load pose data from file
+    pose = np.load(pth)
+
+    # Convert quaternion poses to transformation matrices
+    # poses = pose_from_quaternion(pose)  # Input: (N,7), Output: (N,3,4) w2c matrices
+    # poses = poses.cpu().numpy()
+    if isinstance(pose, np.ndarray):
+        if pose.shape[1] == 3:
+            c2w = np.eye(4)
+            c2w = repeat(c2w, "i j -> n i j", n=pose.shape[0])
+            c2w[:, :3] = pose
+            pose = c2w
+        poses = np.linalg.inv(pose)[:, :3]
+    else:
+        poses = np.linalg.inv(pose["data"])[:, :3]
+
+    # Apply scaling to translation part (camera positions) while keeping rotation unchanged
+    # Create scaled copy of poses
+    poses_scaled = poses.copy()
+    poses_scaled[..., :3, 3] = poses[..., :3, 3] * scale_factor
+
+    print(f"Original poses shape: {poses.shape}")
+    print(f"Applied scale factor: {scale_factor}")
+
+    # Generate visualization based on dynamic flag
+    if dynamic:
+        write_html_animated(poses_scaled, file, vis_depth=vis_depth)
+    else:
+        write_html(poses_scaled, file, vis_depth=vis_depth)
+
+
+def vis_to_html(outdir, datas, scale_factor=0.3, dynamic=False, vis_depth=0.2):
+    # Create output directory and process pose files
+    os.makedirs(outdir, exist_ok=True)
+
+    with open(f"{outdir}/visualize.html", "w") as file:
+        for i, pth in enumerate(tqdm(datas, desc="Processing pose files")):
+            if not os.path.exists(pth):
+                print(f"Warning: Path {pth} does not exist, skipping.")
+                continue
+            print(f"Processing: {pth} (#{i+1})")
+            viz_poses(i, pth, file, scale_factor, dynamic, vis_depth)
+
+
+if __name__ == "__main__":
+    # Set up command-line argument parser
+    parser = argparse.ArgumentParser(
+        description="Visualize camera poses with interactive 3D plots",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--datas",
+        type=str,
+        nargs="+",
+        required=True,
+        help="List of pose file paths (.npz format) to visualize.",
+    )
+    parser.add_argument(
+        "--vis_depth",
+        type=float,
+        default=0.2,
+        help="Depth of camera frustum visualization (default: 0.2).",
+    )
+    parser.add_argument(
+        "--scale_factor",
+        type=float,
+        default=0.3,
+        help="Scale factor to reduce distance between cameras - smaller values bring cameras closer together (default: 0.3).",
+    )
+    parser.add_argument(
+        "--outdir",
+        type=str,
+        default="./visualize",
+        help="Output directory to save HTML visualization files (default: ./visualize).",
+    )
+    parser.add_argument(
+        "--dynamic",
+        action="store_true",
+        help="Create animated visualization showing camera trajectory progression over time.",
+    )
+
+    # Parse command-line arguments
+    args = parser.parse_args()
+
+    print(f"Processing {len(args.datas)} pose file(s)...")
+    print(f"Output directory: {args.outdir}")
+    print(f"Visualization type: {'Animated' if args.dynamic else 'Static'}")
+
+    vis_to_html(args.outdir, args.datas, args.scale_factor, args.dynamic, args.vis_depth)
+
+    print(
+        f"Visualization complete! Open {args.outdir}/visualize.html in your browser to view results."
+    )

UCPE/tools/visualize_re10k.py

@@ -0,0 +1,222 @@
+import random
+import argparse
+import numpy as np
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+from matplotlib.patches import Patch
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+import os
+from tqdm.auto import tqdm
+import imageio
+from matplotlib.backends.backend_agg import FigureCanvasAgg
+
+
+class CameraPoseVisualizer:
+    def __init__(self, xlim, ylim, zlim):
+        self.fig = plt.figure(figsize=(7, 7), dpi=300)
+        self.ax = self.fig.add_subplot(projection='3d')
+        self.plotly_data = None  # plotly data traces
+        self.xlim = xlim
+        self.ylim = ylim
+        self.zlim = zlim
+        self.init_ax()
+        print('initialize camera pose visualizer')
+
+    def init_ax(self):
+        self.ax.cla()
+        self.ax.set_aspect("auto")
+        self.ax.set_xlim(self.xlim)
+        self.ax.set_ylim(self.ylim)
+        self.ax.set_zlim(self.zlim)
+        self.ax.set_xlabel('x')
+        self.ax.set_ylabel('y')
+        self.ax.set_zlabel('z')
+
+    def extrinsic2pyramid(self, extrinsic, color_map='red', hw_ratio=9/16, base_xval=1, zval=3):
+        vertex_std = np.array([[0, 0, 0, 1],
+                               [base_xval, -base_xval * hw_ratio, zval, 1],
+                               [base_xval, base_xval * hw_ratio, zval, 1],
+                               [-base_xval, base_xval * hw_ratio, zval, 1],
+                               [-base_xval, -base_xval * hw_ratio, zval, 1]])
+        vertex_transformed = vertex_std @ extrinsic.T
+        meshes = [[vertex_transformed[0, :-1], vertex_transformed[1][:-1], vertex_transformed[2, :-1]],
+                            [vertex_transformed[0, :-1], vertex_transformed[2, :-1], vertex_transformed[3, :-1]],
+                            [vertex_transformed[0, :-1], vertex_transformed[3, :-1], vertex_transformed[4, :-1]],
+                            [vertex_transformed[0, :-1], vertex_transformed[4, :-1], vertex_transformed[1, :-1]],
+                            [vertex_transformed[1, :-1], vertex_transformed[2, :-1], vertex_transformed[3, :-1], vertex_transformed[4, :-1]]]
+
+        color = color_map if isinstance(color_map, str) else plt.cm.rainbow(color_map)
+
+        self.ax.add_collection3d(
+            Poly3DCollection(meshes, facecolors=color, linewidths=0.3, edgecolors=color, alpha=0.35))
+
+    def customize_legend(self, list_label):
+        list_handle = []
+        for idx, label in enumerate(list_label):
+            color = plt.cm.rainbow(idx / len(list_label))
+            patch = Patch(color=color, label=label)
+            list_handle.append(patch)
+        plt.legend(loc='right', bbox_to_anchor=(1.8, 0.5), handles=list_handle)
+
+    def colorbar(self, max_frame_length):
+        cmap = mpl.cm.rainbow
+        norm = mpl.colors.Normalize(vmin=0, vmax=max_frame_length)
+        self.fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap), ax=self.ax, orientation='vertical', label='Frame Number')
+
+    def show(self, out_file_path):
+        plt.title('Extrinsic Parameters')
+        os.makedirs('debug/visualize_re10k', exist_ok=True)
+        plt.savefig(out_file_path, format='png', dpi=600)
+        plt.show()
+
+    def draw(self):
+        canvas = FigureCanvasAgg(self.fig)
+        canvas.draw()
+        buf = canvas.buffer_rgba()
+        img = np.asarray(buf, dtype=np.uint8)
+        return img[:, :, :3]
+
+    def vis_pose(self, c2ws, out_file_path, hw_ratio, base_xval, zval):
+        num_frames = len(c2ws)
+        self.colorbar(num_frames)
+        self.init_ax()
+
+        for frame_idx, c2w in enumerate(c2ws):
+            self.extrinsic2pyramid(
+                c2w,
+                frame_idx / num_frames,
+                hw_ratio=hw_ratio,
+                base_xval=base_xval,
+                zval=zval
+            )
+
+        self.show(out_file_path)
+
+    def anim_pose(self, c2ws, out_file_path, hw_ratio, base_xval, zval, fps, keyframe_interval):
+        num_frames = len(c2ws)
+        self.colorbar(num_frames)
+        writer = imageio.get_writer(out_file_path, fps=fps)
+
+        keyframes = []
+        for frame_idx, c2w in enumerate(c2ws):
+            self.init_ax()
+            if keyframes:
+                for kf in keyframes:
+                    self.extrinsic2pyramid(
+                        c2ws[kf],
+                        kf / num_frames,
+                        hw_ratio=hw_ratio,
+                        base_xval=base_xval,
+                        zval=zval
+                    )
+            if frame_idx % keyframe_interval == 0:
+                keyframes.append(frame_idx)
+            self.extrinsic2pyramid(
+                c2w,
+                frame_idx / num_frames,
+                hw_ratio=hw_ratio,
+                base_xval=base_xval,
+                zval=zval
+            )
+            img = self.draw()
+            writer.append_data(img)
+
+        writer.close()
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--pose_file_path', required=True, help='path to the trajectory txt file')
+    parser.add_argument('--filter_file', required=True, help='path to the filter txt file')
+    parser.add_argument('--out_path', required=True, help='path to save the visualization results')
+    parser.add_argument('--num_videos', type=int, default=150, help='number of videos to visualize')
+    parser.add_argument('--hw_ratio', default=2/3, type=float, help='the height over width of the film plane')
+    parser.add_argument('--sample_stride', type=int, default=1)
+    parser.add_argument('--num_frames', type=int, default=81)
+    parser.add_argument('--all_frames', action='store_true')
+    parser.add_argument('--base_xval', type=float, default=0.25)
+    parser.add_argument('--zval', type=float, default=0.5)
+    parser.add_argument('--use_exact_fx', action='store_true')
+    parser.add_argument('--relative_c2w', action='store_true')
+    parser.add_argument('--x_min', type=float, default=-2)
+    parser.add_argument('--x_max', type=float, default=2)
+    parser.add_argument('--y_min', type=float, default=-2)
+    parser.add_argument('--y_max', type=float, default=2)
+    parser.add_argument('--z_min', type=float, default=-2)
+    parser.add_argument('--z_max', type=float, default=2)
+    parser.add_argument('--animate_camera', action='store_true')
+    parser.add_argument('--fps', type=int, default=16)
+    parser.add_argument('--keyframe_interval', type=int, default=10)
+    return parser.parse_args()
+
+
+def get_c2w(w2cs, transform_matrix, relative_c2w):
+    if relative_c2w:
+        target_cam_c2w = np.array([
+            [1, 0, 0, 0],
+            [0, 1, 0, 0],
+            [0, 0, 1, 0],
+            [0, 0, 0, 1]
+        ])
+        abs2rel = target_cam_c2w @ w2cs[0]
+        ret_poses = [target_cam_c2w, ] + [abs2rel @ np.linalg.inv(w2c) for w2c in w2cs[1:]]
+    else:
+        ret_poses = [np.linalg.inv(w2c) for w2c in w2cs]
+    ret_poses = [transform_matrix @ x for x in ret_poses]
+    return np.array(ret_poses, dtype=np.float32)
+
+
+if __name__ == '__main__':
+    args = get_args()
+    os.makedirs(args.out_path, exist_ok=True)
+    with open(args.filter_file, 'r') as f:
+        video_ids = f.read().splitlines()
+    video_ids = video_ids[:args.num_videos]
+    for video_id in tqdm(video_ids, desc='Visualizing camera poses'):
+        pose_file_path = os.path.join(args.pose_file_path, f'{video_id}.txt')
+        if not os.path.exists(pose_file_path):
+            print(f'Pose file {pose_file_path} does not exist, skip.')
+            continue
+        print(f'Visualizing {pose_file_path}...')
+
+        with open(pose_file_path, 'r') as f:
+            poses = f.readlines()
+        w2cs = [np.asarray([float(p) for p in pose.strip().split(' ')[7:]]).reshape(3, 4) for pose in poses[1:]]
+        fxs = [float(pose.strip().split(' ')[1]) for pose in poses[1:]]
+        if args.all_frames:
+            args.num_frames = len(fxs)
+            args.sample_stride = 1
+        cropped_length = args.num_frames * args.sample_stride
+        total_frames = len(w2cs)
+        start_frame_ind = 0
+        end_frame_ind = min(start_frame_ind + cropped_length, total_frames)
+        frame_ind = np.linspace(start_frame_ind, end_frame_ind - 1, args.num_frames, dtype=int)
+        w2cs = [w2cs[x] for x in frame_ind]
+        transform_matrix = np.asarray([[1, 0, 0, 0], [0, 0, 1, 0], [0, -1, 0, 0], [0, 0, 0, 1]]).reshape(4, 4)
+        last_row = np.zeros((1, 4))
+        last_row[0, -1] = 1.0
+        w2cs = [np.concatenate((w2c, last_row), axis=0) for w2c in w2cs]
+        c2ws = get_c2w(w2cs, transform_matrix, args.relative_c2w)
+
+        visualizer = CameraPoseVisualizer([args.x_min, args.x_max], [args.y_min, args.y_max], [args.z_min, args.z_max])
+        zval = fxs[0] if args.use_exact_fx else args.zval
+        if args.animate_camera:
+            out_file_path = os.path.join(args.out_path, f'{video_id}.mp4')
+            visualizer.anim_pose(
+                c2ws,
+                out_file_path,
+                args.hw_ratio,
+                args.base_xval,
+                zval,
+                args.fps,
+                args.keyframe_interval,
+            )
+        else:
+            out_file_path = os.path.join(args.out_path, f'{video_id}.png')
+            visualizer.vis_pose(
+                c2ws,
+                out_file_path,
+                args.hw_ratio,
+                args.base_xval,
+                zval
+            )