+
+
+ 
+
+
+ 
+
+
+ We thank VGGT for their excellent gradio implementation!
") + + submit_btn.click( + fn=clear_fields, + inputs=[], + outputs=[reconstruction_output, rgb_video, depth_video], + ).then( + fn=gradio_demo, + inputs=[ + target_dir_output, + ], + outputs=[reconstruction_output, rgb_video, depth_video], + ).then( + fn=lambda: "False", inputs=[], outputs=[is_example] + ) + + input_video.change( + fn=update_gallery_on_upload, + inputs=[input_video, input_images], + outputs=[reconstruction_output, target_dir_output, image_gallery], + ) + input_images.change( + fn=update_gallery_on_upload, + inputs=[input_video, input_images], + outputs=[reconstruction_output, target_dir_output, image_gallery], + ) + + # demo.launch(share=share, server_name=server_name, server_port=server_port) + demo.queue(max_size=20).launch(show_error=True, share=True) + + # We thank VGGT for their excellent gradio implementation diff --git a/examples/video/bungeenerf_colosseum.mp4 b/examples/video/bungeenerf_colosseum.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..333243a8dcc35dbc1c58fac65694fe77ede1b8bf --- /dev/null +++ b/examples/video/bungeenerf_colosseum.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:416b6af945547b5d19476823672de552944c7b5a147d29e9e8243e91a16aee3e +size 329073 diff --git a/examples/video/dtu_scan_106.mp4 b/examples/video/dtu_scan_106.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..60ab7c593788a8b52156d3b0bf7e296ef260d158 --- /dev/null +++ b/examples/video/dtu_scan_106.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:16d7a06325cd368b134908e600a6c0741c7d0d188f1db690532b8ac85d65fef5 +size 352188 diff --git a/examples/video/fillerbuster_hand_hand.mp4 b/examples/video/fillerbuster_hand_hand.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..c7978c090a854d1798284f44f14d7aa5da70a3cc --- /dev/null +++ b/examples/video/fillerbuster_hand_hand.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7b4ca982672bc92342b3e722c171d9d2e4d67a5a8116cd9f346956fbe01e253f +size 319404 diff --git a/examples/video/fillerbuster_ramen.mp4 b/examples/video/fillerbuster_ramen.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..31a2941246aa4c708a488979f39be2a0fdfc4962 --- /dev/null +++ b/examples/video/fillerbuster_ramen.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d60346a64a0a0d6805131d0d57edeeb0dae24f24c3f10560e95df65531221229 +size 660736 diff --git a/examples/video/fox.mp4 b/examples/video/fox.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..a42782fe6c0a29797a2a152c40373bd90d968be1 --- /dev/null +++ b/examples/video/fox.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d3fa2ccff78e5d8085bb58f3def2d482e8df285ced5ef1b56abfe3766f0d90e0 +size 2361921 diff --git a/examples/video/horizongs_hillside_summer.mp4 b/examples/video/horizongs_hillside_summer.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..7e8b209dba1ae99397e91e37db5b6d0f7745e37e --- /dev/null +++ b/examples/video/horizongs_hillside_summer.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e5dff78d9c00b3776bfca3a370061698bddead2ae940fe5a42d082ccf2ca80d1 +size 1606537 diff --git a/examples/video/kitti360.mp4 b/examples/video/kitti360.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..961d93921562cad405c46960d64e9d4e1aa8049d --- /dev/null +++ b/examples/video/kitti360.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c6b13929b2c2aae8b95921d8626f5be06f6afffe05ea4e47940ffeb9906f9fc +size 1843629 diff --git a/examples/video/llff_fortress.mp4 b/examples/video/llff_fortress.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..c858dffb86c1d5bea72b2220ad1545b641e205c4 --- /dev/null +++ b/examples/video/llff_fortress.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:90ea046a0ec78651975529ebe6b9c72b60c19561fe61b15b15b9df0e44d9fe9a +size 196243 diff --git a/examples/video/llff_horns.mp4 b/examples/video/llff_horns.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..5f0717367a7d8750013d0ac4e6d6b3f47978a778 --- /dev/null +++ b/examples/video/llff_horns.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3bc4c443c2a3f889f0c1283e98bd6a7026c36858fb37808bb2e8699ad1a2c1d8 +size 372570 diff --git a/examples/video/matrixcity_street.mp4 b/examples/video/matrixcity_street.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..7d6b95073dfa14a45df403080002ff18ffbf33d3 --- /dev/null +++ b/examples/video/matrixcity_street.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aa415f27177398b4e06f580beb3778701ca55784afade2fd6a058212213febc8 +size 3163684 diff --git a/examples/video/meganerf_rubble.mp4 b/examples/video/meganerf_rubble.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..50797215d0e71fb57268c1811ea51f0d2b2bcd0c --- /dev/null +++ b/examples/video/meganerf_rubble.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3410c759eb73ca2403ab8fe35d5ebabdbc25e3a0e67d8670a89fe17686246ed0 +size 450116 diff --git a/examples/video/re10k_1eca36ec55b88fe4.mp4 b/examples/video/re10k_1eca36ec55b88fe4.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..84ca9af870cbfd343dc32398579b029bb5ef2a1d Binary files /dev/null and b/examples/video/re10k_1eca36ec55b88fe4.mp4 differ diff --git a/examples/video/vrnerf_apartment.mp4 b/examples/video/vrnerf_apartment.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..bc87ca70b8b9053528cb023f9d117f8d10d94d78 --- /dev/null +++ b/examples/video/vrnerf_apartment.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4fdd5f165a4293cd95e3dd88d84b1f370decdd86308aa67a9d3832e01f4d6906 +size 2076392 diff --git a/examples/video/vrnerf_kitchen.mp4 b/examples/video/vrnerf_kitchen.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..4ef4be83a532a1357a7760fea728f81c66c9846e --- /dev/null +++ b/examples/video/vrnerf_kitchen.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3db5d766ec86a7abdfe1f033b252337e6d934ea15035fafb4d0fc0c0e9e9740a +size 775715 diff --git a/examples/video/vrnerf_riverview.mp4 b/examples/video/vrnerf_riverview.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..b4a30cf0165251fa9bb990ebd662b1e3805e03d4 --- /dev/null +++ b/examples/video/vrnerf_riverview.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9b8187936cc49910ef330a37b1bbdab0076096d6c01f33b097c11937184de168 +size 768290 diff --git a/examples/video/vrnerf_workshop.mp4 b/examples/video/vrnerf_workshop.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..afb7f73b9f56b6182f38371a344077b000410bfd --- /dev/null +++ b/examples/video/vrnerf_workshop.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c0f1334acc74bd70086a9be94d0c36838ebd7499af27f942c315e1ba282e285b +size 1718918 diff --git a/examples/vrnerf/riverview/21_DSC0001.jpg b/examples/vrnerf/riverview/21_DSC0001.jpg new file mode 100644 index 0000000000000000000000000000000000000000..802a075a6334bda6007962080b03af5cdbe7a39a --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0001.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7600a24a0725bf42c2260c748f11f39ef495a065f187dd894a6a6b643d209a79 +size 478234 diff --git a/examples/vrnerf/riverview/21_DSC0010.jpg b/examples/vrnerf/riverview/21_DSC0010.jpg new file mode 100644 index 0000000000000000000000000000000000000000..ab75e2ceb80f2f95d4637ce5d0b37ecf674104a4 --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0010.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1621945dbca88d01326e23372993067c7400acf7076c803a3300a159fdb053cc +size 467263 diff --git a/examples/vrnerf/riverview/21_DSC0019.jpg b/examples/vrnerf/riverview/21_DSC0019.jpg new file mode 100644 index 0000000000000000000000000000000000000000..b2f7ba0ee6a8bd8cbf4a52d61d051b2601880359 --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0019.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2309a28deb2797f96f516cf9bb758e97f86b84f78b9bc88dea793cd43f67715f +size 453802 diff --git a/examples/vrnerf/riverview/21_DSC0028.jpg b/examples/vrnerf/riverview/21_DSC0028.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2a025f1b078ddc914bdd75bbac05220f8762aecd --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0028.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bffe2e849bb4bbb32e6e028f966ed1169a0a7d3c6707c90c37a42a471939a1d6 +size 451209 diff --git a/examples/vrnerf/riverview/21_DSC0037.jpg b/examples/vrnerf/riverview/21_DSC0037.jpg new file mode 100644 index 0000000000000000000000000000000000000000..2d5cf7c1eee7763895ed7d88fa0115cdd1e424a1 --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0037.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2a2f68fc6c9036a62b9b601587ae82334c957d1a93ef47ed92cbff6209648591 +size 450246 diff --git a/examples/vrnerf/riverview/21_DSC0046.jpg b/examples/vrnerf/riverview/21_DSC0046.jpg new file mode 100644 index 0000000000000000000000000000000000000000..4e32c91d346d2a832e19c94b4eab31d7a74f36a4 --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0046.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4799bf9781b4e6c0ea0b3bc6bb60a718b14c7343063f7828cc99a934ca2718bf +size 452324 diff --git a/examples/vrnerf/riverview/21_DSC0055.jpg b/examples/vrnerf/riverview/21_DSC0055.jpg new file mode 100644 index 0000000000000000000000000000000000000000..0db079c6359ab986766f91b1edbfd3f9c20d7d4a --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0055.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4f8ae4628e53df1dee87e24eb0e32f475f864c693e7ff6e565478abd327ab746 +size 455865 diff --git a/examples/vrnerf/riverview/21_DSC0064.jpg b/examples/vrnerf/riverview/21_DSC0064.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d7d72edb989d336e1a02a5a377806293c068a8ae --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0064.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a89effe9cba1bb833cf231b3c43b2d2b148194a416e42d46d67d32e57a4b9a13 +size 448712 diff --git a/examples/vrnerf/riverview/21_DSC0073.jpg b/examples/vrnerf/riverview/21_DSC0073.jpg new file mode 100644 index 0000000000000000000000000000000000000000..1f677a9d4bc3a67c2efac76f17ff03ef2eee7e3f --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0073.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6ce299db65e14c3ec67568e58eec6e33b1177251153bb0284609838cb95a40d1 +size 440003 diff --git a/examples/vrnerf/riverview/21_DSC0082.jpg b/examples/vrnerf/riverview/21_DSC0082.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7eafbb06fc3d8a506a80468a88d1759b3363778b --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0082.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6dcf7be5b41a46d2a879163bcbde6cc11dcbab7fde2c437b120c428da97cfe98 +size 432309 diff --git a/examples/vrnerf/riverview/21_DSC0091.jpg b/examples/vrnerf/riverview/21_DSC0091.jpg new file mode 100644 index 0000000000000000000000000000000000000000..768a4497034b1c5e4422843d0a481ed041dacc78 --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0091.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:97229231f32411afd68a25b9bb6ffc67b5b068312a3aea675d2deda4471f6cd9 +size 437419 diff --git a/examples/vrnerf/riverview/21_DSC0100.jpg b/examples/vrnerf/riverview/21_DSC0100.jpg new file mode 100644 index 0000000000000000000000000000000000000000..b3553ef7e126e0a274c91c372da10a782eeab1a2 --- /dev/null +++ b/examples/vrnerf/riverview/21_DSC0100.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ea6bbb2a4dbafdd50d8ea49b23349d64e7677f5a70db88663cbec3e8d2211ba1 +size 436117 diff --git a/inference.py b/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a30790fdedc795621cd01c3e8b4e01121d4048 --- /dev/null +++ b/inference.py @@ -0,0 +1,39 @@ + +from pathlib import Path +import torch +import os +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.misc.image_io import save_interpolated_video +from src.model.ply_export import export_ply +from src.model.model.anysplat import AnySplat +from src.utils.image import process_image + +def main(): + # Load the model from Hugging Face + model = AnySplat.from_pretrained("lhjiang/anysplat") + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = model.to(device) + model.eval() + for param in model.parameters(): + param.requires_grad = False + + # Load Images + image_folder = "examples/vrnerf/riverview" + images = sorted([os.path.join(image_folder, f) for f in os.listdir(image_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]) + images = [process_image(img_path) for img_path in images] + images = torch.stack(images, dim=0).unsqueeze(0).to(device) # [1, K, 3, 448, 448] + b, v, _, h, w = images.shape + + # Run Inference + gaussians, pred_context_pose = model.inference((images+1)*0.5) + + # Save the results + pred_all_extrinsic = pred_context_pose['extrinsic'] + pred_all_intrinsic = pred_context_pose['intrinsic'] + save_interpolated_video(pred_all_extrinsic, pred_all_intrinsic, b, h, w, gaussians, image_folder, model.decoder) + export_ply(gaussians.means[0], gaussians.scales[0], gaussians.rotations[0], gaussians.harmonics[0], gaussians.opacities[0], Path(image_folder) / "gaussians.ply") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..b1c809772e8cab1e73844177229c2c3814db82c6 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,38 @@ +numpy==1.25.0 +wheel +tqdm +lightning +black +ruff +hydra-core +jaxtyping +beartype +wandb +einops +colorama +scikit-image +colorspacious +matplotlib +moviepy +imageio +timm +dacite +lpips +e3nn +plyfile +tabulate +svg.py +scikit-video +opencv-python +Pillow +huggingface_hub +gradio +xformers==0.0.24 +torch_scatter==2.1.2 +moviepy==1.0.3 +pydantic +open3d +einops +safetensors +git+https://github.com/facebookresearch/pytorch3d.git +https://github.com/nerfstudio-project/gsplat/releases/download/v1.4.0/gsplat-1.4.0%2Bpt22cu121-cp310-cp310-linux_x86_64.whl diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000000000000000000000000000000000000..4646a0e79be3cd6801b9fc6bc2ca363eb09dd359 --- /dev/null +++ b/src/config.py @@ -0,0 +1,106 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, Optional, Type, TypeVar + +from dacite import Config, from_dict +from omegaconf import DictConfig, OmegaConf + +from .dataset import DatasetCfgWrapper +from .dataset.data_module import DataLoaderCfg +from .loss import LossCfgWrapper +from .model.decoder import DecoderCfg +from .model.encoder import EncoderCfg +from .model.model_wrapper import OptimizerCfg, TestCfg, TrainCfg + + +@dataclass +class CheckpointingCfg: + load: Optional[str] # Not a path, since it could be something like wandb://... + every_n_train_steps: int + save_top_k: int + save_weights_only: bool + + +@dataclass +class ModelCfg: + decoder: DecoderCfg + encoder: EncoderCfg + + +@dataclass +class TrainerCfg: + max_steps: int + val_check_interval: int | float | None + gradient_clip_val: int | float | None + num_nodes: int = 1 + accumulate_grad_batches: int = 1 + precision: Literal["32", "16-mixed", "bf16-mixed"] = "32" + + +@dataclass +class RootCfg: + wandb: dict + mode: Literal["train", "test"] + dataset: list[DatasetCfgWrapper] + data_loader: DataLoaderCfg + model: ModelCfg + optimizer: OptimizerCfg + checkpointing: CheckpointingCfg + trainer: TrainerCfg + loss: list[LossCfgWrapper] + test: TestCfg + train: TrainCfg + seed: int + + +TYPE_HOOKS = { + Path: Path, +} + + +T = TypeVar("T") + + +def load_typed_config( + cfg: DictConfig, + data_class: Type[T], + extra_type_hooks: dict = {}, +) -> T: + return from_dict( + data_class, + OmegaConf.to_container(cfg), + config=Config(type_hooks={**TYPE_HOOKS, **extra_type_hooks}), + ) + + +def separate_loss_cfg_wrappers(joined: dict) -> list[LossCfgWrapper]: + # The dummy allows the union to be converted. + @dataclass + class Dummy: + dummy: LossCfgWrapper + + return [ + load_typed_config(DictConfig({"dummy": {k: v}}), Dummy).dummy + for k, v in joined.items() + ] + + +def separate_dataset_cfg_wrappers(joined: dict) -> list[DatasetCfgWrapper]: + # The dummy allows the union to be converted. + @dataclass + class Dummy: + dummy: DatasetCfgWrapper + + return [ + load_typed_config(DictConfig({"dummy": {k: v}}), Dummy).dummy + for k, v in joined.items() + ] + + +def load_typed_root_config(cfg: DictConfig) -> RootCfg: + return load_typed_config( + cfg, + RootCfg, + {list[LossCfgWrapper]: separate_loss_cfg_wrappers, + list[DatasetCfgWrapper]: separate_dataset_cfg_wrappers}, + ) diff --git a/src/dataset/__init__.py b/src/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d1d7936add0b5e69799329be24447e536058d72a --- /dev/null +++ b/src/dataset/__init__.py @@ -0,0 +1,98 @@ +from dataclasses import fields +from typing import Callable +from torch.utils.data import Dataset, ConcatDataset +import bisect + +from ..misc.step_tracker import StepTracker +from .types import Stage +from .view_sampler import get_view_sampler +from .dataset_dl3dv import DatasetDL3DV, DatasetDL3DVCfgWrapper +from .dataset_scannetpp import DatasetScannetpp, DatasetScannetppCfgWrapper +from .dataset_co3d import DatasetCo3d, DatasetCo3dCfgWrapper + +DATASETS: dict[str, Dataset] = { + "co3d": DatasetCo3d, + "scannetpp": DatasetScannetpp, + "dl3dv": DatasetDL3DV, +} + +DatasetCfgWrapper = DatasetDL3DVCfgWrapper | DatasetScannetppCfgWrapper | DatasetCo3dCfgWrapper + +class TestDatasetWarpper(Dataset): + def __init__(self, dataset: Dataset): + self.dataset = dataset + + def __getitem__(self, idx): + + return self.dataset[(idx, self.dataset.view_sampler.num_context_views, self.dataset.cfg.input_image_shape[1] // 14)] # fake parameters here, to fit the input of dataset + + def __len__(self): + return len(self.dataset) + + + +class CustomConcatDataset(ConcatDataset): + + def __getitem__(self, idx_tuple): + + if isinstance(idx_tuple, list): + idx_tuple = idx_tuple[0] + + idx = idx_tuple[0] + if idx < 0: + if -idx > len(self): + raise ValueError("absolute value of index should not exceed dataset length") + idx = len(self) + idx + dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) + if dataset_idx == 0: + sample_idx = idx + else: + sample_idx = idx - self.cumulative_sizes[dataset_idx - 1] + return self.datasets[dataset_idx][(sample_idx, idx_tuple[1], idx_tuple[2])] + + +def get_dataset( + cfgs: list[DatasetCfgWrapper], + stage: Stage, + step_tracker: StepTracker | None, + dataset_shim: Callable[[Dataset, str], Dataset] +) -> list[Dataset]: + datasets = [] + if stage != "test": + if stage == "val": + cfgs = [cfgs[0]] + for cfg in cfgs: + (field,) = fields(type(cfg)) + cfg = getattr(cfg, field.name) + + view_sampler = get_view_sampler( + cfg.view_sampler, + stage, + cfg.overfit_to_scene is not None, + cfg.cameras_are_circular, + step_tracker, + ) + dataset = DATASETS[cfg.name](cfg, stage, view_sampler) + dataset = dataset_shim(dataset, stage) + datasets.append(dataset) + + return CustomConcatDataset(datasets), datasets + elif stage == "test": + assert len(cfgs) == 1 + cfg = cfgs[0] + (field,) = fields(type(cfg)) + cfg = getattr(cfg, field.name) + + view_sampler = get_view_sampler( + cfg.view_sampler, + stage, + cfg.overfit_to_scene is not None, + cfg.cameras_are_circular, + step_tracker, + ) + dataset = DATASETS[cfg.name](cfg, stage, view_sampler) + dataset = dataset_shim(dataset, stage) + + return TestDatasetWarpper(dataset) + else: + NotImplementedError(f"Stage {stage} is not supported") diff --git a/src/dataset/data_module.py b/src/dataset/data_module.py new file mode 100644 index 0000000000000000000000000000000000000000..ed4767e3fa087f8ad50ad9e6538a88516e4c98a7 --- /dev/null +++ b/src/dataset/data_module.py @@ -0,0 +1,194 @@ +import random +from dataclasses import dataclass +from typing import Callable + +import numpy as np +import torch +from lightning.pytorch import LightningDataModule +from torch import Generator, nn +from torch.utils.data import DataLoader, Dataset, IterableDataset + +from src.dataset import * +from src.global_cfg import get_cfg + + +from ..misc.step_tracker import StepTracker +from ..misc.utils import get_world_size, get_rank +from . import DatasetCfgWrapper, get_dataset +from .types import DataShim, Stage +from .data_sampler import BatchedRandomSampler, MixedBatchSampler, custom_collate_fn +from .validation_wrapper import ValidationWrapper + +def get_data_shim(encoder: nn.Module) -> DataShim: + """Get functions that modify the batch. It's sometimes necessary to modify batches + outside the data loader because GPU computations are required to modify the batch or + because the modification depends on something outside the data loader. + """ + + shims: list[DataShim] = [] + if hasattr(encoder, "get_data_shim"): + shims.append(encoder.get_data_shim()) + + def combined_shim(batch): + for shim in shims: + batch = shim(batch) + return batch + + return combined_shim + +# the training ratio of datasets (example) +prob_mapping = {DatasetScannetpp: 0.5, + DatasetDL3DV: 0.5, + DatasetCo3d: 0.5} + +@dataclass +class DataLoaderStageCfg: + batch_size: int + num_workers: int + persistent_workers: bool + seed: int | None + + +@dataclass +class DataLoaderCfg: + train: DataLoaderStageCfg + test: DataLoaderStageCfg + val: DataLoaderStageCfg + + +DatasetShim = Callable[[Dataset, Stage], Dataset] + + +def worker_init_fn(worker_id: int) -> None: + random.seed(int(torch.utils.data.get_worker_info().seed) % (2**32 - 1)) + np.random.seed(int(torch.utils.data.get_worker_info().seed) % (2**32 - 1)) + + +class DataModule(LightningDataModule): + dataset_cfgs: list[DatasetCfgWrapper] + data_loader_cfg: DataLoaderCfg + step_tracker: StepTracker | None + dataset_shim: DatasetShim + global_rank: int + + def __init__( + self, + dataset_cfgs: list[DatasetCfgWrapper], + data_loader_cfg: DataLoaderCfg, + step_tracker: StepTracker | None = None, + dataset_shim: DatasetShim = lambda dataset, _: dataset, + global_rank: int = 0, + ) -> None: + super().__init__() + self.dataset_cfgs = dataset_cfgs + self.data_loader_cfg = data_loader_cfg + self.step_tracker = step_tracker + self.dataset_shim = dataset_shim + self.global_rank = global_rank + + def get_persistent(self, loader_cfg: DataLoaderStageCfg) -> bool | None: + return None if loader_cfg.num_workers == 0 else loader_cfg.persistent_workers + + def get_generator(self, loader_cfg: DataLoaderStageCfg) -> torch.Generator | None: + if loader_cfg.seed is None: + return None + generator = Generator() + generator.manual_seed(loader_cfg.seed + self.global_rank) + self.generator = generator + return self.generator + + def train_dataloader(self): + dataset, datasets_ls = get_dataset(self.dataset_cfgs, "train", self.step_tracker, self.dataset_shim) + world_size = get_world_size() + rank = get_rank() + # breakpoint() + prob_ls = [prob_mapping[type(dataset)] for dataset in datasets_ls] + # we assume all the dataset share the same num_context_views + + if len(datasets_ls) > 1: + prob = prob_ls + context_num_views = [dataset.cfg.view_sampler.num_context_views for dataset in datasets_ls] + else: + prob = None + dataset_key = next(iter(get_cfg()["dataset"])) + dataset_cfg = get_cfg()["dataset"][dataset_key] + context_num_views = dataset_cfg['view_sampler']['num_context_views'] + + sampler = MixedBatchSampler(datasets_ls, + batch_size=self.data_loader_cfg.train.batch_size, # Not used here! + num_context_views=context_num_views, + world_size=world_size, + rank=rank, + prob=prob, + generator=self.get_generator(self.data_loader_cfg.train)) + sampler.set_epoch(0) + self.train_loader = DataLoader( + dataset, + # self.data_loader_cfg.train.batch_size, + # shuffle=not isinstance(dataset, IterableDataset), + batch_sampler=sampler, + num_workers=self.data_loader_cfg.train.num_workers, + generator=self.generator, + worker_init_fn=worker_init_fn, + # collate_fn=custom_collate_fn, + persistent_workers=self.get_persistent(self.data_loader_cfg.train), + ) + # breakpoint() + # Set epoch for train and validation loaders (if applicable) + if hasattr(self.train_loader, "dataset") and hasattr(self.train_loader.dataset, "set_epoch"): + print("Training: Set Epoch in DataModule") + self.train_loader.dataset.set_epoch(0) + if hasattr(self.train_loader, "sampler") and hasattr(self.train_loader.sampler, "set_epoch"): + print("Training: Set Epoch in DataModule") + self.train_loader.sampler.set_epoch(0) + + return self.train_loader + + def val_dataloader(self): + dataset, datasets_ls = get_dataset(self.dataset_cfgs, "val", self.step_tracker, self.dataset_shim) + world_size = get_world_size() + rank = get_rank() + # here, we random select one dataset for val + dataset_key = next(iter(get_cfg()["dataset"])) + dataset_cfg = get_cfg()["dataset"][dataset_key] + if len(datasets_ls) > 1: + prob = [0.5] * len(datasets_ls) + else: + prob = None + sampler = MixedBatchSampler(datasets_ls, + batch_size=self.data_loader_cfg.train.batch_size, + num_context_views=dataset_cfg['view_sampler']['num_context_views'], + world_size=world_size, + rank=rank, + prob=prob, + generator=self.get_generator(self.data_loader_cfg.train)) + sampler.set_epoch(0) + self.val_loader = DataLoader( + dataset, + self.data_loader_cfg.val.batch_size, + num_workers=self.data_loader_cfg.val.num_workers, + sampler=sampler, + generator=self.get_generator(self.data_loader_cfg.val), + worker_init_fn=worker_init_fn, + persistent_workers=self.get_persistent(self.data_loader_cfg.val), + ) + if hasattr(self.val_loader, "dataset") and hasattr(self.val_loader.dataset, "set_epoch"): + print("Validation: Set Epoch in DataModule") + self.val_loader.dataset.set_epoch(0) + if hasattr(self.val_loader, "sampler") and hasattr(self.val_loader.sampler, "set_epoch"): + print("Validation: Set Epoch in DataModule") + self.val_loader.sampler.set_epoch(0) + return self.val_loader + + def test_dataloader(self): + dataset = get_dataset(self.dataset_cfgs, "test", self.step_tracker, self.dataset_shim) + data_loader = DataLoader( + dataset, + self.data_loader_cfg.test.batch_size, + num_workers=self.data_loader_cfg.test.num_workers, + generator=self.get_generator(self.data_loader_cfg.test), + worker_init_fn=worker_init_fn, + persistent_workers=self.get_persistent(self.data_loader_cfg.test), + ) + + return data_loader \ No newline at end of file diff --git a/src/dataset/data_sampler.py b/src/dataset/data_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..5f5dda411ddad56a66c552540957691f93062240 --- /dev/null +++ b/src/dataset/data_sampler.py @@ -0,0 +1,374 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# Copyright (C) 2024-present Naver Corporation. All rights reserved. +# Licensed under CC BY-NC-SA 4.0 (non-commercial use only). +# +# -------------------------------------------------------- +# Random sampling under a constraint +# -------------------------------------------------------- +import numpy as np +import torch +from typing import Callable, Iterable, Optional +from torch.utils.data import DataLoader, Dataset, DistributedSampler, IterableDataset, Sampler, BatchSampler +import random + +def custom_collate_fn(batch): + """ + Custom collate function to handle variable batch sizes + + Args: + batch: A list where each element could be either: + - A single tuple (idx, num_images, ...) + - A list of tuples [(idx1, num_images1, ...), (idx2, num_images2, ...)] + """ + # If batch contains lists (variable batch size case) + breakpoint() + if isinstance(batch[0], list): + # Flatten the batch + flattened = [] + for item in batch: + flattened.extend(item) + batch = flattened + + # Now batch is a list of tuples, process normally + return torch.utils.data.default_collate(batch) + +class BatchedRandomSampler: + """Random sampling under a constraint: each sample in the batch has the same feature, + which is chosen randomly from a known pool of 'features' for each batch. + + For instance, the 'feature' could be the image aspect-ratio. + + The index returned is a tuple (sample_idx, feat_idx). + This sampler ensures that each series of `batch_size` indices has the same `feat_idx`. + """ + + def __init__( + self, dataset, batch_size, num_context_views, min_patch_num=20, max_patch_num=32, world_size=1, rank=0, drop_last=True + ): + self.batch_size = batch_size + self.num_context_views = num_context_views + + self.len_dataset = N = len(dataset) + self.total_size = round_by(N, batch_size * world_size) if drop_last else N + self.min_patch_num = min_patch_num + self.max_patch_num = max_patch_num + assert ( + world_size == 1 or drop_last + ), "must drop the last batch in distributed mode" + + # distributed sampler + self.world_size = world_size + self.rank = rank + self.epoch = None + + def __len__(self): + + + return self.total_size // self.world_size + + def set_epoch(self, epoch): + self.epoch = epoch + + def __iter__(self): + # prepare RNG + if self.epoch is None: + assert ( + self.world_size == 1 and self.rank == 0 + ), "use set_epoch() if distributed mode is used" + seed = int(torch.empty((), dtype=torch.int64).random_().item()) + else: + seed = self.epoch + 777 + rng = np.random.default_rng(seed=seed) + + # random indices (will restart from 0 if not drop_last) + sample_idxs = np.arange(self.total_size) + rng.shuffle(sample_idxs) + + # random feat_idxs (same across each batch) + n_batches = (self.total_size + self.batch_size - 1) // self.batch_size + num_imgs = rng.integers(low=2, high=self.num_context_views, size=n_batches) + # num_imgs = (np.ones(n_batches) * self.num_context_views).astype(np.int64) # same number of context views for each batch + num_imgs = np.broadcast_to(num_imgs[:, None], (n_batches, self.batch_size)) + num_imgs = num_imgs.ravel()[: self.total_size] + + # put them together + idxs = np.c_[sample_idxs, num_imgs] # shape = (total_size, 2) + + # Distributed sampler: we select a subset of batches + # make sure the slice for each node is aligned with batch_size + size_per_proc = self.batch_size * ( + (self.total_size + self.world_size * self.batch_size - 1) + // (self.world_size * self.batch_size) + ) + idxs = idxs[self.rank * size_per_proc : (self.rank + 1) * size_per_proc] + + yield from (tuple(idx) for idx in idxs) + +class DynamicBatchSampler(Sampler): + """ + A custom batch sampler that dynamically adjusts batch size, aspect ratio, and image number + for each sample. Batches within a sample share the same aspect ratio and image number. + """ + def __init__(self, + sampler, + image_num_range, + h_range, + epoch=0, + seed=42, + max_img_per_gpu=48): + """ + Initializes the dynamic batch sampler. + + Args: + sampler: Instance of DynamicDistributedSampler. + aspect_ratio_range: List containing [min_aspect_ratio, max_aspect_ratio]. + image_num_range: List containing [min_images, max_images] per sample. + epoch: Current epoch number. + seed: Random seed for reproducibility. + max_img_per_gpu: Maximum number of images to fit in GPU memory. + """ + self.sampler = sampler + self.image_num_range = image_num_range + self.h_range = h_range + self.rng = random.Random() + + # Uniformly sample from the range of possible image numbers + # For any image number, the weight is 1.0 (uniform sampling). You can set any different weights here. + self.image_num_weights = {num_images: float(num_images**2) for num_images in range(image_num_range[0], image_num_range[1]+1)} + + # Possible image numbers, e.g., [2, 3, 4, ..., 24] + self.possible_nums = np.array([n for n in self.image_num_weights.keys() + if self.image_num_range[0] <= n <= self.image_num_range[1]]) + + # Normalize weights for sampling + weights = [self.image_num_weights[n] for n in self.possible_nums] + self.normalized_weights = np.array(weights) / sum(weights) + + # Maximum image number per GPU + self.max_img_per_gpu = max_img_per_gpu + + # Set the epoch for the sampler + self.set_epoch(epoch + seed) + + def set_epoch(self, epoch): + """ + Sets the epoch for this sampler, affecting the random sequence. + + Args: + epoch: The epoch number. + """ + self.sampler.set_epoch(epoch) + self.epoch = epoch + self.rng.seed(epoch * 100) + + def __iter__(self): + """ + Yields batches of samples with synchronized dynamic parameters. + + Returns: + Iterator yielding batches of indices with associated parameters. + """ + sampler_iterator = iter(self.sampler) + + while True: + try: + # Sample random image number and aspect ratio + random_image_num = int(np.random.choice(self.possible_nums, p=self.normalized_weights)) + random_ps_h = np.random.randint(low=(self.h_range[0] // 14), high=(self.h_range[1] // 14)+1) + + # Update sampler parameters + self.sampler.update_parameters( + image_num=random_image_num, + ps_h=random_ps_h + ) + + # Calculate batch size based on max images per GPU and current image number + batch_size = self.max_img_per_gpu / random_image_num + batch_size = np.floor(batch_size).astype(int) + batch_size = max(1, batch_size) # Ensure batch size is at least 1 + + # Collect samples for the current batch + current_batch = [] + for _ in range(batch_size): + try: + item = next(sampler_iterator) # item is (idx, aspect_ratio, image_num) + current_batch.append(item) + except StopIteration: + break # No more samples + + if not current_batch: + break # No more data to yield + + yield current_batch + + except StopIteration: + break # End of sampler's iterator + + def __len__(self): + # Return a large dummy length + return 1000000 + + +class DynamicDistributedSampler(DistributedSampler): + """ + Extends PyTorch's DistributedSampler to include dynamic aspect_ratio and image_num + parameters, which can be passed into the dataset's __getitem__ method. + """ + def __init__( + self, + dataset, + num_replicas: Optional[int] = None, + rank: Optional[int] = None, + shuffle: bool = False, + seed: int = 0, + drop_last: bool = False, + ): + super().__init__( + dataset, + num_replicas=num_replicas, + rank=rank, + shuffle=shuffle, + seed=seed, + drop_last=drop_last + ) + self.image_num = None + self.ps_h = None + + def __iter__(self): + """ + Yields a sequence of (index, image_num, aspect_ratio). + Relies on the parent class's logic for shuffling/distributing + the indices across replicas, then attaches extra parameters. + """ + indices_iter = super().__iter__() + + for idx in indices_iter: + yield (idx, self.image_num, self.ps_h, ) + + def update_parameters(self, image_num, ps_h): + """ + Updates dynamic parameters for each new epoch or iteration. + + Args: + aspect_ratio: The aspect ratio to set. + image_num: The number of images to set. + """ + self.image_num = image_num + self.ps_h = ps_h + +class MixedBatchSampler(BatchSampler): + """Sample one batch from a selected dataset with given probability. + Compatible with datasets at different resolution + """ + + def __init__( + self, src_dataset_ls, batch_size, num_context_views, world_size=1, rank=0, prob=None, sampler=None, generator=None + ): + self.base_sampler = None + self.batch_size = batch_size + self.num_context_views = num_context_views + self.world_size = world_size + self.rank = rank + self.drop_last = True + self.generator = generator + + self.src_dataset_ls = src_dataset_ls + self.n_dataset = len(self.src_dataset_ls) + + # Dataset length + self.dataset_length = [len(ds) for ds in self.src_dataset_ls] + self.cum_dataset_length = [ + sum(self.dataset_length[:i]) for i in range(self.n_dataset) + ] # cumulative dataset length + + # BatchSamplers for each source dataset + self.src_batch_samplers = [] + for ds in self.src_dataset_ls: + sampler = DynamicDistributedSampler(ds, num_replicas=self.world_size, rank=self.rank, seed=42, shuffle=True) + sampler.set_epoch(0) + + if hasattr(ds, "epoch"): + ds.epoch = 0 + if hasattr(ds, "set_epoch"): + ds.set_epoch(0) + batch_sampler = DynamicBatchSampler( + sampler, + [2, ds.cfg.view_sampler.num_context_views], + ds.cfg.input_image_shape, + seed=42, + max_img_per_gpu=ds.cfg.view_sampler.max_img_per_gpu + ) + self.src_batch_samplers.append(batch_sampler) + + # self.src_batch_samplers = [ + # BatchedRandomSampler( + # ds, + # num_context_views=ds.cfg.view_sampler.num_context_views, + # world_size=self.world_size, + # rank=self.rank, + # batch_size=self.batch_size, + # drop_last=self.drop_last, + # ) + # for ds in self.src_dataset_ls + # ] + # set epoch here + print("Setting epoch for all underlying BatchedRandomSamplers") + # for sampler in self.src_batch_samplers: + # sampler.set_epoch(0) + self.raw_batches = [ + list(bs) for bs in self.src_batch_samplers + ] # index in original dataset + self.n_batches = [len(b) for b in self.raw_batches] + self.n_total_batch = sum(self.n_batches) + # print("Total batch num is ", self.n_total_batch) + # sampling probability + if prob is None: + # if not given, decide by dataset length + self.prob = torch.tensor(self.n_batches) / self.n_total_batch + else: + self.prob = torch.as_tensor(prob) + + def __iter__(self): + """Yields batches of indices in the format of (sample_idx, feat_idx) tuples, + where indices correspond to ConcatDataset of src_dataset_ls + """ + for _ in range(self.n_total_batch): + idx_ds = torch.multinomial( + self.prob, 1, replacement=True, generator=self.generator + ).item() + + if 0 == len(self.raw_batches[idx_ds]): + self.raw_batches[idx_ds] = list(self.src_batch_samplers[idx_ds]) + + # get a batch from list - this is already in (sample_idx, feat_idx) format + batch_raw = self.raw_batches[idx_ds].pop() + + # shift only the sample_idx by cumulative dataset length, keep feat_idx unchanged + shift = self.cum_dataset_length[idx_ds] + processed_batch = [] + + for item in batch_raw: + # item[0] is the sample index, item[1] is the number of images + processed_item = (item[0] + shift, item[1], item[2]) + processed_batch.append(processed_item) + yield processed_batch + + def set_epoch(self, epoch): + """Set epoch for all underlying BatchedRandomSamplers""" + for sampler in self.src_batch_samplers: + sampler.set_epoch(epoch) + # Reset raw_batches after setting new epoch + self.raw_batches = [list(bs) for bs in self.src_batch_samplers] + + def __len__(self): + return self.n_total_batch + +def round_by(total, multiple, up=False): + if up: + total = total + multiple - 1 + return (total // multiple) * multiple \ No newline at end of file diff --git a/src/dataset/dataset.py b/src/dataset/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..18a29582cd56d920447fdf613b3d0491f15cd0ee --- /dev/null +++ b/src/dataset/dataset.py @@ -0,0 +1,13 @@ +from dataclasses import dataclass + +from .view_sampler import ViewSamplerCfg + + +@dataclass +class DatasetCfgCommon: + original_image_shape: list[int] + input_image_shape: list[int] + background_color: list[float] + cameras_are_circular: bool + overfit_to_scene: str | None + view_sampler: ViewSamplerCfg diff --git a/src/dataset/dataset_co3d.py b/src/dataset/dataset_co3d.py new file mode 100644 index 0000000000000000000000000000000000000000..cf96da135b9279e1e2b08ad031624cd9a06cf610 --- /dev/null +++ b/src/dataset/dataset_co3d.py @@ -0,0 +1,380 @@ +from concurrent.futures import ThreadPoolExecutor, as_completed +import json +from dataclasses import dataclass +from functools import cached_property +from pathlib import Path +import random +from typing import Literal +import os +import numpy as np +import torch +import torchvision.transforms as tf +from einops import rearrange, repeat +from jaxtyping import Float, UInt8 +from PIL import Image +from torch import Tensor +from torch.utils.data import Dataset +import os.path as osp +import cv2 +from ..geometry.projection import get_fov +from .dataset import DatasetCfgCommon +from .shims.augmentation_shim import apply_augmentation_shim +from .shims.crop_shim import apply_crop_shim +from .types import Stage +from .view_sampler import ViewSampler +from ..misc.cam_utils import camera_normalization + +from .shims.geometry_shim import depthmap_to_absolute_camera_coordinates + +CATEGORY = {'train': + ["backpack", "ball", "banana", "baseballbat", "baseballglove", + "bench", "bicycle", "book", "bottle", "bowl", "broccoli", "cake", "car", "carrot", + "cellphone", "chair", "couch", "cup", "donut", "frisbee", "hairdryer", "handbag", + "hotdog", "hydrant", "keyboard", "kite", "laptop", "microwave", + "motorcycle", + "mouse", "orange", "parkingmeter", "pizza", "plant", "remote", "sandwich", + "skateboard", "stopsign", + "suitcase", "teddybear", "toaster", "toilet", "toybus", + "toyplane", "toytrain", "toytruck", "tv", + "umbrella", "vase", "wineglass",], + 'test': ['teddybear']} + +@dataclass +class DatasetCo3dCfg(DatasetCfgCommon): + name: str + roots: list[Path] + baseline_min: float + baseline_max: float + max_fov: float + make_baseline_1: bool + augment: bool + relative_pose: bool + skip_bad_shape: bool + normalize_by_pts3d: bool + intr_augment: bool + rescale_to_1cube: bool + mask_bg: Literal['rand', True, False] = True + +@dataclass +class DatasetCo3dCfgWrapper: + co3d: DatasetCo3dCfg + + +class DatasetCo3d(Dataset): + cfg: DatasetCo3dCfg + stage: Stage + view_sampler: ViewSampler + + to_tensor: tf.ToTensor + chunks: list[Path] + near: float = 0.1 + far: float = 100.0 + + def __init__( + self, + cfg: DatasetCo3dCfg, + stage: Stage, + view_sampler: ViewSampler, + ) -> None: + super().__init__() + self.cfg = cfg + self.stage = stage + self.view_sampler = view_sampler + self.to_tensor = tf.ToTensor() + + self.root = cfg.roots[0] + self.mask_bg = cfg.mask_bg + assert self.mask_bg in ('rand', True, False) + + # load all scenes + self.categories = CATEGORY[self.data_stage] + self.scene_seq_dict = {} + self.scene_ids = [] + for category in self.categories: + with open(osp.join(self.root, f"{category}/valid_seq.json"), "r") as f: + scene_seq_dict = json.load(f) + for scene, seqs in scene_seq_dict.items(): + self.scene_seq_dict[f"{category}/{scene}"] = seqs + self.scene_ids.append(f"{category}/{scene}") + + print(f"CO3Dv2 {self.stage}: loaded {len(self.scene_seq_dict)} scenes") + + def load_frames(self, scene_id, frame_ids): + with ThreadPoolExecutor(max_workers=32) as executor: + # Create a list to store futures with their original indices + futures_with_idx = [] + for idx, frame_id in enumerate(frame_ids): + file_path = os.path.join(self.root, f"{scene_id}/images/frame{frame_id:06d}.jpg") + futures_with_idx.append( + ( + idx, + executor.submit( + lambda p: self.to_tensor(Image.open(p).convert("RGB")), + file_path, + ), + ) + ) + + # Pre-allocate list with correct size to maintain order + torch_images = [None] * len(frame_ids) + for idx, future in futures_with_idx: + torch_images[idx] = future.result() + # Check if all images have the same size + sizes = set(img.shape for img in torch_images) + if len(sizes) == 1: + torch_images = torch.stack(torch_images) + # Return as list if images have different sizes + return torch_images + + def load_npz(self, scene_id, frame_id): + npzpath = os.path.join(self.root, f"{scene_id}/images/frame{frame_id:06d}.npz") + imgpath = os.path.join(self.root, f"{scene_id}/images/frame{frame_id:06d}.jpg") + img = Image.open(imgpath) + # breakpoint() + W, H = img.size + npzdata = np.load(npzpath) + intri = npzdata['camera_intrinsics'] + extri = npzdata['camera_pose'] + intri[0, 0] /= float(W) + intri[1, 1] /= float(H) + intri[0, 2] /= float(W) + intri[1, 2] /= float(H) + md = npzdata['maximum_depth'] + return intri, extri, md + + def load_depth(self, scene_id, frame_ids, mds): + torch_depths = [] + for frame_id in frame_ids: + depthpath = os.path.join(self.root, f"{scene_id}/depths/frame{frame_id:06d}.jpg.geometric.png") + depth = cv2.imread(depthpath, cv2.IMREAD_UNCHANGED)/65535*np.nan_to_num(mds[frame_id]) + depth = np.nan_to_num(depth) + torch_depths.append(torch.from_numpy(depth)) + return torch_depths + + def load_masks(self, scene_id, frame_ids): + masks = [] + for frame_id in frame_ids: + maskpath = os.path.join(self.root, f"{scene_id}/masks/frame{frame_id:06d}.png") + maskmap = cv2.imread(maskpath, cv2.IMREAD_UNCHANGED).astype(np.float32) + maskmap = (maskmap / 255.0) > 0.1 + masks.append(torch.from_numpy(maskmap)) + return masks + + def getitem(self, index: int, num_context_views: int, patchsize: tuple) -> dict: + scene_id = self.scene_ids[index] + seq = self.scene_seq_dict[scene_id] + + extrinsics = [] + intrinsics = [] + frame_ids = [] + mds = {} + for frame_id in seq: + intri, extri, md = self.load_npz(scene_id, frame_id) + extrinsics.append(extri) + intrinsics.append(intri) + frame_ids.append(frame_id) + mds[frame_id] = md + + extrinsics = np.array(extrinsics) + intrinsics = np.array(intrinsics) + extrinsics = torch.tensor(extrinsics, dtype=torch.float32) + intrinsics = torch.tensor(intrinsics, dtype=torch.float32) + + num_views = extrinsics.shape[0] + context_indices = torch.tensor(random.sample(range(num_views), num_context_views)) + remaining_indices = torch.tensor([i for i in range(num_views) if i not in context_indices]) + target_indices = torch.tensor(random.sample(remaining_indices.tolist(), self.view_sampler.num_target_views)) + + # Skip the example if the field of view is too wide. + if (get_fov(intrinsics).rad2deg() > self.cfg.max_fov).any(): + raise Exception("Field of view too wide") + + input_frames = [frame_ids[i] for i in context_indices] + target_frame = [frame_ids[i] for i in target_indices] + + context_images = self.load_frames(scene_id, input_frames) + target_images = self.load_frames(scene_id, target_frame) + context_depths = self.load_depth(scene_id, input_frames, mds) + target_depths = self.load_depth(scene_id, target_frame, mds) + + mask_bg = (self.mask_bg == True) or (self.mask_bg == "rand" and np.random.random() < 0.5) + if mask_bg: + context_masks = self.load_masks(scene_id, input_frames) + target_mask = self.load_masks(scene_id, target_frame) + + # update the depthmap with mask + context_depths = [depth * mask for depth, mask in zip(context_depths, context_masks)] + target_depths = [depth * mask for depth, mask in zip(target_depths, target_mask)] + + + # Resize the world to make the baseline 1. + context_extrinsics = extrinsics[context_indices] + if self.cfg.make_baseline_1: + a, b = context_extrinsics[0, :3, 3], context_extrinsics[-1, :3, 3] + scale = (a - b).norm() + if scale < self.cfg.baseline_min or scale > self.cfg.baseline_max: + print( + f"Skipped {scene_id} because of baseline out of range: " + f"{scale:.6f}" + ) + raise Exception("baseline out of range") + extrinsics[:, :3, 3] /= scale + else: + scale = 1 + + if self.cfg.relative_pose: + extrinsics = camera_normalization(extrinsics[context_indices][0:1], extrinsics) + + # self.cfg.rescale_to_1cube = True + if self.cfg.rescale_to_1cube: + scene_scale = torch.max(torch.abs(extrinsics[context_indices][:, :3, 3])) # target pose is not included + # all_extrinsics = torch.cat([extrinsics[context_indices], extrinsics[target_indices]], dim=0) # [N, 4, 4] + # scene_scale = torch.max(torch.abs(all_extrinsics[:, :3, 3])) + rescale_factor = 1 * scene_scale + extrinsics[:, :3, 3] /= rescale_factor + + example = { + "context": { + "extrinsics": extrinsics[context_indices], + "intrinsics": intrinsics[context_indices], + "image": context_images, + "depth": context_depths, + "near": self.get_bound("near", len(context_indices)), + "far": self.get_bound("far", len(context_indices)), + "index": context_indices, + # "overlap": overlap, + }, + "target": { + "extrinsics": extrinsics[target_indices], + "intrinsics": intrinsics[target_indices], + "image": target_images, + "depth": target_depths, + "near": self.get_bound("near", len(target_indices)), + "far": self.get_bound("far", len(target_indices)), + "index": target_indices, + }, + "scene": f"CO3Dv2 {scene_id}", + } + + if self.stage == "train" and self.cfg.intr_augment: + intr_aug = True + else: + intr_aug = False + + example = apply_crop_shim(example, (patchsize[0] * 14, patchsize[1] * 14), intr_aug=intr_aug) + + if self.stage == "train" and self.cfg.augment: + example = apply_augmentation_shim(example) + + # example_1 = copy.deepcopy(example) + # world pts + image_size = example["context"]["image"].shape[2:] + context_intrinsics = example["context"]["intrinsics"].clone().detach().numpy() + context_intrinsics[:, 0] = context_intrinsics[:, 0] * image_size[1] + context_intrinsics[:, 1] = context_intrinsics[:, 1] * image_size[0] + + target_intrinsics = example["target"]["intrinsics"].clone().detach().numpy() + target_intrinsics[:, 0] = target_intrinsics[:, 0] * image_size[1] + target_intrinsics[:, 1] = target_intrinsics[:, 1] * image_size[0] + + context_pts3d_list, context_valid_mask_list = [], [] + target_pts3d_list, target_valid_mask_list = [], [] + + for i in range(len(example["context"]["depth"])): + context_pts3d, context_valid_mask = depthmap_to_absolute_camera_coordinates(example["context"]["depth"][i].numpy(), context_intrinsics[i], example["context"]["extrinsics"][i].numpy()) + context_pts3d_list.append(torch.from_numpy(context_pts3d).to(torch.float32)) + context_valid_mask_list.append(torch.from_numpy(context_valid_mask)) + + context_pts3d = torch.stack(context_pts3d_list, dim=0) + context_valid_mask = torch.stack(context_valid_mask_list, dim=0) + + for i in range(len(example["target"]["depth"])): + target_pts3d, target_valid_mask = depthmap_to_absolute_camera_coordinates(example["target"]["depth"][i].numpy(), target_intrinsics[i], example["target"]["extrinsics"][i].numpy()) + target_pts3d_list.append(torch.from_numpy(target_pts3d).to(torch.float32)) + target_valid_mask_list.append(torch.from_numpy(target_valid_mask)) + + target_pts3d = torch.stack(target_pts3d_list, dim=0) + target_valid_mask = torch.stack(target_valid_mask_list, dim=0) + + # normalize by context pts3d + if self.cfg.normalize_by_pts3d: + transformed_pts3d = context_pts3d[context_valid_mask] + scene_factor = transformed_pts3d.norm(dim=-1).mean().clip(min=1e-8) + + context_pts3d /= scene_factor + example["context"]["depth"] /= scene_factor + example["context"]["extrinsics"][:, :3, 3] /= scene_factor + + target_pts3d /= scene_factor + example["target"]["depth"] /= scene_factor + example["target"]["extrinsics"][:, :3, 3] /= scene_factor + + example["context"]["pts3d"] = context_pts3d + example["target"]["pts3d"] = target_pts3d + example["context"]["valid_mask"] = context_valid_mask + example["target"]["valid_mask"] = target_valid_mask + + if torch.isnan(example["context"]["depth"]).any() or torch.isinf(example["context"]["depth"]).any() or \ + torch.isnan(example["context"]["extrinsics"]).any() or torch.isinf(example["context"]["extrinsics"]).any() or \ + torch.isnan(example["context"]["pts3d"]).any() or torch.isinf(example["context"]["pts3d"]).any() or \ + torch.isnan(example["context"]["intrinsics"]).any() or torch.isinf(example["context"]["intrinsics"]).any() or \ + torch.isnan(example["target"]["depth"]).any() or torch.isinf(example["target"]["depth"]).any() or \ + torch.isnan(example["target"]["extrinsics"]).any() or torch.isinf(example["target"]["extrinsics"]).any() or \ + torch.isnan(example["target"]["pts3d"]).any() or torch.isinf(example["target"]["pts3d"]).any() or \ + torch.isnan(example["target"]["intrinsics"]).any() or torch.isinf(example["target"]["intrinsics"]).any(): + raise Exception("encounter nan or inf in context depth") + + for key in ["context", "target"]: + example[key]["valid_mask"] = (torch.ones_like(example[key]["valid_mask"]) * -1).type(torch.int32) + + return example + + + def __getitem__(self, index_tuple: tuple) -> dict: + index, num_context_views, patchsize_h = index_tuple + patchsize_w = (self.cfg.input_image_shape[1] // 14) + try: + return self.getitem(index, num_context_views, (patchsize_h, patchsize_w)) + except Exception as e: + print(f"Error: {e}") + index = np.random.randint(len(self)) + return self.__getitem__((index, num_context_views, patchsize_h)) + + def get_bound( + self, + bound: Literal["near", "far"], + num_views: int, + ) -> Float[Tensor, " view"]: + value = torch.tensor(getattr(self, bound), dtype=torch.float32) + return repeat(value, "-> v", v=num_views) + + @property + def data_stage(self) -> Stage: + if self.cfg.overfit_to_scene is not None: + return "test" + if self.stage == "val": + return "test" + return self.stage + + @cached_property + def index(self) -> dict[str, Path]: + merged_index = {} + data_stages = [self.data_stage] + if self.cfg.overfit_to_scene is not None: + data_stages = ("test", "train") + for data_stage in data_stages: + for root in self.cfg.roots: + # Load the root's index. + with (root / data_stage / "index.json").open("r") as f: + index = json.load(f) + index = {k: Path(root / data_stage / v) for k, v in index.items()} + + # The constituent datasets should have unique keys. + assert not (set(merged_index.keys()) & set(index.keys())) + + # Merge the root's index into the main index. + merged_index = {**merged_index, **index} + return merged_index + + def __len__(self) -> int: + return len(self.scene_ids) \ No newline at end of file diff --git a/src/dataset/dataset_dl3dv.py b/src/dataset/dataset_dl3dv.py new file mode 100644 index 0000000000000000000000000000000000000000..8da706db5c0bcde903a11442d144c5318717dd75 --- /dev/null +++ b/src/dataset/dataset_dl3dv.py @@ -0,0 +1,421 @@ +from concurrent.futures import ThreadPoolExecutor, as_completed +import json +from dataclasses import dataclass +from functools import cached_property +from io import BytesIO +from pathlib import Path +from typing import Literal +import os +import numpy as np +import torch +import torchvision.transforms as tf +from einops import rearrange, repeat +from jaxtyping import Float, UInt8 +from PIL import Image +from torch import Tensor +from torch.utils.data import Dataset +import torch.nn.functional as F + +from ..geometry.projection import get_fov +from .dataset import DatasetCfgCommon +from .shims.augmentation_shim import apply_augmentation_shim +from .shims.crop_shim import apply_crop_shim +from .types import Stage +from .view_sampler import ViewSampler +from ..misc.cam_utils import camera_normalization + + +@dataclass +class DatasetDl3dvCfg(DatasetCfgCommon): + name: str + roots: list[Path] + baseline_min: float + baseline_max: float + max_fov: float + make_baseline_1: bool + augment: bool + relative_pose: bool + skip_bad_shape: bool + avg_pose: bool + rescale_to_1cube: bool + intr_augment: bool + normalize_by_pts3d: bool + rescale_to_1cube: bool + + +@dataclass +class DatasetDL3DVCfgWrapper: + dl3dv: DatasetDl3dvCfg + + + +class DatasetDL3DV(Dataset): + cfg: DatasetDl3dvCfg + stage: Stage + view_sampler: ViewSampler + + to_tensor: tf.ToTensor + chunks: list[Path] + near: float = 0.1 + far: float = 100.0 + + def __init__( + self, + cfg: DatasetDl3dvCfg, + stage: Stage, + view_sampler: ViewSampler, + ) -> None: + super().__init__() + self.cfg = cfg + self.stage = stage + self.view_sampler = view_sampler + self.to_tensor = tf.ToTensor() + + # load data + self.data_root = cfg.roots[0] + self.data_list = [] + with open(f"{self.data_root}/{self.data_stage}_index.json", "r") as file: + data_index = json.load(file) + + self.data_list = [ + os.path.join(self.data_root, item) for item in data_index + ] # train: 9900 test: 140 + + self.scene_ids = {} + self.scenes = {} + index = 0 + with ThreadPoolExecutor(max_workers=32) as executor: + futures = [executor.submit(self.load_jsons, scene_path) for scene_path in self.data_list] + for future in as_completed(futures): + scene_frames, scene_id = future.result() + self.scenes[scene_id] = scene_frames + self.scene_ids[index] = scene_id + index += 1 + print(f"DL3DV: {self.stage}: loaded {len(self.scene_ids)} scenes") + + def convert_intrinsics(self, meta_data): + store_h, store_w = meta_data["h"], meta_data["w"] + fx, fy, cx, cy = ( + meta_data["fl_x"], + meta_data["fl_y"], + meta_data["cx"], + meta_data["cy"], + ) + intrinsics = np.eye(3, dtype=np.float32) + intrinsics[0, 0] = float(fx) / float(store_w) + intrinsics[1, 1] = float(fy) / float(store_h) + intrinsics[0, 2] = float(cx) / float(store_w) + intrinsics[1, 2] = float(cy) / float(store_h) + return intrinsics + + def blender2opencv_c2w(self, pose): + blender2opencv = np.array( + [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]] + ) + opencv_c2w = np.array(pose) @ blender2opencv + return opencv_c2w.tolist() + + def load_jsons(self, scene_path): + json_path = os.path.join(scene_path, "transforms.json") + with open(json_path, "r") as f: + data = json.load(f) + + scene_frames = [] + scene_id = scene_path.split("/")[-1].split(".")[0] + for i, frame in enumerate(data["frames"]): + frame_tmp = {} + frame_tmp["file_path"] = os.path.join(scene_path, frame["file_path"]) + frame_tmp["intrinsics"] = self.convert_intrinsics(data).tolist() + frame_tmp["extrinsics"] = self.blender2opencv_c2w(frame["transform_matrix"]) + scene_frames.append(frame_tmp) + return scene_frames, scene_id + + def load_frames(self, frames): + with ThreadPoolExecutor(max_workers=32) as executor: + # Create a list to store futures with their original indices + futures_with_idx = [] + for idx, file_path in enumerate(frames): + file_path = file_path["file_path"].replace("images", "images_4") + futures_with_idx.append( + ( + idx, + executor.submit( + lambda p: self.to_tensor(Image.open(p).convert("RGB")), + file_path, + ), + ) + ) + + # Pre-allocate list with correct size to maintain order + torch_images = [None] * len(frames) + for idx, future in futures_with_idx: + torch_images[idx] = future.result() + # Check if all images have the same size + sizes = set(img.shape for img in torch_images) + if len(sizes) == 1: + torch_images = torch.stack(torch_images) + # Return as list if images have different sizes + return torch_images + + def load_depth(self, frames): + depth_list = [] + for frame_name in frames: + depth_path = frame_name.replace("images", "depth").replace("jpg", "npy") + depth = torch.from_numpy(np.load(depth_path)) + positive_depths = depth[depth > 0] + if len(positive_depths) > 1000000: # If more than 1M points, sample randomly + indices = torch.randperm(len(positive_depths))[:1000000] + positive_depths = positive_depths[indices] + percentile_95 = torch.quantile(positive_depths, 0.95) + # Set depth values greater than the 95th percentile to 0 + depth[depth > percentile_95] = 0 + depth_list.append(depth) + + return torch.stack(depth_list) + + def shuffle(self, lst: list) -> list: + indices = torch.randperm(len(lst)) + return [lst[x] for x in indices] + + def getitem(self, index: int, num_context_views: int, patchsize: tuple) -> dict: + + scene = self.scene_ids[index] + + example = self.scenes[scene] + # load poses + extrinsics = [] + intrinsics = [] + for frame in example: + extrinsic = frame["extrinsics"] + intrinsic = frame["intrinsics"] + extrinsics.append(extrinsic) + intrinsics.append(intrinsic) + + extrinsics = np.array(extrinsics) + intrinsics = np.array(intrinsics) + extrinsics = torch.tensor(extrinsics, dtype=torch.float32) + intrinsics = torch.tensor(intrinsics, dtype=torch.float32) + + try: + context_indices, target_indices, overlap = self.view_sampler.sample( + scene, + num_context_views, + extrinsics, + intrinsics, + ) + except ValueError: + # Skip because the example doesn't have enough frames. + raise Exception("Not enough frames") + + # Skip the example if the field of view is too wide. + if (get_fov(intrinsics).rad2deg() > self.cfg.max_fov).any(): + raise Exception("Field of view too wide") + + # Load the images. + input_frames = [example[i] for i in context_indices] + target_frame = [example[i] for i in target_indices] + + context_images = self.load_frames(input_frames) + target_images = self.load_frames(target_frame) + + # context_depth = self.load_depth(input_frames) + # target_depth = self.load_depth(target_frame) + context_depth = torch.ones_like(context_images)[:, 0] + target_depth = torch.ones_like(target_images)[:, 0] + + # Skip the example if the images don't have the right shape. + context_image_invalid = context_images.shape[1:] != (3, *self.cfg.original_image_shape) + target_image_invalid = target_images.shape[1:] != (3, *self.cfg.original_image_shape) + if self.cfg.skip_bad_shape and (context_image_invalid or target_image_invalid): + print( + f"Skipped bad example {example['key']}. Context shape was " + f"{context_images.shape} and target shape was " + f"{target_images.shape}." + ) + raise Exception("Bad example image shape") + + # Resize the world to make the baseline 1. + context_extrinsics = extrinsics[context_indices] + if self.cfg.make_baseline_1: + a, b = context_extrinsics[0, :3, 3], context_extrinsics[-1, :3, 3] + scale = (a - b).norm() + if scale < self.cfg.baseline_min or scale > self.cfg.baseline_max: + print( + f"Skipped {scene} because of baseline out of range: " + f"{scale:.6f}" + ) + raise Exception("baseline out of range") + extrinsics[:, :3, 3] /= scale + else: + scale = 1 + + if self.cfg.relative_pose: + extrinsics = camera_normalization(extrinsics[context_indices][0:1], extrinsics) + + if self.cfg.rescale_to_1cube: + scene_scale = torch.max(torch.abs(extrinsics[context_indices][:, :3, 3])) # target pose is not included + rescale_factor = 1 * scene_scale + extrinsics[:, :3, 3] /= rescale_factor + + if torch.isnan(extrinsics).any() or torch.isinf(extrinsics).any(): + raise Exception("encounter nan or inf in input poses") + + example = { + "context": { + "extrinsics": extrinsics[context_indices], + "intrinsics": intrinsics[context_indices], + "image": context_images, + "depth": context_depth, + "near": self.get_bound("near", len(context_indices)) / scale, + "far": self.get_bound("far", len(context_indices)) / scale, + "index": context_indices, + # "overlap": overlap, + }, + "target": { + "extrinsics": extrinsics[target_indices], + "intrinsics": intrinsics[target_indices], + "image": target_images, + "depth": target_depth, + "near": self.get_bound("near", len(target_indices)) / scale, + "far": self.get_bound("far", len(target_indices)) / scale, + "index": target_indices, + }, + "scene": "dl3dv_"+scene, + } + if self.stage == "train" and self.cfg.augment: + example = apply_augmentation_shim(example) + + if self.stage == "train" and self.cfg.intr_augment: + intr_aug = True + else: + intr_aug = False + + example = apply_crop_shim(example, (patchsize[0] * 14, patchsize[1] * 14), intr_aug=intr_aug) + + image_size = example["context"]["image"].shape[2:] + context_intrinsics = example["context"]["intrinsics"].clone().detach().numpy() + context_intrinsics[:, 0] = context_intrinsics[:, 0] * image_size[1] + context_intrinsics[:, 1] = context_intrinsics[:, 1] * image_size[0] + + target_intrinsics = example["target"]["intrinsics"].clone().detach().numpy() + target_intrinsics[:, 0] = target_intrinsics[:, 0] * image_size[1] + target_intrinsics[:, 1] = target_intrinsics[:, 1] * image_size[0] + + context_pts3d_list, context_valid_mask_list = [], [] + target_pts3d_list, target_valid_mask_list = [], [] + + # for i in range(len(example["context"]["depth"])): + # context_pts3d, context_valid_mask = depthmap_to_absolute_camera_coordinates(example["context"]["depth"][i].numpy(), context_intrinsics[i], example["context"]["extrinsics"][i].numpy()) + # context_pts3d_list.append(torch.from_numpy(context_pts3d).to(torch.float32)) + # context_valid_mask_list.append(torch.from_numpy(context_valid_mask)) + + # context_pts3d = torch.stack(context_pts3d_list, dim=0) + # context_valid_mask = torch.stack(context_valid_mask_list, dim=0) + + context_pts3d = torch.ones_like(example["context"]["image"]).permute(0, 2, 3, 1) # [N, H, W, 3] + context_valid_mask = torch.ones_like(example["context"]["image"])[:, 0].bool() # [N, H, W] + + target_pts3d = torch.ones_like(target_images).permute(0, 2, 3, 1) # [N, H, W, 3] + target_valid_mask = torch.ones_like(target_images)[:, 0].bool() # [N, H, W] + + # normalize by context pts3d + if self.cfg.normalize_by_pts3d: + transformed_pts3d = context_pts3d[context_valid_mask] + scene_factor = transformed_pts3d.norm(dim=-1).mean().clip(min=1e-8) + + context_pts3d /= scene_factor + example["context"]["depth"] /= scene_factor + example["context"]["extrinsics"][:, :3, 3] /= scene_factor + + target_pts3d /= scene_factor + example["target"]["depth"] /= scene_factor + example["target"]["extrinsics"][:, :3, 3] /= scene_factor + + example["context"]["pts3d"] = context_pts3d + example["target"]["pts3d"] = target_pts3d + example["context"]["valid_mask"] = context_valid_mask * -1 + example["target"]["valid_mask"] = target_valid_mask * -1 + + return example + + def __getitem__(self, index_tuple: tuple) -> dict: + index, num_context_views, patchsize_h = index_tuple + patchsize_w = (self.cfg.input_image_shape[1] // 14) + try: + return self.getitem(index, num_context_views, (patchsize_h, patchsize_w)) + except Exception as e: + print(f"Error: {e}") + index = np.random.randint(len(self)) + return self.__getitem__((index, num_context_views, patchsize_h)) + + def convert_poses( + self, + poses: Float[Tensor, "batch 18"], + ) -> tuple[ + Float[Tensor, "batch 4 4"], # extrinsics + Float[Tensor, "batch 3 3"], # intrinsics + ]: + b, _ = poses.shape + + # Convert the intrinsics to a 3x3 normalized K matrix. + intrinsics = torch.eye(3, dtype=torch.float32) + intrinsics = repeat(intrinsics, "h w -> b h w", b=b).clone() + fx, fy, cx, cy = poses[:, :4].T + intrinsics[:, 0, 0] = fx + intrinsics[:, 1, 1] = fy + intrinsics[:, 0, 2] = cx + intrinsics[:, 1, 2] = cy + + # Convert the extrinsics to a 4x4 OpenCV-style W2C matrix. + w2c = repeat(torch.eye(4, dtype=torch.float32), "h w -> b h w", b=b).clone() + w2c[:, :3] = rearrange(poses[:, 6:], "b (h w) -> b h w", h=3, w=4) + return w2c.inverse(), intrinsics + + def convert_images( + self, + images: list[UInt8[Tensor, "..."]], + ) -> Float[Tensor, "batch 3 height width"]: + torch_images = [] + for image in images: + image = Image.open(BytesIO(image.numpy().tobytes())) + torch_images.append(self.to_tensor(image)) + return torch.stack(torch_images) + + def get_bound( + self, + bound: Literal["near", "far"], + num_views: int, + ) -> Float[Tensor, " view"]: + value = torch.tensor(getattr(self, bound), dtype=torch.float32) + return repeat(value, "-> v", v=num_views) + + @property + def data_stage(self) -> Stage: + if self.cfg.overfit_to_scene is not None: + return "test" + if self.stage == "val": + return "test" + return self.stage + + @cached_property + def index(self) -> dict[str, Path]: + merged_index = {} + data_stages = [self.data_stage] + if self.cfg.overfit_to_scene is not None: + data_stages = ("test", "train") + for data_stage in data_stages: + for root in self.cfg.roots: + # Load the root's index. + with (root / data_stage / "index.json").open("r") as f: + index = json.load(f) + index = {k: Path(root / data_stage / v) for k, v in index.items()} + + # The constituent datasets should have unique keys. + assert not (set(merged_index.keys()) & set(index.keys())) + + # Merge the root's index into the main index. + merged_index = {**merged_index, **index} + return merged_index + + def __len__(self) -> int: + return len(self.scene_ids) diff --git a/src/dataset/dataset_scannetpp.py b/src/dataset/dataset_scannetpp.py new file mode 100644 index 0000000000000000000000000000000000000000..050a0943a682ca92481b1d194785fcf5e96b6429 --- /dev/null +++ b/src/dataset/dataset_scannetpp.py @@ -0,0 +1,443 @@ +import json +from dataclasses import dataclass +from functools import cached_property +from io import BytesIO +from pathlib import Path +import random +from typing import Literal +import os +import cv2 +import numpy as np +import torch +import torchvision.transforms as tf +from einops import rearrange, repeat +from jaxtyping import Float, UInt8 +from PIL import Image +import torchvision +from torch import Tensor +from torch.utils.data import Dataset +from concurrent.futures import ThreadPoolExecutor, as_completed +import copy +from .shims.geometry_shim import depthmap_to_absolute_camera_coordinates + +from .shims.load_shim import imread_cv2 + +from ..geometry.projection import get_fov +from .dataset import DatasetCfgCommon +from .shims.augmentation_shim import apply_augmentation_shim +from .shims.crop_shim import apply_crop_shim +from .types import Stage +from .view_sampler import ViewSampler +from ..misc.cam_utils import camera_normalization + +@dataclass +class DatasetScannetppCfg(DatasetCfgCommon): + name: str + roots: list[Path] + baseline_min: float + baseline_max: float + max_fov: float + make_baseline_1: bool + augment: bool + relative_pose: bool + skip_bad_shape: bool + metric_thre: float + intr_augment: bool + make_baseline_1: bool + rescale_to_1cube: bool + normalize_by_pts3d: bool + +@dataclass +class DatasetScannetppCfgWrapper: + scannetpp: DatasetScannetppCfg + + +class DatasetScannetpp(Dataset): + cfg: DatasetScannetppCfgWrapper + stage: Stage + view_sampler: ViewSampler + + to_tensor: tf.ToTensor + chunks: list[Path] + near: float = 0.1 + far: float = 100.0 + + def __init__( + self, + cfg: DatasetScannetppCfgWrapper, + stage: Stage, + view_sampler: ViewSampler, + ) -> None: + super().__init__() + self.cfg = cfg + self.stage = stage + self.view_sampler = view_sampler + self.to_tensor = tf.ToTensor() + + # load data + self.data_root = cfg.roots[0] + self.data_list = [] # we use dslr rather than iphone + data_index = os.listdir(f"{self.data_root}") # we train all the scenes + + if self.stage != "train": + with open(f"{self.data_root}/valid.json", "r") as file: + data_index = json.load(file)[:10] + data_index = data_index * 100 + random.shuffle(data_index) + else: + with open(f"{self.data_root}/valid.json", "r") as file: + data_index = json.load(file)[10:] + + self.data_list = [ + os.path.join(self.data_root, item) for item in data_index + ] + + self.scene_ids = {} + self.scenes = {} + index = 0 + with ThreadPoolExecutor(max_workers=32) as executor: + futures = [executor.submit(self.load_metadata, scene_path) for scene_path in self.data_list] + for future in as_completed(futures): + scene_frames, scene_id = future.result() + self.scenes[scene_id] = scene_frames + self.scene_ids[index] = scene_id + index += 1 + + # if self.stage != "train": + # self.scene_ids = self.scene_ids + # random.shuffle(self.scene_ids) + print(f"Scannetpp: {self.stage}: loaded {len(self.scene_ids)} scenes") + + def shuffle(self, lst: list) -> list: + indices = torch.randperm(len(lst)) + return [lst[x] for x in indices] + + def load_metadata(self, scene_path): + metadata_path = os.path.join(scene_path, "scene_metadata.npz") + metadata = np.load(metadata_path, allow_pickle=True) + intrinsics = metadata["intrinsics"] + trajectories = metadata["trajectories"] + images = metadata["images"] + + scene_id = scene_path.split("/")[-1].split(".")[0] + scene_frames = [ + { + "file_path": os.path.join(scene_path, "images", images[i].split(".")[0] + ".jpg"), + "depth_path": os.path.join(scene_path, "depth", images[i].split(".")[0] + ".png"), + "intrinsics": self.convert_intrinsics(intrinsics[i]), + "extrinsics": trajectories[i], + } + for i in range(len(images)) + ] + scene_frames.sort(key=lambda x: x["file_path"]) # sort by file path to ensure correct order + return scene_frames, scene_id + + def convert_intrinsics(self, intrinsics): + w = intrinsics[0, 2] * 2 + h = intrinsics[1, 2] * 2 + intrinsics[0, 0] = intrinsics[0, 0] / w + intrinsics[1, 1] = intrinsics[1, 1] / h + intrinsics[0, 2] = intrinsics[0, 2] / w + intrinsics[1, 2] = intrinsics[1, 2] / h + return intrinsics + + def blender2opencv_c2w(self, pose): + blender2opencv = np.array( + [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]] + ) + opencv_c2w = np.array(pose) @ blender2opencv + return opencv_c2w.tolist() + + def load_frames(self, frames): + with ThreadPoolExecutor(max_workers=1) as executor: + # Create a list to store futures with their original indices + futures_with_idx = [] + for idx, file_path in enumerate(frames): + file_path = file_path["file_path"] + futures_with_idx.append( + ( + idx, + executor.submit( + lambda p: self.to_tensor(Image.open(p).convert("RGB")), + file_path, + ), + ) + ) + + # Pre-allocate list with correct size to maintain order + torch_images = [None] * len(frames) + for idx, future in futures_with_idx: + torch_images[idx] = future.result() + # Check if all images have the same size + sizes = set(img.shape for img in torch_images) + if len(sizes) == 1: + torch_images = torch.stack(torch_images) + # Return as list if images have different sizes + return torch_images + + def load_depths(self, frames): + torch_depths = [] + for idx, frame in enumerate(frames): + depthmap = imread_cv2(frame["depth_path"], cv2.IMREAD_UNCHANGED) + depthmap = depthmap.astype(np.float32) / 1000 + depthmap[~np.isfinite(depthmap)] = 0 + torch_depths.append(torch.from_numpy(depthmap)) + return torch.stack(torch_depths) # [N, H, W] + + + def getitem(self, index: int, num_context_views: int, patchsize: tuple) -> dict: + # import time + # start_time = time.time() + + scene = self.scene_ids[index] + example = self.scenes[scene] + # load poses + extrinsics = [] + intrinsics = [] + for frame in example: + extrinsic = frame["extrinsics"] + intrinsic = frame["intrinsics"] + extrinsics.append(extrinsic) + intrinsics.append(intrinsic) + + extrinsics = np.array(extrinsics) + intrinsics = np.array(intrinsics) + extrinsics = torch.tensor(extrinsics, dtype=torch.float32) + intrinsics = torch.tensor(intrinsics, dtype=torch.float32) + + try: + context_indices, target_indices, overlap = self.view_sampler.sample( + "scannetpp_"+scene, + num_context_views, + extrinsics, + intrinsics, + ) + except ValueError: + # Skip because the example doesn't have enough frames. + raise Exception("Not enough frames") + + # Skip the example if the field of view is too wide. + if (get_fov(intrinsics).rad2deg() > self.cfg.max_fov).any(): + raise Exception("Field of view too wide") + + # Load the images. + input_frames = [example[i] for i in context_indices] + target_frame = [example[i] for i in target_indices] + + context_images = self.load_frames(input_frames) + target_images = self.load_frames(target_frame) + + context_depths = self.load_depths(input_frames) + target_depths = self.load_depths(target_frame) + + # Skip the example if the images don't have the right shape. + context_image_invalid = context_images.shape[1:] != (3, *self.cfg.original_image_shape) + target_image_invalid = target_images.shape[1:] != (3, *self.cfg.original_image_shape) + if self.cfg.skip_bad_shape and (context_image_invalid or target_image_invalid): + print( + f"Skipped bad example {example['key']}. Context shape was " + f"{context_images.shape} and target shape was " + f"{target_images.shape}." + ) + raise Exception("Bad example image shape") + + context_extrinsics = extrinsics[context_indices] + + if self.cfg.make_baseline_1: + a, b = context_extrinsics[0, :3, 3], context_extrinsics[-1, :3, 3] + scale = (a - b).norm() + if scale < self.cfg.baseline_min or scale > self.cfg.baseline_max: + print( + f"Skipped {scene} because of baseline out of range: " + f"{scale:.6f}" + ) + raise Exception("baseline out of range") + extrinsics[:, :3, 3] /= scale + else: + scale = 1 + + if self.cfg.relative_pose: + extrinsics = camera_normalization(extrinsics[context_indices][0:1], extrinsics) + + if self.cfg.rescale_to_1cube: + scene_scale = torch.max(torch.abs(extrinsics[context_indices][:, :3, 3])) # target pose is not included + rescale_factor = 1 * scene_scale + extrinsics[:, :3, 3] /= rescale_factor + + if torch.isnan(extrinsics).any() or torch.isinf(extrinsics).any(): + raise Exception("encounter nan or inf in input poses") + + example = { + "context": { + "extrinsics": extrinsics[context_indices], + "intrinsics": intrinsics[context_indices], + "image": context_images, + "depth": context_depths, + "near": self.get_bound("near", len(context_indices)) / scale, + "far": self.get_bound("far", len(context_indices)) / scale, + "index": context_indices, + "overlap": overlap, + }, + "target": { + "extrinsics": extrinsics[target_indices], + "intrinsics": intrinsics[target_indices], + "image": target_images, + "depth": target_depths, + "near": self.get_bound("near", len(target_indices)) / scale, + "far": self.get_bound("far", len(target_indices)) / scale, + "index": target_indices, + }, + "scene": f"Scannetpp {scene}", + } + if self.stage == "train" and self.cfg.augment: + example = apply_augmentation_shim(example) + + if self.stage == "train" and self.cfg.intr_augment: + intr_aug = True + else: + intr_aug = False + + example = apply_crop_shim(example, (patchsize[0] * 14, patchsize[1] * 14), intr_aug=intr_aug) + + # world pts + image_size = example["context"]["image"].shape[2:] + context_intrinsics = example["context"]["intrinsics"].clone().detach().numpy() + context_intrinsics[:, 0] = context_intrinsics[:, 0] * image_size[1] + context_intrinsics[:, 1] = context_intrinsics[:, 1] * image_size[0] + + target_intrinsics = example["target"]["intrinsics"].clone().detach().numpy() + target_intrinsics[:, 0] = target_intrinsics[:, 0] * image_size[1] + target_intrinsics[:, 1] = target_intrinsics[:, 1] * image_size[0] + + context_pts3d_list, context_valid_mask_list = [], [] + target_pts3d_list, target_valid_mask_list = [], [] + + for i in range(len(example["context"]["depth"])): + context_pts3d, context_valid_mask = depthmap_to_absolute_camera_coordinates(example["context"]["depth"][i].numpy(), context_intrinsics[i], example["context"]["extrinsics"][i].numpy()) + context_pts3d_list.append(torch.from_numpy(context_pts3d).to(torch.float32)) + context_valid_mask_list.append(torch.from_numpy(context_valid_mask)) + + context_pts3d = torch.stack(context_pts3d_list, dim=0) + context_valid_mask = torch.stack(context_valid_mask_list, dim=0) + + for i in range(len(example["target"]["depth"])): + target_pts3d, target_valid_mask = depthmap_to_absolute_camera_coordinates(example["target"]["depth"][i].numpy(), target_intrinsics[i], example["target"]["extrinsics"][i].numpy()) + target_pts3d_list.append(torch.from_numpy(target_pts3d).to(torch.float32)) + target_valid_mask_list.append(torch.from_numpy(target_valid_mask)) + + target_pts3d = torch.stack(target_pts3d_list, dim=0) + target_valid_mask = torch.stack(target_valid_mask_list, dim=0) + + # normalize by context pts3d + if self.cfg.normalize_by_pts3d: + transformed_pts3d = context_pts3d[context_valid_mask] + scene_factor = transformed_pts3d.norm(dim=-1).mean().clip(min=1e-8) + context_pts3d /= scene_factor + example["context"]["depth"] /= scene_factor + example["context"]["extrinsics"][:, :3, 3] /= scene_factor + + target_pts3d /= scene_factor + example["target"]["depth"] /= scene_factor + example["target"]["extrinsics"][:, :3, 3] /= scene_factor + + example["context"]["pts3d"] = context_pts3d + example["target"]["pts3d"] = target_pts3d + example["context"]["valid_mask"] = context_valid_mask + example["target"]["valid_mask"] = target_valid_mask + + if torch.isnan(example["context"]["depth"]).any() or torch.isinf(example["context"]["depth"]).any() or \ + torch.isnan(example["context"]["extrinsics"]).any() or torch.isinf(example["context"]["extrinsics"]).any() or \ + torch.isnan(example["context"]["intrinsics"]).any() or torch.isinf(example["context"]["intrinsics"]).any() or \ + torch.isnan(example["target"]["depth"]).any() or torch.isinf(example["target"]["depth"]).any() or \ + torch.isnan(example["target"]["extrinsics"]).any() or torch.isinf(example["target"]["extrinsics"]).any() or \ + torch.isnan(example["target"]["intrinsics"]).any() or torch.isinf(example["target"]["intrinsics"]).any(): + raise Exception("encounter nan or inf in context depth") + + for key in ["context", "target"]: + example[key]["valid_mask"] = (torch.ones_like(example[key]["valid_mask"]) * -1).type(torch.int32) + + return example + + + def __getitem__(self, index_tuple: tuple) -> dict: + index, num_context_views, patchsize_h = index_tuple + # generate a random patch size + patchsize_w = (self.cfg.input_image_shape[1] // 14) + try: + return self.getitem(index, num_context_views, (patchsize_h, patchsize_w)) + except Exception as e: + print(f"Error: {e}") + index = np.random.randint(len(self)) + return self.__getitem__((index, num_context_views, patchsize_h)) + + def convert_poses( + self, + poses: Float[Tensor, "batch 18"], + ) -> tuple[ + Float[Tensor, "batch 4 4"], # extrinsics + Float[Tensor, "batch 3 3"], # intrinsics + ]: + b, _ = poses.shape + + # Convert the intrinsics to a 3x3 normalized K matrix. + intrinsics = torch.eye(3, dtype=torch.float32) + intrinsics = repeat(intrinsics, "h w -> b h w", b=b).clone() + fx, fy, cx, cy = poses[:, :4].T + intrinsics[:, 0, 0] = fx + intrinsics[:, 1, 1] = fy + intrinsics[:, 0, 2] = cx + intrinsics[:, 1, 2] = cy + + # Convert the extrinsics to a 4x4 OpenCV-style W2C matrix. + w2c = repeat(torch.eye(4, dtype=torch.float32), "h w -> b h w", b=b).clone() + w2c[:, :3] = rearrange(poses[:, 6:], "b (h w) -> b h w", h=3, w=4) + return w2c.inverse(), intrinsics + + def convert_images( + self, + images: list[UInt8[Tensor, "..."]], + ) -> Float[Tensor, "batch 3 height width"]: + torch_images = [] + for image in images: + image = Image.open(BytesIO(image.numpy().tobytes())) + torch_images.append(self.to_tensor(image)) + return torch.stack(torch_images) + + def get_bound( + self, + bound: Literal["near", "far"], + num_views: int, + ) -> Float[Tensor, " view"]: + value = torch.tensor(getattr(self, bound), dtype=torch.float32) + return repeat(value, "-> v", v=num_views) + + @property + def data_stage(self) -> Stage: + if self.cfg.overfit_to_scene is not None: + return "test" + if self.stage == "val": + return "test" + return self.stage + + @cached_property + def index(self) -> dict[str, Path]: + merged_index = {} + data_stages = [self.data_stage] + if self.cfg.overfit_to_scene is not None: + data_stages = ("test", "train") + for data_stage in data_stages: + for root in self.cfg.roots: + # Load the root's index. + with (root / data_stage / "index.json").open("r") as f: + index = json.load(f) + index = {k: Path(root / data_stage / v) for k, v in index.items()} + + # The constituent datasets should have unique keys. + assert not (set(merged_index.keys()) & set(index.keys())) + + # Merge the root's index into the main index. + merged_index = {**merged_index, **index} + return merged_index + + def __len__(self) -> int: + return len(self.data_list) \ No newline at end of file diff --git a/src/dataset/shims/augmentation_shim.py b/src/dataset/shims/augmentation_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..b4e6e7dd0d4ffe77476a34768c5b0f9d91867a57 --- /dev/null +++ b/src/dataset/shims/augmentation_shim.py @@ -0,0 +1,219 @@ +import copy +import random +import numpy as np +import torch +from jaxtyping import Float +from torch import Tensor + +from ..types import AnyExample, AnyViews + + +def reflect_extrinsics( + extrinsics: Float[Tensor, "*batch 4 4"], +) -> Float[Tensor, "*batch 4 4"]: + reflect = torch.eye(4, dtype=torch.float32, device=extrinsics.device) + reflect[0, 0] = -1 + return reflect @ extrinsics @ reflect + + +def reflect_views(views: AnyViews) -> AnyViews: + if "depth" in views.keys(): + return { + **views, + "image": views["image"].flip(-1), + "extrinsics": reflect_extrinsics(views["extrinsics"]), + "depth": views["depth"].flip(-1), + } + else: + return { + **views, + "image": views["image"].flip(-1), + "extrinsics": reflect_extrinsics(views["extrinsics"]), + } + + +def apply_augmentation_shim( + example: AnyExample, + generator: torch.Generator | None = None, +) -> AnyExample: + """Randomly augment the training images.""" + # Do not augment with 50% chance. + if torch.rand(tuple(), generator=generator) < 0.5: + return example + + return { + **example, + "context": reflect_views(example["context"]), + "target": reflect_views(example["target"]), + } + +def rotate_90_degrees( + image: torch.Tensor, depth_map: torch.Tensor | None, extri_opencv: torch.Tensor, intri_opencv: torch.Tensor, clockwise=True +): + """ + Rotates the input image, depth map, and camera parameters by 90 degrees. + + Applies one of two 90-degree rotations: + - Clockwise + - Counterclockwise (if clockwise=False) + + The extrinsic and intrinsic matrices are adjusted accordingly to maintain + correct camera geometry. + + Args: + image (torch.Tensor): + Input image tensor of shape (C, H, W). + depth_map (torch.Tensor or None): + Depth map tensor of shape (H, W), or None if not available. + extri_opencv (torch.Tensor): + Extrinsic matrix (3x4) in OpenCV convention. + intri_opencv (torch.Tensor): + Intrinsic matrix (3x3). + clockwise (bool): + If True, rotates the image 90 degrees clockwise; else 90 degrees counterclockwise. + + Returns: + tuple: + ( + rotated_image, + rotated_depth_map, + new_extri_opencv, + new_intri_opencv + ) + + Where each is the updated version after the rotation. + """ + image_height, image_width = image.shape[-2:] + + # Rotate the image and depth map + rotated_image, rotated_depth_map = rotate_image_and_depth_rot90(image, depth_map, clockwise) + # Adjust the intrinsic matrix + new_intri_opencv = adjust_intrinsic_matrix_rot90(intri_opencv, image_width, image_height, clockwise) + # Adjust the extrinsic matrix + new_extri_opencv = adjust_extrinsic_matrix_rot90(extri_opencv, clockwise) + + return ( + rotated_image, + rotated_depth_map, + new_extri_opencv, + new_intri_opencv, + ) + + +def rotate_image_and_depth_rot90(image: torch.Tensor, depth_map: torch.Tensor | None, clockwise: bool): + """ + Rotates the given image and depth map by 90 degrees (clockwise or counterclockwise). + + Args: + image (torch.Tensor): + Input image tensor of shape (C, H, W). + depth_map (torch.Tensor or None): + Depth map tensor of shape (H, W), or None if not available. + clockwise (bool): + If True, rotate 90 degrees clockwise; else 90 degrees counterclockwise. + + Returns: + tuple: + (rotated_image, rotated_depth_map) + """ + rotated_depth_map = None + if clockwise: + rotated_image = torch.rot90(image, k=-1, dims=[-2, -1]) + if depth_map is not None: + rotated_depth_map = torch.rot90(depth_map, k=-1, dims=[-2, -1]) + else: + rotated_image = torch.rot90(image, k=1, dims=[-2, -1]) + if depth_map is not None: + rotated_depth_map = torch.rot90(depth_map, k=1, dims=[-2, -1]) + return rotated_image, rotated_depth_map + + +def adjust_extrinsic_matrix_rot90(extri_opencv: torch.Tensor, clockwise: bool): + """ + Adjusts the extrinsic matrix (3x4) for a 90-degree rotation of the image. + + The rotation is in the image plane. This modifies the camera orientation + accordingly. The function applies either a clockwise or counterclockwise + 90-degree rotation. + + Args: + extri_opencv (torch.Tensor): + Extrinsic matrix (3x4) in OpenCV convention. + clockwise (bool): + If True, rotate extrinsic for a 90-degree clockwise image rotation; + otherwise, counterclockwise. + + Returns: + torch.Tensor: + A new 3x4 extrinsic matrix after the rotation. + """ + R = extri_opencv[:3, :3] + t = extri_opencv[:3, 3] + + if clockwise: + R_rotation = torch.tensor([ + [0, -1, 0], + [1, 0, 0], + [0, 0, 1] + ], dtype=extri_opencv.dtype, device=extri_opencv.device) + else: + R_rotation = torch.tensor([ + [0, 1, 0], + [-1, 0, 0], + [0, 0, 1] + ], dtype=extri_opencv.dtype, device=extri_opencv.device) + + new_R = torch.matmul(R_rotation, R) + new_t = torch.matmul(R_rotation, t) + new_extri_opencv = torch.cat((new_R, new_t.reshape(-1, 1)), dim=1) + new_extri_opencv = torch.cat((new_extri_opencv, + torch.tensor([[0, 0, 0, 1]], + dtype=extri_opencv.dtype, device=extri_opencv.device)), dim=0) + return new_extri_opencv + + +def adjust_intrinsic_matrix_rot90(intri_opencv: torch.Tensor, image_width: int, image_height: int, clockwise: bool): + """ + Adjusts the intrinsic matrix (3x3) for a 90-degree rotation of the image in the image plane. + + Args: + intri_opencv (torch.Tensor): + Intrinsic matrix (3x3). + image_width (int): + Original width of the image. + image_height (int): + Original height of the image. + clockwise (bool): + If True, rotate 90 degrees clockwise; else 90 degrees counterclockwise. + + Returns: + torch.Tensor: + A new 3x3 intrinsic matrix after the rotation. + """ + intri_opencv = copy.deepcopy(intri_opencv) + intri_opencv[0, :] *= image_width + intri_opencv[1, :] *= image_height + + fx, fy, cx, cy = ( + intri_opencv[0, 0], + intri_opencv[1, 1], + intri_opencv[0, 2], + intri_opencv[1, 2], + ) + + new_intri_opencv = torch.eye(3, dtype=intri_opencv.dtype, device=intri_opencv.device) + if clockwise: + new_intri_opencv[0, 0] = fy + new_intri_opencv[1, 1] = fx + new_intri_opencv[0, 2] = image_height - cy + new_intri_opencv[1, 2] = cx + else: + new_intri_opencv[0, 0] = fy + new_intri_opencv[1, 1] = fx + new_intri_opencv[0, 2] = cy + new_intri_opencv[1, 2] = image_width - cx + + new_intri_opencv[0, :] /= image_height + new_intri_opencv[1, :] /= image_width + + return new_intri_opencv diff --git a/src/dataset/shims/bounds_shim.py b/src/dataset/shims/bounds_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..f699867554de265bace335d22a90dfb259560650 --- /dev/null +++ b/src/dataset/shims/bounds_shim.py @@ -0,0 +1,80 @@ +import torch +from einops import einsum, reduce, repeat +from jaxtyping import Float +from torch import Tensor + +from ..types import BatchedExample + + +def compute_depth_for_disparity( + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + image_shape: tuple[int, int], + disparity: float, + delta_min: float = 1e-6, # This prevents motionless scenes from lacking depth. +) -> Float[Tensor, " batch"]: + """Compute the depth at which moving the maximum distance between cameras + corresponds to the specified disparity (in pixels). + """ + + # Use the furthest distance between cameras as the baseline. + origins = extrinsics[:, :, :3, 3] + deltas = (origins[:, None, :, :] - origins[:, :, None, :]).norm(dim=-1) + deltas = deltas.clip(min=delta_min) + baselines = reduce(deltas, "b v ov -> b", "max") + + # Compute a single pixel's size at depth 1. + h, w = image_shape + pixel_size = 1 / torch.tensor((w, h), dtype=torch.float32, device=extrinsics.device) + pixel_size = einsum( + intrinsics[..., :2, :2].inverse(), pixel_size, "... i j, j -> ... i" + ) + + # This wouldn't make sense with non-square pixels, but then again, non-square pixels + # don't make much sense anyway. + mean_pixel_size = reduce(pixel_size, "b v xy -> b", "mean") + + return baselines / (disparity * mean_pixel_size) + + +def apply_bounds_shim( + batch: BatchedExample, + near_disparity: float, + far_disparity: float, +) -> BatchedExample: + """Compute reasonable near and far planes (lower and upper bounds on depth). This + assumes that all of an example's views are of roughly the same thing. + """ + + context = batch["context"] + _, cv, _, h, w = context["image"].shape + + # Compute near and far planes using the context views. + near = compute_depth_for_disparity( + context["extrinsics"], + context["intrinsics"], + (h, w), + near_disparity, + ) + far = compute_depth_for_disparity( + context["extrinsics"], + context["intrinsics"], + (h, w), + far_disparity, + ) + + target = batch["target"] + _, tv, _, _, _ = target["image"].shape + return { + **batch, + "context": { + **context, + "near": repeat(near, "b -> b v", v=cv), + "far": repeat(far, "b -> b v", v=cv), + }, + "target": { + **target, + "near": repeat(near, "b -> b v", v=tv), + "far": repeat(far, "b -> b v", v=tv), + }, + } diff --git a/src/dataset/shims/crop_shim.py b/src/dataset/shims/crop_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..25b50f4b7359cb96aca86615dfcb6091f76df3a8 --- /dev/null +++ b/src/dataset/shims/crop_shim.py @@ -0,0 +1,196 @@ +import random +import numpy as np +import torch +from einops import rearrange +from jaxtyping import Float +from PIL import Image +from torch import Tensor +import torchvision.transforms.functional as F +import cv2 + +from ..types import AnyExample, AnyViews + + +def rescale( + image: Float[Tensor, "3 h_in w_in"], + shape: tuple[int, int], +) -> Float[Tensor, "3 h_out w_out"]: + h, w = shape + image_new = (image * 255).clip(min=0, max=255).type(torch.uint8) + image_new = rearrange(image_new, "c h w -> h w c").detach().cpu().numpy() + image_new = Image.fromarray(image_new) + image_new = image_new.resize((w, h), Image.LANCZOS) + image_new = np.array(image_new) / 255 + image_new = torch.tensor(image_new, dtype=image.dtype, device=image.device) + return rearrange(image_new, "h w c -> c h w") + +def rescale_depth( + depth: Float[Tensor, "1 h w"], + shape: tuple[int, int], +) -> Float[Tensor, "1 h_out w_out"]: + h, w = shape + depth_new = depth.detach().cpu().numpy() + depth_new = cv2.resize(depth_new, (w,h), interpolation=cv2.INTER_NEAREST) + depth_new = torch.from_numpy(depth_new).to(depth.device) + return depth_new + +def center_crop( + images: Float[Tensor, "*#batch c h w"], + intrinsics: Float[Tensor, "*#batch 3 3"], + shape: tuple[int, int], + depths: Float[Tensor, "*#batch 1 h w"] | None = None, +) -> tuple[ + Float[Tensor, "*#batch c h_out w_out"], # updated images + Float[Tensor, "*#batch 3 3"], # updated intrinsics + Float[Tensor, "*#batch 1 h_out w_out"] | None, # updated depths +]: + *_, h_in, w_in = images.shape + h_out, w_out = shape + + # Note that odd input dimensions induce half-pixel misalignments. + row = (h_in - h_out) // 2 + col = (w_in - w_out) // 2 + + # Center-crop the image. + images = images[..., :, row : row + h_out, col : col + w_out] + + if depths is not None: + depths = depths[..., row : row + h_out, col : col + w_out] + + # Adjust the intrinsics to account for the cropping. + intrinsics = intrinsics.clone() + intrinsics[..., 0, 0] *= w_in / w_out # fx + intrinsics[..., 1, 1] *= h_in / h_out # fy + + + if depths is not None: + return images, intrinsics, depths + else: + return images, intrinsics + + +def rescale_and_crop( + images: Float[Tensor, "*#batch c h w"], + intrinsics: Float[Tensor, "*#batch 3 3"], + shape: tuple[int, int], + intr_aug: bool = False, + scale_range: tuple[float, float] = (0.77, 1.0), + depths: Float[Tensor, "*#batch 1 h w"] | None = None, +) -> tuple[ + Float[Tensor, "*#batch c h_out w_out"], # updated images + Float[Tensor, "*#batch 3 3"], # updated intrinsics + Float[Tensor, "*#batch 1 h_out w_out"] | None, # updated depths +]: + if type(images) == list: + images_new = [] + intrinsics_new = [] + for i in range(len(images)): + image = images[i] + intrinsic = intrinsics[i] + + *_, h_in, w_in = image.shape + h_out, w_out = shape + + scale_factor = max(h_out / h_in, w_out / w_in) + h_scaled = round(h_in * scale_factor) + w_scaled = round(w_in * scale_factor) + image = F.resize(image, (h_scaled, w_scaled)) + image = F.center_crop(image, (h_out, w_out)) + images_new.append(image) + + intrinsic_new = intrinsic.clone() + intrinsic_new[..., 0, 0] *= w_scaled / w_in # fx + intrinsic_new[..., 1, 1] *= h_scaled / h_in # fy + intrinsics_new.append(intrinsic_new) + + if depths is not None: + depths_new = [] + for i in range(len(depths)): + depth = depths[i] + depth = rescale_depth(depth, (h_out, w_out)) + depth = F.center_crop(depth, (h_out, w_out)) + depths_new.append(depth) + return torch.stack(images_new), torch.stack(intrinsics_new), torch.stack(depths_new) + else: + return torch.stack(images_new), torch.stack(intrinsics_new) + + else: + # we only support intr_aug for clean datasets + *_, h_in, w_in = images.shape + h_out, w_out = shape + # assert h_out <= h_in and w_out <= w_in # to avoid the case that the image is too small, like co3d + + if intr_aug: + scale = random.uniform(*scale_range) + h_scale = round(h_out * scale) + w_scale = round(w_out * scale) + else: + h_scale = h_out + w_scale = w_out + + scale_factor = max(h_scale / h_in, w_scale / w_in) + h_scaled = round(h_in * scale_factor) + w_scaled = round(w_in * scale_factor) + assert h_scaled == h_scale or w_scaled == w_scale + + # Reshape the images to the correct size. Assume we don't have to worry about + # changing the intrinsics based on how the images are rounded. + *batch, c, h, w = images.shape + images = images.reshape(-1, c, h, w) + images = torch.stack([rescale(image, (h_scaled, w_scaled)) for image in images]) + images = images.reshape(*batch, c, h_scaled, w_scaled) + + if depths is not None: + if type(depths) == list: + depths_new = [] + for i in range(len(depths)): + depth = depths[i] + depth = rescale_depth(depth, (h_scaled, w_scaled)) + depths_new.append(depth) + depths = torch.stack(depths_new) + else: + depths = depths.reshape(-1, h, w) + depths = torch.stack([rescale_depth(depth, (h_scaled, w_scaled)) for depth in depths]) + depths = depths.reshape(*batch, h_scaled, w_scaled) + + images, intrinsics, depths = center_crop(images, intrinsics, (h_scale, w_scale), depths) + + if intr_aug: + images = F.resize(images, size=(h_out, w_out), interpolation=F.InterpolationMode.BILINEAR) + depths = F.resize(depths, size=(h_out, w_out), interpolation=F.InterpolationMode.NEAREST) + + return images, intrinsics, depths + else: + images, intrinsics = center_crop(images, intrinsics, (h_scale, w_scale)) + + if intr_aug: + images = F.resize(images, size=(h_out, w_out)) + + return images, intrinsics + + +def apply_crop_shim_to_views(views: AnyViews, shape: tuple[int, int], intr_aug: bool = False) -> AnyViews: + if "depth" in views.keys(): + images, intrinsics, depths = rescale_and_crop(views["image"], views["intrinsics"], shape, depths=views["depth"], intr_aug=intr_aug) + return { + **views, + "image": images, + "intrinsics": intrinsics, + "depth": depths, + } + else: + images, intrinsics = rescale_and_crop(views["image"], views["intrinsics"], shape, intr_aug) + return { + **views, + "image": images, + "intrinsics": intrinsics, + } + + +def apply_crop_shim(example: AnyExample, shape: tuple[int, int], intr_aug: bool = False) -> AnyExample: + """Crop images in the example.""" + return { + **example, + "context": apply_crop_shim_to_views(example["context"], shape, intr_aug), + "target": apply_crop_shim_to_views(example["target"], shape, intr_aug), + } diff --git a/src/dataset/shims/geometry_shim.py b/src/dataset/shims/geometry_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..e5a1bbceaac16e18914678162d7f55b27c4d7561 --- /dev/null +++ b/src/dataset/shims/geometry_shim.py @@ -0,0 +1,383 @@ +# Copyright (C) 2024-present Naver Corporation. All rights reserved. +# Licensed under CC BY-NC-SA 4.0 (non-commercial use only). +# +# -------------------------------------------------------- +# geometry utilitary functions +# -------------------------------------------------------- +import torch +import numpy as np +from scipy.spatial import cKDTree as KDTree + +def invalid_to_nans(arr, valid_mask, ndim=999): + if valid_mask is not None: + arr = arr.clone() + arr[~valid_mask] = float('nan') + if arr.ndim > ndim: + arr = arr.flatten(-2 - (arr.ndim - ndim), -2) + return arr + + +def invalid_to_zeros(arr, valid_mask, ndim=999): + if valid_mask is not None: + arr = arr.clone() + arr[~valid_mask] = 0 + nnz = valid_mask.view(len(valid_mask), -1).sum(1) + else: + nnz = arr.numel() // len(arr) if len(arr) else 0 # number of point per image + if arr.ndim > ndim: + arr = arr.flatten(-2 - (arr.ndim - ndim), -2) + return arr, nnz + + +def xy_grid(W, H, device=None, origin=(0, 0), unsqueeze=None, cat_dim=-1, homogeneous=False, **arange_kw): + """ Output a (H,W,2) array of int32 + with output[j,i,0] = i + origin[0] + output[j,i,1] = j + origin[1] + """ + if device is None: + # numpy + arange, meshgrid, stack, ones = np.arange, np.meshgrid, np.stack, np.ones + else: + # torch + arange = lambda *a, **kw: torch.arange(*a, device=device, **kw) + meshgrid, stack = torch.meshgrid, torch.stack + ones = lambda *a: torch.ones(*a, device=device) + + tw, th = [arange(o, o + s, **arange_kw) for s, o in zip((W, H), origin)] + grid = meshgrid(tw, th, indexing='xy') + if homogeneous: + grid = grid + (ones((H, W)),) + if unsqueeze is not None: + grid = (grid[0].unsqueeze(unsqueeze), grid[1].unsqueeze(unsqueeze)) + if cat_dim is not None: + grid = stack(grid, cat_dim) + return grid + + +def geotrf(Trf, pts, ncol=None, norm=False): + """ Apply a geometric transformation to a list of 3-D points. + + H: 3x3 or 4x4 projection matrix (typically a Homography) + p: numpy/torch/tuple of coordinates. Shape must be (...,2) or (...,3) + + ncol: int. number of columns of the result (2 or 3) + norm: float. if != 0, the resut is projected on the z=norm plane. + + Returns an array of projected 2d points. + """ + assert Trf.ndim >= 2 + if isinstance(Trf, np.ndarray): + pts = np.asarray(pts) + elif isinstance(Trf, torch.Tensor): + pts = torch.as_tensor(pts, dtype=Trf.dtype) + + # adapt shape if necessary + output_reshape = pts.shape[:-1] + ncol = ncol or pts.shape[-1] + + # optimized code + if (isinstance(Trf, torch.Tensor) and isinstance(pts, torch.Tensor) and + Trf.ndim == 3 and pts.ndim == 4): + d = pts.shape[3] + if Trf.shape[-1] == d: + pts = torch.einsum("bij, bhwj -> bhwi", Trf, pts) + elif Trf.shape[-1] == d + 1: + pts = torch.einsum("bij, bhwj -> bhwi", Trf[:, :d, :d], pts) + Trf[:, None, None, :d, d] + else: + raise ValueError(f'bad shape, not ending with 3 or 4, for {pts.shape=}') + else: + if Trf.ndim >= 3: + n = Trf.ndim - 2 + assert Trf.shape[:n] == pts.shape[:n], 'batch size does not match' + Trf = Trf.reshape(-1, Trf.shape[-2], Trf.shape[-1]) + + if pts.ndim > Trf.ndim: + # Trf == (B,d,d) & pts == (B,H,W,d) --> (B, H*W, d) + pts = pts.reshape(Trf.shape[0], -1, pts.shape[-1]) + elif pts.ndim == 2: + # Trf == (B,d,d) & pts == (B,d) --> (B, 1, d) + pts = pts[:, None, :] + + if pts.shape[-1] + 1 == Trf.shape[-1]: + Trf = Trf.swapaxes(-1, -2) # transpose Trf + pts = pts @ Trf[..., :-1, :] + Trf[..., -1:, :] + elif pts.shape[-1] == Trf.shape[-1]: + Trf = Trf.swapaxes(-1, -2) # transpose Trf + pts = pts @ Trf + else: + pts = Trf @ pts.T + if pts.ndim >= 2: + pts = pts.swapaxes(-1, -2) + + if norm: + pts = pts / pts[..., -1:] # DONT DO /= BECAUSE OF WEIRD PYTORCH BUG + if norm != 1: + pts *= norm + + res = pts[..., :ncol].reshape(*output_reshape, ncol) + return res + + +def inv(mat): + """ Invert a torch or numpy matrix + """ + if isinstance(mat, torch.Tensor): + return torch.linalg.inv(mat) + if isinstance(mat, np.ndarray): + return np.linalg.inv(mat) + raise ValueError(f'bad matrix type = {type(mat)}') + + +def depthmap_to_pts3d(depth, pseudo_focal, pp=None, **_): + """ + Args: + - depthmap (BxHxW array): + - pseudo_focal: [B,H,W] ; [B,2,H,W] or [B,1,H,W] + Returns: + pointmap of absolute coordinates (BxHxWx3 array) + """ + + if len(depth.shape) == 4: + B, H, W, n = depth.shape + else: + B, H, W = depth.shape + n = None + + if len(pseudo_focal.shape) == 3: # [B,H,W] + pseudo_focalx = pseudo_focaly = pseudo_focal + elif len(pseudo_focal.shape) == 4: # [B,2,H,W] or [B,1,H,W] + pseudo_focalx = pseudo_focal[:, 0] + if pseudo_focal.shape[1] == 2: + pseudo_focaly = pseudo_focal[:, 1] + else: + pseudo_focaly = pseudo_focalx + else: + raise NotImplementedError("Error, unknown input focal shape format.") + + assert pseudo_focalx.shape == depth.shape[:3] + assert pseudo_focaly.shape == depth.shape[:3] + grid_x, grid_y = xy_grid(W, H, cat_dim=0, device=depth.device)[:, None] + + # set principal point + if pp is None: + grid_x = grid_x - (W - 1) / 2 + grid_y = grid_y - (H - 1) / 2 + else: + grid_x = grid_x.expand(B, -1, -1) - pp[:, 0, None, None] + grid_y = grid_y.expand(B, -1, -1) - pp[:, 1, None, None] + + if n is None: + pts3d = torch.empty((B, H, W, 3), device=depth.device) + pts3d[..., 0] = depth * grid_x / pseudo_focalx + pts3d[..., 1] = depth * grid_y / pseudo_focaly + pts3d[..., 2] = depth + else: + pts3d = torch.empty((B, H, W, 3, n), device=depth.device) + pts3d[..., 0, :] = depth * (grid_x / pseudo_focalx)[..., None] + pts3d[..., 1, :] = depth * (grid_y / pseudo_focaly)[..., None] + pts3d[..., 2, :] = depth + return pts3d + + +def depthmap_to_camera_coordinates(depthmap, camera_intrinsics, pseudo_focal=None): + """ + Args: + - depthmap (HxW array): + - camera_intrinsics: a 3x3 matrix + Returns: + pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels. + """ + camera_intrinsics = np.float32(camera_intrinsics) + H, W = depthmap.shape + + # Compute 3D ray associated with each pixel + # Strong assumption: there are no skew terms + assert camera_intrinsics[0, 1] == 0.0 + assert camera_intrinsics[1, 0] == 0.0 + if pseudo_focal is None: + fu = camera_intrinsics[0, 0] + fv = camera_intrinsics[1, 1] + else: + assert pseudo_focal.shape == (H, W) + fu = fv = pseudo_focal + cu = camera_intrinsics[0, 2] + cv = camera_intrinsics[1, 2] + + u, v = np.meshgrid(np.arange(W), np.arange(H)) + z_cam = depthmap + x_cam = (u - cu) * z_cam / fu + y_cam = (v - cv) * z_cam / fv + X_cam = np.stack((x_cam, y_cam, z_cam), axis=-1).astype(np.float32) + + # Mask for valid coordinates + valid_mask = (depthmap > 0.0) + return X_cam, valid_mask + + +def depthmap_to_absolute_camera_coordinates(depthmap, camera_intrinsics, camera_pose=None, **kw): + """ + Args: + - depthmap (HxW array): + - camera_intrinsics: a 3x3 matrix + - camera_pose: a 4x3 or 4x4 cam2world matrix + Returns: + pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels.""" + X_cam, valid_mask = depthmap_to_camera_coordinates(depthmap, camera_intrinsics) + + X_world = X_cam # default + if camera_pose is not None: + # R_cam2world = np.float32(camera_params["R_cam2world"]) + # t_cam2world = np.float32(camera_params["t_cam2world"]).squeeze() + R_cam2world = camera_pose[:3, :3] + t_cam2world = camera_pose[:3, 3] + + # Express in absolute coordinates (invalid depth values) + X_world = np.einsum("ik, vuk -> vui", R_cam2world, X_cam) + t_cam2world[None, None, :] + + return X_world, valid_mask + + +def colmap_to_opencv_intrinsics(K): + """ + Modify camera intrinsics to follow a different convention. + Coordinates of the center of the top-left pixels are by default: + - (0.5, 0.5) in Colmap + - (0,0) in OpenCV + """ + K = K.copy() + K[0, 2] -= 0.5 + K[1, 2] -= 0.5 + return K + + +def opencv_to_colmap_intrinsics(K): + """ + Modify camera intrinsics to follow a different convention. + Coordinates of the center of the top-left pixels are by default: + - (0.5, 0.5) in Colmap + - (0,0) in OpenCV + """ + K = K.copy() + K[0, 2] += 0.5 + K[1, 2] += 0.5 + return K + + +def normalize_pointcloud(pts1, pts2, norm_mode='avg_dis', valid1=None, valid2=None, ret_factor=False): + """ renorm pointmaps pts1, pts2 with norm_mode + """ + assert pts1.ndim >= 3 and pts1.shape[-1] == 3 + assert pts2 is None or (pts2.ndim >= 3 and pts2.shape[-1] == 3) + norm_mode, dis_mode = norm_mode.split('_') + + if norm_mode == 'avg': + # gather all points together (joint normalization) + nan_pts1, nnz1 = invalid_to_zeros(pts1, valid1, ndim=3) + nan_pts2, nnz2 = invalid_to_zeros(pts2, valid2, ndim=3) if pts2 is not None else (None, 0) + all_pts = torch.cat((nan_pts1, nan_pts2), dim=1) if pts2 is not None else nan_pts1 + + # compute distance to origin + all_dis = all_pts.norm(dim=-1) + if dis_mode == 'dis': + pass # do nothing + elif dis_mode == 'log1p': + all_dis = torch.log1p(all_dis) + elif dis_mode == 'warp-log1p': + # actually warp input points before normalizing them + log_dis = torch.log1p(all_dis) + warp_factor = log_dis / all_dis.clip(min=1e-8) + H1, W1 = pts1.shape[1:-1] + pts1 = pts1 * warp_factor[:, :W1 * H1].view(-1, H1, W1, 1) + if pts2 is not None: + H2, W2 = pts2.shape[1:-1] + pts2 = pts2 * warp_factor[:, W1 * H1:].view(-1, H2, W2, 1) + all_dis = log_dis # this is their true distance afterwards + else: + raise ValueError(f'bad {dis_mode=}') + + norm_factor = all_dis.sum(dim=1) / (nnz1 + nnz2 + 1e-8) + else: + # gather all points together (joint normalization) + nan_pts1 = invalid_to_nans(pts1, valid1, ndim=3) + nan_pts2 = invalid_to_nans(pts2, valid2, ndim=3) if pts2 is not None else None + all_pts = torch.cat((nan_pts1, nan_pts2), dim=1) if pts2 is not None else nan_pts1 + + # compute distance to origin + all_dis = all_pts.norm(dim=-1) + + if norm_mode == 'avg': + norm_factor = all_dis.nanmean(dim=1) + elif norm_mode == 'median': + norm_factor = all_dis.nanmedian(dim=1).values.detach() + elif norm_mode == 'sqrt': + norm_factor = all_dis.sqrt().nanmean(dim=1)**2 + else: + raise ValueError(f'bad {norm_mode=}') + + norm_factor = norm_factor.clip(min=1e-8) + while norm_factor.ndim < pts1.ndim: + norm_factor.unsqueeze_(-1) + + res = pts1 / norm_factor + if pts2 is not None: + res = (res, pts2 / norm_factor) + if ret_factor: + res = res + (norm_factor,) + return res + + +@torch.no_grad() +def get_joint_pointcloud_depth(z1, z2, valid_mask1, valid_mask2=None, quantile=0.5): + # set invalid points to NaN + _z1 = invalid_to_nans(z1, valid_mask1).reshape(len(z1), -1) + _z2 = invalid_to_nans(z2, valid_mask2).reshape(len(z2), -1) if z2 is not None else None + _z = torch.cat((_z1, _z2), dim=-1) if z2 is not None else _z1 + + # compute median depth overall (ignoring nans) + if quantile == 0.5: + shift_z = torch.nanmedian(_z, dim=-1).values + else: + shift_z = torch.nanquantile(_z, quantile, dim=-1) + return shift_z # (B,) + + +@torch.no_grad() +def get_joint_pointcloud_center_scale(pts1, pts2, valid_mask1=None, valid_mask2=None, z_only=False, center=True): + # set invalid points to NaN + _pts1 = invalid_to_nans(pts1, valid_mask1).reshape(len(pts1), -1, 3) + _pts2 = invalid_to_nans(pts2, valid_mask2).reshape(len(pts2), -1, 3) if pts2 is not None else None + _pts = torch.cat((_pts1, _pts2), dim=1) if pts2 is not None else _pts1 + + # compute median center + _center = torch.nanmedian(_pts, dim=1, keepdim=True).values # (B,1,3) + if z_only: + _center[..., :2] = 0 # do not center X and Y + + # compute median norm + _norm = ((_pts - _center) if center else _pts).norm(dim=-1) + scale = torch.nanmedian(_norm, dim=1).values + return _center[:, None, :, :], scale[:, None, None, None] + + +def find_reciprocal_matches(P1, P2): + """ + returns 3 values: + 1 - reciprocal_in_P2: a boolean array of size P2.shape[0], a "True" value indicates a match + 2 - nn2_in_P1: a int array of size P2.shape[0], it contains the indexes of the closest points in P1 + 3 - reciprocal_in_P2.sum(): the number of matches + """ + tree1 = KDTree(P1) + tree2 = KDTree(P2) + + _, nn1_in_P2 = tree2.query(P1, workers=8) + _, nn2_in_P1 = tree1.query(P2, workers=8) + + reciprocal_in_P1 = (nn2_in_P1[nn1_in_P2] == np.arange(len(nn1_in_P2))) + reciprocal_in_P2 = (nn1_in_P2[nn2_in_P1] == np.arange(len(nn2_in_P1))) + assert reciprocal_in_P1.sum() == reciprocal_in_P2.sum() + return reciprocal_in_P2, nn2_in_P1, reciprocal_in_P2.sum() + + +def get_med_dist_between_poses(poses): + from scipy.spatial.distance import pdist + return np.median(pdist([p[:3, 3].detach().cpu().numpy() for p in poses])) \ No newline at end of file diff --git a/src/dataset/shims/load_shim.py b/src/dataset/shims/load_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..d94f272f06adddc5be565e17e0275770bb2023df --- /dev/null +++ b/src/dataset/shims/load_shim.py @@ -0,0 +1,12 @@ +import cv2 + +def imread_cv2(path, options=cv2.IMREAD_COLOR): + """Open an image or a depthmap with opencv-python.""" + if path.endswith((".exr", "EXR")): + options = cv2.IMREAD_ANYDEPTH + img = cv2.imread(path, options) + if img is None: + raise IOError(f"Could not load image={path} with {options=}") + if img.ndim == 3: + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + return img \ No newline at end of file diff --git a/src/dataset/shims/normalize_shim.py b/src/dataset/shims/normalize_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..5650a51ccc0cd90b8db0f2a3c45eeb5c22b5c52c --- /dev/null +++ b/src/dataset/shims/normalize_shim.py @@ -0,0 +1,27 @@ +import torch +from einops import einsum, reduce, repeat +from jaxtyping import Float +from torch import Tensor + +from ..types import BatchedExample + + +def inverse_normalize_image(tensor, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)): + mean = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1) + std = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1) + return tensor * std + mean + + +def normalize_image(tensor, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)): + mean = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1) + std = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1) + return (tensor - mean) / std + + +def apply_normalize_shim( + batch: BatchedExample, + mean: tuple[float, float, float] = (0.5, 0.5, 0.5), + std: tuple[float, float, float] = (0.5, 0.5, 0.5), +) -> BatchedExample: + batch["context"]["image"] = normalize_image(batch["context"]["image"], mean, std) + return batch diff --git a/src/dataset/shims/patch_shim.py b/src/dataset/shims/patch_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..60dc944995332726b060e7eac9bf914b63c3c840 --- /dev/null +++ b/src/dataset/shims/patch_shim.py @@ -0,0 +1,38 @@ +from ..types import BatchedExample, BatchedViews + + +def apply_patch_shim_to_views(views: BatchedViews, patch_size: int) -> BatchedViews: + _, _, _, h, w = views["image"].shape + + # Image size must be even so that naive center-cropping does not cause misalignment. + assert h % 2 == 0 and w % 2 == 0 + + h_new = (h // patch_size) * patch_size + row = (h - h_new) // 2 + w_new = (w // patch_size) * patch_size + col = (w - w_new) // 2 + + # Center-crop the image. + image = views["image"][:, :, :, row : row + h_new, col : col + w_new] + + # Adjust the intrinsics to account for the cropping. + intrinsics = views["intrinsics"].clone() + intrinsics[:, :, 0, 0] *= w / w_new # fx + intrinsics[:, :, 1, 1] *= h / h_new # fy + + return { + **views, + "image": image, + "intrinsics": intrinsics, + } + + +def apply_patch_shim(batch: BatchedExample, patch_size: int) -> BatchedExample: + """Crop images in the batch so that their dimensions are cleanly divisible by the + specified patch size. + """ + return { + **batch, + "context": apply_patch_shim_to_views(batch["context"], patch_size), + "target": apply_patch_shim_to_views(batch["target"], patch_size), + } diff --git a/src/dataset/types.py b/src/dataset/types.py new file mode 100644 index 0000000000000000000000000000000000000000..76736c58b2daed4ca317a1031b6767cda6eed4f0 --- /dev/null +++ b/src/dataset/types.py @@ -0,0 +1,51 @@ +from typing import Callable, Literal, TypedDict + +from jaxtyping import Float, Int64 +from torch import Tensor + +Stage = Literal["train", "val", "test"] + + +# The following types mainly exist to make type-hinted keys show up in VS Code. Some +# dimensions are annotated as "_" because either: +# 1. They're expected to change as part of a function call (e.g., resizing the dataset). +# 2. They're expected to vary within the same function call (e.g., the number of views, +# which differs between context and target BatchedViews). + + +class BatchedViews(TypedDict, total=False): + extrinsics: Float[Tensor, "batch _ 4 4"] # batch view 4 4 + intrinsics: Float[Tensor, "batch _ 3 3"] # batch view 3 3 + image: Float[Tensor, "batch _ _ _ _"] # batch view channel height width + near: Float[Tensor, "batch _"] # batch view + far: Float[Tensor, "batch _"] # batch view + index: Int64[Tensor, "batch _"] # batch view + overlap: Float[Tensor, "batch _"] # batch view + + +class BatchedExample(TypedDict, total=False): + target: BatchedViews + context: BatchedViews + scene: list[str] + + +class UnbatchedViews(TypedDict, total=False): + extrinsics: Float[Tensor, "_ 4 4"] + intrinsics: Float[Tensor, "_ 3 3"] + image: Float[Tensor, "_ 3 height width"] + near: Float[Tensor, " _"] + far: Float[Tensor, " _"] + index: Int64[Tensor, " _"] + + +class UnbatchedExample(TypedDict, total=False): + target: UnbatchedViews + context: UnbatchedViews + scene: str + + +# A data shim modifies the example after it's been returned from the data loader. +DataShim = Callable[[BatchedExample], BatchedExample] + +AnyExample = BatchedExample | UnbatchedExample +AnyViews = BatchedViews | UnbatchedViews diff --git a/src/dataset/validation_wrapper.py b/src/dataset/validation_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..2c3fef9ffdfa52c1fa1e75c52aaecd79bf147e43 --- /dev/null +++ b/src/dataset/validation_wrapper.py @@ -0,0 +1,34 @@ +from typing import Iterator, Optional + +import torch +from torch.utils.data import Dataset, IterableDataset + + +class ValidationWrapper(Dataset): + """Wraps a dataset so that PyTorch Lightning's validation step can be turned into a + visualization step. + """ + + dataset: Dataset + dataset_iterator: Optional[Iterator] + length: int + + def __init__(self, dataset: Dataset, length: int) -> None: + super().__init__() + self.dataset = dataset + self.length = length + self.dataset_iterator = None + + def __len__(self): + return self.length + + def __getitem__(self, index: tuple): + if isinstance(self.dataset, IterableDataset): + if self.dataset_iterator is None: + self.dataset_iterator = iter(self.dataset) + return next(self.dataset_iterator) + + random_index = torch.randint(0, len(self.dataset), tuple()) + random_context_num = torch.randint(2, self.dataset.view_sampler.num_context_views + 1, tuple()) + # breakpoint() + return self.dataset[random_index.item(), random_context_num.item()] diff --git a/src/dataset/view_sampler/__init__.py b/src/dataset/view_sampler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1c09a072e743d1f06e73770295d12d7ada57702b --- /dev/null +++ b/src/dataset/view_sampler/__init__.py @@ -0,0 +1,40 @@ +from typing import Any + +from ...misc.step_tracker import StepTracker +from ..types import Stage +from .view_sampler import ViewSampler +from .view_sampler_all import ViewSamplerAll, ViewSamplerAllCfg +from .view_sampler_arbitrary import ViewSamplerArbitrary, ViewSamplerArbitraryCfg +from .view_sampler_bounded import ViewSamplerBounded, ViewSamplerBoundedCfg +from .view_sampler_evaluation import ViewSamplerEvaluation, ViewSamplerEvaluationCfg +from .view_sampler_rank import ViewSamplerRank, ViewSamplerRankCfg +VIEW_SAMPLERS: dict[str, ViewSampler[Any]] = { + "all": ViewSamplerAll, + "arbitrary": ViewSamplerArbitrary, + "bounded": ViewSamplerBounded, + "evaluation": ViewSamplerEvaluation, + "rank": ViewSamplerRank, +} + +ViewSamplerCfg = ( + ViewSamplerArbitraryCfg + | ViewSamplerBoundedCfg + | ViewSamplerEvaluationCfg + | ViewSamplerAllCfg + | ViewSamplerRankCfg +) + +def get_view_sampler( + cfg: ViewSamplerCfg, + stage: Stage, + overfit: bool, + cameras_are_circular: bool, + step_tracker: StepTracker | None, +) -> ViewSampler[Any]: + return VIEW_SAMPLERS[cfg.name]( + cfg, + stage, + overfit, + cameras_are_circular, + step_tracker, + ) diff --git a/src/dataset/view_sampler/three_view_hack.py b/src/dataset/view_sampler/three_view_hack.py new file mode 100644 index 0000000000000000000000000000000000000000..3a47cb624a1472651bc6611de2c69b04fdc085ad --- /dev/null +++ b/src/dataset/view_sampler/three_view_hack.py @@ -0,0 +1,10 @@ +import torch +from jaxtyping import Int +from torch import Tensor + + +def add_third_context_index( + indices: Int[Tensor, "*batch 2"] +) -> Int[Tensor, "*batch 3"]: + left, right = indices.unbind(dim=-1) + return torch.stack((left, (left + right) // 2, right), dim=-1) diff --git a/src/dataset/view_sampler/view_sampler.py b/src/dataset/view_sampler/view_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..45a6276945363ad7a8691295edcb2d5ba4dd0134 --- /dev/null +++ b/src/dataset/view_sampler/view_sampler.py @@ -0,0 +1,61 @@ +from abc import ABC, abstractmethod +from typing import Generic, TypeVar + +import torch +from jaxtyping import Float, Int64 +from torch import Tensor + +from ...misc.step_tracker import StepTracker +from ..types import Stage + +T = TypeVar("T") + + +class ViewSampler(ABC, Generic[T]): + cfg: T + stage: Stage + is_overfitting: bool + cameras_are_circular: bool + step_tracker: StepTracker | None + + def __init__( + self, + cfg: T, + stage: Stage, + is_overfitting: bool, + cameras_are_circular: bool, + step_tracker: StepTracker | None, + ) -> None: + self.cfg = cfg + self.stage = stage + self.is_overfitting = is_overfitting + self.cameras_are_circular = cameras_are_circular + self.step_tracker = step_tracker + + @abstractmethod + def sample( + self, + scene: str, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + Float[Tensor, " overlap"], # overlap + ]: + pass + + @property + @abstractmethod + def num_target_views(self) -> int: + pass + + @property + @abstractmethod + def num_context_views(self) -> int: + pass + + @property + def global_step(self) -> int: + return 0 if self.step_tracker is None else self.step_tracker.get_step() diff --git a/src/dataset/view_sampler/view_sampler_all.py b/src/dataset/view_sampler/view_sampler_all.py new file mode 100644 index 0000000000000000000000000000000000000000..91ac3dcc5c9f8cfdec7fc495f3162008cf5942f2 --- /dev/null +++ b/src/dataset/view_sampler/view_sampler_all.py @@ -0,0 +1,37 @@ +from dataclasses import dataclass +from typing import Literal + +import torch +from jaxtyping import Float, Int64 +from torch import Tensor + +from .view_sampler import ViewSampler + + +@dataclass +class ViewSamplerAllCfg: + name: Literal["all"] + + +class ViewSamplerAll(ViewSampler[ViewSamplerAllCfg]): + def sample( + self, + scene: str, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + ]: + v, _, _ = extrinsics.shape + all_frames = torch.arange(v, device=device) + return all_frames, all_frames + + @property + def num_context_views(self) -> int: + return 0 + + @property + def num_target_views(self) -> int: + return 0 diff --git a/src/dataset/view_sampler/view_sampler_arbitrary.py b/src/dataset/view_sampler/view_sampler_arbitrary.py new file mode 100644 index 0000000000000000000000000000000000000000..44e68d051bf48629b0b51df187844a9f91f7d493 --- /dev/null +++ b/src/dataset/view_sampler/view_sampler_arbitrary.py @@ -0,0 +1,77 @@ +from dataclasses import dataclass +from typing import Literal + +import torch +from jaxtyping import Float, Int64 +from torch import Tensor + +from .three_view_hack import add_third_context_index +from .view_sampler import ViewSampler + + +@dataclass +class ViewSamplerArbitraryCfg: + name: Literal["arbitrary"] + num_context_views: int + num_target_views: int + context_views: list[int] | None + target_views: list[int] | None + + +class ViewSamplerArbitrary(ViewSampler[ViewSamplerArbitraryCfg]): + def sample( + self, + scene: str, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + Float[Tensor, " overlap"], # overlap + ]: + """Arbitrarily sample context and target views.""" + num_views, _, _ = extrinsics.shape + + index_context = torch.randint( + 0, + num_views, + size=(self.cfg.num_context_views,), + device=device, + ) + + # Allow the context views to be fixed. + if self.cfg.context_views is not None: + index_context = torch.tensor( + self.cfg.context_views, dtype=torch.int64, device=device + ) + + if self.cfg.num_context_views == 3 and len(self.cfg.context_views) == 2: + index_context = add_third_context_index(index_context) + else: + assert len(self.cfg.context_views) == self.cfg.num_context_views + index_target = torch.randint( + 0, + num_views, + size=(self.cfg.num_target_views,), + device=device, + ) + + # Allow the target views to be fixed. + if self.cfg.target_views is not None: + assert len(self.cfg.target_views) == self.cfg.num_target_views + index_target = torch.tensor( + self.cfg.target_views, dtype=torch.int64, device=device + ) + + overlap = torch.tensor([0.5], dtype=torch.float32, device=device) # dummy + + return index_context, index_target, overlap + + @property + def num_context_views(self) -> int: + return self.cfg.num_context_views + + @property + def num_target_views(self) -> int: + return self.cfg.num_target_views diff --git a/src/dataset/view_sampler/view_sampler_bounded.py b/src/dataset/view_sampler/view_sampler_bounded.py new file mode 100644 index 0000000000000000000000000000000000000000..ca4b9cef21dfc9da95ba6b1c3e600a3f51392d80 --- /dev/null +++ b/src/dataset/view_sampler/view_sampler_bounded.py @@ -0,0 +1,151 @@ +from dataclasses import dataclass +from typing import Literal + +import torch +from jaxtyping import Float, Int64 +from torch import Tensor + +from .view_sampler import ViewSampler + + +@dataclass +class ViewSamplerBoundedCfg: + name: Literal["bounded"] + num_context_views: int + num_target_views: int + min_distance_between_context_views: int + max_distance_between_context_views: int + min_distance_to_context_views: int + warm_up_steps: int + initial_min_distance_between_context_views: int + initial_max_distance_between_context_views: int + max_img_per_gpu: int + min_gap_multiplier: int + max_gap_multiplier: int + +class ViewSamplerBounded(ViewSampler[ViewSamplerBoundedCfg]): + def schedule(self, initial: int, final: int) -> int: + fraction = self.global_step / self.cfg.warm_up_steps + return min(initial + int((final - initial) * fraction), final) + + def sample( + self, + scene: str, + num_context_views: int, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + Float[Tensor, " overlap"], # overlap + ]: + num_views, _, _ = extrinsics.shape + + # Compute the context view spacing based on the current global step. + if self.stage == "test": + # When testing, always use the full gap. + max_gap = self.cfg.max_distance_between_context_views + min_gap = self.cfg.max_distance_between_context_views + # elif self.cfg.warm_up_steps > 0: + # max_gap = self.schedule( + # self.cfg.initial_max_distance_between_context_views, + # self.cfg.max_distance_between_context_views, + # ) + # min_gap = self.schedule( + # self.cfg.initial_min_distance_between_context_views, + # self.cfg.min_distance_between_context_views, + # ) + # else: + # max_gap = self.cfg.max_distance_between_context_views + # min_gap = self.cfg.min_distance_between_context_views + + min_gap, max_gap = self.num_ctxt_gap_mapping[num_context_views] + max_gap = min(max_gap, num_views-1) + # Pick the gap between the context views. + if not self.cameras_are_circular: + max_gap = min(num_views - 1, max_gap) + min_gap = max(2 * self.cfg.min_distance_to_context_views, min_gap) + if max_gap < min_gap: + raise ValueError("Example does not have enough frames!") + context_gap = torch.randint( + min_gap, + max_gap + 1, + size=tuple(), + device=device, + ).item() + + # Pick the left and right context indices. + index_context_left = torch.randint( + num_views if self.cameras_are_circular else num_views - context_gap, + size=tuple(), + device=device, + ).item() + if self.stage == "test": + index_context_left = index_context_left * 0 + index_context_right = index_context_left + context_gap + + if self.is_overfitting: + index_context_left *= 0 + index_context_right *= 0 + index_context_right += max_gap + + # Pick the target view indices. + if self.stage == "test": + # When testing, pick all. + index_target = torch.arange( + index_context_left, + index_context_right + 1, + device=device, + ) + else: + # When training or validating (visualizing), pick at random. + index_target = torch.randint( + index_context_left + self.cfg.min_distance_to_context_views, + index_context_right + 1 - self.cfg.min_distance_to_context_views, + size=(self.cfg.num_target_views,), + device=device, + ) + + # Apply modulo for circular datasets. + if self.cameras_are_circular: + index_target %= num_views + index_context_right %= num_views + + # If more than two context views are desired, pick extra context views between + # the left and right ones. + if num_context_views > 2: + num_extra_views = num_context_views - 2 + extra_views = [] + while len(set(extra_views)) != num_extra_views: + extra_views = torch.randint( + index_context_left + 1, + index_context_right, + (num_extra_views,), + ).tolist() + else: + extra_views = [] + + overlap = torch.tensor([0.5], dtype=torch.float32, device=device) # dummy + + return ( + torch.tensor((index_context_left, *extra_views, index_context_right)), + index_target, + overlap + ) + + @property + def num_context_views(self) -> int: + return self.cfg.num_context_views + + @property + def num_target_views(self) -> int: + return self.cfg.num_target_views + + @property + def num_ctxt_gap_mapping(self) -> dict: + mapping = dict() + for num_ctxt in range(2, self.cfg.num_context_views + 1): + mapping[num_ctxt] = [min(num_ctxt * self.cfg.min_gap_multiplier, self.cfg.min_distance_between_context_views), + min(max(num_ctxt * self.cfg.max_gap_multiplier, num_ctxt ** 2), self.cfg.max_distance_between_context_views)] + return mapping diff --git a/src/dataset/view_sampler/view_sampler_evaluation.py b/src/dataset/view_sampler/view_sampler_evaluation.py new file mode 100644 index 0000000000000000000000000000000000000000..ad0b91e5c5527db2432dd11eca33b52a183d6213 --- /dev/null +++ b/src/dataset/view_sampler/view_sampler_evaluation.py @@ -0,0 +1,70 @@ +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import torch +from dacite import Config, from_dict +from jaxtyping import Float, Int64 +from torch import Tensor + +from ...evaluation.evaluation_index_generator import IndexEntry +from ...global_cfg import get_cfg +from ...misc.step_tracker import StepTracker +from ..types import Stage +from .three_view_hack import add_third_context_index +from .view_sampler import ViewSampler + + +@dataclass +class ViewSamplerEvaluationCfg: + name: Literal["evaluation"] + index_path: Path + num_context_views: int + + +class ViewSamplerEvaluation(ViewSampler[ViewSamplerEvaluationCfg]): + index: dict[str, IndexEntry | None] + + def __init__( + self, + cfg: ViewSamplerEvaluationCfg, + stage: Stage, + is_overfitting: bool, + cameras_are_circular: bool, + step_tracker: StepTracker | None, + ) -> None: + super().__init__(cfg, stage, is_overfitting, cameras_are_circular, step_tracker) + + dacite_config = Config(cast=[tuple]) + with cfg.index_path.open("r") as f: + self.index = { + k: None if v is None else from_dict(IndexEntry, v, dacite_config) + for k, v in json.load(f).items() + } + + def sample( + self, + scene: str, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + **kwargs, + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + ]: + entry = self.index.get(scene) + if entry is None: + raise ValueError(f"No indices available for scene {scene}.") + context_indices = torch.tensor(entry.context, dtype=torch.int64, device=device) + target_indices = torch.tensor(entry.target, dtype=torch.int64, device=device) + return context_indices, target_indices, torch.zeros(1) + + @property + def num_context_views(self) -> int: + return 0 + + @property + def num_target_views(self) -> int: + return 0 diff --git a/src/dataset/view_sampler/view_sampler_rank.py b/src/dataset/view_sampler/view_sampler_rank.py new file mode 100644 index 0000000000000000000000000000000000000000..6be507f07620c32eba589ff26a779c2456348827 --- /dev/null +++ b/src/dataset/view_sampler/view_sampler_rank.py @@ -0,0 +1,263 @@ + +from dataclasses import dataclass +from typing import Literal + +import torch +import copy +from jaxtyping import Float, Int64 +from torch import Tensor +import random +from .view_sampler import ViewSampler + + +@dataclass +class ViewSamplerRankCfg: + name: Literal["rank"] + num_context_views: int # max number of context views + num_target_views: int + min_distance_between_context_views: int + max_distance_between_context_views: int + min_distance_to_context_views: int + warm_up_steps: int + initial_min_distance_between_context_views: int + initial_max_distance_between_context_views: int + max_img_per_gpu: int + + +def rotation_angle(R1, R2): + # R1 and R2 are 3x3 rotation matrices + R = R1.T @ R2 + # Numerical stability: clamp values into [-1,1] + val = (torch.trace(R) - 1) / 2 + val = torch.clamp(val, -1.0, 1.0) + angle_rad = torch.acos(val) + angle_deg = angle_rad * 180 / torch.pi # Convert radians to degrees + return angle_deg + +def extrinsic_distance(extrinsic1, extrinsic2, lambda_t=1.0): + R1, t1 = extrinsic1[:3, :3], extrinsic1[:3, 3] + R2, t2 = extrinsic2[:3, :3], extrinsic2[:3, 3] + rot_diff = rotation_angle(R1, R2) / 180 + + center_diff = torch.norm(t1 - t2) + return rot_diff + lambda_t * center_diff + +def rotation_angle_batch(R1, R2): + # R1, R2: shape (N, 3, 3) + # We want a matrix of rotation angles for all pairs. + # We'll get R1^T R2 for each pair. + # Expand dimensions to broadcast: + # R1^T: (N,3,3) -> (N,1,3,3) + # R2: (N,3,3) -> (1,N,3,3) + R1_t = R1.transpose(-2, -1)[:, None, :, :] # shape (N,1,3,3) + R2_b = R2[None, :, :, :] # shape (1,N,3,3) + R_mult = torch.matmul(R1_t, R2_b) # shape (N,N,3,3) + # trace(R) for each pair + trace_vals = R_mult[..., 0, 0] + R_mult[..., 1, 1] + R_mult[..., 2, 2] # (N,N) + val = (trace_vals - 1) / 2 + val = torch.clamp(val, -1.0, 1.0) + angle_rad = torch.acos(val) + angle_deg = angle_rad * 180 / torch.pi + return angle_deg / 180.0 # normalized rotation difference + +def extrinsic_distance_batch(extrinsics, lambda_t=1.0): + # extrinsics: (N,4,4) + # Extract rotation and translation + R = extrinsics[:, :3, :3] # (N,3,3) + t = extrinsics[:, :3, 3] # (N,3) + # Compute all pairwise rotation differences + rot_diff = rotation_angle_batch(R, R) # (N,N) + # Compute all pairwise translation differences + # For t, shape (N,3). We want all pair differences: t[i] - t[j]. + # t_i: (N,1,3), t_j: (1,N,3) + t_i = t[:, None, :] # (N,1,3) + t_j = t[None, :, :] # (1,N,3) + trans_diff = torch.norm(t_i - t_j, dim=2) # (N,N) + dists = rot_diff + lambda_t * trans_diff + return dists + + +def compute_ranking(extrinsics, lambda_t=1.0, normalize=True, batched=True): + + if normalize: + extrinsics = copy.deepcopy(extrinsics) + camera_center = copy.deepcopy(extrinsics[:, :3, 3]) + camera_center_scale = torch.norm(camera_center, dim=1) + avg_scale = torch.mean(camera_center_scale) + extrinsics[:, :3, 3] = extrinsics[:, :3, 3] / avg_scale + + + if batched: + dists = extrinsic_distance_batch(extrinsics, lambda_t=lambda_t) + else: + N = extrinsics.shape[0] + dists = torch.zeros((N, N), device=extrinsics.device) + for i in range(N): + for j in range(N): + dists[i,j] = extrinsic_distance(extrinsics[i], extrinsics[j], lambda_t=lambda_t) + ranking = torch.argsort(dists, dim=1) + return ranking, dists + +# class ViewSamplerRank(ViewSampler[ViewSamplerRankCfg]): +# def schedule(self, initial: int, final: int) -> int: +# fraction = self.global_step / self.cfg.warm_up_steps +# return min(initial + int((final - initial) * fraction), final) + +# def sample( +# self, +# scene: str, +# extrinsics: Float[Tensor, "view 4 4"], +# intrinsics: Float[Tensor, "view 3 3"], +# device: torch.device = torch.device("cpu"), +# ) -> tuple[ +# Int64[Tensor, " context_view"], # indices for context views +# Int64[Tensor, " target_view"], # indices for target views +# Float[Tensor, " overlap"], # overlap +# ]: +# num_views, _, _ = extrinsics.shape +# # breakpoint() +# # Compute the context view spacing based on the current global step. +# ranking, dists = compute_ranking(extrinsics, lambda_t=1.0, normalize=True, batched=True) +# reference_view = random.sample(range(num_views), 1)[0] + +# refview_ranking = ranking[reference_view] +# # if self.cfg.warm_up_steps > 0: +# # max_gap = self.schedule( +# # self.cfg.initial_max_distance_between_context_views, +# # self.cfg.max_distance_between_context_views, +# # ) +# # min_gap = self.schedule( +# # self.cfg.initial_min_distance_between_context_views, +# # self.cfg.min_distance_between_context_views, +# # ) +# # else: +# max_gap = self.cfg.max_distance_between_context_views +# min_gap = self.cfg.min_distance_between_context_views + +# index_context_left = reference_view +# rightmost_index = random.sample(range(min_gap, max_gap + 1), 1)[0] + 1 +# index_context_right = refview_ranking[rightmost_index].item() + +# middle_indices = refview_ranking[1: rightmost_index].tolist() +# index_target = random.sample(middle_indices, self.num_target_views) + +# remaining_indices = [idx for idx in middle_indices if idx not in index_target] + +# # Sample extra context views if needed +# extra_views = [] +# num_extra_views = self.num_context_views - 2 # subtract left and right context views +# if num_extra_views > 0 and remaining_indices: +# extra_views = random.sample(remaining_indices, min(num_extra_views, len(remaining_indices))) +# else: +# extra_views = [] + +# overlap = torch.zeros(1) + +# return ( +# torch.tensor((index_context_left, *extra_views, index_context_right)), +# torch.tensor(index_target), +# overlap +# ) + + +# @property +# def num_context_views(self) -> int: +# return self.cfg.num_context_views + +# @property +# def num_target_views(self) -> int: +# return self.cfg.num_target_views + + +class ViewSamplerRank(ViewSampler[ViewSamplerRankCfg]): + + def sample( + self, + scene: str, + num_context_views: int, + extrinsics: Float[Tensor, "view 4 4"], + intrinsics: Float[Tensor, "view 3 3"], + device: torch.device = torch.device("cpu"), + ) -> tuple[ + Int64[Tensor, " context_view"], # indices for context views + Int64[Tensor, " target_view"], # indices for target views + Float[Tensor, " overlap"], # overlap + ]: + num_views, _, _ = extrinsics.shape + # breakpoint() + extrinsics = extrinsics.clone() + # Compute the context view spacing based on the current global step. + ranking, dists = compute_ranking(extrinsics, lambda_t=1.0, normalize=True, batched=True) + reference_view = random.sample(range(num_views), 1)[0] + + refview_ranking = ranking[reference_view] + # if self.cfg.warm_up_steps > 0: + # max_gap = self.schedule( + # self.cfg.initial_max_distance_between_context_views, + # self.cfg.max_distance_between_context_views, + # ) + # min_gap = self.schedule( + # self.cfg.initial_min_distance_between_context_views, + # self.cfg.min_distance_between_context_views, + # ) + # else: + min_gap, max_gap = self.num_ctxt_gap_mapping[num_context_views] + + # min_gap = self.cfg.min_distance_between_context_views + # max_gap = self.cfg.max_distance_between_context_views + + max_gap = min(max_gap, num_views-1) + # print(f"num_context_views: {num_context_views}, min_gap: {min_gap}, max_gap: {max_gap}") + index_context_left = reference_view + rightmost_index = random.sample(range(min_gap, max_gap + 1), 1)[0] + + # #! hard code for visualization + # rightmost_index = self.cfg.max_distance_between_context_views + + index_context_right = refview_ranking[rightmost_index].item() + + middle_indices = refview_ranking[1: rightmost_index].tolist() + index_target = random.sample(middle_indices, self.num_target_views) + + remaining_indices = [idx for idx in middle_indices if idx not in index_target] + + # Sample extra context views if needed + extra_views = [] + num_extra_views = num_context_views - 2 # subtract left and right context views + if num_extra_views > 0 and remaining_indices: + extra_views = random.sample(remaining_indices, min(num_extra_views, len(remaining_indices))) + else: + extra_views = [] + + overlap = torch.zeros(1) + + return ( + torch.tensor((index_context_left, *extra_views, index_context_right)), + torch.tensor(index_target), + overlap + ) + + + @property + def num_context_views(self) -> int: + return self.cfg.num_context_views + + @property + def num_target_views(self) -> int: + return self.cfg.num_target_views + + @property + def num_ctxt_gap_mapping_target(self) -> dict: + mapping = dict() + for num_ctxt in range(2, self.cfg.num_context_views + 1): + mapping[num_ctxt] = [max(num_ctxt * 2, self.cfg.num_target_views + num_ctxt), max(self.cfg.num_target_views + num_ctxt, min(num_ctxt ** 2, self.cfg.max_distance_between_context_views))] + return mapping + + @property + def num_ctxt_gap_mapping(self) -> dict: + mapping = dict() + for num_ctxt in range(2, self.cfg.num_context_views + 1): + mapping[num_ctxt] = [min(num_ctxt * 3, self.cfg.min_distance_between_context_views), min(max(num_ctxt * 5, num_ctxt ** 2), self.cfg.max_distance_between_context_views)] + return mapping + + diff --git a/src/eval_nvs.py b/src/eval_nvs.py new file mode 100644 index 0000000000000000000000000000000000000000..06725ab45992ec5c7668c1e1f5caad62de2efce0 --- /dev/null +++ b/src/eval_nvs.py @@ -0,0 +1,115 @@ +import os +from pathlib import Path +import sys +import json +import gzip +import argparse +import numpy as np +from PIL import Image + +import torch +import torch.nn as nn +import torchvision +from einops import rearrange + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.evaluation.metrics import compute_lpips, compute_psnr, compute_ssim +from misc.image_io import save_image, save_interpolated_video +from src.utils.image import process_image + +from src.model.model.anysplat import AnySplat +from src.model.encoder.vggt.utils.pose_enc import pose_encoding_to_extri_intri + +def setup_args(): + """Set up command-line arguments for the eval NVS script.""" + parser = argparse.ArgumentParser(description='Test AnySplat on NVS evaluation') + parser.add_argument('--data_dir', type=str, required=True, help='Path to NVS dataset') + parser.add_argument('--llffhold', type=int, default=8, help='LLFF holdout') + parser.add_argument('--output_path', type=str, default="outputs/nvs", help='Path to output directory') + return parser.parse_args() + +def compute_metrics(pred_image, image): + psnr = compute_psnr(pred_image, image) + ssim = compute_ssim(pred_image, image) + lpips = compute_lpips(pred_image, image) + return psnr, ssim, lpips + +def evaluate(args: argparse.Namespace): + model = AnySplat.from_pretrained("lhjiang/anysplat") + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model.to(device) + model.eval() + for param in model.parameters(): + param.requires_grad = False + + os.makedirs(args.output_path, exist_ok=True) + + # load images + image_folder = args.data_dir + image_names = sorted([os.path.join(image_folder, f) for f in os.listdir(image_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]) + images = [process_image(img_path) for img_path in image_names] + ctx_indices = [idx for idx, name in enumerate(image_names) if idx % args.llffhold != 0] + tgt_indices = [idx for idx, name in enumerate(image_names) if idx % args.llffhold == 0] + + ctx_images = torch.stack([images[i] for i in ctx_indices], dim=0).unsqueeze(0).to(device) + tgt_images = torch.stack([images[i] for i in tgt_indices], dim=0).unsqueeze(0).to(device) + ctx_images = (ctx_images+1)*0.5 + tgt_images = (tgt_images+1)*0.5 + b, v, _, h, w = tgt_images.shape + + # run inference + encoder_output = model.encoder( + ctx_images, + global_step=0, + visualization_dump={}, + ) + gaussians, pred_context_pose = encoder_output.gaussians, encoder_output.pred_context_pose + + num_context_view = ctx_images.shape[1] + vggt_input_image = torch.cat((ctx_images, tgt_images), dim=1).to(torch.bfloat16) + with torch.no_grad(), torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16): + aggregated_tokens_list, patch_start_idx = model.encoder.aggregator(vggt_input_image, intermediate_layer_idx=model.encoder.cfg.intermediate_layer_idx) + with torch.cuda.amp.autocast(enabled=False): + fp32_tokens = [token.float() for token in aggregated_tokens_list] + pred_all_pose_enc = model.encoder.camera_head(fp32_tokens)[-1] + pred_all_extrinsic, pred_all_intrinsic = pose_encoding_to_extri_intri(pred_all_pose_enc, vggt_input_image.shape[-2:]) + + extrinsic_padding = torch.tensor([0, 0, 0, 1], device=pred_all_extrinsic.device, dtype=pred_all_extrinsic.dtype).view(1, 1, 1, 4).repeat(b, vggt_input_image.shape[1], 1, 1) + pred_all_extrinsic = torch.cat([pred_all_extrinsic, extrinsic_padding], dim=2).inverse() + + pred_all_intrinsic[:, :, 0] = pred_all_intrinsic[:, :, 0] / w + pred_all_intrinsic[:, :, 1] = pred_all_intrinsic[:, :, 1] / h + pred_all_context_extrinsic, pred_all_target_extrinsic = pred_all_extrinsic[:, :num_context_view], pred_all_extrinsic[:, num_context_view:] + pred_all_context_intrinsic, pred_all_target_intrinsic = pred_all_intrinsic[:, :num_context_view], pred_all_intrinsic[:, num_context_view:] + + scale_factor = pred_context_pose['extrinsic'][:, :, :3, 3].mean() / pred_all_context_extrinsic[:, :, :3, 3].mean() + pred_all_target_extrinsic[..., :3, 3] = pred_all_target_extrinsic[..., :3, 3] * scale_factor + pred_all_context_extrinsic[..., :3, 3] = pred_all_context_extrinsic[..., :3, 3] * scale_factor + print("scale_factor:", scale_factor) + + output = model.decoder.forward( + gaussians, + pred_all_target_extrinsic, + pred_all_target_intrinsic.float(), + torch.ones(1, v, device=device) * 0.01, + torch.ones(1, v, device=device) * 100, + (h, w) + ) + + save_interpolated_video(pred_all_context_extrinsic, pred_all_context_intrinsic, b, h, w, gaussians, args.output_path, model.decoder) + + # Save original images + save_path = Path(args.output_path) + # os.makedirs(save_path, exist_ok=True) + for idx, (gt_image, pred_image) in enumerate(zip(tgt_images[0], output.color[0])): + save_image(gt_image, save_path / "gt" / f"{idx:0>6}.jpg") + save_image(pred_image, save_path / "pred" / f"{idx:0>6}.jpg") + + # compute metrics + psnr, ssim, lpips = compute_metrics(output.color[0], tgt_images[0]) + print(f"PSNR: {psnr.mean():.2f}, SSIM: {ssim.mean():.3f}, LPIPS: {lpips.mean():.3f}") + +if __name__ == "__main__": + args = setup_args() + evaluate(args) diff --git a/src/eval_pose.py b/src/eval_pose.py new file mode 100644 index 0000000000000000000000000000000000000000..068626e0f9b3818065d717e07cd3051843fa9f27 --- /dev/null +++ b/src/eval_pose.py @@ -0,0 +1,339 @@ +import os +import sys +import json +import gzip +import argparse +import numpy as np +from PIL import Image + +import torch +import torch.nn as nn +import torchvision +from einops import rearrange +from lpips import LPIPS + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from src.model.model.anysplat import AnySplat +from src.model.encoder.vggt.utils.pose_enc import pose_encoding_to_extri_intri +from src.model.encoder.vggt.utils.load_fn import load_and_preprocess_images +from src.utils.pose import align_to_first_camera, calculate_auc_np, convert_pt3d_RT_to_opencv, se3_to_relative_pose_error +from src.misc.cam_utils import camera_normalization, pose_auc, rotation_6d_to_matrix, update_pose, get_pnp_pose + +def setup_args(): + """Set up command-line arguments for the CO3D evaluation script.""" + parser = argparse.ArgumentParser(description='Test AnySplat on CO3D dataset') + parser.add_argument('--debug', action='store_true', help='Enable debug mode (only test on specific category)') + parser.add_argument('--use_ba', action='store_true', default=False, help='Enable bundle adjustment') + parser.add_argument('--fast_eval', action='store_true', default=False, help='Only evaluate 10 sequences per category') + parser.add_argument('--min_num_images', type=int, default=50, help='Minimum number of images for a sequence') + parser.add_argument('--num_frames', type=int, default=10, help='Number of frames to use for testing') + parser.add_argument('--co3d_dir', type=str, required=True, help='Path to CO3D dataset') + parser.add_argument('--co3d_anno_dir', type=str, required=True, help='Path to CO3D annotations') + parser.add_argument('--seed', type=int, default=0, help='Random seed for reproducibility') + return parser.parse_args() + +lpips = LPIPS(net="vgg") + +def rendering_loss(pred_image, image): + lpips_loss = lpips.forward(rearrange(pred_image, "b v c h w -> (b v) c h w"), rearrange(image, "b v c h w -> (b v) c h w"), normalize=True) + delta = pred_image - (image + 1) / 2 + mse_loss = (delta**2).mean() + return mse_loss + 0.05 * lpips_loss.mean() + +def process_sequence(model, seq_name, seq_data, category, co3d_dir, min_num_images, num_frames, use_ba, device, dtype): + """ + Process a single sequence and compute pose errors. + + Args: + model: AnySplat model + seq_name: Sequence name + seq_data: Sequence data + category: Category name + co3d_dir: CO3D dataset directory + min_num_images: Minimum number of images required + num_frames: Number of frames to sample + use_ba: Whether to use bundle adjustment + device: Device to run on + dtype: Data type for model inference + + Returns: + rError: Rotation errors + tError: Translation errors + """ + if len(seq_data) < min_num_images: + return None, None + + metadata = [] + for data in seq_data: + # Make sure translations are not ridiculous + if data["T"][0] + data["T"][1] + data["T"][2] > 1e5: + return None, None + + extri_opencv = convert_pt3d_RT_to_opencv(data["R"], data["T"]) + metadata.append({ + "filepath": data["filepath"], + "extri": extri_opencv, + }) + + ids = np.random.choice(len(metadata), num_frames, replace=False) + image_names = [os.path.join(co3d_dir, metadata[i]["filepath"]) for i in ids] + gt_extri = [np.array(metadata[i]["extri"]) for i in ids] + gt_extri = np.stack(gt_extri, axis=0) + + max_size = max(Image.open(image_names[0]).size) + if max_size < 448: + return None, None + images = load_and_preprocess_images(image_names)[None].to(device) + + batch = { + "context": { + "image": images*2.0-1, + "image_names": image_names, + "index": ids, + }, + "scene": "co3d" + } + + if use_ba: + try: + encoder_output = model.encoder( + batch, + global_step=0, + visualization_dump={}, + ) + gaussians, pred_context_pose = encoder_output.gaussians, encoder_output.pred_context_pose + pred_extrinsic = pred_context_pose['extrinsic'] + pred_intrinsic = pred_context_pose['intrinsic'] + # rendering ba + b, v, _, h, w = images.shape + with torch.set_grad_enabled(True), torch.cuda.amp.autocast(enabled=False, dtype=torch.float32): + cam_rot_delta = nn.Parameter(torch.zeros([b, v, 6], requires_grad=True, device=pred_extrinsic.device, dtype=torch.float32)) + cam_trans_delta = nn.Parameter(torch.zeros([b, v, 3], requires_grad=True, device=pred_extrinsic.device, dtype=torch.float32)) + opt_params = [] + model.register_buffer("identity", torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0], dtype=torch.float32).to(pred_extrinsic.device)) + opt_params.append( + { + "params": [cam_rot_delta], + "lr": 0.005, + } + ) + opt_params.append( + { + "params": [cam_trans_delta], + "lr": 0.005, + } + ) + pose_optimizer = torch.optim.Adam(opt_params) + extrinsics = pred_extrinsic.clone().float() + + for i in range(100): + pose_optimizer.zero_grad() + dx, drot = cam_trans_delta, cam_rot_delta + rot = rotation_6d_to_matrix( + drot + model.identity.expand(b, v, -1) + ) # (..., 3, 3) + + transform = torch.eye(4, device=extrinsics.device).repeat((b, v, 1, 1)) + transform[..., :3, :3] = rot + transform[..., :3, 3] = dx + + new_extrinsics = torch.matmul(extrinsics, transform) + # breakpoint() + output = model.decoder.forward( + gaussians, + new_extrinsics, + pred_intrinsic.float(), + 0.1, + 100.0, + (h, w), + # cam_rot_delta=cam_rot_delta, + # cam_trans_delta=cam_trans_delta, + ) + # export_ply(gaussians.means[0], gaussians.scales[0], gaussians.rotations[0], gaussians.harmonics[0], gaussians.opacities[0], Path(f"gaussians_co3d.ply")) + rendering_loss = rendering_loss(output.color, images*2.0-1) + torchvision.utils.save_image(output.color[0], f"outputs/vis/output_co3d_{i}.png") + print(f"Rendering loss: {rendering_loss.item()}") + # print(f"Rendering loss: {rendering_loss.item()}") + + rendering_loss.backward() + pose_optimizer.step() + torchvision.utils.save_image(images[0], f"outputs/vis/gt_co3d.png") + pred_extrinsic = new_extrinsics.inverse()[0][:,:-1,:] + + except Exception as e: + print(f"BA failed with error: {e}. Falling back to standard VGGT inference.") + with torch.no_grad(), torch.cuda.amp.autocast(dtype=dtype): + aggregated_tokens_list, patch_start_idx = model.encoder.aggregator(images, intermediate_layer_idx=model.encoder.cfg.intermediate_layer_idx) + with torch.cuda.amp.autocast(dtype=torch.float32): + fp32_tokens = [token.float() for token in aggregated_tokens_list] + pred_all_pose_enc = model.encoder.camera_head(fp32_tokens)[-1] + pred_all_extrinsic, pred_all_intrinsic = pose_encoding_to_extri_intri(pred_all_pose_enc, images.shape[-2:]) + pred_extrinsic = pred_all_extrinsic[0] + else: + with torch.no_grad(), torch.cuda.amp.autocast(dtype=dtype): + aggregated_tokens_list, patch_start_idx = model.encoder.aggregator(images, intermediate_layer_idx=model.encoder.cfg.intermediate_layer_idx) + with torch.cuda.amp.autocast(dtype=torch.float32): + fp32_tokens = [token.float() for token in aggregated_tokens_list] + pred_all_pose_enc = model.encoder.camera_head(fp32_tokens)[-1] + pred_all_extrinsic, pred_all_intrinsic = pose_encoding_to_extri_intri(pred_all_pose_enc, images.shape[-2:]) + pred_extrinsic = pred_all_extrinsic[0] + + with torch.cuda.amp.autocast(dtype=torch.float32): + gt_extrinsic = torch.from_numpy(gt_extri).to(device) + add_row = torch.tensor([0, 0, 0, 1], device=device).expand(pred_extrinsic.size(0), 1, 4) + + pred_se3 = torch.cat((pred_extrinsic, add_row), dim=1) + gt_se3 = torch.cat((gt_extrinsic, add_row), dim=1) + + # Set the coordinate of the first camera as the coordinate of the world + # NOTE: DO NOT REMOVE THIS UNLESS YOU KNOW WHAT YOU ARE DOING + # pred_se3 = align_to_first_camera(pred_se3) + gt_se3 = align_to_first_camera(gt_se3) + + rel_rangle_deg, rel_tangle_deg = se3_to_relative_pose_error(pred_se3, gt_se3, num_frames) + print(f"{category} sequence {seq_name} Rot Error: {rel_rangle_deg.mean().item():.4f}") + print(f"{category} sequence {seq_name} Trans Error: {rel_tangle_deg.mean().item():.4f}") + + return rel_rangle_deg.cpu().numpy(), rel_tangle_deg.cpu().numpy() + +def evaluate(args: argparse.Namespace): + model = AnySplat.from_pretrained("lhjiang/anysplat") + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model.to(device) + model.eval() + for param in model.parameters(): + param.requires_grad = False + + # CO3D evaluation + SEEN_CATEGORIES = [ + "apple", "backpack", "banana", "baseballbat", "baseballglove", + "bench", "bicycle", "bottle", "bowl", "broccoli", + "cake", "car", "carrot", "cellphone", "chair", + "cup", "donut", "hairdryer", "handbag", "hydrant", + "keyboard", "laptop", "microwave", "motorcycle", "mouse", + "orange", "parkingmeter", "pizza", "plant", "stopsign", + "teddybear", "toaster", "toilet", "toybus", "toyplane", + "toytrain", "toytruck", "tv", "umbrella", "vase", "wineglass", + ] + + if args.debug: + SEEN_CATEGORIES = ["apple"] + + per_category_results = {} + + for category in SEEN_CATEGORIES: + print(f"Loading annotation for {category} test set") + annotation_file = os.path.join(args.co3d_anno_dir, f"{category}_test.jgz") + + try: + with gzip.open(annotation_file, "r") as fin: + annotation = json.loads(fin.read()) + except FileNotFoundError: + print(f"Annotation file not found for {category}, skipping") + continue + + rError = [] + tError = [] + + for seq_name, seq_data in annotation.items(): + print("-" * 50) + + print(f"Processing {seq_name} for {category} test set") + if args.debug and not os.path.exists(os.path.join(args.co3d_dir, category, seq_name)): + print(f"Skipping {seq_name} (not found)") + continue + + seq_rError, seq_tError = process_sequence( + model, seq_name, seq_data, category, args.co3d_dir, + args.min_num_images, args.num_frames, args.use_ba, device, torch.bfloat16 + ) + + print("-" * 50) + + if seq_rError is not None and seq_tError is not None: + rError.extend(seq_rError) + tError.extend(seq_tError) + + if not rError: + print(f"No valid sequences found for {category}, skipping") + continue + + rError = np.array(rError) + tError = np.array(tError) + + thresholds = [5, 10, 20, 30] + Aucs = {} + + for threshold in thresholds: + Auc, _ = calculate_auc_np(rError, tError, max_threshold=threshold) + Aucs[threshold] = Auc + + print("="*80) + print(f"AUC of {category} test set: {Aucs[30]:.4f}") + print("="*80) + + per_category_results[category] = { + "rError": rError, + "tError": tError, + "Auc_5": Aucs[5], + "Auc_10": Aucs[10], + "Auc_20": Aucs[20], + "Auc_30": Aucs[30], + } + + # Print summary results + print("\nSummary of AUC results:") + print("-"*50) + for category in sorted(per_category_results.keys()): + print(f"{category:<15} AUC_5: {per_category_results[category]['Auc_5']:.4f}") + print(f"{category:<15} AUC_30: {per_category_results[category]['Auc_30']:.4f}") + print(f"{category:<15} AUC_20: {per_category_results[category]['Auc_20']:.4f}") + print(f"{category:<15} AUC_10: {per_category_results[category]['Auc_10']:.4f}") + + if per_category_results: + mean_AUC_30 = np.mean([per_category_results[category]["Auc_30"] for category in per_category_results]) + mean_AUC_20 = np.mean([per_category_results[category]["Auc_20"] for category in per_category_results]) + mean_AUC_10 = np.mean([per_category_results[category]["Auc_10"] for category in per_category_results]) + mean_AUC_5 = np.mean([per_category_results[category]["Auc_5"] for category in per_category_results]) + print("-"*50) + print(f"Mean AUC_5: {mean_AUC_5:.4f}") + print(f"Mean AUC_30: {mean_AUC_30:.4f}") + print(f"Mean AUC_20: {mean_AUC_20:.4f}") + print(f"Mean AUC_10: {mean_AUC_10:.4f}") + + # Generate a random index to avoid overwriting previous results + # random_index = torch.randint(0, 10000, (1,)).item() + # Use timestamp as index instead of random number + import datetime + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + random_index = timestamp + results_file = f"co3d_results_{random_index}.txt" + + with open(results_file, "w") as f: + f.write("CO3D Evaluation Results\n") + f.write("=" * 50 + "\n\n") + + f.write("Per-category results:\n") + f.write("-" * 50 + "\n") + for category in sorted(per_category_results.keys()): + f.write(f"{category:<15} AUC_30: {per_category_results[category]['Auc_30']:.4f}\n") + f.write(f"{category:<15} AUC_20: {per_category_results[category]['Auc_20']:.4f}\n") + f.write(f"{category:<15} AUC_10: {per_category_results[category]['Auc_10']:.4f}\n") + f.write(f"{category:<15} AUC_5: {per_category_results[category]['Auc_5']:.4f}\n") + f.write("\n") + + if per_category_results: + f.write("-" * 50 + "\n") + f.write(f"Mean AUC_30: {mean_AUC_30:.4f}\n") + f.write(f"Mean AUC_20: {mean_AUC_20:.4f}\n") + f.write(f"Mean AUC_10: {mean_AUC_10:.4f}\n") + f.write(f"Mean AUC_5: {mean_AUC_5:.4f}\n") + f.write("\n" + "=" * 50 + "\n") + + print(f"Results saved to {results_file}") + + +if __name__ == "__main__": + args = setup_args() + evaluate(args) diff --git a/src/evaluation/evaluation_cfg.py b/src/evaluation/evaluation_cfg.py new file mode 100644 index 0000000000000000000000000000000000000000..1a0245c787f2941b2c79b960c3691fcabe700842 --- /dev/null +++ b/src/evaluation/evaluation_cfg.py @@ -0,0 +1,24 @@ +from dataclasses import dataclass +from pathlib import Path + + +@dataclass +class MethodCfg: + name: str + key: str + path: Path + + +@dataclass +class SceneCfg: + scene: str + target_index: int + + +@dataclass +class EvaluationCfg: + methods: list[MethodCfg] + side_by_side_path: Path | None + output_metrics_path: Path + animate_side_by_side: bool + highlighted: list[SceneCfg] diff --git a/src/evaluation/evaluation_index_generator.py b/src/evaluation/evaluation_index_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..a3d037a442a76adcbeb20c0ed16c075f9a3c81e1 --- /dev/null +++ b/src/evaluation/evaluation_index_generator.py @@ -0,0 +1,160 @@ +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Optional + +import torch +from einops import rearrange +from lightning.pytorch import LightningModule +from tqdm import tqdm + +from ..geometry.epipolar_lines import project_rays +from ..geometry.projection import get_world_rays, sample_image_grid +from ..misc.image_io import save_image +from ..visualization.annotation import add_label +from ..visualization.layout import add_border, hcat + + +@dataclass +class EvaluationIndexGeneratorCfg: + num_target_views: int + min_distance: int + max_distance: int + min_overlap: float + max_overlap: float + output_path: Path + save_previews: bool + seed: int + + +@dataclass +class IndexEntry: + context: tuple[int, ...] + target: tuple[int, ...] + overlap: Optional[str | float] = None # choose from ["small", "medium", "large"] or a float number indicates the overlap ratio + + +class EvaluationIndexGenerator(LightningModule): + generator: torch.Generator + cfg: EvaluationIndexGeneratorCfg + index: dict[str, IndexEntry | None] + + def __init__(self, cfg: EvaluationIndexGeneratorCfg) -> None: + super().__init__() + self.cfg = cfg + self.generator = torch.Generator() + self.generator.manual_seed(cfg.seed) + self.index = {} + + def test_step(self, batch, batch_idx): + b, v, _, h, w = batch["target"]["image"].shape + assert b == 1 + extrinsics = batch["target"]["extrinsics"][0] + intrinsics = batch["target"]["intrinsics"][0] + scene = batch["scene"][0] + + context_indices = torch.randperm(v, generator=self.generator) + for context_index in tqdm(context_indices, "Finding context pair"): + xy, _ = sample_image_grid((h, w), self.device) + context_origins, context_directions = get_world_rays( + rearrange(xy, "h w xy -> (h w) xy"), + extrinsics[context_index], + intrinsics[context_index], + ) + + # Step away from context view until the minimum overlap threshold is met. + valid_indices = [] + for step in (1, -1): + min_distance = self.cfg.min_distance + max_distance = self.cfg.max_distance + current_index = context_index + step * min_distance + + while 0 <= current_index.item() < v: + # Compute overlap. + current_origins, current_directions = get_world_rays( + rearrange(xy, "h w xy -> (h w) xy"), + extrinsics[current_index], + intrinsics[current_index], + ) + projection_onto_current = project_rays( + context_origins, + context_directions, + extrinsics[current_index], + intrinsics[current_index], + ) + projection_onto_context = project_rays( + current_origins, + current_directions, + extrinsics[context_index], + intrinsics[context_index], + ) + overlap_a = projection_onto_context["overlaps_image"].float().mean() + overlap_b = projection_onto_current["overlaps_image"].float().mean() + + overlap = min(overlap_a, overlap_b) + delta = (current_index - context_index).abs() + + min_overlap = self.cfg.min_overlap + max_overlap = self.cfg.max_overlap + if min_overlap <= overlap <= max_overlap: + valid_indices.append( + (current_index.item(), overlap_a, overlap_b) + ) + + # Stop once the camera has panned away too much. + if overlap < min_overlap or delta > max_distance: + break + + current_index += step + + if valid_indices: + # Pick a random valid view. Index the resulting views. + num_options = len(valid_indices) + chosen = torch.randint( + 0, num_options, size=tuple(), generator=self.generator + ) + chosen, overlap_a, overlap_b = valid_indices[chosen] + + context_left = min(chosen, context_index.item()) + context_right = max(chosen, context_index.item()) + delta = context_right - context_left + + # Pick non-repeated random target views. + while True: + target_views = torch.randint( + context_left, + context_right + 1, + (self.cfg.num_target_views,), + generator=self.generator, + ) + if (target_views.unique(return_counts=True)[1] == 1).all(): + break + + target = tuple(sorted(target_views.tolist())) + self.index[scene] = IndexEntry( + context=(context_left, context_right), + target=target, + ) + + # Optionally, save a preview. + if self.cfg.save_previews: + preview_path = self.cfg.output_path / "previews" + preview_path.mkdir(exist_ok=True, parents=True) + a = batch["target"]["image"][0, chosen] + a = add_label(a, f"Overlap: {overlap_a * 100:.1f}%") + b = batch["target"]["image"][0, context_index] + b = add_label(b, f"Overlap: {overlap_b * 100:.1f}%") + vis = add_border(add_border(hcat(a, b)), 1, 0) + vis = add_label(vis, f"Distance: {delta} frames") + save_image(add_border(vis), preview_path / f"{scene}.png") + break + else: + # This happens if no starting frame produces a valid evaluation example. + self.index[scene] = None + + def save_index(self) -> None: + self.cfg.output_path.mkdir(exist_ok=True, parents=True) + with (self.cfg.output_path / "evaluation_index.json").open("w") as f: + json.dump( + {k: None if v is None else asdict(v) for k, v in self.index.items()}, f + ) diff --git a/src/evaluation/metrics.py b/src/evaluation/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..75894fa795653dfcd4e0594aecfb221a2c079e86 --- /dev/null +++ b/src/evaluation/metrics.py @@ -0,0 +1,134 @@ +from functools import cache + +import torch +from einops import reduce +from jaxtyping import Float +from lpips import LPIPS +from skimage.metrics import structural_similarity +from torch import Tensor + + +@torch.no_grad() +def compute_psnr( + ground_truth: Float[Tensor, "batch channel height width"], + predicted: Float[Tensor, "batch channel height width"], +) -> Float[Tensor, " batch"]: + ground_truth = ground_truth.clip(min=0, max=1) + predicted = predicted.clip(min=0, max=1) + mse = reduce((ground_truth - predicted) ** 2, "b c h w -> b", "mean") + return -10 * mse.log10() + + +@cache +def get_lpips(device: torch.device) -> LPIPS: + return LPIPS(net="vgg").to(device) + + +@torch.no_grad() +def compute_lpips( + ground_truth: Float[Tensor, "batch channel height width"], + predicted: Float[Tensor, "batch channel height width"], +) -> Float[Tensor, " batch"]: + value = get_lpips(predicted.device).forward(ground_truth, predicted, normalize=True) + return value[:, 0, 0, 0] + + +@torch.no_grad() +def compute_ssim( + ground_truth: Float[Tensor, "batch channel height width"], + predicted: Float[Tensor, "batch channel height width"], +) -> Float[Tensor, " batch"]: + ssim = [ + structural_similarity( + gt.detach().cpu().numpy(), + hat.detach().cpu().numpy(), + win_size=11, + gaussian_weights=True, + channel_axis=0, + data_range=1.0, + ) + for gt, hat in zip(ground_truth, predicted) + ] + return torch.tensor(ssim, dtype=predicted.dtype, device=predicted.device) + + +def compute_geodesic_distance_from_two_matrices(m1, m2): + batch = m1.shape[0] + m = torch.bmm(m1, m2.transpose(1, 2)) # batch*3*3 + + cos = (m[:, 0, 0] + m[:, 1, 1] + m[:, 2, 2] - 1) / 2 + cos = torch.min(cos, torch.autograd.Variable(torch.ones(batch).to(m1.device))) + cos = torch.max(cos, torch.autograd.Variable(torch.ones(batch).to(m1.device)) * -1) + + theta = torch.acos(cos) + + # theta = torch.min(theta, 2*np.pi - theta) + + return theta + + +def angle_error_mat(R1, R2): + cos = (torch.trace(torch.mm(R1.T, R2)) - 1) / 2 + cos = torch.clamp(cos, -1.0, 1.0) # numerical errors can make it out of bounds + return torch.rad2deg(torch.abs(torch.acos(cos))) + + +def angle_error_vec(v1, v2): + n = torch.norm(v1) * torch.norm(v2) + cos_theta = torch.dot(v1, v2) / n + cos_theta = torch.clamp(cos_theta, -1.0, 1.0) # numerical errors can make it out of bounds + return torch.rad2deg(torch.acos(cos_theta)) + + +def compute_translation_error(t1, t2): + return torch.norm(t1 - t2) + + +@torch.no_grad() +def compute_pose_error(pose_gt, pose_pred): + R_gt = pose_gt[:3, :3] + t_gt = pose_gt[:3, 3] + + R = pose_pred[:3, :3] + t = pose_pred[:3, 3] + + error_t = angle_error_vec(t, t_gt) + error_t = torch.minimum(error_t, 180 - error_t) # ambiguity of E estimation + error_t_scale = compute_translation_error(t, t_gt) + error_R = angle_error_mat(R, R_gt) + return error_t, error_t_scale, error_R + +@torch.no_grad() +def abs_relative_difference(output, target, valid_mask=None): + actual_output = output + actual_target = target + abs_relative_diff = torch.abs(actual_output - actual_target) / actual_target + if valid_mask is not None: + abs_relative_diff[~valid_mask] = 0 + n = valid_mask.sum((-1, -2)) + else: + n = output.shape[-1] * output.shape[-2] + abs_relative_diff = torch.sum(abs_relative_diff, (-1, -2)) / n + return abs_relative_diff.mean() + +# adapt from: https://github.com/imran3180/depth-map-prediction/blob/master/main.py +@torch.no_grad() +def threshold_percentage(output, target, threshold_val, valid_mask=None): + d1 = output / target + d2 = target / output + max_d1_d2 = torch.max(d1, d2) + zero = torch.zeros_like(output) + one = torch.ones_like(output) + bit_mat = torch.where(max_d1_d2 < threshold_val, one, zero) + if valid_mask is not None: + bit_mat[~valid_mask] = 0 + n = valid_mask.sum((-1, -2)) + else: + n = output.shape[-1] * output.shape[-2] + count_mat = torch.sum(bit_mat, (-1, -2)) + threshold_mat = count_mat / n + return threshold_mat.mean() + +@torch.no_grad() +def delta1_acc(pred, gt, valid_mask): + return threshold_percentage(pred, gt, 1.25, valid_mask) \ No newline at end of file diff --git a/src/geometry/camera_emb.py b/src/geometry/camera_emb.py new file mode 100644 index 0000000000000000000000000000000000000000..39beeae7fc9b78743d70db702185d3cb07562184 --- /dev/null +++ b/src/geometry/camera_emb.py @@ -0,0 +1,29 @@ +from einops import rearrange + +from .projection import sample_image_grid, get_local_rays +from ..misc.sht import rsh_cart_2, rsh_cart_4, rsh_cart_6, rsh_cart_8 + + +def get_intrinsic_embedding(context, degree=0, downsample=1, merge_hw=False): + assert degree in [0, 2, 4, 8] + + b, v, _, h, w = context["image"].shape + device = context["image"].device + tgt_h, tgt_w = h // downsample, w // downsample + xy_ray, _ = sample_image_grid((tgt_h, tgt_w), device) + xy_ray = xy_ray[None, None, ...].expand(b, v, -1, -1, -1) # [b, v, h, w, 2] + directions = get_local_rays(xy_ray, rearrange(context["intrinsics"], "b v i j -> b v () () i j"),) + + if degree == 2: + directions = rsh_cart_2(directions) + elif degree == 4: + directions = rsh_cart_4(directions) + elif degree == 8: + directions = rsh_cart_8(directions) + + if merge_hw: + directions = rearrange(directions, "b v h w d -> b v (h w) d") + else: + directions = rearrange(directions, "b v h w d -> b v d h w") + + return directions diff --git a/src/geometry/epipolar_lines.py b/src/geometry/epipolar_lines.py new file mode 100644 index 0000000000000000000000000000000000000000..d4ea46bfbdfefb77851e747dd16d1eade1970d2c --- /dev/null +++ b/src/geometry/epipolar_lines.py @@ -0,0 +1,292 @@ +import itertools +from typing import Iterable, Literal, Optional, TypedDict + +import torch +from einops import einsum, repeat +from jaxtyping import Bool, Float +from torch import Tensor +from torch.utils.data.dataloader import default_collate + +from .projection import ( + get_world_rays, + homogenize_points, + homogenize_vectors, + intersect_rays, + project_camera_space, +) + + +def _is_in_bounds( + xy: Float[Tensor, "*batch 2"], + epsilon: float = 1e-6, +) -> Bool[Tensor, " *batch"]: + """Check whether the specified XY coordinates are within the normalized image plane, + which has a range from 0 to 1 in each direction. + """ + return (xy >= -epsilon).all(dim=-1) & (xy <= 1 + epsilon).all(dim=-1) + + +def _is_in_front_of_camera( + xyz: Float[Tensor, "*batch 3"], + epsilon: float = 1e-6, +) -> Bool[Tensor, " *batch"]: + """Check whether the specified points in camera space are in front of the camera.""" + return xyz[..., -1] > -epsilon + + +def _is_positive_t( + t: Float[Tensor, " *batch"], + epsilon: float = 1e-6, +) -> Bool[Tensor, " *batch"]: + """Check whether the specified t value is positive.""" + return t > -epsilon + + +class PointProjection(TypedDict): + t: Float[Tensor, " *batch"] # ray parameter, as in xyz = origin + t * direction + xy: Float[Tensor, "*batch 2"] # image-space xy (normalized to 0 to 1) + + # A "valid" projection satisfies two conditions: + # 1. It is in front of the camera (i.e., its 3D Z coordinate is positive). + # 2. It is within the image frame (i.e., its 2D coordinates are between 0 and 1). + valid: Bool[Tensor, " *batch"] + + +def _intersect_image_coordinate( + intrinsics: Float[Tensor, "*#batch 3 3"], + origins: Float[Tensor, "*#batch 3"], + directions: Float[Tensor, "*#batch 3"], + dimension: Literal["x", "y"], + coordinate_value: float, +) -> PointProjection: + """Compute the intersection of the projection of a camera-space ray with a line + that's parallel to the image frame, either horizontally or vertically. + """ + + # Define shorthands. + dim = "xy".index(dimension) + other_dim = 1 - dim + fs = intrinsics[..., dim, dim] # focal length, same coordinate + fo = intrinsics[..., other_dim, other_dim] # focal length, other coordinate + cs = intrinsics[..., dim, 2] # principal point, same coordinate + co = intrinsics[..., other_dim, 2] # principal point, other coordinate + os = origins[..., dim] # ray origin, same coordinate + oo = origins[..., other_dim] # ray origin, other coordinate + ds = directions[..., dim] # ray direction, same coordinate + do = directions[..., other_dim] # ray direction, other coordinate + oz = origins[..., 2] # ray origin, z coordinate + dz = directions[..., 2] # ray direction, z coordinate + c = (coordinate_value - cs) / fs # coefficient (computed once and factored out) + + # Compute the value of t at the intersection. + # Note: Infinite values of t are fine. No need to handle division by zero. + t_numerator = c * oz - os + t_denominator = ds - c * dz + t = t_numerator / t_denominator + + # Compute the value of the other coordinate at the intersection. + # Note: Infinite coordinate values are fine. No need to handle division by zero. + coordinate_numerator = fo * (oo * (c * dz - ds) + do * (os - c * oz)) + coordinate_denominator = dz * os - ds * oz + coordinate_other = co + coordinate_numerator / coordinate_denominator + coordinate_same = torch.ones_like(coordinate_other) * coordinate_value + xy = [coordinate_same] + xy.insert(other_dim, coordinate_other) + xy = torch.stack(xy, dim=-1) + xyz = origins + t[..., None] * directions + + # These will all have exactly the same batch shape (no broadcasting necessary). In + # terms of jaxtyping annotations, they all match *batch, not just *#batch. + return { + "t": t, + "xy": xy, + "valid": _is_in_bounds(xy) & _is_in_front_of_camera(xyz) & _is_positive_t(t), + } + + +def _compare_projections( + intersections: Iterable[PointProjection], + reduction: Literal["min", "max"], +) -> PointProjection: + intersections = {k: v.clone() for k, v in default_collate(intersections).items()} + t = intersections["t"] + xy = intersections["xy"] + valid = intersections["valid"] + + # Make sure out-of-bounds values are not chosen. + lowest_priority = { + "min": torch.inf, + "max": -torch.inf, + }[reduction] + t[~valid] = lowest_priority + + # Run the reduction (either t.min() or t.max()). + reduced, selector = getattr(t, reduction)(dim=0) + + # Index the results. + return { + "t": reduced, + "xy": xy.gather(0, repeat(selector, "... -> () ... xy", xy=2))[0], + "valid": valid.gather(0, selector[None])[0], + } + + +def _compute_point_projection( + xyz: Float[Tensor, "*#batch 3"], + t: Float[Tensor, "*#batch"], + intrinsics: Float[Tensor, "*#batch 3 3"], +) -> PointProjection: + xy = project_camera_space(xyz, intrinsics) + return { + "t": t, + "xy": xy, + "valid": _is_in_bounds(xy) & _is_in_front_of_camera(xyz) & _is_positive_t(t), + } + + +class RaySegmentProjection(TypedDict): + t_min: Float[Tensor, " *batch"] # ray parameter + t_max: Float[Tensor, " *batch"] # ray parameter + xy_min: Float[Tensor, "*batch 2"] # image-space xy (normalized to 0 to 1) + xy_max: Float[Tensor, "*batch 2"] # image-space xy (normalized to 0 to 1) + + # Whether the segment overlaps the image. If not, the above values are meaningless. + overlaps_image: Bool[Tensor, " *batch"] + + +def project_rays( + origins: Float[Tensor, "*#batch 3"], + directions: Float[Tensor, "*#batch 3"], + extrinsics: Float[Tensor, "*#batch 4 4"], + intrinsics: Float[Tensor, "*#batch 3 3"], + near: Optional[Float[Tensor, "*#batch"]] = None, + far: Optional[Float[Tensor, "*#batch"]] = None, + epsilon: float = 1e-6, +) -> RaySegmentProjection: + # Transform the rays into camera space. + world_to_cam = torch.linalg.inv(extrinsics) + origins = homogenize_points(origins) + origins = einsum(world_to_cam, origins, "... i j, ... j -> ... i") + directions = homogenize_vectors(directions) + directions = einsum(world_to_cam, directions, "... i j, ... j -> ... i") + origins = origins[..., :3] + directions = directions[..., :3] + + # Compute intersections with the image's frame. + frame_intersections = ( + _intersect_image_coordinate(intrinsics, origins, directions, "x", 0.0), + _intersect_image_coordinate(intrinsics, origins, directions, "x", 1.0), + _intersect_image_coordinate(intrinsics, origins, directions, "y", 0.0), + _intersect_image_coordinate(intrinsics, origins, directions, "y", 1.0), + ) + frame_intersection_min = _compare_projections(frame_intersections, "min") + frame_intersection_max = _compare_projections(frame_intersections, "max") + + if near is None: + # Compute the ray's projection at zero depth. If an origin's depth (z value) is + # within epsilon of zero, this can mean one of two things: + # 1. The origin is at the camera's position. In this case, use the direction + # instead (the ray is probably coming from the camera). + # 2. The origin isn't at the camera's position, and randomly happens to be on + # the plane at zero depth. In this case, its projection is outside the image + # plane, and is thus marked as invalid. + origins_for_projection = origins.clone() + mask_depth_zero = origins_for_projection[..., -1] < epsilon + mask_at_camera = origins_for_projection.norm(dim=-1) < epsilon + origins_for_projection[mask_at_camera] = directions[mask_at_camera] + projection_at_zero = _compute_point_projection( + origins_for_projection, + torch.zeros_like(frame_intersection_min["t"]), + intrinsics, + ) + projection_at_zero["valid"][mask_depth_zero & ~mask_at_camera] = False + else: + # If a near plane is specified, use it instead. + t_near = near.broadcast_to(frame_intersection_min["t"].shape) + projection_at_zero = _compute_point_projection( + origins + near[..., None] * directions, + t_near, + intrinsics, + ) + + if far is None: + # Compute the ray's projection at infinite depth. Using the projection function + # with directions (vectors) instead of points may seem wonky, but is equivalent + # to projecting the point at (origins + infinity * directions). + projection_at_infinity = _compute_point_projection( + directions, + torch.ones_like(frame_intersection_min["t"]) * torch.inf, + intrinsics, + ) + else: + # If a far plane is specified, use it instead. + t_far = far.broadcast_to(frame_intersection_min["t"].shape) + projection_at_infinity = _compute_point_projection( + origins + far[..., None] * directions, + t_far, + intrinsics, + ) + + # Build the result by handling cases for ray intersection. + result = { + "t_min": torch.empty_like(projection_at_zero["t"]), + "t_max": torch.empty_like(projection_at_infinity["t"]), + "xy_min": torch.empty_like(projection_at_zero["xy"]), + "xy_max": torch.empty_like(projection_at_infinity["xy"]), + "overlaps_image": torch.empty_like(projection_at_zero["valid"]), + } + + for min_valid, max_valid in itertools.product([True, False], [True, False]): + min_mask = projection_at_zero["valid"] ^ (not min_valid) + max_mask = projection_at_infinity["valid"] ^ (not max_valid) + mask = min_mask & max_mask + min_value = projection_at_zero if min_valid else frame_intersection_min + max_value = projection_at_infinity if max_valid else frame_intersection_max + result["t_min"][mask] = min_value["t"][mask] + result["t_max"][mask] = max_value["t"][mask] + result["xy_min"][mask] = min_value["xy"][mask] + result["xy_max"][mask] = max_value["xy"][mask] + result["overlaps_image"][mask] = (min_value["valid"] & max_value["valid"])[mask] + + return result + + +class RaySegmentProjection(TypedDict): + t_min: Float[Tensor, " *batch"] # ray parameter + t_max: Float[Tensor, " *batch"] # ray parameter + xy_min: Float[Tensor, "*batch 2"] # image-space xy (normalized to 0 to 1) + xy_max: Float[Tensor, "*batch 2"] # image-space xy (normalized to 0 to 1) + + # Whether the segment overlaps the image. If not, the above values are meaningless. + overlaps_image: Bool[Tensor, " *batch"] + + +def lift_to_3d( + origins: Float[Tensor, "*#batch 3"], + directions: Float[Tensor, "*#batch 3"], + xy: Float[Tensor, "*#batch 2"], + extrinsics: Float[Tensor, "*#batch 4 4"], + intrinsics: Float[Tensor, "*#batch 3 3"], +) -> Float[Tensor, "*batch 3"]: + """Calculate the 3D positions that correspond to the specified 2D points on the + epipolar lines defined by the origins and directions. The extrinsics and intrinsics + are for the images the 2D points lie on. + """ + + xy_origins, xy_directions = get_world_rays(xy, extrinsics, intrinsics) + return intersect_rays(origins, directions, xy_origins, xy_directions) + + +def get_depth( + origins: Float[Tensor, "*#batch 3"], + directions: Float[Tensor, "*#batch 3"], + xy: Float[Tensor, "*#batch 2"], + extrinsics: Float[Tensor, "*#batch 4 4"], + intrinsics: Float[Tensor, "*#batch 3 3"], +) -> Float[Tensor, " *batch"]: + """Calculate the depths that correspond to the specified 2D points on the epipolar + lines defined by the origins and directions. The extrinsics and intrinsics are for + the images the 2D points lie on. + """ + xyz = lift_to_3d(origins, directions, xy, extrinsics, intrinsics) + return (xyz - origins).norm(dim=-1) diff --git a/src/geometry/projection.py b/src/geometry/projection.py new file mode 100644 index 0000000000000000000000000000000000000000..c7fcf6f4cc439f4110f8f1ca8fb1c653231cd024 --- /dev/null +++ b/src/geometry/projection.py @@ -0,0 +1,261 @@ +from math import prod + +import torch +from einops import einsum, rearrange, reduce, repeat +from jaxtyping import Bool, Float, Int64 +from torch import Tensor + + +def homogenize_points( + points: Float[Tensor, "*batch dim"], +) -> Float[Tensor, "*batch dim+1"]: + """Convert batched points (xyz) to (xyz1).""" + return torch.cat([points, torch.ones_like(points[..., :1])], dim=-1) + + +def homogenize_vectors( + vectors: Float[Tensor, "*batch dim"], +) -> Float[Tensor, "*batch dim+1"]: + """Convert batched vectors (xyz) to (xyz0).""" + return torch.cat([vectors, torch.zeros_like(vectors[..., :1])], dim=-1) + + +def transform_rigid( + homogeneous_coordinates: Float[Tensor, "*#batch dim"], + transformation: Float[Tensor, "*#batch dim dim"], +) -> Float[Tensor, "*batch dim"]: + """Apply a rigid-body transformation to points or vectors.""" + return einsum(transformation, homogeneous_coordinates, "... i j, ... j -> ... i") + + +def transform_cam2world( + homogeneous_coordinates: Float[Tensor, "*#batch dim"], + extrinsics: Float[Tensor, "*#batch dim dim"], +) -> Float[Tensor, "*batch dim"]: + """Transform points from 3D camera coordinates to 3D world coordinates.""" + return transform_rigid(homogeneous_coordinates, extrinsics) + + +def transform_world2cam( + homogeneous_coordinates: Float[Tensor, "*#batch dim"], + extrinsics: Float[Tensor, "*#batch dim dim"], +) -> Float[Tensor, "*batch dim"]: + """Transform points from 3D world coordinates to 3D camera coordinates.""" + return transform_rigid(homogeneous_coordinates, extrinsics.inverse()) + + +def project_camera_space( + points: Float[Tensor, "*#batch dim"], + intrinsics: Float[Tensor, "*#batch dim dim"], + epsilon: float = torch.finfo(torch.float32).eps, + infinity: float = 1e8, +) -> Float[Tensor, "*batch dim-1"]: + points = points / (points[..., -1:] + epsilon) + points = points.nan_to_num(posinf=infinity, neginf=-infinity) + points = einsum(intrinsics, points, "... i j, ... j -> ... i") + return points[..., :-1] + + +def project( + points: Float[Tensor, "*#batch dim"], + extrinsics: Float[Tensor, "*#batch dim+1 dim+1"], + intrinsics: Float[Tensor, "*#batch dim dim"], + epsilon: float = torch.finfo(torch.float32).eps, +) -> tuple[ + Float[Tensor, "*batch dim-1"], # xy coordinates + Bool[Tensor, " *batch"], # whether points are in front of the camera +]: + points = homogenize_points(points) + points = transform_world2cam(points, extrinsics)[..., :-1] + in_front_of_camera = points[..., -1] >= 0 + return project_camera_space(points, intrinsics, epsilon=epsilon), in_front_of_camera + + +def unproject( + coordinates: Float[Tensor, "*#batch dim"], + z: Float[Tensor, "*#batch"], + intrinsics: Float[Tensor, "*#batch dim+1 dim+1"], +) -> Float[Tensor, "*batch dim+1"]: + """Unproject 2D camera coordinates with the given Z values.""" + + # Apply the inverse intrinsics to the coordinates. + coordinates = homogenize_points(coordinates) + ray_directions = einsum( + intrinsics.inverse(), coordinates, "... i j, ... j -> ... i" + ) + + # Apply the supplied depth values. + return ray_directions * z[..., None] + + +def get_world_rays( + coordinates: Float[Tensor, "*#batch dim"], + extrinsics: Float[Tensor, "*#batch dim+2 dim+2"], + intrinsics: Float[Tensor, "*#batch dim+1 dim+1"], +) -> tuple[ + Float[Tensor, "*batch dim+1"], # origins + Float[Tensor, "*batch dim+1"], # directions +]: + # Get camera-space ray directions. + directions = unproject( + coordinates, + torch.ones_like(coordinates[..., 0]), + intrinsics, + ) + directions = directions / directions.norm(dim=-1, keepdim=True) + + # Transform ray directions to world coordinates. + directions = homogenize_vectors(directions) + directions = transform_cam2world(directions, extrinsics)[..., :-1] + + # Tile the ray origins to have the same shape as the ray directions. + origins = extrinsics[..., :-1, -1].broadcast_to(directions.shape) + + return origins, directions + + +def get_local_rays( + coordinates: Float[Tensor, "*#batch dim"], + intrinsics: Float[Tensor, "*#batch dim+1 dim+1"], +) -> Float[Tensor, "*batch dim+1"]: + # Get camera-space ray directions. + directions = unproject( + coordinates, + torch.ones_like(coordinates[..., 0]), + intrinsics, + ) + directions = directions / directions.norm(dim=-1, keepdim=True) + return directions + + +def sample_image_grid( + shape: tuple[int, ...], + device: torch.device = torch.device("cpu"), +) -> tuple[ + Float[Tensor, "*shape dim"], # float coordinates (xy indexing) + Int64[Tensor, "*shape dim"], # integer indices (ij indexing) +]: + """Get normalized (range 0 to 1) coordinates and integer indices for an image.""" + + # Each entry is a pixel-wise integer coordinate. In the 2D case, each entry is a + # (row, col) coordinate. + indices = [torch.arange(length, device=device) for length in shape] + stacked_indices = torch.stack(torch.meshgrid(*indices, indexing="ij"), dim=-1) + + # Each entry is a floating-point coordinate in the range (0, 1). In the 2D case, + # each entry is an (x, y) coordinate. + coordinates = [(idx + 0.5) / length for idx, length in zip(indices, shape)] + coordinates = reversed(coordinates) + coordinates = torch.stack(torch.meshgrid(*coordinates, indexing="xy"), dim=-1) + + return coordinates, stacked_indices + + +def sample_training_rays( + image: Float[Tensor, "batch view channel ..."], + intrinsics: Float[Tensor, "batch view dim dim"], + extrinsics: Float[Tensor, "batch view dim+1 dim+1"], + num_rays: int, +) -> tuple[ + Float[Tensor, "batch ray dim"], # origins + Float[Tensor, "batch ray dim"], # directions + Float[Tensor, "batch ray 3"], # sampled color +]: + device = extrinsics.device + b, v, _, *grid_shape = image.shape + + # Generate all possible target rays. + xy, _ = sample_image_grid(tuple(grid_shape), device) + origins, directions = get_world_rays( + rearrange(xy, "... d -> ... () () d"), + extrinsics, + intrinsics, + ) + origins = rearrange(origins, "... b v xy -> b (v ...) xy", b=b, v=v) + directions = rearrange(directions, "... b v xy -> b (v ...) xy", b=b, v=v) + pixels = rearrange(image, "b v c ... -> b (v ...) c") + + # Sample random rays. + num_possible_rays = v * prod(grid_shape) + ray_indices = torch.randint(num_possible_rays, (b, num_rays), device=device) + batch_indices = repeat(torch.arange(b, device=device), "b -> b n", n=num_rays) + + return ( + origins[batch_indices, ray_indices], + directions[batch_indices, ray_indices], + pixels[batch_indices, ray_indices], + ) + + +def intersect_rays( + origins_x: Float[Tensor, "*#batch 3"], + directions_x: Float[Tensor, "*#batch 3"], + origins_y: Float[Tensor, "*#batch 3"], + directions_y: Float[Tensor, "*#batch 3"], + eps: float = 1e-5, + inf: float = 1e10, +) -> Float[Tensor, "*batch 3"]: + """Compute the least-squares intersection of rays. Uses the math from here: + https://math.stackexchange.com/a/1762491/286022 + """ + + # Broadcast the rays so their shapes match. + shape = torch.broadcast_shapes( + origins_x.shape, + directions_x.shape, + origins_y.shape, + directions_y.shape, + ) + origins_x = origins_x.broadcast_to(shape) + directions_x = directions_x.broadcast_to(shape) + origins_y = origins_y.broadcast_to(shape) + directions_y = directions_y.broadcast_to(shape) + + # Detect and remove batch elements where the directions are parallel. + parallel = einsum(directions_x, directions_y, "... xyz, ... xyz -> ...") > 1 - eps + origins_x = origins_x[~parallel] + directions_x = directions_x[~parallel] + origins_y = origins_y[~parallel] + directions_y = directions_y[~parallel] + + # Stack the rays into (2, *shape). + origins = torch.stack([origins_x, origins_y], dim=0) + directions = torch.stack([directions_x, directions_y], dim=0) + dtype = origins.dtype + device = origins.device + + # Compute n_i * n_i^T - eye(3) from the equation. + n = einsum(directions, directions, "r b i, r b j -> r b i j") + n = n - torch.eye(3, dtype=dtype, device=device).broadcast_to((2, 1, 3, 3)) + + # Compute the left-hand side of the equation. + lhs = reduce(n, "r b i j -> b i j", "sum") + + # Compute the right-hand side of the equation. + rhs = einsum(n, origins, "r b i j, r b j -> r b i") + rhs = reduce(rhs, "r b i -> b i", "sum") + + # Left-matrix-multiply both sides by the pseudo-inverse of lhs to find p. + result = torch.linalg.lstsq(lhs, rhs).solution + + # Handle the case of parallel lines by setting depth to infinity. + result_all = torch.ones(shape, dtype=dtype, device=device) * inf + result_all[~parallel] = result + return result_all + + +def get_fov(intrinsics: Float[Tensor, "batch 3 3"]) -> Float[Tensor, "batch 2"]: + intrinsics_inv = intrinsics.inverse() + + def process_vector(vector): + vector = torch.tensor(vector, dtype=torch.float32, device=intrinsics.device) + vector = einsum(intrinsics_inv, vector, "b i j, j -> b i") + return vector / vector.norm(dim=-1, keepdim=True) + + left = process_vector([0, 0.5, 1]) + right = process_vector([1, 0.5, 1]) + top = process_vector([0.5, 0, 1]) + bottom = process_vector([0.5, 1, 1]) + fov_x = (left * right).sum(dim=-1).acos() + fov_y = (top * bottom).sum(dim=-1).acos() + return torch.stack((fov_x, fov_y), dim=-1) diff --git a/src/geometry/ptc_geometry.py b/src/geometry/ptc_geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..4cef303162fce19b7a9d2e7bb2ad520ad7c2cd96 --- /dev/null +++ b/src/geometry/ptc_geometry.py @@ -0,0 +1,385 @@ +# Copyright (C) 2024-present Naver Corporation. All rights reserved. +# Licensed under CC BY-NC-SA 4.0 (non-commercial use only). +# +# -------------------------------------------------------- +# geometry utilitary functions +# -------------------------------------------------------- +import torch +import numpy as np +from scipy.spatial import cKDTree as KDTree + +from ..model.encoder.backbone.croco.misc import invalid_to_zeros, invalid_to_nans + +# from dust3r.utils.device import to_numpy + + +def xy_grid(W, H, device=None, origin=(0, 0), unsqueeze=None, cat_dim=-1, homogeneous=False, **arange_kw): + """ Output a (H,W,2) array of int32 + with output[j,i,0] = i + origin[0] + output[j,i,1] = j + origin[1] + """ + if device is None: + # numpy + arange, meshgrid, stack, ones = np.arange, np.meshgrid, np.stack, np.ones + else: + # torch + arange = lambda *a, **kw: torch.arange(*a, device=device, **kw) + meshgrid, stack = torch.meshgrid, torch.stack + ones = lambda *a: torch.ones(*a, device=device) + + tw, th = [arange(o, o+s, **arange_kw) for s, o in zip((W, H), origin)] + grid = meshgrid(tw, th, indexing='xy') + if homogeneous: + grid = grid + (ones((H, W)),) + if unsqueeze is not None: + grid = (grid[0].unsqueeze(unsqueeze), grid[1].unsqueeze(unsqueeze)) + if cat_dim is not None: + grid = stack(grid, cat_dim) + return grid + + +def geotrf(Trf, pts, ncol=None, norm=False): + """ Apply a geometric transformation to a list of 3-D points. + + H: 3x3 or 4x4 projection matrix (typically a Homography) + p: numpy/torch/tuple of coordinates. Shape must be (...,2) or (...,3) + + ncol: int. number of columns of the result (2 or 3) + norm: float. if != 0, the resut is projected on the z=norm plane. + + Returns an array of projected 2d points. + """ + assert Trf.ndim >= 2 + if isinstance(Trf, np.ndarray): + pts = np.asarray(pts) + elif isinstance(Trf, torch.Tensor): + pts = torch.as_tensor(pts, dtype=Trf.dtype) + + # adapt shape if necessary + output_reshape = pts.shape[:-1] + ncol = ncol or pts.shape[-1] + + # optimized code + if (isinstance(Trf, torch.Tensor) and isinstance(pts, torch.Tensor) and + Trf.ndim == 3 and pts.ndim == 4): + d = pts.shape[3] + if Trf.shape[-1] == d: + pts = torch.einsum("bij, bhwj -> bhwi", Trf, pts) + elif Trf.shape[-1] == d+1: + pts = torch.einsum("bij, bhwj -> bhwi", Trf[:, :d, :d], pts) + Trf[:, None, None, :d, d] + else: + raise ValueError(f'bad shape, not ending with 3 or 4, for {pts.shape=}') + else: + if Trf.ndim >= 3: + n = Trf.ndim-2 + assert Trf.shape[:n] == pts.shape[:n], 'batch size does not match' + Trf = Trf.reshape(-1, Trf.shape[-2], Trf.shape[-1]) + + if pts.ndim > Trf.ndim: + # Trf == (B,d,d) & pts == (B,H,W,d) --> (B, H*W, d) + pts = pts.reshape(Trf.shape[0], -1, pts.shape[-1]) + elif pts.ndim == 2: + # Trf == (B,d,d) & pts == (B,d) --> (B, 1, d) + pts = pts[:, None, :] + + if pts.shape[-1]+1 == Trf.shape[-1]: + Trf = Trf.swapaxes(-1, -2) # transpose Trf + pts = pts @ Trf[..., :-1, :] + Trf[..., -1:, :] + elif pts.shape[-1] == Trf.shape[-1]: + Trf = Trf.swapaxes(-1, -2) # transpose Trf + pts = pts @ Trf + else: + pts = Trf @ pts.T + if pts.ndim >= 2: + pts = pts.swapaxes(-1, -2) + + if norm: + pts = pts / pts[..., -1:] # DONT DO /= BECAUSE OF WEIRD PYTORCH BUG + if norm != 1: + pts *= norm + + res = pts[..., :ncol].reshape(*output_reshape, ncol) + return res + + +def inv(mat): + """ Invert a torch or numpy matrix + """ + if isinstance(mat, torch.Tensor): + return torch.linalg.inv(mat) + if isinstance(mat, np.ndarray): + return np.linalg.inv(mat) + raise ValueError(f'bad matrix type = {type(mat)}') + + +def depthmap_to_pts3d(depth, pseudo_focal, pp=None, **_): + """ + Args: + - depthmap (BxHxW array): + - pseudo_focal: [B,H,W] ; [B,2,H,W] or [B,1,H,W] + Returns: + pointmap of absolute coordinates (BxHxWx3 array) + """ + + if len(depth.shape) == 4: + B, H, W, n = depth.shape + else: + B, H, W = depth.shape + n = None + + if len(pseudo_focal.shape) == 3: # [B,H,W] + pseudo_focalx = pseudo_focaly = pseudo_focal + elif len(pseudo_focal.shape) == 4: # [B,2,H,W] or [B,1,H,W] + pseudo_focalx = pseudo_focal[:, 0] + if pseudo_focal.shape[1] == 2: + pseudo_focaly = pseudo_focal[:, 1] + else: + pseudo_focaly = pseudo_focalx + else: + raise NotImplementedError("Error, unknown input focal shape format.") + + assert pseudo_focalx.shape == depth.shape[:3] + assert pseudo_focaly.shape == depth.shape[:3] + grid_x, grid_y = xy_grid(W, H, cat_dim=0, device=depth.device)[:, None] + + # set principal point + if pp is None: + grid_x = grid_x - (W-1)/2 + grid_y = grid_y - (H-1)/2 + else: + grid_x = grid_x.expand(B, -1, -1) - pp[:, 0, None, None] + grid_y = grid_y.expand(B, -1, -1) - pp[:, 1, None, None] + + if n is None: + pts3d = torch.empty((B, H, W, 3), device=depth.device) + pts3d[..., 0] = depth * grid_x / pseudo_focalx + pts3d[..., 1] = depth * grid_y / pseudo_focaly + pts3d[..., 2] = depth + else: + pts3d = torch.empty((B, H, W, 3, n), device=depth.device) + pts3d[..., 0, :] = depth * (grid_x / pseudo_focalx)[..., None] + pts3d[..., 1, :] = depth * (grid_y / pseudo_focaly)[..., None] + pts3d[..., 2, :] = depth + return pts3d + + +def depthmap_to_camera_coordinates(depthmap, camera_intrinsics, pseudo_focal=None): + """ + Args: + - depthmap (HxW array): + - camera_intrinsics: a 3x3 matrix + Returns: + pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels. + """ + camera_intrinsics = np.float32(camera_intrinsics) + H, W = depthmap.shape + + # Compute 3D ray associated with each pixel + # Strong assumption: there are no skew terms + assert camera_intrinsics[0, 1] == 0.0 + assert camera_intrinsics[1, 0] == 0.0 + if pseudo_focal is None: + fu = camera_intrinsics[0, 0] + fv = camera_intrinsics[1, 1] + else: + assert pseudo_focal.shape == (H, W) + fu = fv = pseudo_focal + cu = camera_intrinsics[0, 2] + cv = camera_intrinsics[1, 2] + + u, v = np.meshgrid(np.arange(W), np.arange(H)) + z_cam = depthmap + x_cam = (u - cu) * z_cam / fu + y_cam = (v - cv) * z_cam / fv + X_cam = np.stack((x_cam, y_cam, z_cam), axis=-1).astype(np.float32) + + # Mask for valid coordinates + valid_mask = (depthmap > 0.0) + return X_cam, valid_mask + + +def depthmap_to_absolute_camera_coordinates(depthmap, camera_intrinsics, camera_pose, **kw): + """ + Args: + - depthmap (HxW array): + - camera_intrinsics: a 3x3 matrix + - camera_pose: a 4x3 or 4x4 cam2world matrix + Returns: + pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels.""" + X_cam, valid_mask = depthmap_to_camera_coordinates(depthmap, camera_intrinsics) + + # R_cam2world = np.float32(camera_params["R_cam2world"]) + # t_cam2world = np.float32(camera_params["t_cam2world"]).squeeze() + R_cam2world = camera_pose[:3, :3] + t_cam2world = camera_pose[:3, 3] + + # Express in absolute coordinates (invalid depth values) + X_world = np.einsum("ik, vuk -> vui", R_cam2world, X_cam) + t_cam2world[None, None, :] + return X_world, valid_mask + + +def colmap_to_opencv_intrinsics(K): + """ + Modify camera intrinsics to follow a different convention. + Coordinates of the center of the top-left pixels are by default: + - (0.5, 0.5) in Colmap + - (0,0) in OpenCV + """ + K = K.copy() + K[0, 2] -= 0.5 + K[1, 2] -= 0.5 + return K + + +def opencv_to_colmap_intrinsics(K): + """ + Modify camera intrinsics to follow a different convention. + Coordinates of the center of the top-left pixels are by default: + - (0.5, 0.5) in Colmap + - (0,0) in OpenCV + """ + K = K.copy() + K[0, 2] += 0.5 + K[1, 2] += 0.5 + return K + + +def obtain_pointcloud_center(pts, valid_mask=None): + """ + Args: + - pts (BxNx3 torch.array): pointmap of absolute coordinates + - valid_mask (BxN torch.array): mask specifying valid pixels. + Returns: + center of the point cloud (3 torch.array) + """ + depth = pts[..., 2] + # only choose the 0.02-0.98 quantile as the valid depth range + valid_depth_mask = ((depth > depth.quantile(0.02, dim=1, keepdim=True)) & + (depth < depth.quantile(0.98, dim=1, keepdim=True))) + if valid_mask is not None: + valid_mask = valid_depth_mask & valid_mask + else: + valid_mask = valid_depth_mask + + # pts: (B, N, 3), valid_mask: (B, N) + all_pts, nnz = invalid_to_zeros(pts, valid_mask, ndim=3) + all_dis = all_pts.norm(dim=-1) + center_dis = all_dis.sum(dim=1) / (nnz + 1e-8) + return center_dis # (B,) + +def normalize_pointcloud(pts1, pts2, norm_mode='avg_dis', valid1=None, valid2=None): + """ renorm pointmaps pts1, pts2 with norm_mode + """ + assert pts1.ndim >= 3 and pts1.shape[-1] == 3 + assert pts2 is None or (pts2.ndim >= 3 and pts2.shape[-1] == 3) + norm_mode, dis_mode = norm_mode.split('_') + + if norm_mode == 'avg': + # gather all points together (joint normalization) + nan_pts1, nnz1 = invalid_to_zeros(pts1, valid1, ndim=3) + nan_pts2, nnz2 = invalid_to_zeros(pts2, valid2, ndim=3) if pts2 is not None else (None, 0) + all_pts = torch.cat((nan_pts1, nan_pts2), dim=1) if pts2 is not None else nan_pts1 + + # compute distance to origin + all_dis = all_pts.norm(dim=-1) + if dis_mode == 'dis': + pass # do nothing + elif dis_mode == 'log1p': + all_dis = torch.log1p(all_dis) + elif dis_mode == 'warp-log1p': + # actually warp input points before normalizing them + log_dis = torch.log1p(all_dis) + warp_factor = log_dis / all_dis.clip(min=1e-8) + H1, W1 = pts1.shape[1:-1] + pts1 = pts1 * warp_factor[:, :W1*H1].view(-1, H1, W1, 1) + if pts2 is not None: + H2, W2 = pts2.shape[1:-1] + pts2 = pts2 * warp_factor[:, W1*H1:].view(-1, H2, W2, 1) + all_dis = log_dis # this is their true distance afterwards + else: + raise ValueError(f'bad {dis_mode=}') + + norm_factor = all_dis.sum(dim=1) / (nnz1 + nnz2 + 1e-8) + else: + # gather all points together (joint normalization) + nan_pts1 = invalid_to_nans(pts1, valid1, ndim=3) + nan_pts2 = invalid_to_nans(pts2, valid2, ndim=3) if pts2 is not None else None + all_pts = torch.cat((nan_pts1, nan_pts2), dim=1) if pts2 is not None else nan_pts1 + + # compute distance to origin + all_dis = all_pts.norm(dim=-1) + + if norm_mode == 'avg': + norm_factor = all_dis.nanmean(dim=1) + elif norm_mode == 'median': + norm_factor = all_dis.nanmedian(dim=1).values.detach() + elif norm_mode == 'sqrt': + norm_factor = all_dis.sqrt().nanmean(dim=1)**2 + else: + raise ValueError(f'bad {norm_mode=}') + + norm_factor = norm_factor.clip(min=1e-8) + while norm_factor.ndim < pts1.ndim: + norm_factor.unsqueeze_(-1) + + res = pts1 / norm_factor + if pts2 is not None: + res = (res, pts2 / norm_factor) + return res + + +@torch.no_grad() +def get_joint_pointcloud_depth(z1, z2, valid_mask1, valid_mask2=None, quantile=0.5): + # set invalid points to NaN + _z1 = invalid_to_nans(z1, valid_mask1).reshape(len(z1), -1) + _z2 = invalid_to_nans(z2, valid_mask2).reshape(len(z2), -1) if z2 is not None else None + _z = torch.cat((_z1, _z2), dim=-1) if z2 is not None else _z1 + + # compute median depth overall (ignoring nans) + if quantile == 0.5: + shift_z = torch.nanmedian(_z, dim=-1).values + else: + shift_z = torch.nanquantile(_z, quantile, dim=-1) + return shift_z # (B,) + + +@torch.no_grad() +def get_joint_pointcloud_center_scale(pts1, pts2, valid_mask1=None, valid_mask2=None, z_only=False, center=True): + # set invalid points to NaN + _pts1 = invalid_to_nans(pts1, valid_mask1).reshape(len(pts1), -1, 3) + _pts2 = invalid_to_nans(pts2, valid_mask2).reshape(len(pts2), -1, 3) if pts2 is not None else None + _pts = torch.cat((_pts1, _pts2), dim=1) if pts2 is not None else _pts1 + + # compute median center + _center = torch.nanmedian(_pts, dim=1, keepdim=True).values # (B,1,3) + if z_only: + _center[..., :2] = 0 # do not center X and Y + + # compute median norm + _norm = ((_pts - _center) if center else _pts).norm(dim=-1) + scale = torch.nanmedian(_norm, dim=1).values + return _center[:, None, :, :], scale[:, None, None, None] + + +def find_reciprocal_matches(P1, P2): + """ + returns 3 values: + 1 - reciprocal_in_P2: a boolean array of size P2.shape[0], a "True" value indicates a match + 2 - nn2_in_P1: a int array of size P2.shape[0], it contains the indexes of the closest points in P1 + 3 - reciprocal_in_P2.sum(): the number of matches + """ + tree1 = KDTree(P1) + tree2 = KDTree(P2) + + _, nn1_in_P2 = tree2.query(P1, workers=8) + _, nn2_in_P1 = tree1.query(P2, workers=8) + + reciprocal_in_P1 = (nn2_in_P1[nn1_in_P2] == np.arange(len(nn1_in_P2))) + reciprocal_in_P2 = (nn1_in_P2[nn2_in_P1] == np.arange(len(nn2_in_P1))) + assert reciprocal_in_P1.sum() == reciprocal_in_P2.sum() + return reciprocal_in_P2, nn2_in_P1, reciprocal_in_P2.sum() + + +def get_med_dist_between_poses(poses): + from scipy.spatial.distance import pdist + return np.median(pdist([to_numpy(p[:3, 3]) for p in poses])) diff --git a/src/global_cfg.py b/src/global_cfg.py new file mode 100644 index 0000000000000000000000000000000000000000..fc8571b0adc3dbc3c41ddb0b364698782ff94ced --- /dev/null +++ b/src/global_cfg.py @@ -0,0 +1,19 @@ +from typing import Optional + +from omegaconf import DictConfig + +cfg: Optional[DictConfig] = None + + +def get_cfg() -> DictConfig: + global cfg + return cfg + + +def set_cfg(new_cfg: DictConfig) -> None: + global cfg + cfg = new_cfg + + +def get_seed() -> int: + return cfg.seed diff --git a/src/loss/__init__.py b/src/loss/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d7f511ba3e4aa6d6f381a4bcf5738e9166e2bb5c --- /dev/null +++ b/src/loss/__init__.py @@ -0,0 +1,26 @@ +from .loss import Loss +from .loss_depth import LossDepth, LossDepthCfgWrapper +from .loss_lpips import LossLpips, LossLpipsCfgWrapper +from .loss_mse import LossMse, LossMseCfgWrapper +from .loss_opacity import LossOpacity, LossOpacityCfgWrapper +from .loss_depth_gt import LossDepthGT, LossDepthGTCfgWrapper +from .loss_lod import LossLOD, LossLODCfgWrapper +from .loss_depth_consis import LossDepthConsis, LossDepthConsisCfgWrapper +from .loss_normal_consis import LossNormalConsis, LossNormalConsisCfgWrapper +from .loss_chamfer_distance import LossChamferDistance, LossChamferDistanceCfgWrapper +LOSSES = { + LossDepthCfgWrapper: LossDepth, + LossLpipsCfgWrapper: LossLpips, + LossMseCfgWrapper: LossMse, + LossOpacityCfgWrapper: LossOpacity, + LossDepthGTCfgWrapper: LossDepthGT, + LossLODCfgWrapper: LossLOD, + LossDepthConsisCfgWrapper: LossDepthConsis, + LossNormalConsisCfgWrapper: LossNormalConsis, + LossChamferDistanceCfgWrapper: LossChamferDistance, +} + +LossCfgWrapper = LossDepthCfgWrapper | LossLpipsCfgWrapper | LossMseCfgWrapper | LossOpacityCfgWrapper | LossDepthGTCfgWrapper | LossLODCfgWrapper | LossDepthConsisCfgWrapper | LossNormalConsisCfgWrapper | LossChamferDistanceCfgWrapper + +def get_losses(cfgs: list[LossCfgWrapper]) -> list[Loss]: + return [LOSSES[type(cfg)](cfg) for cfg in cfgs] diff --git a/src/loss/loss.py b/src/loss/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..42690697b6a6134930771f382caa9424eb6c13c8 --- /dev/null +++ b/src/loss/loss.py @@ -0,0 +1,37 @@ +from abc import ABC, abstractmethod +from dataclasses import fields +from typing import Generic, TypeVar + +from jaxtyping import Float +from torch import Tensor, nn + +from src.dataset.types import BatchedExample +from src.model.decoder.decoder import DecoderOutput +from src.model.types import Gaussians + +T_cfg = TypeVar("T_cfg") +T_wrapper = TypeVar("T_wrapper") + + +class Loss(nn.Module, ABC, Generic[T_cfg, T_wrapper]): + cfg: T_cfg + name: str + + def __init__(self, cfg: T_wrapper) -> None: + super().__init__() + + # Extract the configuration from the wrapper. + (field,) = fields(type(cfg)) + self.cfg = getattr(cfg, field.name) + self.name = field.name + + @abstractmethod + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + depth_dict: dict, + global_step: int, + ) -> Float[Tensor, ""]: + pass diff --git a/src/loss/loss_chamfer_distance.py b/src/loss/loss_chamfer_distance.py new file mode 100644 index 0000000000000000000000000000000000000000..17f0bc2e07824ac50e5308e342b3fc626c7d8222 --- /dev/null +++ b/src/loss/loss_chamfer_distance.py @@ -0,0 +1,75 @@ +from dataclasses import dataclass + +import torch +from einops import reduce +from jaxtyping import Float +from torch import Tensor + +from src.dataset.types import BatchedExample +from src.model.decoder.decoder import DecoderOutput +from src.model.types import Gaussians +from .loss import Loss +from typing import Generic, TypeVar +from dataclasses import fields +import torch.nn.functional as F +import sys +from pytorch3d.loss import chamfer_distance +import os +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +# from src.loss.depth_anything.dpt import DepthAnything +from src.misc.utils import vis_depth_map + +T_cfg = TypeVar("T_cfg") +T_wrapper = TypeVar("T_wrapper") + + +@dataclass +class LossChamferDistanceCfg: + weight: float + down_sample_ratio: float + sigma_image: float | None + + +@dataclass +class LossChamferDistanceCfgWrapper: + chamfer_distance: LossChamferDistanceCfg + +class LossChamferDistance(Loss[LossChamferDistanceCfg, LossChamferDistanceCfgWrapper]): + def __init__(self, cfg: T_wrapper) -> None: + super().__init__(cfg) + + # Extract the configuration from the wrapper. + (field,) = fields(type(cfg)) + self.cfg = getattr(cfg, field.name) + self.name = field.name + + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + depth_dict: dict, + global_step: int, + ) -> Float[Tensor, ""]: + # Scale the depth between the near and far planes. + b, v, h, w, _ = depth_dict['distill_infos']['pts_all'].shape + pred_pts = depth_dict['distill_infos']['pts_all'].flatten(0, 1) + + conf_mask = depth_dict['distill_infos']['conf_mask'] + gaussian_meas = gaussians.means + + pred_pts = pred_pts.view(b, v, h, w, -1) + conf_mask = conf_mask.view(b, v, h, w) + + pts_mask = torch.abs(gaussian_meas[..., -1]) < 1e2 # + # conf_mask = conf_mask & pts_mask + + cd_losses = 0.0 + for b_idx in range(b): + batch_pts, batch_conf, batch_gaussian = pred_pts[b_idx], conf_mask[b_idx], gaussian_meas[b_idx][pts_mask[b_idx]] + batch_pts = batch_pts[batch_conf] + batch_pts = batch_pts[torch.randperm(batch_pts.shape[0])[:int(batch_pts.shape[0] * self.cfg.down_sample_ratio)]] + batch_gaussian = batch_gaussian[torch.randperm(batch_gaussian.shape[0])[:int(batch_gaussian.shape[0] * self.cfg.down_sample_ratio)]] + cd_loss = chamfer_distance(batch_pts.unsqueeze(0), batch_gaussian.unsqueeze(0))[0] + cd_losses = cd_losses + cd_loss + return self.cfg.weight * torch.nan_to_num(cd_losses / b, nan=0.0) \ No newline at end of file diff --git a/src/loss/loss_depth.py b/src/loss/loss_depth.py new file mode 100644 index 0000000000000000000000000000000000000000..81929387682d0e65386b3f913ee10d1363c3cf8f --- /dev/null +++ b/src/loss/loss_depth.py @@ -0,0 +1,118 @@ +from dataclasses import dataclass + +import torch +from einops import reduce +from jaxtyping import Float +from torch import Tensor + +from src.dataset.types import BatchedExample +from src.model.decoder.decoder import DecoderOutput +from src.model.types import Gaussians +from .loss import Loss +from typing import Generic, TypeVar +from dataclasses import fields +import torch.nn.functional as F +import sys +import os +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +# from src.loss.depth_anything.dpt import DepthAnything +from src.misc.utils import vis_depth_map + +T_cfg = TypeVar("T_cfg") +T_wrapper = TypeVar("T_wrapper") + + +@dataclass +class LossDepthCfg: + weight: float + sigma_image: float | None + use_second_derivative: bool + + +@dataclass +class LossDepthCfgWrapper: + depth: LossDepthCfg + + +class LossDepth(Loss[LossDepthCfg, LossDepthCfgWrapper]): + def __init__(self, cfg: T_wrapper) -> None: + super().__init__(cfg) + + # Extract the configuration from the wrapper. + (field,) = fields(type(cfg)) + self.cfg = getattr(cfg, field.name) + self.name = field.name + + model_configs = { + 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]}, + 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]}, + 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]} + } + encoder = 'vits' # or 'vitb', 'vits' + depth_anything = DepthAnything(model_configs[encoder]) + depth_anything.load_state_dict(torch.load(f'src/loss/depth_anything/depth_anything_{encoder}14.pth')) + + self.depth_anything = depth_anything + for param in self.depth_anything.parameters(): + param.requires_grad = False + + def disp_rescale(self, disp: Float[Tensor, "B H W"]): + disp = disp.flatten(1, 2) + disp_median = torch.median(disp, dim=-1, keepdim=True)[0] # (B, V, 1) + disp_var = (disp - disp_median).abs().mean(dim=-1, keepdim=True) # (B, V, 1) + disp = (disp - disp_median) / (disp_var + 1e-6) + return disp + + def smooth_l1_loss(self, pred, target, beta=1.0, reduction='none'): + diff = pred - target + abs_diff = torch.abs(diff) + + loss = torch.where(abs_diff < beta, 0.5 * diff ** 2 / beta, abs_diff - 0.5 * beta) + + if reduction == 'mean': + return loss.mean() + elif reduction == 'sum': + return loss.sum() + elif reduction == 'none': + return loss + else: + raise ValueError("Invalid reduction type. Choose from 'mean', 'sum', or 'none'.") + + def ctx_depth_loss(self, + depth_map: Float[Tensor, "B V H W C"], + depth_conf: Float[Tensor, "B V H W"], + batch: BatchedExample, + cxt_depth_weight: float = 0.01, + alpha: float = 0.2): + B, V, _, H, W = batch["context"]["image"].shape + ctx_imgs = batch["context"]["image"].view(B * V, 3, H, W).float() + da_output = self.depth_anything(ctx_imgs) + da_output = self.disp_rescale(da_output) + + disp_context = 1.0 / depth_map.flatten(0, 1).squeeze(-1).clamp(1e-3) # (B * V, H, W) + context_output = self.disp_rescale(disp_context) + + depth_conf = depth_conf.flatten(0, 1).flatten(1, 2) # (B * V) + + return cxt_depth_weight * (self.smooth_l1_loss(context_output*depth_conf, da_output*depth_conf, reduction='none') - alpha * torch.log(depth_conf)).mean() + + + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + global_step: int, + ) -> Float[Tensor, ""]: + # Scale the depth between the near and far planes. + target_imgs = batch["target"]["image"] + B, V, _, H, W = target_imgs.shape + target_imgs = target_imgs.view(B * V, 3, H, W) + da_output = self.depth_anything(target_imgs.float()) + da_output = self.disp_rescale(da_output) + + disp_gs = 1.0 / prediction.depth.flatten(0, 1).clamp(1e-3).float() + gs_output = self.disp_rescale(disp_gs) + + + return self.cfg.weight * torch.nan_to_num(F.smooth_l1_loss(da_output, gs_output), nan=0.0) \ No newline at end of file diff --git a/src/loss/loss_depth_consis.py b/src/loss/loss_depth_consis.py new file mode 100644 index 0000000000000000000000000000000000000000..cbd9eada7dcb547994a44d1e9be2aad6e9f25414 --- /dev/null +++ b/src/loss/loss_depth_consis.py @@ -0,0 +1,140 @@ +from dataclasses import dataclass + +import torch +from einops import reduce +from jaxtyping import Float +from torch import Tensor + +from src.dataset.types import BatchedExample +from src.model.decoder.decoder import DecoderOutput +from src.model.types import Gaussians +from .loss import Loss +from typing import Generic, Literal, Optional, TypeVar +from dataclasses import fields +import torch.nn.functional as F +import sys +import os +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +# from src.loss.depth_anything.dpt import DepthAnything +from src.misc.utils import vis_depth_map + +T_cfg = TypeVar("T_cfg") +T_wrapper = TypeVar("T_wrapper") + + +@dataclass +class LossDepthConsisCfg: + weight: float + sigma_image: float | None + use_second_derivative: bool + loss_type: Literal['MSE', 'EdgeAwareLogL1', 'PearsonDepth'] = 'MSE' + detach: bool = False + conf: bool = False + not_use_valid_mask: bool = False + apply_after_step: int = 0 + +@dataclass +class LossDepthConsisCfgWrapper: + depth_consis: LossDepthConsisCfg + + +class LogL1(torch.nn.Module): + """Log-L1 loss""" + + def __init__( + self, implementation: Literal["scalar", "per-pixel"] = "scalar", **kwargs + ): + super().__init__() + self.implementation = implementation + + def forward(self, pred, gt): + if self.implementation == "scalar": + return torch.log(1 + torch.abs(pred - gt)).mean() + else: + return torch.log(1 + torch.abs(pred - gt)) + +class EdgeAwareLogL1(torch.nn.Module): + """Gradient aware Log-L1 loss""" + + def __init__( + self, implementation: Literal["scalar", "per-pixel"] = "scalar", **kwargs + ): + super().__init__() + self.implementation = implementation + self.logl1 = LogL1(implementation="per-pixel") + + def forward(self, pred: Tensor, gt: Tensor, rgb: Tensor, mask: Optional[Tensor]): + logl1 = self.logl1(pred, gt) + + grad_img_x = torch.mean( + torch.abs(rgb[..., :, :-1, :] - rgb[..., :, 1:, :]), -1, keepdim=True + ) + grad_img_y = torch.mean( + torch.abs(rgb[..., :-1, :, :] - rgb[..., 1:, :, :]), -1, keepdim=True + ) + lambda_x = torch.exp(-grad_img_x) + lambda_y = torch.exp(-grad_img_y) + + loss_x = lambda_x * logl1[..., :, :-1, :] + loss_y = lambda_y * logl1[..., :-1, :, :] + + if self.implementation == "per-pixel": + if mask is not None: + loss_x[~mask[..., :, :-1, :]] = 0 + loss_y[~mask[..., :-1, :, :]] = 0 + return loss_x[..., :-1, :, :] + loss_y[..., :, :-1, :] + + if mask is not None: + assert mask.shape[:2] == pred.shape[:2] + loss_x = loss_x[mask[..., :, :-1, :]] + loss_y = loss_y[mask[..., :-1, :, :]] + + if self.implementation == "scalar": + return loss_x.mean() + loss_y.mean() + +class LossDepthConsis(Loss[LossDepthConsisCfg, LossDepthConsisCfgWrapper]): + def __init__(self, cfg: T_wrapper) -> None: + super().__init__(cfg) + + # Extract the configuration from the wrapper. + (field,) = fields(type(cfg)) + self.cfg = getattr(cfg, field.name) + self.name = field.name + + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + depth_dict: dict, + global_step: int, + ) -> Float[Tensor, ""]: + + # Before the specified step, don't apply the loss. + if global_step < self.cfg.apply_after_step: + return torch.tensor(0.0, dtype=torch.float32, device=prediction.depth.device) + + # Scale the depth between the near and far planes. + # conf_valid_mask = depth_dict['conf_valid_mask'] + rendered_depth = prediction.depth + gt_rgb = (batch["context"]["image"] + 1) / 2 + valid_mask = depth_dict["distill_infos"]['conf_mask'] + + if batch['context']['valid_mask'].sum() > 0: + valid_mask = batch['context']['valid_mask'] + # if self.cfg.conf: + # valid_mask = valid_mask & conf_valid_mask + if self.cfg.not_use_valid_mask: + valid_mask = torch.ones_like(valid_mask, device=valid_mask.device) + pred_depth = depth_dict['depth'].squeeze(-1) + if self.cfg.detach: + pred_depth = pred_depth.detach() + if self.cfg.loss_type == 'MSE': + depth_loss = F.mse_loss(rendered_depth, pred_depth, reduction='none')[valid_mask].mean() + elif self.cfg.loss_type == 'EdgeAwareLogL1': + rendered_depth = rendered_depth.flatten(0, 1).unsqueeze(-1) + pred_depth = pred_depth.flatten(0, 1).unsqueeze(-1) + gt_rgb = gt_rgb.flatten(0, 1).permute(0, 2, 3, 1) + valid_mask = valid_mask.flatten(0, 1).unsqueeze(-1) + depth_loss = EdgeAwareLogL1()(rendered_depth, pred_depth, gt_rgb, valid_mask) + return self.cfg.weight * torch.nan_to_num(depth_loss, nan=0.0) \ No newline at end of file diff --git a/src/loss/loss_depth_gt.py b/src/loss/loss_depth_gt.py new file mode 100644 index 0000000000000000000000000000000000000000..5a4caeae8b1dfe19da9e1f4729096806ca46bc2d --- /dev/null +++ b/src/loss/loss_depth_gt.py @@ -0,0 +1,93 @@ +from dataclasses import dataclass + +import torch +from einops import reduce +from jaxtyping import Float +from torch import Tensor + +from src.dataset.types import BatchedExample +from src.model.decoder.decoder import DecoderOutput +from src.model.types import Gaussians +from .loss import Loss +from typing import Generic, Literal, TypeVar +from dataclasses import fields +import torch.nn.functional as F +import sys +import os +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +# from src.loss.depth_anything.dpt import DepthAnything +from src.misc.utils import vis_depth_map + +T_cfg = TypeVar("T_cfg") +T_wrapper = TypeVar("T_wrapper") + + +@dataclass +class LossDepthGTCfg: + weight: float + type: Literal["l1", "mse", "silog", "gradient", "l1+gradient"] | None + +@dataclass +class LossDepthGTCfgWrapper: + depthgt: LossDepthGTCfg + + +class LossDepthGT(Loss[LossDepthGTCfg, LossDepthGTCfgWrapper]): + def gradient_loss(self, gs_depth, target_depth, target_valid_mask): + diff = gs_depth - target_depth + + grad_x_diff = diff[:, :, :, 1:] - diff[:, :, :, :-1] + grad_y_diff = diff[:, :, 1:, :] - diff[:, :, :-1, :] + + mask_x = target_valid_mask[:, :, :, 1:] * target_valid_mask[:, :, :, :-1] + mask_y = target_valid_mask[:, :, 1:, :] * target_valid_mask[:, :, :-1, :] + + grad_x_diff = grad_x_diff * mask_x + grad_y_diff = grad_y_diff * mask_y + + grad_x_diff = grad_x_diff.clamp(min=-100, max=100) + grad_y_diff = grad_y_diff.clamp(min=-100, max=100) + + loss_x = grad_x_diff.abs().sum() + loss_y = grad_y_diff.abs().sum() + num_valid = mask_x.sum() + mask_y.sum() + + if num_valid == 0: + gradient_loss = 0 + else: + gradient_loss = (loss_x + loss_y) / (num_valid + 1e-6) + + return gradient_loss + + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + global_step: int, + ) -> Float[Tensor, ""]: + # Scale the depth between the near and far planes. + + # prediction: B, H, W, C + # target: B, H, W, C + # mask: B, H, W + + target_depth = batch["target"]["depth"] + target_valid_mask = batch["target"]["valid_mask"] + gs_depth = prediction.depth.clamp(1e-3) + + if self.cfg.type == "l1": + depth_loss = torch.abs(target_depth[target_valid_mask] - gs_depth[target_valid_mask]).mean() + elif self.cfg.type == "mse": + depth_loss = F.mse_loss(target_depth[target_valid_mask], gs_depth[target_valid_mask]) + elif self.cfg.type == "silog": + depth_loss = torch.log(gs_depth[target_valid_mask]) ** 2 + (gs_depth[target_valid_mask] - target_depth[target_valid_mask]) ** 2 - 0.5 + depth_loss = depth_loss.mean() + elif self.cfg.type == "gradient": + depth_loss = self.gradient_loss(gs_depth, target_depth, target_valid_mask) + elif self.cfg.type == "l1+gradient": + depth_loss_l1 = torch.abs(target_depth[target_valid_mask] - gs_depth[target_valid_mask]).mean() + depth_loss_gradient = self.gradient_loss(gs_depth, target_depth, target_valid_mask) + depth_loss = depth_loss_l1 + depth_loss_gradient + + return self.cfg.weight * torch.nan_to_num(depth_loss, nan=0.0, posinf=0.0, neginf=0.0) \ No newline at end of file diff --git a/src/loss/loss_distill.py b/src/loss/loss_distill.py new file mode 100644 index 0000000000000000000000000000000000000000..3c32da2af139f5dad0758cec0d73bad64e9831d1 --- /dev/null +++ b/src/loss/loss_distill.py @@ -0,0 +1,151 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from copy import copy, deepcopy + +from src.model.encoder.vggt.utils.pose_enc import pose_encoding_to_extri_intri +from src.model.encoder.vggt.utils.rotation import mat_to_quat +from src.utils.point import get_normal_map + +def extri_intri_to_pose_encoding( + extrinsics, + intrinsics, + image_size_hw=None, # e.g., (256, 512) + pose_encoding_type="absT_quaR_FoV", +): + """Convert camera extrinsics and intrinsics to a compact pose encoding. + + This function transforms camera parameters into a unified pose encoding format, + which can be used for various downstream tasks like pose prediction or representation. + + Args: + extrinsics (torch.Tensor): Camera extrinsic parameters with shape BxSx3x4, + where B is batch size and S is sequence length. + In OpenCV coordinate system (x-right, y-down, z-forward), representing camera from world transformation. + The format is [R|t] where R is a 3x3 rotation matrix and t is a 3x1 translation vector. + intrinsics (torch.Tensor): Camera intrinsic parameters with shape BxSx3x3. + Defined in pixels, with format: + [[fx, 0, cx], + [0, fy, cy], + [0, 0, 1]] + where fx, fy are focal lengths and (cx, cy) is the principal point + image_size_hw (tuple): Tuple of (height, width) of the image in pixels. + Required for computing field of view values. For example: (256, 512). + pose_encoding_type (str): Type of pose encoding to use. Currently only + supports "absT_quaR_FoV" (absolute translation, quaternion rotation, field of view). + + Returns: + torch.Tensor: Encoded camera pose parameters with shape BxSx9. + For "absT_quaR_FoV" type, the 9 dimensions are: + - [:3] = absolute translation vector T (3D) + - [3:7] = rotation as quaternion quat (4D) + - [7:] = field of view (2D) + """ + + # extrinsics: BxSx3x4 + # intrinsics: BxSx3x3 + + if pose_encoding_type == "absT_quaR_FoV": + R = extrinsics[:, :, :3, :3] # BxSx3x3 + T = extrinsics[:, :, :3, 3] # BxSx3 + + quat = mat_to_quat(R) + # Note the order of h and w here + # H, W = image_size_hw + # fov_h = 2 * torch.atan((H / 2) / intrinsics[..., 1, 1]) + # fov_w = 2 * torch.atan((W / 2) / intrinsics[..., 0, 0]) + fov_h = 2 * torch.atan(0.5 / intrinsics[..., 1, 1]) + fov_w = 2 * torch.atan(0.5 / intrinsics[..., 0, 0]) + pose_encoding = torch.cat([T, quat, fov_h[..., None], fov_w[..., None]], dim=-1).float() + else: + raise NotImplementedError + + return pose_encoding + +def huber_loss(x, y, delta=1.0): + """Calculate element-wise Huber loss between x and y""" + diff = x - y + abs_diff = diff.abs() + flag = (abs_diff <= delta).to(diff.dtype) + return flag * 0.5 * diff**2 + (1 - flag) * delta * (abs_diff - 0.5 * delta) + +class DistillLoss(nn.Module): + def __init__(self, delta=1.0, gamma=0.6, weight_pose=1.0, weight_depth=1.0, weight_normal=1.0): + super().__init__() + self.delta = delta + self.gamma = gamma + self.weight_pose = weight_pose + self.weight_depth = weight_depth + self.weight_normal = weight_normal + + def camera_loss_single(self, cur_pred_pose_enc, gt_pose_encoding, loss_type="l1"): + if loss_type == "l1": + loss_T = (cur_pred_pose_enc[..., :3] - gt_pose_encoding[..., :3]).abs() + loss_R = (cur_pred_pose_enc[..., 3:7] - gt_pose_encoding[..., 3:7]).abs() + loss_fl = (cur_pred_pose_enc[..., 7:] - gt_pose_encoding[..., 7:]).abs() + elif loss_type == "l2": + loss_T = (cur_pred_pose_enc[..., :3] - gt_pose_encoding[..., :3]).norm(dim=-1, keepdim=True) + loss_R = (cur_pred_pose_enc[..., 3:7] - gt_pose_encoding[..., 3:7]).norm(dim=-1) + loss_fl = (cur_pred_pose_enc[..., 7:] - gt_pose_encoding[..., 7:]).norm(dim=-1) + elif loss_type == "huber": + loss_T = huber_loss(cur_pred_pose_enc[..., :3], gt_pose_encoding[..., :3]) + loss_R = huber_loss(cur_pred_pose_enc[..., 3:7], gt_pose_encoding[..., 3:7]) + loss_fl = huber_loss(cur_pred_pose_enc[..., 7:], gt_pose_encoding[..., 7:]) + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + loss_T = torch.nan_to_num(loss_T, nan=0.0, posinf=0.0, neginf=0.0) + loss_R = torch.nan_to_num(loss_R, nan=0.0, posinf=0.0, neginf=0.0) + loss_fl = torch.nan_to_num(loss_fl, nan=0.0, posinf=0.0, neginf=0.0) + + loss_T = torch.clamp(loss_T, min=-100, max=100) + loss_R = torch.clamp(loss_R, min=-100, max=100) + loss_fl = torch.clamp(loss_fl, min=-100, max=100) + + loss_T = loss_T.mean() + loss_R = loss_R.mean() + loss_fl = loss_fl.mean() + + return loss_T, loss_R, loss_fl + + def forward(self, distill_infos, pred_pose_enc_list, prediction, batch): + loss_pose = 0.0 + + if pred_pose_enc_list is not None: + num_predictions = len(pred_pose_enc_list) + pesudo_gt_pose_enc = distill_infos['pred_pose_enc_list'] + for i in range(num_predictions): + i_weight = self.gamma ** (num_predictions - i - 1) + cur_pred_pose_enc = pred_pose_enc_list[i] + cur_pesudo_gt_pose_enc = pesudo_gt_pose_enc[i] + loss_pose += i_weight * huber_loss(cur_pred_pose_enc, cur_pesudo_gt_pose_enc).mean() + loss_pose = loss_pose / num_predictions + loss_pose = torch.nan_to_num(loss_pose, nan=0.0, posinf=0.0, neginf=0.0) + + pred_depth = prediction.depth.flatten(0, 1) + pesudo_gt_depth = distill_infos['depth_map'].flatten(0, 1).squeeze(-1) + conf_mask = distill_infos['conf_mask'].flatten(0, 1) + + if batch['context']['valid_mask'].sum() > 0: + conf_mask = batch['context']['valid_mask'].flatten(0, 1) + + loss_depth = F.mse_loss(pred_depth[conf_mask], pesudo_gt_depth[conf_mask], reduction='none').mean() + + render_normal = get_normal_map(pred_depth, batch["context"]["intrinsics"].flatten(0, 1)) + pred_normal = get_normal_map(pesudo_gt_depth, batch["context"]["intrinsics"].flatten(0, 1)) + + alpha1_loss = (1 - (render_normal[conf_mask] * pred_normal[conf_mask]).sum(-1)).mean() + alpha2_loss = F.l1_loss(render_normal[conf_mask], pred_normal[conf_mask], reduction='mean') + loss_normal = (alpha1_loss + alpha2_loss) / 2 + + loss_distill = loss_pose * self.weight_pose + loss_depth * self.weight_depth + loss_normal * self.weight_normal + loss_distill = torch.nan_to_num(loss_distill, nan=0.0, posinf=0.0, neginf=0.0) + + loss_dict = { + "loss_distill": loss_distill, + "loss_pose": loss_pose * self.weight_pose, + "loss_depth": loss_depth * self.weight_depth, + "loss_normal": loss_normal * self.weight_normal + } + + return loss_dict diff --git a/src/loss/loss_huber.py b/src/loss/loss_huber.py new file mode 100644 index 0000000000000000000000000000000000000000..3bb69c3b10e7b89b2edcdd562a17a7a4124de589 --- /dev/null +++ b/src/loss/loss_huber.py @@ -0,0 +1,171 @@ +import torch +import torch.nn as nn +from copy import copy, deepcopy + +from src.model.encoder.vggt.utils.pose_enc import pose_encoding_to_extri_intri +from src.model.encoder.vggt.utils.rotation import mat_to_quat + +def extri_intri_to_pose_encoding( + extrinsics, + intrinsics, + image_size_hw=None, # e.g., (256, 512) + pose_encoding_type="absT_quaR_FoV", +): + """Convert camera extrinsics and intrinsics to a compact pose encoding. + + This function transforms camera parameters into a unified pose encoding format, + which can be used for various downstream tasks like pose prediction or representation. + + Args: + extrinsics (torch.Tensor): Camera extrinsic parameters with shape BxSx3x4, + where B is batch size and S is sequence length. + In OpenCV coordinate system (x-right, y-down, z-forward), representing camera from world transformation. + The format is [R|t] where R is a 3x3 rotation matrix and t is a 3x1 translation vector. + intrinsics (torch.Tensor): Camera intrinsic parameters with shape BxSx3x3. + Defined in pixels, with format: + [[fx, 0, cx], + [0, fy, cy], + [0, 0, 1]] + where fx, fy are focal lengths and (cx, cy) is the principal point + image_size_hw (tuple): Tuple of (height, width) of the image in pixels. + Required for computing field of view values. For example: (256, 512). + pose_encoding_type (str): Type of pose encoding to use. Currently only + supports "absT_quaR_FoV" (absolute translation, quaternion rotation, field of view). + + Returns: + torch.Tensor: Encoded camera pose parameters with shape BxSx9. + For "absT_quaR_FoV" type, the 9 dimensions are: + - [:3] = absolute translation vector T (3D) + - [3:7] = rotation as quaternion quat (4D) + - [7:] = field of view (2D) + """ + + # extrinsics: BxSx3x4 + # intrinsics: BxSx3x3 + + if pose_encoding_type == "absT_quaR_FoV": + R = extrinsics[:, :, :3, :3] # BxSx3x3 + T = extrinsics[:, :, :3, 3] # BxSx3 + + quat = mat_to_quat(R) + # Note the order of h and w here + # H, W = image_size_hw + # fov_h = 2 * torch.atan((H / 2) / intrinsics[..., 1, 1]) + # fov_w = 2 * torch.atan((W / 2) / intrinsics[..., 0, 0]) + fov_h = 2 * torch.atan(0.5 / intrinsics[..., 1, 1]) + fov_w = 2 * torch.atan(0.5 / intrinsics[..., 0, 0]) + pose_encoding = torch.cat([T, quat, fov_h[..., None], fov_w[..., None]], dim=-1).float() + else: + raise NotImplementedError + + return pose_encoding + +def huber_loss(x, y, delta=1.0): + """Calculate element-wise Huber loss between x and y""" + diff = x - y + abs_diff = diff.abs() + flag = (abs_diff <= delta).to(diff.dtype) + return flag * 0.5 * diff**2 + (1 - flag) * delta * (abs_diff - 0.5 * delta) + +class HuberLoss(nn.Module): + def __init__(self, alpha=1.0, delta=1.0, gamma=0.6, weight_T=1.0, weight_R=1.0, weight_fl=0.5): + super().__init__() + self.alpha = alpha + self.delta = delta + self.gamma = gamma + self.weight_T = weight_T + self.weight_R = weight_R + self.weight_fl = weight_fl + + def camera_loss_single(self, cur_pred_pose_enc, gt_pose_encoding, loss_type="l1"): + if loss_type == "l1": + loss_T = (cur_pred_pose_enc[..., :3] - gt_pose_encoding[..., :3]).abs() + loss_R = (cur_pred_pose_enc[..., 3:7] - gt_pose_encoding[..., 3:7]).abs() + loss_fl = (cur_pred_pose_enc[..., 7:] - gt_pose_encoding[..., 7:]).abs() + elif loss_type == "l2": + loss_T = (cur_pred_pose_enc[..., :3] - gt_pose_encoding[..., :3]).norm(dim=-1, keepdim=True) + loss_R = (cur_pred_pose_enc[..., 3:7] - gt_pose_encoding[..., 3:7]).norm(dim=-1) + loss_fl = (cur_pred_pose_enc[..., 7:] - gt_pose_encoding[..., 7:]).norm(dim=-1) + elif loss_type == "huber": + loss_T = huber_loss(cur_pred_pose_enc[..., :3], gt_pose_encoding[..., :3]) + loss_R = huber_loss(cur_pred_pose_enc[..., 3:7], gt_pose_encoding[..., 3:7]) + loss_fl = huber_loss(cur_pred_pose_enc[..., 7:], gt_pose_encoding[..., 7:]) + else: + raise ValueError(f"Unknown loss type: {loss_type}") + + loss_T = torch.nan_to_num(loss_T, nan=0.0, posinf=0.0, neginf=0.0) + loss_R = torch.nan_to_num(loss_R, nan=0.0, posinf=0.0, neginf=0.0) + loss_fl = torch.nan_to_num(loss_fl, nan=0.0, posinf=0.0, neginf=0.0) + + loss_T = torch.clamp(loss_T, min=-100, max=100) + loss_R = torch.clamp(loss_R, min=-100, max=100) + loss_fl = torch.clamp(loss_fl, min=-100, max=100) + + loss_T = loss_T.mean() + loss_R = loss_R.mean() + loss_fl = loss_fl.mean() + + return loss_T, loss_R, loss_fl + + def forward(self, pred_pose_enc_list, batch): + context_extrinsics = batch["context"]["extrinsics"] + context_intrinsics = batch["context"]["intrinsics"] + image_size_hw = batch["context"]["image"].shape[-2:] + + # transform extrinsics and intrinsics to pose_enc + GT_pose_enc = extri_intri_to_pose_encoding(context_extrinsics, context_intrinsics, image_size_hw) + num_predictions = len(pred_pose_enc_list) + loss_T = loss_R = loss_fl = 0 + + for i in range(num_predictions): + i_weight = self.gamma ** (num_predictions - i - 1) + + cur_pred_pose_enc = pred_pose_enc_list[i] + + loss_T_i, loss_R_i, loss_fl_i = self.camera_loss_single(cur_pred_pose_enc.clone(), GT_pose_enc.clone(), loss_type="huber") + loss_T += i_weight * loss_T_i + loss_R += i_weight * loss_R_i + loss_fl += i_weight * loss_fl_i + + loss_T = loss_T / num_predictions + loss_R = loss_R / num_predictions + loss_fl = loss_fl / num_predictions + loss_camera = loss_T * self.weight_T + loss_R * self.weight_R + loss_fl * self.weight_fl + + loss_dict = { + "loss_camera": loss_camera, + "loss_T": loss_T, + "loss_R": loss_R, + "loss_fl": loss_fl + } + + # with torch.no_grad(): + # # compute auc + # last_pred_pose_enc = pred_pose_enc_list[-1] + + # last_pred_extrinsic, _ = pose_encoding_to_extri_intri(last_pred_pose_enc.detach(), image_size_hw, pose_encoding_type='absT_quaR_FoV', build_intrinsics=False) + + # rel_rangle_deg, rel_tangle_deg = camera_to_rel_deg(last_pred_extrinsic.float(), context_extrinsics.float(), context_extrinsics.device) + + # if rel_rangle_deg.numel() == 0 and rel_tangle_deg.numel() == 0: + # rel_rangle_deg = torch.FloatTensor([0]).to(context_extrinsics.device).to(context_extrinsics.dtype) + # rel_tangle_deg = torch.FloatTensor([0]).to(context_extrinsics.device).to(context_extrinsics.dtype) + + # thresholds = [5, 15] + # for threshold in thresholds: + # loss_dict[f"Rac_{threshold}"] = (rel_rangle_deg < threshold).float().mean() + # loss_dict[f"Tac_{threshold}"] = (rel_tangle_deg < threshold).float().mean() + + # _, normalized_histogram = calculate_auc( + # rel_rangle_deg, rel_tangle_deg, max_threshold=30, return_list=True + # ) + + # auc_thresholds = [30, 10, 5, 3] + # for auc_threshold in auc_thresholds: + # cur_auc = torch.cumsum( + # normalized_histogram[:auc_threshold], dim=0 + # ).mean() + # loss_dict[f"Auc_{auc_threshold}"] = cur_auc + + return loss_dict + diff --git a/src/loss/loss_lod.py b/src/loss/loss_lod.py new file mode 100644 index 0000000000000000000000000000000000000000..2af7af4cd768c9e2b828c2ba206d51c9d50e305e --- /dev/null +++ b/src/loss/loss_lod.py @@ -0,0 +1,64 @@ +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from einops import rearrange +from jaxtyping import Float +from lpips import LPIPS +from torch import Tensor + +from src.dataset.types import BatchedExample +from src.misc.nn_module_tools import convert_to_buffer +from src.model.decoder.decoder import DecoderOutput +from src.model.types import Gaussians +from .loss import Loss + + +@dataclass +class LossLODCfg: + mse_weight: float + lpips_weight: float + +@dataclass +class LossLODCfgWrapper: + lod: LossLODCfg + +WEIGHT_LEVEL_MAPPING = {0: 0.1, 1: 0.1, 2: 0.2, 3: 0.6} + +class LossLOD(Loss[LossLODCfg, LossLODCfgWrapper]): + lpips: LPIPS + + def __init__(self, cfg: LossLODCfgWrapper) -> None: + super().__init__(cfg) + + self.lpips = LPIPS(net="vgg") + convert_to_buffer(self.lpips, persistent=False) + + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + global_step: int, + ) -> Float[Tensor, ""]: + image = batch["target"]["image"] + # breakpoint() + def mse_loss(x, y): + delta = x - y + return torch.nan_to_num((delta**2).mean().mean(), nan=0.0, posinf=0.0, neginf=0.0) + # Before the specified step, don't apply the loss. + lod_rendering = prediction.lod_rendering + loss = 0.0 + for level in lod_rendering.keys(): + # level_weight + # breakpoint() + # if level != 3: + # continue + rendered_imgs = lod_rendering[level]['rendered_imgs'].flatten(0, 1) + _h, _w = rendered_imgs.shape[2:] + resized_image = F.interpolate(image.clone().flatten(0, 1), size=(_h, _w), mode='bilinear', align_corners=False) + level_mse_loss = mse_loss(rendered_imgs, resized_image) + level_lpips_loss = self.lpips.forward(rendered_imgs, resized_image, normalize=True).mean() + + loss += (torch.nan_to_num(level_mse_loss, nan=0.0, posinf=0.0, neginf=0.0) * self.cfg.mse_weight + torch.nan_to_num(level_lpips_loss, nan=0.0, posinf=0.0, neginf=0.0) * self.cfg.lpips_weight) + return loss / len(lod_rendering.keys()) diff --git a/src/loss/loss_lpips.py b/src/loss/loss_lpips.py new file mode 100644 index 0000000000000000000000000000000000000000..13411fa94c3b40d9875c1f48e97d69929d070766 --- /dev/null +++ b/src/loss/loss_lpips.py @@ -0,0 +1,76 @@ +from dataclasses import dataclass + +import torch +from einops import rearrange +from jaxtyping import Float +from lpips import LPIPS +from torch import Tensor + +from src.dataset.types import BatchedExample +from src.misc.nn_module_tools import convert_to_buffer +from src.model.decoder.decoder import DecoderOutput +from src.model.types import Gaussians +from .loss import Loss + + +@dataclass +class LossLpipsCfg: + weight: float + apply_after_step: int + conf: bool = False + alpha: bool = False + mask: bool = False + + +@dataclass +class LossLpipsCfgWrapper: + lpips: LossLpipsCfg + + +class LossLpips(Loss[LossLpipsCfg, LossLpipsCfgWrapper]): + lpips: LPIPS + + def __init__(self, cfg: LossLpipsCfgWrapper) -> None: + super().__init__(cfg) + + self.lpips = LPIPS(net="vgg") + convert_to_buffer(self.lpips, persistent=False) + + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + depth_dict: dict | None, + global_step: int, + ) -> Float[Tensor, ""]: + image = (batch["context"]["image"] + 1) / 2 + + # Before the specified step, don't apply the loss. + if global_step < self.cfg.apply_after_step: + return torch.tensor(0, dtype=torch.float32, device=image.device) + + if self.cfg.mask or self.cfg.alpha or self.cfg.conf: + if self.cfg.mask: + mask = batch["context"]["valid_mask"] + elif self.cfg.alpha: + mask = prediction.alpha + elif self.cfg.conf: + mask = depth_dict['conf_valid_mask'] + b, v, c, h, w = prediction.color.shape + expanded_mask = mask.unsqueeze(2).expand(-1, -1, c, -1, -1) + masked_pred = prediction.color * expanded_mask + masked_img = image * expanded_mask + + loss = self.lpips.forward( + rearrange(masked_pred, "b v c h w -> (b v) c h w"), + rearrange(masked_img, "b v c h w -> (b v) c h w"), + normalize=True, + ) + else: + loss = self.lpips.forward( + rearrange(prediction.color, "b v c h w -> (b v) c h w"), + rearrange(image, "b v c h w -> (b v) c h w"), + normalize=True, + ) + return self.cfg.weight * torch.nan_to_num(loss.mean(), nan=0.0, posinf=0.0, neginf=0.0) diff --git a/src/loss/loss_mse.py b/src/loss/loss_mse.py new file mode 100644 index 0000000000000000000000000000000000000000..c72cb8735b0f88ff560085066d3eb7f86152ca5f --- /dev/null +++ b/src/loss/loss_mse.py @@ -0,0 +1,59 @@ +from dataclasses import dataclass + +from jaxtyping import Float +from torch import Tensor +import torch +from src.dataset.types import BatchedExample +from src.model.decoder.decoder import DecoderOutput +from src.model.types import Gaussians +from .loss import Loss + + +@dataclass +class LossMseCfg: + weight: float + conf: bool = False + mask: bool = False + alpha: bool = False + + +@dataclass +class LossMseCfgWrapper: + mse: LossMseCfg + + +class LossMse(Loss[LossMseCfg, LossMseCfgWrapper]): + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + depth_dict: dict | None, + global_step: int, + ) -> Float[Tensor, ""]: + # Get alpha and valid mask from inputs + alpha = prediction.alpha + # valid_mask = torch.ones_like(alpha, device=alpha.device).bool() + valid_mask = batch['context']['valid_mask'] + + # # only for objaverse + # if batch['context']['valid_mask'].sum() > 0: + # valid_mask = batch['context']['valid_mask'] + + # Determine which mask to use based on config + if self.cfg.mask: + mask = valid_mask + elif self.cfg.alpha: + mask = alpha + elif self.cfg.conf: + mask = depth_dict['conf_valid_mask'] + else: + mask = torch.ones_like(alpha, device=alpha.device).bool() + + # Rearrange and mask predicted and ground truth images + pred_img = prediction.color.permute(0, 1, 3, 4, 2)[mask] + gt_img = ((batch["context"]["image"][:, batch["using_index"]] + 1) / 2).permute(0, 1, 3, 4, 2)[mask] + + delta = pred_img - gt_img + + return self.cfg.weight * torch.nan_to_num((delta**2).mean(), nan=0.0, posinf=0.0, neginf=0.0) diff --git a/src/loss/loss_normal_consis.py b/src/loss/loss_normal_consis.py new file mode 100644 index 0000000000000000000000000000000000000000..3b8933bbbc1291fc7fdcc53256078805038f8b34 --- /dev/null +++ b/src/loss/loss_normal_consis.py @@ -0,0 +1,134 @@ +from dataclasses import dataclass + +import torch +from einops import reduce +from jaxtyping import Float +from torch import Tensor + +from src.dataset.types import BatchedExample +from src.model.decoder.decoder import DecoderOutput +from src.model.types import Gaussians +from .loss import Loss +from typing import Generic, TypeVar +from dataclasses import fields +import torch.nn.functional as F +import sys +import os +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +# from src.loss.depth_anything.dpt import DepthAnything +from src.misc.utils import vis_depth_map +import open3d as o3d +T_cfg = TypeVar("T_cfg") +T_wrapper = TypeVar("T_wrapper") + +@dataclass +class LossNormalConsisCfg: + normal_weight: float + smooth_weight: float + sigma_image: float | None + use_second_derivative: bool + detach: bool = False + conf: bool = False + not_use_valid_mask: bool = False + +@dataclass +class LossNormalConsisCfgWrapper: + normal_consis: LossNormalConsisCfg + +class TVLoss(torch.nn.Module): + """TV loss""" + + def __init__(self): + super().__init__() + + def forward(self, pred): + """ + Args: + pred: [batch, H, W, 3] + + Returns: + tv_loss: [batch] + """ + h_diff = pred[..., :, :-1, :] - pred[..., :, 1:, :] + w_diff = pred[..., :-1, :, :] - pred[..., 1:, :, :] + return torch.mean(torch.abs(h_diff)) + torch.mean(torch.abs(w_diff)) + + +class LossNormalConsis(Loss[LossNormalConsisCfg, LossNormalConsisCfgWrapper]): + def __init__(self, cfg: T_wrapper) -> None: + super().__init__(cfg) + + # Extract the configuration from the wrapper. + (field,) = fields(type(cfg)) + self.cfg = getattr(cfg, field.name) + self.name = field.name + + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + depth_dict: dict, + global_step: int, + ) -> Float[Tensor, ""]: + # Scale the depth between the near and far planes. + conf_valid_mask = depth_dict['conf_valid_mask'].flatten(0, 1) + valid_mask = batch["context"]["valid_mask"][:, batch["using_index"]].flatten(0, 1) + if self.cfg.conf: + valid_mask = valid_mask & conf_valid_mask + if self.cfg.not_use_valid_mask: + valid_mask = torch.ones_like(valid_mask, device=valid_mask.device) + render_normal = self.get_normal_map(prediction.depth.flatten(0, 1), batch["context"]["intrinsics"].flatten(0, 1)) + pred_normal = self.get_normal_map(depth_dict['depth'].flatten(0, 1).squeeze(-1), batch["context"]["intrinsics"].flatten(0, 1)) + if self.cfg.detach: + pred_normal = pred_normal.detach() + alpha1_loss = (1 - (render_normal * pred_normal).sum(-1)).mean() + alpha2_loss = F.l1_loss(render_normal, pred_normal, reduction='mean') + normal_smooth_loss = TVLoss()(render_normal) + normal_loss = (alpha1_loss + alpha2_loss) / 2 + return self.cfg.normal_weight * torch.nan_to_num(normal_loss, nan=0.0) + self.cfg.smooth_weight * torch.nan_to_num(normal_smooth_loss, nan=0.0) + + def get_normal_map(self, depth_map: torch.Tensor, intrinsic: torch.Tensor) -> torch.Tensor: + """ + Convert a depth map to camera coordinates. + + Args: + depth_map (torch.Tensor): Depth map of shape (H, W). + intrinsic (torch.Tensor): Camera intrinsic matrix of shape (3, 3). + + Returns: + tuple[torch.Tensor, torch.Tensor]: Camera coordinates (H, W, 3) + """ + B, H, W = depth_map.shape + assert intrinsic.shape == (B, 3, 3), "Intrinsic matrix must be Bx3x3" + assert (intrinsic[:, 0, 1] == 0).all() and (intrinsic[:, 1, 0] == 0).all(), "Intrinsic matrix must have zero skew" + + # Intrinsic parameters + fu = intrinsic[:, 0, 0] * W # (B,) + fv = intrinsic[:, 1, 1] * H # (B,) + cu = intrinsic[:, 0, 2] * W # (B,) + cv = intrinsic[:, 1, 2] * H # (B,) + + # Generate grid of pixel coordinates + u = torch.arange(W, device=depth_map.device)[None, None, :].expand(B, H, W) + v = torch.arange(H, device=depth_map.device)[None, :, None].expand(B, H, W) + + # Unproject to camera coordinates (B, H, W) + x_cam = (u - cu[:, None, None]) * depth_map / fu[:, None, None] + y_cam = (v - cv[:, None, None]) * depth_map / fv[:, None, None] + z_cam = depth_map + + # Stack to form camera coordinates (B, H, W, 3) + cam_coords = torch.stack((x_cam, y_cam, z_cam), dim=-1).to(dtype=torch.float32) + + output = torch.zeros_like(cam_coords) + # Calculate dx using batch dimension (B, H-2, W-2, 3) + dx = cam_coords[:, 2:, 1:-1] - cam_coords[:, :-2, 1:-1] + # Calculate dy using batch dimension (B, H-2, W-2, 3) + dy = cam_coords[:, 1:-1, 2:] - cam_coords[:, 1:-1, :-2] + # Cross product and normalization (B, H-2, W-2, 3) + normal_map = torch.nn.functional.normalize(torch.cross(dx, dy, dim=-1), dim=-1) + # Assign the computed normal map to the output tensor + output[:, 1:-1, 1:-1, :] = normal_map + + return output \ No newline at end of file diff --git a/src/loss/loss_opacity.py b/src/loss/loss_opacity.py new file mode 100644 index 0000000000000000000000000000000000000000..0e3aa78a5091b03784d803e1ca2fea986721e5a5 --- /dev/null +++ b/src/loss/loss_opacity.py @@ -0,0 +1,43 @@ +from dataclasses import dataclass +from typing import Literal + +from jaxtyping import Float +from torch import Tensor +import torch +import torch.nn.functional as F +from src.dataset.types import BatchedExample +from src.model.decoder.decoder import DecoderOutput +from src.model.types import Gaussians +from .loss import Loss + + +@dataclass +class LossOpacityCfg: + weight: float + type: Literal["exp", "mean", "exp+mean"] = "exp+mean" + + +@dataclass +class LossOpacityCfgWrapper: + opacity: LossOpacityCfg + + +class LossOpacity(Loss[LossOpacityCfg, LossOpacityCfgWrapper]): + def forward( + self, + prediction: DecoderOutput, + batch: BatchedExample, + gaussians: Gaussians, + depth_dict: dict | None, + global_step: int, + ) -> Float[Tensor, ""]: + alpha = prediction.alpha + valid_mask = batch['context']['valid_mask'].float() + opacity_loss = F.mse_loss(alpha, valid_mask, reduction='none').mean() + # if self.cfg.type == "exp": + # opacity_loss = torch.exp(-(gaussians.opacities - 0.5) ** 2 / 0.05).mean() + # elif self.cfg.type == "mean": + # opacity_loss = gaussians.opacities.mean() + # elif self.cfg.type == "exp+mean": + # opacity_loss = 0.5 * torch.exp(-(gaussians.opacities - 0.5) ** 2 / 0.05).mean() + gaussians.opacities.mean() + return self.cfg.weight * torch.nan_to_num(opacity_loss, nan=0.0, posinf=0.0, neginf=0.0) diff --git a/src/loss/loss_point.py b/src/loss/loss_point.py new file mode 100644 index 0000000000000000000000000000000000000000..d0c6792c86913381c7197e39733299ebd9205078 --- /dev/null +++ b/src/loss/loss_point.py @@ -0,0 +1,358 @@ +import torch +import torch.nn as nn +from copy import copy, deepcopy + +from src.geometry.ptc_geometry import geotrf, inv, normalize_pointcloud, depthmap_to_pts3d +# from torchmetrics.functional.regression import pearson_corrcoef +# from pytorch3d.loss import chamfer_distance + + +def get_pred_pts3d(gt, pred, use_pose=False): + if 'depth' in pred and 'pseudo_focal' in pred: + try: + pp = gt['camera_intrinsics'][..., :2, 2] + except KeyError: + pp = None + pts3d = depthmap_to_pts3d(**pred, pp=pp) + + elif 'pts3d' in pred: + # pts3d from my camera + pts3d = pred['pts3d'] + + elif 'pts3d_in_other_view' in pred: + # pts3d from the other camera, already transformed + assert use_pose is True + return pred['pts3d_in_other_view'] # return! + + if use_pose: + camera_pose = pred.get('camera_pose') + assert camera_pose is not None + pts3d = geotrf(camera_pose, pts3d) + + return pts3d + + +class LLoss (nn.Module): + """ L-norm loss + """ + + def __init__(self, reduction='mean'): + super().__init__() + self.reduction = reduction + + def forward(self, a, b): + assert a.shape == b.shape and a.ndim >= 2 and 1 <= a.shape[-1] <= 3, f'Bad shape = {a.shape}' + dist = self.distance(a, b) + assert dist.ndim == a.ndim-1 # one dimension less + if self.reduction == 'none': + return dist + if self.reduction == 'sum': + return dist.sum() + if self.reduction == 'mean': + return dist.mean() if dist.numel() > 0 else dist.new_zeros(()) + raise ValueError(f'bad {self.reduction=} mode') + + def distance(self, a, b): + raise NotImplementedError() + + +class L21Loss (LLoss): + """ Euclidean distance between 3d points """ + + def distance(self, a, b): + return torch.norm(a - b, dim=-1) # normalized L2 distance + + +class MultiLoss (nn.Module): + """ Easily combinable losses (also keep track of individual loss values): + loss = MyLoss1() + 0.1*MyLoss2() + Usage: + Inherit from this class and override get_name() and compute_loss() + """ + + def __init__(self): + super().__init__() + self._alpha = 1 + self._loss2 = None + + def compute_loss(self, *args, **kwargs): + raise NotImplementedError() + + def get_name(self): + raise NotImplementedError() + + def __mul__(self, alpha): + assert isinstance(alpha, (int, float)) + res = copy(self) + res._alpha = alpha + return res + __rmul__ = __mul__ # same + + def __add__(self, loss2): + assert isinstance(loss2, MultiLoss) + res = cur = copy(self) + # find the end of the chain + while cur._loss2 is not None: + cur = cur._loss2 + cur._loss2 = loss2 + return res + + def __repr__(self): + name = self.get_name() + if self._alpha != 1: + name = f'{self._alpha:g}*{name}' + if self._loss2: + name = f'{name} + {self._loss2}' + return name + + def forward(self, *args, **kwargs): + loss = self.compute_loss(*args, **kwargs) + if isinstance(loss, tuple): + loss, details = loss + elif loss.ndim == 0: + details = {self.get_name(): float(loss)} + else: + details = {} + loss = loss * self._alpha + + if self._loss2: + loss2, details2 = self._loss2(*args, **kwargs) + loss = loss + loss2 + details |= details2 + + return loss, details + + +class Criterion (nn.Module): + def __init__(self, criterion=None): + super().__init__() + assert isinstance(criterion, LLoss), f'{criterion} is not a proper criterion!'+bb() + self.criterion = copy(criterion) + + def get_name(self): + return f'{type(self).__name__}({self.criterion})' + + def with_reduction(self, mode): + res = loss = deepcopy(self) + while loss is not None: + assert isinstance(loss, Criterion) + loss.criterion.reduction = 'none' # make it return the loss for each sample + loss = loss._loss2 # we assume loss is a Multiloss + return res + + +class ConfLoss (MultiLoss): + """ Weighted regression by learned confidence. + Assuming the input pixel_loss is a pixel-level regression loss. + + Principle: + high-confidence means high conf = 0.1 ==> conf_loss = x / 10 + alpha*log(10) + low confidence means low conf = 10 ==> conf_loss = x * 10 - alpha*log(10) + + alpha: hyperparameter + """ + + def __init__(self, pixel_loss, alpha=1): + super().__init__() + assert alpha > 0 + self.alpha = alpha + self.pixel_loss = pixel_loss.with_reduction('none') + + def get_name(self): + return f'ConfLoss({self.pixel_loss})' + + def get_conf_log(self, x): + return x, torch.log(x) + + def compute_loss(self, gt1, gt2, pred1, pred2, **kw): + # compute per-pixel loss + ((loss1, msk1), (loss2, msk2)), details = self.pixel_loss(gt1, gt2, pred1, pred2, **kw) + if loss1.numel() == 0: + print('NO VALID POINTS in img1', force=True) + if loss2.numel() == 0: + print('NO VALID POINTS in img2', force=True) + + # weight by confidence + conf1, log_conf1 = self.get_conf_log(pred1['conf'][msk1]) + conf2, log_conf2 = self.get_conf_log(pred2['conf'][msk2]) + conf_loss1 = loss1 * conf1 - self.alpha * log_conf1 + conf_loss2 = loss2 * conf2 - self.alpha * log_conf2 + + # average + nan protection (in case of no valid pixels at all) + conf_loss1 = conf_loss1.mean() if conf_loss1.numel() > 0 else 0 + conf_loss2 = conf_loss2.mean() if conf_loss2.numel() > 0 else 0 + + return conf_loss1 + conf_loss2, dict(conf_loss_1=float(conf_loss1), conf_loss2=float(conf_loss2), **details) + + +class Regr3D(nn.Module): + """ Ensure that all 3D points are correct. + Asymmetric loss: view1 is supposed to be the anchor. + + P1 = RT1 @ D1 + P2 = RT2 @ D2 + loss1 = (I @ pred_D1) - (RT1^-1 @ RT1 @ D1) + loss2 = (RT21 @ pred_D2) - (RT1^-1 @ P2) + = (RT21 @ pred_D2) - (RT1^-1 @ RT2 @ D2) + """ + + def __init__(self, norm_mode='avg_dis', alpha=0.2, gt_scale=False): + super().__init__() + self.norm_mode = norm_mode + self.alpha = alpha + self.gt_scale = gt_scale + + def get_conf_log(self, x): + return x, torch.log(x) + + def forward(self, gt_pts1, gt_pts2, pr_pts1, pr_pts2, conf1=None, conf2=None, dist_clip=None, disable_view1=False): + valid1 = valid2 = torch.ones_like(conf1, dtype=torch.bool) + if dist_clip is not None: + # points that are too far-away == invalid + dis1 = gt_pts1.norm(dim=-1) # (B, H, W) + dis2 = gt_pts2.norm(dim=-1) # (B, H, W) + valid1 = (dis1 <= dist_clip) + valid2 = (dis2 <= dist_clip) + else: + dis1 = gt_pts1.norm(dim=-1) # (B, H, W) + dis2 = gt_pts2.norm(dim=-1) # (B, H, W) + + # only keep the points norm whithin the range of 1% to 99% of each batch + # Flatten along the H and W dimensions + dis1_flat = dis1.view(dis1.shape[0], -1) + dis2_flat = dis2.view(dis2.shape[0], -1) + + # Compute the 0.1% and 99.9% quantiles for each batch + # quantiles_1 = torch.quantile(dis1_flat, torch.tensor([0.01, 0.99]).to(dis1_flat.device), dim=1) + # quantiles_2 = torch.quantile(dis2_flat, torch.tensor([0.01, 0.99]).to(dis2_flat.device), dim=1) + quantiles_1 = torch.quantile(dis1_flat, torch.tensor([0.002, 0.998]).to(dis1_flat.device), dim=1) + quantiles_2 = torch.quantile(dis2_flat, torch.tensor([0.002, 0.998]).to(dis2_flat.device), dim=1) + + # Create masks based on the quantiles + valid1 = (dis1 >= quantiles_1[0].view(-1, 1, 1)) & (dis1 <= quantiles_1[1].view(-1, 1, 1)) + valid2 = (dis2 >= quantiles_2[0].view(-1, 1, 1)) & (dis2 <= quantiles_2[1].view(-1, 1, 1)) + + # set min confidence to 3 + valid1 = valid1 & (conf1 >= 3) + valid2 = valid2 & (conf2 >= 3) + + # normalize 3d points + if self.norm_mode: + pr_pts1, pr_pts2 = normalize_pointcloud(pr_pts1, pr_pts2, self.norm_mode, valid1, valid2) + if self.norm_mode and not self.gt_scale: + gt_pts1, gt_pts2 = normalize_pointcloud(gt_pts1, gt_pts2, self.norm_mode, valid1, valid2) + + loss1 = torch.norm(pr_pts1 - gt_pts1, dim=-1) + loss2 = torch.norm(pr_pts2 - gt_pts2, dim=-1) + # loss1 = (pr_pts1[..., -1] - gt_pts1[..., -1]).abs() + # loss2 = (pr_pts2[..., -1] - gt_pts2[..., -1]).abs() + + loss1, loss2 = loss1[valid1], loss2[valid2] + + if disable_view1: + return loss2.mean() + return loss1.mean() + loss2.mean() + + # conf1, conf2 = conf1[valid1], conf2[valid2] + # conf1, conf2 = conf1.softmax(dim=-1), conf2.softmax(dim=-1) + # loss1 = (loss1 * conf1).sum() + # loss2 = (loss2 * conf2).sum() + # return loss1 + loss2 + # + # # weight by confidence + # conf1, log_conf1 = self.get_conf_log(conf1[valid1]) + # conf2, log_conf2 = self.get_conf_log(conf2[valid2]) + # conf_loss1 = loss1 * conf1 - self.alpha * log_conf1 + # conf_loss2 = loss2 * conf2 - self.alpha * log_conf2 + # + # # average + nan protection (in case of no valid pixels at all) + # conf_loss1 = conf_loss1.mean() if conf_loss1.numel() > 0 else 0 + # conf_loss2 = conf_loss2.mean() if conf_loss2.numel() > 0 else 0 + # + # return conf_loss1 + conf_loss2 + + # def forward(self, gt_pts1, gt_pts2, pr_pts1, pr_pts2, conf1=None, conf2=None, dist_clip=None, disable_view1=False): + # # valid1 = valid2 = torch.ones_like(conf1, dtype=torch.bool) + # if dist_clip is not None: + # # points that are too far-away == invalid + # dis1 = gt_pts1.norm(dim=-1) # (B, H, W) + # dis2 = gt_pts2.norm(dim=-1) # (B, H, W) + # valid1 = (dis1 <= dist_clip) + # valid2 = (dis2 <= dist_clip) + # else: + # dis1 = gt_pts1.norm(dim=-1) # (B, H, W) + # dis2 = gt_pts2.norm(dim=-1) # (B, H, W) + # + # # only keep the points norm whithin the range of 1% to 99% of each batch + # # Flatten along the H and W dimensions + # dis1_flat = dis1.view(dis1.shape[0], -1) + # dis2_flat = dis2.view(dis2.shape[0], -1) + # + # # Compute the 0.1% and 99.9% quantiles for each batch + # quantiles_1 = torch.quantile(dis1_flat, torch.tensor([0.1, 0.9]).to(dis1_flat.device), dim=1) + # quantiles_2 = torch.quantile(dis2_flat, torch.tensor([0.1, 0.9]).to(dis2_flat.device), dim=1) + # # quantiles_1 = torch.quantile(dis1_flat, torch.tensor([0.002, 0.998]).to(dis1_flat.device), dim=1) + # # quantiles_2 = torch.quantile(dis2_flat, torch.tensor([0.002, 0.998]).to(dis2_flat.device), dim=1) + # + # # Create masks based on the quantiles + # valid1 = (dis1 >= quantiles_1[0].view(-1, 1, 1)) & (dis1 <= quantiles_1[1].view(-1, 1, 1)) + # valid2 = (dis2 >= quantiles_2[0].view(-1, 1, 1)) & (dis2 <= quantiles_2[1].view(-1, 1, 1)) + # + # # set min opacity to 3 + # valid1 = valid1 & (conf1 >= 0.2) + # valid2 = valid2 & (conf2 >= 0.2) + # + # # # normalize 3d points + # # if self.norm_mode: + # # pr_pts1, pr_pts2 = normalize_pointcloud(pr_pts1, pr_pts2, self.norm_mode, valid1, valid2) + # # if self.norm_mode and not self.gt_scale: + # # gt_pts1, gt_pts2 = normalize_pointcloud(gt_pts1, gt_pts2, self.norm_mode, valid1, valid2) + # + # # L1 loss + # # loss1 = (pr_pts1[..., -1] - gt_pts1[..., -1]).abs() + # # loss2 = (pr_pts2[..., -1] - gt_pts2[..., -1]).abs() + # + # # L2 loss + # loss1 = torch.norm(pr_pts1 - gt_pts1, dim=-1) + # loss2 = torch.norm(pr_pts2 - gt_pts2, dim=-1) + # loss1, loss2 = loss1[valid1], loss2[valid2] + # + # # Pearson correlation coefficient loss + # # pr_pts1, pr_pts2 = pr_pts1[valid1], pr_pts2[valid2] + # # gt_pts1, gt_pts2 = gt_pts1[valid1], gt_pts2[valid2] + # # loss1 = 1 - pearson_corrcoef(pr_pts1.view(-1, 3), gt_pts1.view(-1, 3)) + # # loss2 = 1 - pearson_corrcoef(pr_pts2.view(-1, 3), gt_pts2.view(-1, 3)) + # + # # # Chamfer distance loss + # # pr_pts = torch.cat([pr_pts1.flatten(1, 2), pr_pts2.flatten(1, 2)], dim=1) + # # gt_pts = torch.cat([gt_pts1.flatten(1, 2), gt_pts2.flatten(1, 2)], dim=1) + # # valid_mask = torch.cat([valid1.flatten(1, 2), valid2.flatten(1, 2)], dim=1) + # # nan_pts_pr, nnz = invalid_to_zeros(pr_pts, valid_mask, ndim=3) + # # nan_pts_gt, nnz = invalid_to_zeros(gt_pts, valid_mask, ndim=3) + # # + # # loss, _ = chamfer_distance(nan_pts_pr, nan_pts_gt, batch_reduction=None, point_reduction=None) + # # loss1, loss2 = loss[0], loss[1] + # # return loss1.sum() / valid_mask.sum() + # + # if disable_view1: + # return loss2.mean() + # return loss1.mean() + loss2.mean() + # + # # conf1, conf2 = conf1[valid1], conf2[valid2] + # # conf1, conf2 = conf1.softmax(dim=-1), conf2.softmax(dim=-1) + # # loss1 = (loss1 * conf1).sum() + # # loss2 = (loss2 * conf2).sum() + # # return loss1 + loss2 + # # + # # # weight by confidence + # # conf1, log_conf1 = self.get_conf_log(conf1[valid1]) + # # conf2, log_conf2 = self.get_conf_log(conf2[valid2]) + # # conf_loss1 = loss1 * conf1 - self.alpha * log_conf1 + # # conf_loss2 = loss2 * conf2 - self.alpha * log_conf2 + # # + # # # average + nan protection (in case of no valid pixels at all) + # # conf_loss1 = conf_loss1.mean() if conf_loss1.numel() > 0 else 0 + # # conf_loss2 = conf_loss2.mean() if conf_loss2.numel() > 0 else 0 + # # + # # return conf_loss1 + conf_loss2 + diff --git a/src/loss/loss_ssim.py b/src/loss/loss_ssim.py new file mode 100644 index 0000000000000000000000000000000000000000..7f9e475b346a5d69bc9b2aeca56ae4a783c8747f --- /dev/null +++ b/src/loss/loss_ssim.py @@ -0,0 +1,357 @@ +# Copyright 2020 by Gongfan Fang, Zhejiang University. +# All rights reserved. +# Modified by Botao Ye from https://github.com/VainF/pytorch-msssim/blob/master/pytorch_msssim/ssim.py. +import warnings +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import Tensor + + +def _fspecial_gauss_1d(size: int, sigma: float) -> Tensor: + r"""Create 1-D gauss kernel + Args: + size (int): the size of gauss kernel + sigma (float): sigma of normal distribution + Returns: + torch.Tensor: 1D kernel (1 x 1 x size) + """ + coords = torch.arange(size, dtype=torch.float) + coords -= size // 2 + + g = torch.exp(-(coords ** 2) / (2 * sigma ** 2)) + g /= g.sum() + + return g.unsqueeze(0).unsqueeze(0) + + +def gaussian_filter(input: Tensor, win: Tensor) -> Tensor: + r""" Blur input with 1-D kernel + Args: + input (torch.Tensor): a batch of tensors to be blurred + window (torch.Tensor): 1-D gauss kernel + Returns: + torch.Tensor: blurred tensors + """ + assert all([ws == 1 for ws in win.shape[1:-1]]), win.shape + if len(input.shape) == 4: + conv = F.conv2d + elif len(input.shape) == 5: + conv = F.conv3d + else: + raise NotImplementedError(input.shape) + + C = input.shape[1] + out = input + for i, s in enumerate(input.shape[2:]): + if s >= win.shape[-1]: + out = conv(out, weight=win.transpose(2 + i, -1), stride=1, padding=0, groups=C) + else: + warnings.warn( + f"Skipping Gaussian Smoothing at dimension 2+{i} for input: {input.shape} and win size: {win.shape[-1]}" + ) + + return out + + +def _ssim( + X: Tensor, + Y: Tensor, + data_range: float, + win: Tensor, + size_average: bool = True, + K: Union[Tuple[float, float], List[float]] = (0.01, 0.03), + retrun_seprate: bool = False, +) -> Tuple[Tensor, Tensor, Tensor | None, Tensor | None, Tensor | None]: + r""" Calculate ssim index for X and Y + + Args: + X (torch.Tensor): images + Y (torch.Tensor): images + data_range (float or int): value range of input images. (usually 1.0 or 255) + win (torch.Tensor): 1-D gauss kernel + size_average (bool, optional): if size_average=True, ssim of all images will be averaged as a scalar + retrun_seprate (bool, optional): if True, return brightness, contrast, and structure similarity maps as well + + Returns: + Tuple[torch.Tensor, torch.Tensor]: ssim results. + """ + K1, K2 = K + # batch, channel, [depth,] height, width = X.shape + compensation = 1.0 + + C1 = (K1 * data_range) ** 2 + C2 = (K2 * data_range) ** 2 + + win = win.to(X.device, dtype=X.dtype) + + mu1 = gaussian_filter(X, win) + mu2 = gaussian_filter(Y, win) + + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + + sigma1_sq = compensation * (gaussian_filter(X * X, win) - mu1_sq) + sigma2_sq = compensation * (gaussian_filter(Y * Y, win) - mu2_sq) + sigma12 = compensation * (gaussian_filter(X * Y, win) - mu1_mu2) + + cs_map = (2 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2) # set alpha=beta=gamma=1 + ssim_map = ((2 * mu1_mu2 + C1) / (mu1_sq + mu2_sq + C1)) * cs_map + ssim_per_channel = torch.flatten(ssim_map, 2).mean(-1) + cs = torch.flatten(cs_map, 2).mean(-1) + + brightness = contrast = structure = torch.zeros_like(ssim_per_channel) + if retrun_seprate: + epsilon = torch.finfo(torch.float32).eps**2 + sigma1_sq = sigma1_sq.clamp(min=epsilon) + sigma2_sq = sigma2_sq.clamp(min=epsilon) + sigma12 = torch.sign(sigma12) * torch.minimum( + torch.sqrt(sigma1_sq * sigma2_sq), torch.abs(sigma12)) + + C3 = C2 / 2 + sigma1_sigma2 = torch.sqrt(sigma1_sq) * torch.sqrt(sigma2_sq) + brightness_map = (2 * mu1_mu2 + C1) / (mu1_sq + mu2_sq + C1) + contrast_map = (2 * sigma1_sigma2 + C2) / (sigma1_sq + sigma2_sq + C2) + structure_map = (sigma12 + C3) / (sigma1_sigma2 + C3) + + contrast_map = contrast_map.clamp(max=0.98) + structure_map = structure_map.clamp(max=0.98) + + brightness = brightness_map.flatten(2).mean(-1) + contrast = contrast_map.flatten(2).mean(-1) + structure = structure_map.flatten(2).mean(-1) + + return ssim_per_channel, cs, brightness, contrast, structure + + +def ssim( + X: Tensor, + Y: Tensor, + data_range: float = 255, + size_average: bool = True, + win_size: int = 11, + win_sigma: float = 1.5, + win: Optional[Tensor] = None, + K: Union[Tuple[float, float], List[float]] = (0.01, 0.03), + nonnegative_ssim: bool = False, + retrun_seprate: bool = False, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + r""" interface of ssim + Args: + X (torch.Tensor): a batch of images, (N,C,H,W) + Y (torch.Tensor): a batch of images, (N,C,H,W) + data_range (float or int, optional): value range of input images. (usually 1.0 or 255) + size_average (bool, optional): if size_average=True, ssim of all images will be averaged as a scalar + win_size: (int, optional): the size of gauss kernel + win_sigma: (float, optional): sigma of normal distribution + win (torch.Tensor, optional): 1-D gauss kernel. if None, a new kernel will be created according to win_size and win_sigma + K (list or tuple, optional): scalar constants (K1, K2). Try a larger K2 constant (e.g. 0.4) if you get a negative or NaN results. + nonnegative_ssim (bool, optional): force the ssim response to be nonnegative with relu + retrun_seprate (bool, optional): if True, return brightness, contrast, and structure similarity maps as well + + Returns: + torch.Tensor: ssim results + """ + if not X.shape == Y.shape: + raise ValueError(f"Input images should have the same dimensions, but got {X.shape} and {Y.shape}.") + + for d in range(len(X.shape) - 1, 1, -1): + X = X.squeeze(dim=d) + Y = Y.squeeze(dim=d) + + if len(X.shape) not in (4, 5): + raise ValueError(f"Input images should be 4-d or 5-d tensors, but got {X.shape}") + + #if not X.type() == Y.type(): + # raise ValueError(f"Input images should have the same dtype, but got {X.type()} and {Y.type()}.") + + if win is not None: # set win_size + win_size = win.shape[-1] + + if not (win_size % 2 == 1): + raise ValueError("Window size should be odd.") + + if win is None: + win = _fspecial_gauss_1d(win_size, win_sigma) + win = win.repeat([X.shape[1]] + [1] * (len(X.shape) - 1)) + + ssim_per_channel, cs, brightness, contrast, structure \ + = _ssim(X, Y, data_range=data_range, win=win, size_average=False, K=K, retrun_seprate=retrun_seprate) + + if nonnegative_ssim: + ssim_per_channel = torch.relu(ssim_per_channel) + + if size_average: + return ssim_per_channel.mean(), brightness.mean(), contrast.mean(), structure.mean() + else: + return ssim_per_channel.mean(1), brightness.mean(1), contrast.mean(1), structure.mean(1) + + +def ms_ssim( + X: Tensor, + Y: Tensor, + data_range: float = 255, + size_average: bool = True, + win_size: int = 11, + win_sigma: float = 1.5, + win: Optional[Tensor] = None, + weights: Optional[List[float]] = None, + K: Union[Tuple[float, float], List[float]] = (0.01, 0.03) +) -> Tensor: + r""" interface of ms-ssim + Args: + X (torch.Tensor): a batch of images, (N,C,[T,]H,W) + Y (torch.Tensor): a batch of images, (N,C,[T,]H,W) + data_range (float or int, optional): value range of input images. (usually 1.0 or 255) + size_average (bool, optional): if size_average=True, ssim of all images will be averaged as a scalar + win_size: (int, optional): the size of gauss kernel + win_sigma: (float, optional): sigma of normal distribution + win (torch.Tensor, optional): 1-D gauss kernel. if None, a new kernel will be created according to win_size and win_sigma + weights (list, optional): weights for different levels + K (list or tuple, optional): scalar constants (K1, K2). Try a larger K2 constant (e.g. 0.4) if you get a negative or NaN results. + Returns: + torch.Tensor: ms-ssim results + """ + if not X.shape == Y.shape: + raise ValueError(f"Input images should have the same dimensions, but got {X.shape} and {Y.shape}.") + + for d in range(len(X.shape) - 1, 1, -1): + X = X.squeeze(dim=d) + Y = Y.squeeze(dim=d) + + #if not X.type() == Y.type(): + # raise ValueError(f"Input images should have the same dtype, but got {X.type()} and {Y.type()}.") + + if len(X.shape) == 4: + avg_pool = F.avg_pool2d + elif len(X.shape) == 5: + avg_pool = F.avg_pool3d + else: + raise ValueError(f"Input images should be 4-d or 5-d tensors, but got {X.shape}") + + if win is not None: # set win_size + win_size = win.shape[-1] + + if not (win_size % 2 == 1): + raise ValueError("Window size should be odd.") + + smaller_side = min(X.shape[-2:]) + assert smaller_side > (win_size - 1) * ( + 2 ** 4 + ), "Image size should be larger than %d due to the 4 downsamplings in ms-ssim" % ((win_size - 1) * (2 ** 4)) + + if weights is None: + weights = [0.0448, 0.2856, 0.3001, 0.2363, 0.1333] + weights_tensor = X.new_tensor(weights) + + if win is None: + win = _fspecial_gauss_1d(win_size, win_sigma) + win = win.repeat([X.shape[1]] + [1] * (len(X.shape) - 1)) + + levels = weights_tensor.shape[0] + mcs = [] + for i in range(levels): + ssim_per_channel, cs = _ssim(X, Y, win=win, data_range=data_range, size_average=False, K=K) + + if i < levels - 1: + mcs.append(torch.relu(cs)) + padding = [s % 2 for s in X.shape[2:]] + X = avg_pool(X, kernel_size=2, padding=padding) + Y = avg_pool(Y, kernel_size=2, padding=padding) + + ssim_per_channel = torch.relu(ssim_per_channel) # type: ignore # (batch, channel) + mcs_and_ssim = torch.stack(mcs + [ssim_per_channel], dim=0) # (level, batch, channel) + ms_ssim_val = torch.prod(mcs_and_ssim ** weights_tensor.view(-1, 1, 1), dim=0) + + if size_average: + return ms_ssim_val.mean() + else: + return ms_ssim_val.mean(1) + + +class SSIM(torch.nn.Module): + def __init__( + self, + data_range: float = 255, + size_average: bool = True, + win_size: int = 11, + win_sigma: float = 1.5, + channel: int = 3, + spatial_dims: int = 2, + K: Union[Tuple[float, float], List[float]] = (0.01, 0.03), + nonnegative_ssim: bool = False, + ) -> None: + r""" class for ssim + Args: + data_range (float or int, optional): value range of input images. (usually 1.0 or 255) + size_average (bool, optional): if size_average=True, ssim of all images will be averaged as a scalar + win_size: (int, optional): the size of gauss kernel + win_sigma: (float, optional): sigma of normal distribution + channel (int, optional): input channels (default: 3) + K (list or tuple, optional): scalar constants (K1, K2). Try a larger K2 constant (e.g. 0.4) if you get a negative or NaN results. + nonnegative_ssim (bool, optional): force the ssim response to be nonnegative with relu. + """ + + super(SSIM, self).__init__() + self.win_size = win_size + self.win = _fspecial_gauss_1d(win_size, win_sigma).repeat([channel, 1] + [1] * spatial_dims) + self.size_average = size_average + self.data_range = data_range + self.K = K + self.nonnegative_ssim = nonnegative_ssim + + def forward(self, X: Tensor, Y: Tensor) -> Tensor: + return ssim( + X, + Y, + data_range=self.data_range, + size_average=self.size_average, + win=self.win, + K=self.K, + nonnegative_ssim=self.nonnegative_ssim, + ) + + +class MS_SSIM(torch.nn.Module): + def __init__( + self, + data_range: float = 255, + size_average: bool = True, + win_size: int = 11, + win_sigma: float = 1.5, + channel: int = 3, + spatial_dims: int = 2, + weights: Optional[List[float]] = None, + K: Union[Tuple[float, float], List[float]] = (0.01, 0.03), + ) -> None: + r""" class for ms-ssim + Args: + data_range (float or int, optional): value range of input images. (usually 1.0 or 255) + size_average (bool, optional): if size_average=True, ssim of all images will be averaged as a scalar + win_size: (int, optional): the size of gauss kernel + win_sigma: (float, optional): sigma of normal distribution + channel (int, optional): input channels (default: 3) + weights (list, optional): weights for different levels + K (list or tuple, optional): scalar constants (K1, K2). Try a larger K2 constant (e.g. 0.4) if you get a negative or NaN results. + """ + + super(MS_SSIM, self).__init__() + self.win_size = win_size + self.win = _fspecial_gauss_1d(win_size, win_sigma).repeat([channel, 1] + [1] * spatial_dims) + self.size_average = size_average + self.data_range = data_range + self.weights = weights + self.K = K + + def forward(self, X: Tensor, Y: Tensor) -> Tensor: + return ms_ssim( + X, + Y, + data_range=self.data_range, + size_average=self.size_average, + win=self.win, + weights=self.weights, + K=self.K, + ) diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000000000000000000000000000000000000..975fe992020ca5dd379a18c34f3f016bdc802625 --- /dev/null +++ b/src/main.py @@ -0,0 +1,158 @@ +import os +from pathlib import Path + +import hydra +import torch +import wandb +import random +from colorama import Fore +from jaxtyping import install_import_hook +from lightning.pytorch import Trainer +from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint +from lightning.pytorch.loggers.wandb import WandbLogger +from lightning.pytorch.plugins.environments import SLURMEnvironment +from lightning.pytorch.strategies import DeepSpeedStrategy +from omegaconf import DictConfig, OmegaConf +from hydra.core.hydra_config import HydraConfig + +import sys +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from src.model.model import get_model +from src.misc.weight_modify import checkpoint_filter_fn + +import warnings +warnings.filterwarnings("ignore") + +# Configure beartype and jaxtyping. +with install_import_hook( + ("src",), + ("beartype", "beartype"), +): + from src.config import load_typed_root_config + from src.dataset.data_module import DataModule + from src.global_cfg import set_cfg + from src.loss import get_losses + from src.misc.LocalLogger import LocalLogger + from src.misc.step_tracker import StepTracker + from src.misc.wandb_tools import update_checkpoint_path + from src.model.decoder import get_decoder + from src.model.encoder import get_encoder + from src.model.model_wrapper import ModelWrapper + + +def cyan(text: str) -> str: + return f"{Fore.CYAN}{text}{Fore.RESET}" + + +@hydra.main( + version_base=None, + config_path="../config", + config_name="main", +) +def train(cfg_dict: DictConfig): + cfg = load_typed_root_config(cfg_dict) + set_cfg(cfg_dict) + + # Set up the output directory. + output_dir = Path( + hydra.core.hydra_config.HydraConfig.get()["runtime"]["output_dir"] + ) + output_dir.mkdir(parents=True, exist_ok=True) + print(cyan(f"Saving outputs to {output_dir}.")) + + cfg.train.output_path = output_dir + + # Set up logging with wandb. + callbacks = [] + if cfg_dict.wandb.mode != "disabled": + logger = WandbLogger( + project=cfg_dict.wandb.project, + mode=cfg_dict.wandb.mode, + name=f"{cfg_dict.wandb.name} ({output_dir.parent.name}/{output_dir.name})", + tags=cfg_dict.wandb.get("tags", None), + log_model=False, + save_dir=output_dir, + config=OmegaConf.to_container(cfg_dict), + ) + callbacks.append(LearningRateMonitor("step", True)) + + # On rank != 0, wandb.run is None. + if wandb.run is not None: + wandb.run.log_code("src") + else: + logger = LocalLogger() + + # Set up checkpointing. + callbacks.append( + ModelCheckpoint( + output_dir / "checkpoints", + every_n_train_steps=cfg.checkpointing.every_n_train_steps, + save_top_k=cfg.checkpointing.save_top_k, + save_weights_only=cfg.checkpointing.save_weights_only, + monitor="info/global_step", + mode="max", + ) + ) + callbacks[-1].CHECKPOINT_EQUALS_CHAR = '_' + + # Prepare the checkpoint for loading. + checkpoint_path = update_checkpoint_path(cfg.checkpointing.load, cfg.wandb) + + # This allows the current step to be shared with the data loader processes. + step_tracker = StepTracker() + + trainer = Trainer( + max_epochs=-1, + num_nodes=cfg.trainer.num_nodes, + # num_sanity_val_steps=0, + accelerator="gpu", + logger=logger, + devices="auto", + strategy=( + "ddp_find_unused_parameters_true" + if torch.cuda.device_count() > 1 + else "auto" + ), + # strategy="deepspeed_stage_1", + callbacks=callbacks, + val_check_interval=cfg.trainer.val_check_interval, + check_val_every_n_epoch=None, + enable_progress_bar=False, + gradient_clip_val=cfg.trainer.gradient_clip_val, + max_steps=cfg.trainer.max_steps, + precision=cfg.trainer.precision, + accumulate_grad_batches=cfg.trainer.accumulate_grad_batches, + # plugins=[SLURMEnvironment(requeue_signal=signal.SIGUSR1)], # Uncomment for SLURM auto resubmission. + inference_mode=False if (cfg.mode == "test" and cfg.test.align_pose) else True, + ) + torch.manual_seed(cfg_dict.seed + trainer.global_rank) + + model = get_model(cfg.model.encoder, cfg.model.decoder) + + model_wrapper = ModelWrapper( + cfg.optimizer, + cfg.test, + cfg.train, + model, + get_losses(cfg.loss), + step_tracker + ) + data_module = DataModule( + cfg.dataset, + cfg.data_loader, + step_tracker, + global_rank=trainer.global_rank, + ) + + if cfg.mode == "train": + trainer.fit(model_wrapper, datamodule=data_module, ckpt_path=checkpoint_path) + else: + trainer.test( + model_wrapper, + datamodule=data_module, + ckpt_path=checkpoint_path, + ) + + +if __name__ == "__main__": + train() diff --git a/src/misc/LocalLogger.py b/src/misc/LocalLogger.py new file mode 100644 index 0000000000000000000000000000000000000000..8761a7aa04c24f10a9ef40e66e7d5cbb0a2591b8 --- /dev/null +++ b/src/misc/LocalLogger.py @@ -0,0 +1,48 @@ +import os +from pathlib import Path +from typing import Any, Optional + +from lightning.pytorch.loggers.logger import Logger +from lightning.pytorch.utilities import rank_zero_only +from PIL import Image + +LOG_PATH = Path("outputs/local") + + +class LocalLogger(Logger): + def __init__(self) -> None: + super().__init__() + self.experiment = None + os.system(f"rm -r {LOG_PATH}") + + @property + def name(self): + return "LocalLogger" + + @property + def version(self): + return 0 + + @rank_zero_only + def log_hyperparams(self, params): + pass + + @rank_zero_only + def log_metrics(self, metrics, step): + pass + + @rank_zero_only + def log_image( + self, + key: str, + images: list[Any], + step: Optional[int] = None, + **kwargs, + ): + # The function signature is the same as the wandb logger's, but the step is + # actually required. + assert step is not None + for index, image in enumerate(images): + path = LOG_PATH / f"{key}/{index:0>2}_{step:0>6}.jpg" + path.parent.mkdir(exist_ok=True, parents=True) + Image.fromarray(image).save(path) diff --git a/src/misc/benchmarker.py b/src/misc/benchmarker.py new file mode 100644 index 0000000000000000000000000000000000000000..3cad5bcd413d08de482b58046187e4f887f3a970 --- /dev/null +++ b/src/misc/benchmarker.py @@ -0,0 +1,37 @@ +import json +from collections import defaultdict +from contextlib import contextmanager +from pathlib import Path +from time import time + +import numpy as np +import torch + + +class Benchmarker: + def __init__(self): + self.execution_times = defaultdict(list) + + @contextmanager + def time(self, tag: str, num_calls: int = 1): + try: + start_time = time() + yield + finally: + end_time = time() + for _ in range(num_calls): + self.execution_times[tag].append((end_time - start_time) / num_calls) + + def dump(self, path: Path) -> None: + path.parent.mkdir(exist_ok=True, parents=True) + with path.open("w") as f: + json.dump(dict(self.execution_times), f) + + def dump_memory(self, path: Path) -> None: + path.parent.mkdir(exist_ok=True, parents=True) + with path.open("w") as f: + json.dump(torch.cuda.memory_stats()["allocated_bytes.all.peak"], f) + + def summarize(self) -> None: + for tag, times in self.execution_times.items(): + print(f"{tag}: {len(times)} calls, avg. {np.mean(times)} seconds per call") diff --git a/src/misc/cam_utils.py b/src/misc/cam_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e194236f38e381a3b813643105b38c5f494cff5d --- /dev/null +++ b/src/misc/cam_utils.py @@ -0,0 +1,218 @@ +import cv2 +import numpy as np +import torch +from jaxtyping import Float +from torch import Tensor +import torch.nn.functional as F + + +def decompose_extrinsic_RT(E: torch.Tensor): + """ + Decompose the standard extrinsic matrix into RT. + Batched I/O. + """ + return E[:, :3, :] + + +def compose_extrinsic_RT(RT: torch.Tensor): + """ + Compose the standard form extrinsic matrix from RT. + Batched I/O. + """ + return torch.cat([ + RT, + torch.tensor([[[0, 0, 0, 1]]], dtype=RT.dtype, device=RT.device).repeat(RT.shape[0], 1, 1) + ], dim=1) + + +def camera_normalization(pivotal_pose: torch.Tensor, poses: torch.Tensor): + # [1, 4, 4], [N, 4, 4] + + canonical_camera_extrinsics = torch.tensor([[ + [1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1], + ]], dtype=torch.float32, device=pivotal_pose.device) + pivotal_pose_inv = torch.inverse(pivotal_pose) + camera_norm_matrix = torch.bmm(canonical_camera_extrinsics, pivotal_pose_inv) + + # normalize all views + poses = torch.bmm(camera_norm_matrix.repeat(poses.shape[0], 1, 1), poses) + + return poses + + +####### Pose update from delta + +def rt2mat(R, T): + mat = np.eye(4) + mat[0:3, 0:3] = R + mat[0:3, 3] = T + return mat + + +def skew_sym_mat(x): + device = x.device + dtype = x.dtype + ssm = torch.zeros(3, 3, device=device, dtype=dtype) + ssm[0, 1] = -x[2] + ssm[0, 2] = x[1] + ssm[1, 0] = x[2] + ssm[1, 2] = -x[0] + ssm[2, 0] = -x[1] + ssm[2, 1] = x[0] + return ssm + + +def SO3_exp(theta): + device = theta.device + dtype = theta.dtype + + W = skew_sym_mat(theta) + W2 = W @ W + angle = torch.norm(theta) + I = torch.eye(3, device=device, dtype=dtype) + if angle < 1e-5: + return I + W + 0.5 * W2 + else: + return ( + I + + (torch.sin(angle) / angle) * W + + ((1 - torch.cos(angle)) / (angle**2)) * W2 + ) + + +def V(theta): + dtype = theta.dtype + device = theta.device + I = torch.eye(3, device=device, dtype=dtype) + W = skew_sym_mat(theta) + W2 = W @ W + angle = torch.norm(theta) + if angle < 1e-5: + V = I + 0.5 * W + (1.0 / 6.0) * W2 + else: + V = ( + I + + W * ((1.0 - torch.cos(angle)) / (angle**2)) + + W2 * ((angle - torch.sin(angle)) / (angle**3)) + ) + return V + + +def SE3_exp(tau): + dtype = tau.dtype + device = tau.device + + rho = tau[:3] + theta = tau[3:] + R = SO3_exp(theta) + t = V(theta) @ rho + + T = torch.eye(4, device=device, dtype=dtype) + T[:3, :3] = R + T[:3, 3] = t + return T + + +def update_pose(cam_trans_delta: Float[Tensor, "batch 3"], + cam_rot_delta: Float[Tensor, "batch 3"], + extrinsics: Float[Tensor, "batch 4 4"], + # original_rot: Float[Tensor, "batch 3 3"], + # original_trans: Float[Tensor, "batch 3"], + # converged_threshold: float = 1e-4 + ): + # extrinsics is c2w, here we need w2c as input, so we need to invert it + bs = cam_trans_delta.shape[0] + + tau = torch.cat([cam_trans_delta, cam_rot_delta], dim=-1) + T_w2c = extrinsics.inverse() + + new_w2c_list = [] + for i in range(bs): + new_w2c = SE3_exp(tau[i]) @ T_w2c[i] + new_w2c_list.append(new_w2c) + + new_w2c = torch.stack(new_w2c_list, dim=0) + return new_w2c.inverse() + + # converged = tau.norm() < converged_threshold + # camera.update_RT(new_R, new_T) + # + # camera.cam_rot_delta.data.fill_(0) + # camera.cam_trans_delta.data.fill_(0) + # return converged + + +####### Pose estimation +def inv(mat): + """ Invert a torch or numpy matrix + """ + if isinstance(mat, torch.Tensor): + return torch.linalg.inv(mat) + if isinstance(mat, np.ndarray): + return np.linalg.inv(mat) + raise ValueError(f'bad matrix type = {type(mat)}') + + +def get_pnp_pose(pts3d, opacity, K, H, W, opacity_threshold=0.3): + pixels = np.mgrid[:W, :H].T.astype(np.float32) + pts3d = pts3d.cpu().numpy() + opacity = opacity.cpu().numpy() + K = K.cpu().numpy() + + K[0, :] = K[0, :] * W + K[1, :] = K[1, :] * H + + mask = opacity > opacity_threshold + + res = cv2.solvePnPRansac(pts3d[mask], pixels[mask], K, None, + iterationsCount=100, reprojectionError=5, flags=cv2.SOLVEPNP_SQPNP) + success, R, T, inliers = res + + assert success + + R = cv2.Rodrigues(R)[0] # world to cam + pose = inv(np.r_[np.c_[R, T], [(0, 0, 0, 1)]]) # cam to world + + return torch.from_numpy(pose.astype(np.float32)) + + +def pose_auc(errors, thresholds): + sort_idx = np.argsort(errors) + errors = np.array(errors.copy())[sort_idx] + recall = (np.arange(len(errors)) + 1) / len(errors) + errors = np.r_[0.0, errors] + recall = np.r_[0.0, recall] + aucs = [] + for t in thresholds: + last_index = np.searchsorted(errors, t) + r = np.r_[recall[:last_index], recall[last_index - 1]] + e = np.r_[errors[:last_index], t] + aucs.append(np.trapz(r, x=e) / t) + return aucs + + +def rotation_6d_to_matrix(d6): + """ + Converts 6D rotation representation by Zhou et al. [1] to rotation matrix + using Gram--Schmidt orthogonalization per Section B of [1]. Adapted from pytorch3d. + Args: + d6: 6D rotation representation, of size (*, 6) + + Returns: + batch of rotation matrices of size (*, 3, 3) + + [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. + On the Continuity of Rotation Representations in Neural Networks. + IEEE Conference on Computer Vision and Pattern Recognition, 2019. + Retrieved from http://arxiv.org/abs/1812.07035 + """ + + a1, a2 = d6[..., :3], d6[..., 3:] + b1 = F.normalize(a1, dim=-1) + b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1 + b2 = F.normalize(b2, dim=-1) + b3 = torch.cross(b1, b2, dim=-1) + return torch.stack((b1, b2, b3), dim=-2) \ No newline at end of file diff --git a/src/misc/collation.py b/src/misc/collation.py new file mode 100644 index 0000000000000000000000000000000000000000..f65cc4f978c049e249117a09995018f4fb102b1d --- /dev/null +++ b/src/misc/collation.py @@ -0,0 +1,15 @@ +from typing import Callable, Dict, Union + +from torch import Tensor + +Tree = Union[Dict[str, "Tree"], Tensor] + + +def collate(trees: list[Tree], merge_fn: Callable[[list[Tensor]], Tensor]) -> Tree: + """Merge nested dictionaries of tensors.""" + if isinstance(trees[0], Tensor): + return merge_fn(trees) + else: + return { + key: collate([tree[key] for tree in trees], merge_fn) for key in trees[0] + } diff --git a/src/misc/colmap_utils.py b/src/misc/colmap_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c9fcb6f16511934664fa23d2919456c35323d5ca --- /dev/null +++ b/src/misc/colmap_utils.py @@ -0,0 +1,347 @@ +# +# Copyright (C) 2023, Inria +# GRAPHDECO research group, https://team.inria.fr/graphdeco +# All rights reserved. +# +# This software is free for non-commercial, research and evaluation use +# under the terms of the LICENSE.md file. +# +# For inquiries contact george.drettakis@inria.fr +# + +import collections +import struct + +import numpy as np + +CameraModel = collections.namedtuple( + "CameraModel", ["model_id", "model_name", "num_params"] +) +Camera = collections.namedtuple("Camera", ["id", "model", "width", "height", "params"]) +BaseImage = collections.namedtuple( + "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"] +) +Point3D = collections.namedtuple( + "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"] +) +CAMERA_MODELS = { + CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3), + CameraModel(model_id=1, model_name="PINHOLE", num_params=4), + CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4), + CameraModel(model_id=3, model_name="RADIAL", num_params=5), + CameraModel(model_id=4, model_name="OPENCV", num_params=8), + CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8), + CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12), + CameraModel(model_id=7, model_name="FOV", num_params=5), + CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4), + CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), + CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12), +} +CAMERA_MODEL_IDS = dict( + [(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS] +) +CAMERA_MODEL_NAMES = dict( + [(camera_model.model_name, camera_model) for camera_model in CAMERA_MODELS] +) + + +def qvec2rotmat(qvec): + return np.array( + [ + [ + 1 - 2 * qvec[2] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], + 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2], + ], + [ + 2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[3] ** 2, + 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1], + ], + [ + 2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], + 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], + 1 - 2 * qvec[1] ** 2 - 2 * qvec[2] ** 2, + ], + ] + ) + + +def rotmat2qvec(R): + Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat + K = ( + np.array( + [ + [Rxx - Ryy - Rzz, 0, 0, 0], + [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0], + [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0], + [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz], + ] + ) + / 3.0 + ) + eigvals, eigvecs = np.linalg.eigh(K) + qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)] + if qvec[0] < 0: + qvec *= -1 + return qvec + + +class Image(BaseImage): + def qvec2rotmat(self): + return qvec2rotmat(self.qvec) + + +def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): + """Read and unpack the next bytes from a binary file. + :param fid: + :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc. + :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. + :param endian_character: Any of {@, =, <, >, !} + :return: Tuple of read and unpacked values. + """ + data = fid.read(num_bytes) + return struct.unpack(endian_character + format_char_sequence, data) + + +def read_points3D_text(path): + """ + see: src/base/reconstruction.cc + void Reconstruction::ReadPoints3DText(const std::string& path) + void Reconstruction::WritePoints3DText(const std::string& path) + """ + xyzs = None + rgbs = None + errors = None + num_points = 0 + with open(path, "r") as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + num_points += 1 + + xyzs = np.empty((num_points, 3)) + rgbs = np.empty((num_points, 3)) + errors = np.empty((num_points, 1)) + count = 0 + with open(path, "r") as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + xyz = np.array(tuple(map(float, elems[1:4]))) + rgb = np.array(tuple(map(int, elems[4:7]))) + error = np.array(float(elems[7])) + xyzs[count] = xyz + rgbs[count] = rgb + errors[count] = error + count += 1 + + return xyzs, rgbs, errors + + +def read_points3D_binary(path_to_model_file): + """ + see: src/base/reconstruction.cc + void Reconstruction::ReadPoints3DBinary(const std::string& path) + void Reconstruction::WritePoints3DBinary(const std::string& path) + """ + + with open(path_to_model_file, "rb") as fid: + num_points = read_next_bytes(fid, 8, "Q")[0] + + xyzs = np.empty((num_points, 3)) + rgbs = np.empty((num_points, 3)) + errors = np.empty((num_points, 1)) + + for p_id in range(num_points): + binary_point_line_properties = read_next_bytes( + fid, num_bytes=43, format_char_sequence="QdddBBBd" + ) + xyz = np.array(binary_point_line_properties[1:4]) + rgb = np.array(binary_point_line_properties[4:7]) + error = np.array(binary_point_line_properties[7]) + track_length = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[ + 0 + ] + track_elems = read_next_bytes( + fid, + num_bytes=8 * track_length, + format_char_sequence="ii" * track_length, + ) + xyzs[p_id] = xyz + rgbs[p_id] = rgb + errors[p_id] = error + return xyzs, rgbs, errors + + +def read_intrinsics_text(path): + """ + Taken from https://github.com/colmap/colmap/blob/dev/scripts/python/read_write_model.py + """ + cameras = {} + with open(path, "r") as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + camera_id = int(elems[0]) + model = elems[1] + assert ( + model == "PINHOLE" + ), "While the loader support other types, the rest of the code assumes PINHOLE" + width = int(elems[2]) + height = int(elems[3]) + params = np.array(tuple(map(float, elems[4:]))) + cameras[camera_id] = Camera( + id=camera_id, model=model, width=width, height=height, params=params + ) + return cameras + + +def read_extrinsics_binary(path_to_model_file): + """ + see: src/base/reconstruction.cc + void Reconstruction::ReadImagesBinary(const std::string& path) + void Reconstruction::WriteImagesBinary(const std::string& path) + """ + images = {} + with open(path_to_model_file, "rb") as fid: + num_reg_images = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_reg_images): + binary_image_properties = read_next_bytes( + fid, num_bytes=64, format_char_sequence="idddddddi" + ) + image_id = binary_image_properties[0] + qvec = np.array(binary_image_properties[1:5]) + tvec = np.array(binary_image_properties[5:8]) + camera_id = binary_image_properties[8] + image_name = "" + current_char = read_next_bytes(fid, 1, "c")[0] + while current_char != b"\x00": # look for the ASCII 0 entry + image_name += current_char.decode("utf-8") + current_char = read_next_bytes(fid, 1, "c")[0] + num_points2D = read_next_bytes(fid, num_bytes=8, format_char_sequence="Q")[ + 0 + ] + x_y_id_s = read_next_bytes( + fid, + num_bytes=24 * num_points2D, + format_char_sequence="ddq" * num_points2D, + ) + xys = np.column_stack( + [tuple(map(float, x_y_id_s[0::3])), tuple(map(float, x_y_id_s[1::3]))] + ) + point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3]))) + images[image_id] = Image( + id=image_id, + qvec=qvec, + tvec=tvec, + camera_id=camera_id, + name=image_name, + xys=xys, + point3D_ids=point3D_ids, + ) + return images + + +def read_intrinsics_binary(path_to_model_file): + """ + see: src/base/reconstruction.cc + void Reconstruction::WriteCamerasBinary(const std::string& path) + void Reconstruction::ReadCamerasBinary(const std::string& path) + """ + cameras = {} + with open(path_to_model_file, "rb") as fid: + num_cameras = read_next_bytes(fid, 8, "Q")[0] + for _ in range(num_cameras): + camera_properties = read_next_bytes( + fid, num_bytes=24, format_char_sequence="iiQQ" + ) + camera_id = camera_properties[0] + model_id = camera_properties[1] + model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name + width = camera_properties[2] + height = camera_properties[3] + num_params = CAMERA_MODEL_IDS[model_id].num_params + params = read_next_bytes( + fid, num_bytes=8 * num_params, format_char_sequence="d" * num_params + ) + cameras[camera_id] = Camera( + id=camera_id, + model=model_name, + width=width, + height=height, + params=np.array(params), + ) + assert len(cameras) == num_cameras + return cameras + + +def read_extrinsics_text(path): + """ + Taken from https://github.com/colmap/colmap/blob/dev/scripts/python/read_write_model.py + """ + images = {} + with open(path, "r") as fid: + while True: + line = fid.readline() + if not line: + break + line = line.strip() + if len(line) > 0 and line[0] != "#": + elems = line.split() + image_id = int(elems[0]) + qvec = np.array(tuple(map(float, elems[1:5]))) + tvec = np.array(tuple(map(float, elems[5:8]))) + camera_id = int(elems[8]) + image_name = elems[9] + elems = fid.readline().split() + xys = np.column_stack( + [tuple(map(float, elems[0::3])), tuple(map(float, elems[1::3]))] + ) + point3D_ids = np.array(tuple(map(int, elems[2::3]))) + images[image_id] = Image( + id=image_id, + qvec=qvec, + tvec=tvec, + camera_id=camera_id, + name=image_name, + xys=xys, + point3D_ids=point3D_ids, + ) + return images + + +def read_colmap_bin_array(path): + """ + Taken from https://github.com/colmap/colmap/blob/dev/scripts/python/read_dense.py + + :param path: path to the colmap binary file. + :return: nd array with the floating point values in the value + """ + with open(path, "rb") as fid: + width, height, channels = np.genfromtxt( + fid, delimiter="&", max_rows=1, usecols=(0, 1, 2), dtype=int + ) + fid.seek(0) + num_delimiter = 0 + byte = fid.read(1) + while True: + if byte == b"&": + num_delimiter += 1 + if num_delimiter >= 3: + break + byte = fid.read(1) + array = np.fromfile(fid, np.float32) + array = array.reshape((width, height, channels), order="F") + return np.transpose(array, (1, 0, 2)).squeeze() diff --git a/src/misc/discrete_probability_distribution.py b/src/misc/discrete_probability_distribution.py new file mode 100644 index 0000000000000000000000000000000000000000..1b9fcb228098da9a12ce7f7b59fe43c3562083f5 --- /dev/null +++ b/src/misc/discrete_probability_distribution.py @@ -0,0 +1,33 @@ +import torch +from einops import reduce +from jaxtyping import Float, Int64 +from torch import Tensor + + +def sample_discrete_distribution( + pdf: Float[Tensor, "*batch bucket"], + num_samples: int, + eps: float = torch.finfo(torch.float32).eps, +) -> tuple[ + Int64[Tensor, "*batch sample"], # index + Float[Tensor, "*batch sample"], # probability density +]: + *batch, bucket = pdf.shape + normalized_pdf = pdf / (eps + reduce(pdf, "... bucket -> ... ()", "sum")) + cdf = normalized_pdf.cumsum(dim=-1) + samples = torch.rand((*batch, num_samples), device=pdf.device) + index = torch.searchsorted(cdf, samples, right=True).clip(max=bucket - 1) + return index, normalized_pdf.gather(dim=-1, index=index) + + +def gather_discrete_topk( + pdf: Float[Tensor, "*batch bucket"], + num_samples: int, + eps: float = torch.finfo(torch.float32).eps, +) -> tuple[ + Int64[Tensor, "*batch sample"], # index + Float[Tensor, "*batch sample"], # probability density +]: + normalized_pdf = pdf / (eps + reduce(pdf, "... bucket -> ... ()", "sum")) + index = pdf.topk(k=num_samples, dim=-1).indices + return index, normalized_pdf.gather(dim=-1, index=index) diff --git a/src/misc/heterogeneous_pairings.py b/src/misc/heterogeneous_pairings.py new file mode 100644 index 0000000000000000000000000000000000000000..76d9b7fd1ef8676ad47135110ecd0795bf86a363 --- /dev/null +++ b/src/misc/heterogeneous_pairings.py @@ -0,0 +1,43 @@ +import torch +from einops import repeat +from jaxtyping import Int +from torch import Tensor + +Index = Int[Tensor, "n n-1"] + + +def generate_heterogeneous_index( + n: int, + device: torch.device = torch.device("cpu"), +) -> tuple[Index, Index]: + """Generate indices for all pairs except self-pairs.""" + arange = torch.arange(n, device=device) + + # Generate an index that represents the item itself. + index_self = repeat(arange, "h -> h w", w=n - 1) + + # Generate an index that represents the other items. + index_other = repeat(arange, "w -> h w", h=n).clone() + index_other += torch.ones((n, n), device=device, dtype=torch.int64).triu() + index_other = index_other[:, :-1] + + return index_self, index_other + + +def generate_heterogeneous_index_transpose( + n: int, + device: torch.device = torch.device("cpu"), +) -> tuple[Index, Index]: + """Generate an index that can be used to "transpose" the heterogeneous index. + Applying the index a second time inverts the "transpose." + """ + arange = torch.arange(n, device=device) + ones = torch.ones((n, n), device=device, dtype=torch.int64) + + index_self = repeat(arange, "w -> h w", h=n).clone() + index_self = index_self + ones.triu() + + index_other = repeat(arange, "h -> h w", w=n) + index_other = index_other - (1 - ones.triu()) + + return index_self[:, :-1], index_other[:, :-1] diff --git a/src/misc/image_io.py b/src/misc/image_io.py new file mode 100644 index 0000000000000000000000000000000000000000..2a6994df1e3f06c3842b6d13364afe3a75e3df3a --- /dev/null +++ b/src/misc/image_io.py @@ -0,0 +1,211 @@ +import io +import os +from pathlib import Path +from typing import Union + +import cv2 +import numpy as np +import skvideo +import torch +import torchvision.transforms as tf +from einops import rearrange, repeat +from jaxtyping import Float, UInt8 + +from matplotlib import pyplot as plt +from matplotlib.figure import Figure +from PIL import Image +from torch import Tensor + +FloatImage = Union[ + Float[Tensor, "height width"], + Float[Tensor, "channel height width"], + Float[Tensor, "batch channel height width"], +] + + +def fig_to_image( + fig: Figure, + dpi: int = 100, + device: torch.device = torch.device("cpu"), +) -> Float[Tensor, "3 height width"]: + buffer = io.BytesIO() + fig.savefig(buffer, format="raw", dpi=dpi) + buffer.seek(0) + data = np.frombuffer(buffer.getvalue(), dtype=np.uint8) + h = int(fig.bbox.bounds[3]) + w = int(fig.bbox.bounds[2]) + data = rearrange(data, "(h w c) -> c h w", h=h, w=w, c=4) + buffer.close() + return (torch.tensor(data, device=device, dtype=torch.float32) / 255)[:3] + + +def prep_image(image: FloatImage) -> UInt8[np.ndarray, "height width channel"]: + # Handle batched images. + if image.ndim == 4: + image = rearrange(image, "b c h w -> c h (b w)") + + # Handle single-channel images. + if image.ndim == 2: + image = rearrange(image, "h w -> () h w") + + # Ensure that there are 3 or 4 channels. + channel, _, _ = image.shape + if channel == 1: + image = repeat(image, "() h w -> c h w", c=3) + assert image.shape[0] in (3, 4) + + image = (image.detach().clip(min=0, max=1) * 255).type(torch.uint8) + return rearrange(image, "c h w -> h w c").cpu().numpy() + + +def save_image( + image: FloatImage, + path: Union[Path, str], +) -> None: + """Save an image. Assumed to be in range 0-1.""" + + # Create the parent directory if it doesn't already exist. + path = Path(path) + path.parent.mkdir(exist_ok=True, parents=True) + + # Save the image. + Image.fromarray(prep_image(image)).save(path) + + +def load_image( + path: Union[Path, str], +) -> Float[Tensor, "3 height width"]: + return tf.ToTensor()(Image.open(path))[:3] + + +def save_video(tensor, save_path, fps=10): + """ + Save a tensor of shape (N, C, H, W) as a video file using imageio. + Args: + tensor: Tensor of shape (N, C, H, W) in range [0, 1] + save_path: Path to save the video file + fps: Frames per second for the video + """ + # Convert tensor to numpy array and adjust dimensions + video = tensor.cpu().detach().numpy() # (N, C, H, W) + video = np.transpose(video, (0, 2, 3, 1)) # (N, H, W, C) + + # Scale to [0, 255] and convert to uint8 + video = (video * 255).astype(np.uint8) + + # Ensure the directory exists + import os + + os.makedirs(os.path.dirname(save_path), exist_ok=True) + + # Use imageio to write video (handles codec compatibility automatically) + import imageio + + writer = imageio.get_writer(save_path, fps=fps) + + for frame in video: + writer.append_data(frame) + + writer.close() + + +def save_interpolated_video( + pred_extrinsics, pred_intrinsics, b, h, w, gaussians, save_path, decoder_func, t=10 +): + # Interpolate between neighboring frames + # t: Number of extra views to interpolate between each pair + interpolated_extrinsics = [] + interpolated_intrinsics = [] + + # For each pair of neighboring frame + for i in range(pred_extrinsics.shape[1] - 1): + # Add the current frame + interpolated_extrinsics.append(pred_extrinsics[:, i : i + 1]) + interpolated_intrinsics.append(pred_intrinsics[:, i : i + 1]) + + # Interpolate between current and next frame + for j in range(1, t + 1): + alpha = j / (t + 1) + + # Interpolate extrinsics + start_extrinsic = pred_extrinsics[:, i] + end_extrinsic = pred_extrinsics[:, i + 1] + + # Separate rotation and translation + start_rot = start_extrinsic[:, :3, :3] + end_rot = end_extrinsic[:, :3, :3] + start_trans = start_extrinsic[:, :3, 3] + end_trans = end_extrinsic[:, :3, 3] + + # Interpolate translation (linear) + interp_trans = (1 - alpha) * start_trans + alpha * end_trans + + # Interpolate rotation (spherical) + start_rot_flat = start_rot.reshape(b, 9) + end_rot_flat = end_rot.reshape(b, 9) + interp_rot_flat = (1 - alpha) * start_rot_flat + alpha * end_rot_flat + interp_rot = interp_rot_flat.reshape(b, 3, 3) + + # Normalize rotation matrix to ensure it's orthogonal + u, _, v = torch.svd(interp_rot) + interp_rot = torch.bmm(u, v.transpose(1, 2)) + + # Combine interpolated rotation and translation + interp_extrinsic = ( + torch.eye(4, device=pred_extrinsics.device).unsqueeze(0).repeat(b, 1, 1) + ) + interp_extrinsic[:, :3, :3] = interp_rot + interp_extrinsic[:, :3, 3] = interp_trans + + # Interpolate intrinsics (linear) + start_intrinsic = pred_intrinsics[:, i] + end_intrinsic = pred_intrinsics[:, i + 1] + interp_intrinsic = (1 - alpha) * start_intrinsic + alpha * end_intrinsic + + # Add interpolated frame + interpolated_extrinsics.append(interp_extrinsic.unsqueeze(1)) + interpolated_intrinsics.append(interp_intrinsic.unsqueeze(1)) + + # Concatenate all frames + pred_all_extrinsic = torch.cat(interpolated_extrinsics, dim=1) + pred_all_intrinsic = torch.cat(interpolated_intrinsics, dim=1) + + # Add the last frame + interpolated_extrinsics.append(pred_all_extrinsic[:, -1:]) + interpolated_intrinsics.append(pred_all_intrinsic[:, -1:]) + + # Update K to reflect the new number of frames + num_frames = pred_all_extrinsic.shape[1] + + # Render interpolated views + interpolated_output = decoder_func.forward( + gaussians, + pred_all_extrinsic, + pred_all_intrinsic.float(), + torch.ones(1, num_frames, device=pred_all_extrinsic.device) * 0.1, + torch.ones(1, num_frames, device=pred_all_extrinsic.device) * 100, + (h, w), + ) + + # Convert to video format + video = interpolated_output.color[0].clip(min=0, max=1) + depth = interpolated_output.depth[0] + + # Normalize depth for visualization + # to avoid `quantile() input tensor is too large` + num_views = pred_extrinsics.shape[1] + depth_norm = (depth - depth[::num_views].quantile(0.01)) / ( + depth[::num_views].quantile(0.99) - depth[::num_views].quantile(0.01) + ) + depth_norm = plt.cm.turbo(depth_norm.cpu().numpy()) + depth_colored = ( + torch.from_numpy(depth_norm[..., :3]).permute(0, 3, 1, 2).to(depth.device) + ) + depth_colored = depth_colored.clip(min=0, max=1) + + # Save depth video + save_video(depth_colored, os.path.join(save_path, f"depth.mp4")) + # Save video + save_video(video, os.path.join(save_path, f"rgb.mp4")) + + return os.path.join(save_path, f"rgb.mp4"), os.path.join(save_path, f"depth.mp4") diff --git a/src/misc/nn_module_tools.py b/src/misc/nn_module_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..21570e2880349a8c6083118d76183e90a501e617 --- /dev/null +++ b/src/misc/nn_module_tools.py @@ -0,0 +1,16 @@ +from torch import nn + + +def convert_to_buffer(module: nn.Module, persistent: bool = True): + # Recurse over child modules. + for name, child in list(module.named_children()): + convert_to_buffer(child, persistent) + + # Also re-save buffers to change persistence. + for name, parameter_or_buffer in ( + *module.named_parameters(recurse=False), + *module.named_buffers(recurse=False), + ): + value = parameter_or_buffer.detach().clone() + delattr(module, name) + module.register_buffer(name, value, persistent=persistent) diff --git a/src/misc/sh_rotation.py b/src/misc/sh_rotation.py new file mode 100644 index 0000000000000000000000000000000000000000..cdca2dafbe845d1bc7e2a8df2f07bdf933e64b98 --- /dev/null +++ b/src/misc/sh_rotation.py @@ -0,0 +1,111 @@ +from math import isqrt + +import torch +from e3nn.o3 import matrix_to_angles, wigner_D +from einops import einsum +from jaxtyping import Float +from torch import Tensor + + +def rotate_sh( + sh_coefficients: Float[Tensor, "*#batch n"], + rotations: Float[Tensor, "*#batch 3 3"], +) -> Float[Tensor, "*batch n"]: + device = sh_coefficients.device + dtype = sh_coefficients.dtype + + # change the basis from YZX -> XYZ to fit the convention of e3nn + P = torch.tensor([[0, 0, 1], [1, 0, 0], [0, 1, 0]], + dtype=sh_coefficients.dtype, device=sh_coefficients.device) + inversed_P = torch.tensor([[0, 1, 0], [0, 0, 1], [1, 0, 0], ], + dtype=sh_coefficients.dtype, device=sh_coefficients.device) + permuted_rotation_matrix = inversed_P @ rotations @ P + + *_, n = sh_coefficients.shape + alpha, beta, gamma = matrix_to_angles(permuted_rotation_matrix) + result = [] + for degree in range(isqrt(n)): + with torch.device(device): + sh_rotations = wigner_D(degree, alpha, -beta, gamma).type(dtype) + sh_rotated = einsum( + sh_rotations, + sh_coefficients[..., degree**2 : (degree + 1) ** 2], + "... i j, ... j -> ... i", + ) + result.append(sh_rotated) + + return torch.cat(result, dim=-1) + + +# def rotate_sh( +# sh_coefficients: Float[Tensor, "*#batch n"], +# rotations: Float[Tensor, "*#batch 3 3"], +# ) -> Float[Tensor, "*batch n"]: +# device = sh_coefficients.device +# dtype = sh_coefficients.dtype +# +# *_, n = sh_coefficients.shape +# alpha, beta, gamma = matrix_to_angles(rotations) +# result = [] +# for degree in range(isqrt(n)): +# with torch.device(device): +# sh_rotations = wigner_D(degree, alpha, beta, gamma).type(dtype) +# sh_rotated = einsum( +# sh_rotations, +# sh_coefficients[..., degree**2 : (degree + 1) ** 2], +# "... i j, ... j -> ... i", +# ) +# result.append(sh_rotated) +# +# return torch.cat(result, dim=-1) + + +if __name__ == "__main__": + from pathlib import Path + + import matplotlib.pyplot as plt + from e3nn.o3 import spherical_harmonics + from matplotlib import cm + from scipy.spatial.transform.rotation import Rotation as R + + device = torch.device("cuda") + + # Generate random spherical harmonics coefficients. + degree = 4 + coefficients = torch.rand((degree + 1) ** 2, dtype=torch.float32, device=device) + + def plot_sh(sh_coefficients, path: Path) -> None: + phi = torch.linspace(0, torch.pi, 100, device=device) + theta = torch.linspace(0, 2 * torch.pi, 100, device=device) + phi, theta = torch.meshgrid(phi, theta, indexing="xy") + x = torch.sin(phi) * torch.cos(theta) + y = torch.sin(phi) * torch.sin(theta) + z = torch.cos(phi) + xyz = torch.stack([x, y, z], dim=-1) + sh = spherical_harmonics(list(range(degree + 1)), xyz, True) + result = einsum(sh, sh_coefficients, "... n, n -> ...") + result = (result - result.min()) / (result.max() - result.min()) + + # Set the aspect ratio to 1 so our sphere looks spherical + fig = plt.figure(figsize=plt.figaspect(1.0)) + ax = fig.add_subplot(111, projection="3d") + ax.plot_surface( + x.cpu().numpy(), + y.cpu().numpy(), + z.cpu().numpy(), + rstride=1, + cstride=1, + facecolors=cm.seismic(result.cpu().numpy()), + ) + # Turn off the axis planes + ax.set_axis_off() + path.parent.mkdir(exist_ok=True, parents=True) + plt.savefig(path) + + for i, angle in enumerate(torch.linspace(0, 2 * torch.pi, 30)): + rotation = torch.tensor( + R.from_euler("x", angle.item()).as_matrix(), device=device + ) + plot_sh(rotate_sh(coefficients, rotation), Path(f"sh_rotation/{i:0>3}.png")) + + print("Done!") diff --git a/src/misc/sht.py b/src/misc/sht.py new file mode 100644 index 0000000000000000000000000000000000000000..4b89273a8f20b4da5ba296b175c856c974df0984 --- /dev/null +++ b/src/misc/sht.py @@ -0,0 +1,1637 @@ +"""Real spherical harmonics in Cartesian form for PyTorch. + +This is an autogenerated file. See +https://github.com/cheind/torch-spherical-harmonics +for more information. +""" + +import torch + + +def rsh_cart_0(xyz: torch.Tensor): + """Computes all real spherical harmonics up to degree 0. + + This is an autogenerated method. See + https://github.com/cheind/torch-spherical-harmonics + for more information. + + Params: + xyz: (N,...,3) tensor of points on the unit sphere + + Returns: + rsh: (N,...,1) real spherical harmonics + projections of input. Ynm is found at index + `n*(n+1) + m`, with `0 <= n <= degree` and + `-n <= m <= n`. + """ + + return torch.stack( + [ + xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]), + ], + -1, + ) + + +def rsh_cart_1(xyz: torch.Tensor): + """Computes all real spherical harmonics up to degree 1. + + This is an autogenerated method. See + https://github.com/cheind/torch-spherical-harmonics + for more information. + + Params: + xyz: (N,...,3) tensor of points on the unit sphere + + Returns: + rsh: (N,...,4) real spherical harmonics + projections of input. Ynm is found at index + `n*(n+1) + m`, with `0 <= n <= degree` and + `-n <= m <= n`. + """ + x = xyz[..., 0] + y = xyz[..., 1] + z = xyz[..., 2] + + return torch.stack( + [ + xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]), + -0.48860251190292 * y, + 0.48860251190292 * z, + -0.48860251190292 * x, + ], + -1, + ) + + +def rsh_cart_2(xyz: torch.Tensor): + """Computes all real spherical harmonics up to degree 2. + + This is an autogenerated method. See + https://github.com/cheind/torch-spherical-harmonics + for more information. + + Params: + xyz: (N,...,3) tensor of points on the unit sphere + + Returns: + rsh: (N,...,9) real spherical harmonics + projections of input. Ynm is found at index + `n*(n+1) + m`, with `0 <= n <= degree` and + `-n <= m <= n`. + """ + x = xyz[..., 0] + y = xyz[..., 1] + z = xyz[..., 2] + + x2 = x**2 + y2 = y**2 + z2 = z**2 + xy = x * y + xz = x * z + yz = y * z + + return torch.stack( + [ + xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]), + -0.48860251190292 * y, + 0.48860251190292 * z, + -0.48860251190292 * x, + 1.09254843059208 * xy, + -1.09254843059208 * yz, + 0.94617469575756 * z2 - 0.31539156525252, + -1.09254843059208 * xz, + 0.54627421529604 * x2 - 0.54627421529604 * y2, + ], + -1, + ) + + +def rsh_cart_3(xyz: torch.Tensor): + """Computes all real spherical harmonics up to degree 3. + + This is an autogenerated method. See + https://github.com/cheind/torch-spherical-harmonics + for more information. + + Params: + xyz: (N,...,3) tensor of points on the unit sphere + + Returns: + rsh: (N,...,16) real spherical harmonics + projections of input. Ynm is found at index + `n*(n+1) + m`, with `0 <= n <= degree` and + `-n <= m <= n`. + """ + x = xyz[..., 0] + y = xyz[..., 1] + z = xyz[..., 2] + + x2 = x**2 + y2 = y**2 + z2 = z**2 + xy = x * y + xz = x * z + yz = y * z + + return torch.stack( + [ + xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]), + -0.48860251190292 * y, + 0.48860251190292 * z, + -0.48860251190292 * x, + 1.09254843059208 * xy, + -1.09254843059208 * yz, + 0.94617469575756 * z2 - 0.31539156525252, + -1.09254843059208 * xz, + 0.54627421529604 * x2 - 0.54627421529604 * y2, + -0.590043589926644 * y * (3.0 * x2 - y2), + 2.89061144264055 * xy * z, + 0.304697199642977 * y * (1.5 - 7.5 * z2), + 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z, + 0.304697199642977 * x * (1.5 - 7.5 * z2), + 1.44530572132028 * z * (x2 - y2), + -0.590043589926644 * x * (x2 - 3.0 * y2), + ], + -1, + ) + + +def rsh_cart_4(xyz: torch.Tensor): + """Computes all real spherical harmonics up to degree 4. + + This is an autogenerated method. See + https://github.com/cheind/torch-spherical-harmonics + for more information. + + Params: + xyz: (N,...,3) tensor of points on the unit sphere + + Returns: + rsh: (N,...,25) real spherical harmonics + projections of input. Ynm is found at index + `n*(n+1) + m`, with `0 <= n <= degree` and + `-n <= m <= n`. + """ + x = xyz[..., 0] + y = xyz[..., 1] + z = xyz[..., 2] + + x2 = x**2 + y2 = y**2 + z2 = z**2 + xy = x * y + xz = x * z + yz = y * z + x4 = x2**2 + y4 = y2**2 + z4 = z2**2 + + return torch.stack( + [ + xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]), + -0.48860251190292 * y, + 0.48860251190292 * z, + -0.48860251190292 * x, + 1.09254843059208 * xy, + -1.09254843059208 * yz, + 0.94617469575756 * z2 - 0.31539156525252, + -1.09254843059208 * xz, + 0.54627421529604 * x2 - 0.54627421529604 * y2, + -0.590043589926644 * y * (3.0 * x2 - y2), + 2.89061144264055 * xy * z, + 0.304697199642977 * y * (1.5 - 7.5 * z2), + 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z, + 0.304697199642977 * x * (1.5 - 7.5 * z2), + 1.44530572132028 * z * (x2 - y2), + -0.590043589926644 * x * (x2 - 3.0 * y2), + 2.5033429417967 * xy * (x2 - y2), + -1.77013076977993 * yz * (3.0 * x2 - y2), + 0.126156626101008 * xy * (52.5 * z2 - 7.5), + 0.267618617422916 * y * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z), + 1.48099765681286 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 0.952069922236839 * z2 + + 0.317356640745613, + 0.267618617422916 * x * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z), + 0.063078313050504 * (x2 - y2) * (52.5 * z2 - 7.5), + -1.77013076977993 * xz * (x2 - 3.0 * y2), + -3.75501441269506 * x2 * y2 + + 0.625835735449176 * x4 + + 0.625835735449176 * y4, + ], + -1, + ) + + +def rsh_cart_5(xyz: torch.Tensor): + """Computes all real spherical harmonics up to degree 5. + + This is an autogenerated method. See + https://github.com/cheind/torch-spherical-harmonics + for more information. + + Params: + xyz: (N,...,3) tensor of points on the unit sphere + + Returns: + rsh: (N,...,36) real spherical harmonics + projections of input. Ynm is found at index + `n*(n+1) + m`, with `0 <= n <= degree` and + `-n <= m <= n`. + """ + x = xyz[..., 0] + y = xyz[..., 1] + z = xyz[..., 2] + + x2 = x**2 + y2 = y**2 + z2 = z**2 + xy = x * y + xz = x * z + yz = y * z + x4 = x2**2 + y4 = y2**2 + z4 = z2**2 + + return torch.stack( + [ + xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]), + -0.48860251190292 * y, + 0.48860251190292 * z, + -0.48860251190292 * x, + 1.09254843059208 * xy, + -1.09254843059208 * yz, + 0.94617469575756 * z2 - 0.31539156525252, + -1.09254843059208 * xz, + 0.54627421529604 * x2 - 0.54627421529604 * y2, + -0.590043589926644 * y * (3.0 * x2 - y2), + 2.89061144264055 * xy * z, + 0.304697199642977 * y * (1.5 - 7.5 * z2), + 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z, + 0.304697199642977 * x * (1.5 - 7.5 * z2), + 1.44530572132028 * z * (x2 - y2), + -0.590043589926644 * x * (x2 - 3.0 * y2), + 2.5033429417967 * xy * (x2 - y2), + -1.77013076977993 * yz * (3.0 * x2 - y2), + 0.126156626101008 * xy * (52.5 * z2 - 7.5), + 0.267618617422916 * y * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z), + 1.48099765681286 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 0.952069922236839 * z2 + + 0.317356640745613, + 0.267618617422916 * x * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z), + 0.063078313050504 * (x2 - y2) * (52.5 * z2 - 7.5), + -1.77013076977993 * xz * (x2 - 3.0 * y2), + -3.75501441269506 * x2 * y2 + + 0.625835735449176 * x4 + + 0.625835735449176 * y4, + -0.65638205684017 * y * (-10.0 * x2 * y2 + 5.0 * x4 + y4), + 8.30264925952416 * xy * z * (x2 - y2), + 0.00931882475114763 * y * (52.5 - 472.5 * z2) * (3.0 * x2 - y2), + 0.0913054625709205 * xy * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z), + 0.241571547304372 + * y + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ), + -1.24747010616985 * z * (1.5 * z2 - 0.5) + + 1.6840846433293 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 0.498988042467941 * z, + 0.241571547304372 + * x + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ), + 0.0456527312854602 * (x2 - y2) * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z), + 0.00931882475114763 * x * (52.5 - 472.5 * z2) * (x2 - 3.0 * y2), + 2.07566231488104 * z * (-6.0 * x2 * y2 + x4 + y4), + -0.65638205684017 * x * (-10.0 * x2 * y2 + x4 + 5.0 * y4), + ], + -1, + ) + + +def rsh_cart_6(xyz: torch.Tensor): + """Computes all real spherical harmonics up to degree 6. + + This is an autogenerated method. See + https://github.com/cheind/torch-spherical-harmonics + for more information. + + Params: + xyz: (N,...,3) tensor of points on the unit sphere + + Returns: + rsh: (N,...,49) real spherical harmonics + projections of input. Ynm is found at index + `n*(n+1) + m`, with `0 <= n <= degree` and + `-n <= m <= n`. + """ + x = xyz[..., 0] + y = xyz[..., 1] + z = xyz[..., 2] + + x2 = x**2 + y2 = y**2 + z2 = z**2 + xy = x * y + xz = x * z + yz = y * z + x4 = x2**2 + y4 = y2**2 + z4 = z2**2 + + return torch.stack( + [ + xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]), + -0.48860251190292 * y, + 0.48860251190292 * z, + -0.48860251190292 * x, + 1.09254843059208 * xy, + -1.09254843059208 * yz, + 0.94617469575756 * z2 - 0.31539156525252, + -1.09254843059208 * xz, + 0.54627421529604 * x2 - 0.54627421529604 * y2, + -0.590043589926644 * y * (3.0 * x2 - y2), + 2.89061144264055 * xy * z, + 0.304697199642977 * y * (1.5 - 7.5 * z2), + 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z, + 0.304697199642977 * x * (1.5 - 7.5 * z2), + 1.44530572132028 * z * (x2 - y2), + -0.590043589926644 * x * (x2 - 3.0 * y2), + 2.5033429417967 * xy * (x2 - y2), + -1.77013076977993 * yz * (3.0 * x2 - y2), + 0.126156626101008 * xy * (52.5 * z2 - 7.5), + 0.267618617422916 * y * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z), + 1.48099765681286 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 0.952069922236839 * z2 + + 0.317356640745613, + 0.267618617422916 * x * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z), + 0.063078313050504 * (x2 - y2) * (52.5 * z2 - 7.5), + -1.77013076977993 * xz * (x2 - 3.0 * y2), + -3.75501441269506 * x2 * y2 + + 0.625835735449176 * x4 + + 0.625835735449176 * y4, + -0.65638205684017 * y * (-10.0 * x2 * y2 + 5.0 * x4 + y4), + 8.30264925952416 * xy * z * (x2 - y2), + 0.00931882475114763 * y * (52.5 - 472.5 * z2) * (3.0 * x2 - y2), + 0.0913054625709205 * xy * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z), + 0.241571547304372 + * y + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ), + -1.24747010616985 * z * (1.5 * z2 - 0.5) + + 1.6840846433293 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 0.498988042467941 * z, + 0.241571547304372 + * x + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ), + 0.0456527312854602 * (x2 - y2) * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z), + 0.00931882475114763 * x * (52.5 - 472.5 * z2) * (x2 - 3.0 * y2), + 2.07566231488104 * z * (-6.0 * x2 * y2 + x4 + y4), + -0.65638205684017 * x * (-10.0 * x2 * y2 + x4 + 5.0 * y4), + 4.09910463115149 * x**4 * xy + - 13.6636821038383 * xy**3 + + 4.09910463115149 * xy * y**4, + -2.36661916223175 * yz * (-10.0 * x2 * y2 + 5.0 * x4 + y4), + 0.00427144889505798 * xy * (x2 - y2) * (5197.5 * z2 - 472.5), + 0.00584892228263444 + * y + * (3.0 * x2 - y2) + * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z), + 0.0701870673916132 + * xy + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ), + 0.221950995245231 + * y + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ), + -1.48328138624466 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + + 1.86469659985043 + * z + * ( + -1.33333333333333 * z * (1.5 * z2 - 0.5) + + 1.8 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 0.533333333333333 * z + ) + + 0.953538034014426 * z2 + - 0.317846011338142, + 0.221950995245231 + * x + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ), + 0.0350935336958066 + * (x2 - y2) + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ), + 0.00584892228263444 + * x + * (x2 - 3.0 * y2) + * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z), + 0.0010678622237645 * (5197.5 * z2 - 472.5) * (-6.0 * x2 * y2 + x4 + y4), + -2.36661916223175 * xz * (-10.0 * x2 * y2 + x4 + 5.0 * y4), + 0.683184105191914 * x2**3 + + 10.2477615778787 * x2 * y4 + - 10.2477615778787 * x4 * y2 + - 0.683184105191914 * y2**3, + ], + -1, + ) + + +def rsh_cart_7(xyz: torch.Tensor): + """Computes all real spherical harmonics up to degree 7. + + This is an autogenerated method. See + https://github.com/cheind/torch-spherical-harmonics + for more information. + + Params: + xyz: (N,...,3) tensor of points on the unit sphere + + Returns: + rsh: (N,...,64) real spherical harmonics + projections of input. Ynm is found at index + `n*(n+1) + m`, with `0 <= n <= degree` and + `-n <= m <= n`. + """ + x = xyz[..., 0] + y = xyz[..., 1] + z = xyz[..., 2] + + x2 = x**2 + y2 = y**2 + z2 = z**2 + xy = x * y + xz = x * z + yz = y * z + x4 = x2**2 + y4 = y2**2 + z4 = z2**2 + + return torch.stack( + [ + xyz.new_tensor(0.282094791773878).expand(xyz.shape[:-1]), + -0.48860251190292 * y, + 0.48860251190292 * z, + -0.48860251190292 * x, + 1.09254843059208 * xy, + -1.09254843059208 * yz, + 0.94617469575756 * z2 - 0.31539156525252, + -1.09254843059208 * xz, + 0.54627421529604 * x2 - 0.54627421529604 * y2, + -0.590043589926644 * y * (3.0 * x2 - y2), + 2.89061144264055 * xy * z, + 0.304697199642977 * y * (1.5 - 7.5 * z2), + 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z, + 0.304697199642977 * x * (1.5 - 7.5 * z2), + 1.44530572132028 * z * (x2 - y2), + -0.590043589926644 * x * (x2 - 3.0 * y2), + 2.5033429417967 * xy * (x2 - y2), + -1.77013076977993 * yz * (3.0 * x2 - y2), + 0.126156626101008 * xy * (52.5 * z2 - 7.5), + 0.267618617422916 * y * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z), + 1.48099765681286 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 0.952069922236839 * z2 + + 0.317356640745613, + 0.267618617422916 * x * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z), + 0.063078313050504 * (x2 - y2) * (52.5 * z2 - 7.5), + -1.77013076977993 * xz * (x2 - 3.0 * y2), + -3.75501441269506 * x2 * y2 + + 0.625835735449176 * x4 + + 0.625835735449176 * y4, + -0.65638205684017 * y * (-10.0 * x2 * y2 + 5.0 * x4 + y4), + 8.30264925952416 * xy * z * (x2 - y2), + 0.00931882475114763 * y * (52.5 - 472.5 * z2) * (3.0 * x2 - y2), + 0.0913054625709205 * xy * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z), + 0.241571547304372 + * y + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ), + -1.24747010616985 * z * (1.5 * z2 - 0.5) + + 1.6840846433293 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 0.498988042467941 * z, + 0.241571547304372 + * x + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ), + 0.0456527312854602 * (x2 - y2) * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z), + 0.00931882475114763 * x * (52.5 - 472.5 * z2) * (x2 - 3.0 * y2), + 2.07566231488104 * z * (-6.0 * x2 * y2 + x4 + y4), + -0.65638205684017 * x * (-10.0 * x2 * y2 + x4 + 5.0 * y4), + 4.09910463115149 * x**4 * xy + - 13.6636821038383 * xy**3 + + 4.09910463115149 * xy * y**4, + -2.36661916223175 * yz * (-10.0 * x2 * y2 + 5.0 * x4 + y4), + 0.00427144889505798 * xy * (x2 - y2) * (5197.5 * z2 - 472.5), + 0.00584892228263444 + * y + * (3.0 * x2 - y2) + * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z), + 0.0701870673916132 + * xy + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ), + 0.221950995245231 + * y + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ), + -1.48328138624466 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + + 1.86469659985043 + * z + * ( + -1.33333333333333 * z * (1.5 * z2 - 0.5) + + 1.8 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 0.533333333333333 * z + ) + + 0.953538034014426 * z2 + - 0.317846011338142, + 0.221950995245231 + * x + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ), + 0.0350935336958066 + * (x2 - y2) + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ), + 0.00584892228263444 + * x + * (x2 - 3.0 * y2) + * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z), + 0.0010678622237645 * (5197.5 * z2 - 472.5) * (-6.0 * x2 * y2 + x4 + y4), + -2.36661916223175 * xz * (-10.0 * x2 * y2 + x4 + 5.0 * y4), + 0.683184105191914 * x2**3 + + 10.2477615778787 * x2 * y4 + - 10.2477615778787 * x4 * y2 + - 0.683184105191914 * y2**3, + -0.707162732524596 + * y + * (7.0 * x2**3 + 21.0 * x2 * y4 - 35.0 * x4 * y2 - y2**3), + 2.6459606618019 * z * (6.0 * x**4 * xy - 20.0 * xy**3 + 6.0 * xy * y**4), + 9.98394571852353e-5 + * y + * (5197.5 - 67567.5 * z2) + * (-10.0 * x2 * y2 + 5.0 * x4 + y4), + 0.00239614697244565 + * xy + * (x2 - y2) + * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z), + 0.00397356022507413 + * y + * (3.0 * x2 - y2) + * ( + 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z) + + 1063.125 * z2 + - 118.125 + ), + 0.0561946276120613 + * xy + * ( + -4.8 * z * (52.5 * z2 - 7.5) + + 2.6 + * z + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ) + + 48.0 * z + ), + 0.206472245902897 + * y + * ( + -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 2.16666666666667 + * z + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ) + - 10.9375 * z2 + + 2.1875 + ), + 1.24862677781952 * z * (1.5 * z2 - 0.5) + - 1.68564615005635 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 2.02901851395672 + * z + * ( + -1.45833333333333 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + + 1.83333333333333 + * z + * ( + -1.33333333333333 * z * (1.5 * z2 - 0.5) + + 1.8 + * z + * ( + 1.75 + * z + * ( + 1.66666666666667 * z * (1.5 * z2 - 0.5) + - 0.666666666666667 * z + ) + - 1.125 * z2 + + 0.375 + ) + + 0.533333333333333 * z + ) + + 0.9375 * z2 + - 0.3125 + ) + - 0.499450711127808 * z, + 0.206472245902897 + * x + * ( + -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 2.16666666666667 + * z + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ) + - 10.9375 * z2 + + 2.1875 + ), + 0.0280973138060306 + * (x2 - y2) + * ( + -4.8 * z * (52.5 * z2 - 7.5) + + 2.6 + * z + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ) + + 48.0 * z + ), + 0.00397356022507413 + * x + * (x2 - 3.0 * y2) + * ( + 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z) + + 1063.125 * z2 + - 118.125 + ), + 0.000599036743111412 + * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z) + * (-6.0 * x2 * y2 + x4 + y4), + 9.98394571852353e-5 + * x + * (5197.5 - 67567.5 * z2) + * (-10.0 * x2 * y2 + x4 + 5.0 * y4), + 2.6459606618019 * z * (x2**3 + 15.0 * x2 * y4 - 15.0 * x4 * y2 - y2**3), + -0.707162732524596 + * x + * (x2**3 + 35.0 * x2 * y4 - 21.0 * x4 * y2 - 7.0 * y2**3), + ], + -1, + ) + + +# @torch.jit.script +def rsh_cart_8(xyz: torch.Tensor): + """Computes all real spherical harmonics up to degree 8. + + This is an autogenerated method. See + https://github.com/cheind/torch-spherical-harmonics + for more information. + + Params: + xyz: (N,...,3) tensor of points on the unit sphere + + Returns: + rsh: (N,...,81) real spherical harmonics + projections of input. Ynm is found at index + `n*(n+1) + m`, with `0 <= n <= degree` and + `-n <= m <= n`. + """ + x = xyz[..., 0] + y = xyz[..., 1] + z = xyz[..., 2] + + x2 = x**2 + y2 = y**2 + z2 = z**2 + xy = x * y + xz = x * z + yz = y * z + x4 = x2**2 + y4 = y2**2 + # z4 = z2**2 + return torch.stack( + [ + 0.282094791773878 * torch.ones(1, device=xyz.device).expand(xyz.shape[:-1]), + -0.48860251190292 * y, + 0.48860251190292 * z, + -0.48860251190292 * x, + 1.09254843059208 * xy, + -1.09254843059208 * yz, + 0.94617469575756 * z2 - 0.31539156525252, + -1.09254843059208 * xz, + 0.54627421529604 * x2 - 0.54627421529604 * y2, + -0.590043589926644 * y * (3.0 * x2 - y2), + 2.89061144264055 * xy * z, + 0.304697199642977 * y * (1.5 - 7.5 * z2), + 1.24392110863372 * z * (1.5 * z2 - 0.5) - 0.497568443453487 * z, + 0.304697199642977 * x * (1.5 - 7.5 * z2), + 1.44530572132028 * z * (x2 - y2), + -0.590043589926644 * x * (x2 - 3.0 * y2), + 2.5033429417967 * xy * (x2 - y2), + -1.77013076977993 * yz * (3.0 * x2 - y2), + 0.126156626101008 * xy * (52.5 * z2 - 7.5), + 0.267618617422916 * y * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z), + 1.48099765681286 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 0.952069922236839 * z2 + + 0.317356640745613, + 0.267618617422916 * x * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z), + 0.063078313050504 * (x2 - y2) * (52.5 * z2 - 7.5), + -1.77013076977993 * xz * (x2 - 3.0 * y2), + -3.75501441269506 * x2 * y2 + + 0.625835735449176 * x4 + + 0.625835735449176 * y4, + -0.65638205684017 * y * (-10.0 * x2 * y2 + 5.0 * x4 + y4), + 8.30264925952416 * xy * z * (x2 - y2), + 0.00931882475114763 * y * (52.5 - 472.5 * z2) * (3.0 * x2 - y2), + 0.0913054625709205 * xy * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z), + 0.241571547304372 + * y + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ), + -1.24747010616985 * z * (1.5 * z2 - 0.5) + + 1.6840846433293 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 0.498988042467941 * z, + 0.241571547304372 + * x + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ), + 0.0456527312854602 * (x2 - y2) * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z), + 0.00931882475114763 * x * (52.5 - 472.5 * z2) * (x2 - 3.0 * y2), + 2.07566231488104 * z * (-6.0 * x2 * y2 + x4 + y4), + -0.65638205684017 * x * (-10.0 * x2 * y2 + x4 + 5.0 * y4), + 4.09910463115149 * x**4 * xy + - 13.6636821038383 * xy**3 + + 4.09910463115149 * xy * y**4, + -2.36661916223175 * yz * (-10.0 * x2 * y2 + 5.0 * x4 + y4), + 0.00427144889505798 * xy * (x2 - y2) * (5197.5 * z2 - 472.5), + 0.00584892228263444 + * y + * (3.0 * x2 - y2) + * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z), + 0.0701870673916132 + * xy + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ), + 0.221950995245231 + * y + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ), + -1.48328138624466 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + + 1.86469659985043 + * z + * ( + -1.33333333333333 * z * (1.5 * z2 - 0.5) + + 1.8 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 0.533333333333333 * z + ) + + 0.953538034014426 * z2 + - 0.317846011338142, + 0.221950995245231 + * x + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ), + 0.0350935336958066 + * (x2 - y2) + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ), + 0.00584892228263444 + * x + * (x2 - 3.0 * y2) + * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z), + 0.0010678622237645 * (5197.5 * z2 - 472.5) * (-6.0 * x2 * y2 + x4 + y4), + -2.36661916223175 * xz * (-10.0 * x2 * y2 + x4 + 5.0 * y4), + 0.683184105191914 * x2**3 + + 10.2477615778787 * x2 * y4 + - 10.2477615778787 * x4 * y2 + - 0.683184105191914 * y2**3, + -0.707162732524596 + * y + * (7.0 * x2**3 + 21.0 * x2 * y4 - 35.0 * x4 * y2 - y2**3), + 2.6459606618019 * z * (6.0 * x**4 * xy - 20.0 * xy**3 + 6.0 * xy * y**4), + 9.98394571852353e-5 + * y + * (5197.5 - 67567.5 * z2) + * (-10.0 * x2 * y2 + 5.0 * x4 + y4), + 0.00239614697244565 + * xy + * (x2 - y2) + * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z), + 0.00397356022507413 + * y + * (3.0 * x2 - y2) + * ( + 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z) + + 1063.125 * z2 + - 118.125 + ), + 0.0561946276120613 + * xy + * ( + -4.8 * z * (52.5 * z2 - 7.5) + + 2.6 + * z + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ) + + 48.0 * z + ), + 0.206472245902897 + * y + * ( + -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 2.16666666666667 + * z + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ) + - 10.9375 * z2 + + 2.1875 + ), + 1.24862677781952 * z * (1.5 * z2 - 0.5) + - 1.68564615005635 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 2.02901851395672 + * z + * ( + -1.45833333333333 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + + 1.83333333333333 + * z + * ( + -1.33333333333333 * z * (1.5 * z2 - 0.5) + + 1.8 + * z + * ( + 1.75 + * z + * ( + 1.66666666666667 * z * (1.5 * z2 - 0.5) + - 0.666666666666667 * z + ) + - 1.125 * z2 + + 0.375 + ) + + 0.533333333333333 * z + ) + + 0.9375 * z2 + - 0.3125 + ) + - 0.499450711127808 * z, + 0.206472245902897 + * x + * ( + -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 2.16666666666667 + * z + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ) + - 10.9375 * z2 + + 2.1875 + ), + 0.0280973138060306 + * (x2 - y2) + * ( + -4.8 * z * (52.5 * z2 - 7.5) + + 2.6 + * z + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ) + + 48.0 * z + ), + 0.00397356022507413 + * x + * (x2 - 3.0 * y2) + * ( + 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z) + + 1063.125 * z2 + - 118.125 + ), + 0.000599036743111412 + * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z) + * (-6.0 * x2 * y2 + x4 + y4), + 9.98394571852353e-5 + * x + * (5197.5 - 67567.5 * z2) + * (-10.0 * x2 * y2 + x4 + 5.0 * y4), + 2.6459606618019 * z * (x2**3 + 15.0 * x2 * y4 - 15.0 * x4 * y2 - y2**3), + -0.707162732524596 + * x + * (x2**3 + 35.0 * x2 * y4 - 21.0 * x4 * y2 - 7.0 * y2**3), + 5.83141328139864 * xy * (x2**3 + 7.0 * x2 * y4 - 7.0 * x4 * y2 - y2**3), + -2.91570664069932 + * yz + * (7.0 * x2**3 + 21.0 * x2 * y4 - 35.0 * x4 * y2 - y2**3), + 7.87853281621404e-6 + * (1013512.5 * z2 - 67567.5) + * (6.0 * x**4 * xy - 20.0 * xy**3 + 6.0 * xy * y**4), + 5.10587282657803e-5 + * y + * (5.0 * z * (5197.5 - 67567.5 * z2) + 41580.0 * z) + * (-10.0 * x2 * y2 + 5.0 * x4 + y4), + 0.00147275890257803 + * xy + * (x2 - y2) + * ( + 3.75 * z * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z) + - 14293.125 * z2 + + 1299.375 + ), + 0.0028519853513317 + * y + * (3.0 * x2 - y2) + * ( + -7.33333333333333 * z * (52.5 - 472.5 * z2) + + 3.0 + * z + * ( + 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z) + + 1063.125 * z2 + - 118.125 + ) + - 560.0 * z + ), + 0.0463392770473559 + * xy + * ( + -4.125 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + + 2.5 + * z + * ( + -4.8 * z * (52.5 * z2 - 7.5) + + 2.6 + * z + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ) + + 48.0 * z + ) + + 137.8125 * z2 + - 19.6875 + ), + 0.193851103820053 + * y + * ( + 3.2 * z * (1.5 - 7.5 * z2) + - 2.51428571428571 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + + 2.14285714285714 + * z + * ( + -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 2.16666666666667 + * z + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 + * z + * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ) + - 10.9375 * z2 + + 2.1875 + ) + + 5.48571428571429 * z + ), + 1.48417251362228 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.86581687426801 + * z + * ( + -1.33333333333333 * z * (1.5 * z2 - 0.5) + + 1.8 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 0.533333333333333 * z + ) + + 2.1808249179756 + * z + * ( + 1.14285714285714 * z * (1.5 * z2 - 0.5) + - 1.54285714285714 + * z + * ( + 1.75 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + - 1.125 * z2 + + 0.375 + ) + + 1.85714285714286 + * z + * ( + -1.45833333333333 + * z + * (1.66666666666667 * z * (1.5 * z2 - 0.5) - 0.666666666666667 * z) + + 1.83333333333333 + * z + * ( + -1.33333333333333 * z * (1.5 * z2 - 0.5) + + 1.8 + * z + * ( + 1.75 + * z + * ( + 1.66666666666667 * z * (1.5 * z2 - 0.5) + - 0.666666666666667 * z + ) + - 1.125 * z2 + + 0.375 + ) + + 0.533333333333333 * z + ) + + 0.9375 * z2 + - 0.3125 + ) + - 0.457142857142857 * z + ) + - 0.954110901614325 * z2 + + 0.318036967204775, + 0.193851103820053 + * x + * ( + 3.2 * z * (1.5 - 7.5 * z2) + - 2.51428571428571 + * z + * ( + 2.25 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + + 2.14285714285714 + * z + * ( + -2.625 * z * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 2.16666666666667 + * z + * ( + -2.8 * z * (1.5 - 7.5 * z2) + + 2.2 + * z + * ( + 2.25 + * z + * (2.33333333333333 * z * (1.5 - 7.5 * z2) + 4.0 * z) + + 9.375 * z2 + - 1.875 + ) + - 4.8 * z + ) + - 10.9375 * z2 + + 2.1875 + ) + + 5.48571428571429 * z + ), + 0.0231696385236779 + * (x2 - y2) + * ( + -4.125 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + + 2.5 + * z + * ( + -4.8 * z * (52.5 * z2 - 7.5) + + 2.6 + * z + * ( + 2.75 * z * (3.0 * z * (52.5 * z2 - 7.5) - 30.0 * z) + - 91.875 * z2 + + 13.125 + ) + + 48.0 * z + ) + + 137.8125 * z2 + - 19.6875 + ), + 0.0028519853513317 + * x + * (x2 - 3.0 * y2) + * ( + -7.33333333333333 * z * (52.5 - 472.5 * z2) + + 3.0 + * z + * ( + 3.25 * z * (3.66666666666667 * z * (52.5 - 472.5 * z2) + 280.0 * z) + + 1063.125 * z2 + - 118.125 + ) + - 560.0 * z + ), + 0.000368189725644507 + * (-6.0 * x2 * y2 + x4 + y4) + * ( + 3.75 * z * (4.33333333333333 * z * (5197.5 * z2 - 472.5) - 3150.0 * z) + - 14293.125 * z2 + + 1299.375 + ), + 5.10587282657803e-5 + * x + * (5.0 * z * (5197.5 - 67567.5 * z2) + 41580.0 * z) + * (-10.0 * x2 * y2 + x4 + 5.0 * y4), + 7.87853281621404e-6 + * (1013512.5 * z2 - 67567.5) + * (x2**3 + 15.0 * x2 * y4 - 15.0 * x4 * y2 - y2**3), + -2.91570664069932 + * xz + * (x2**3 + 35.0 * x2 * y4 - 21.0 * x4 * y2 - 7.0 * y2**3), + -20.4099464848952 * x2**3 * y2 + - 20.4099464848952 * x2 * y2**3 + + 0.72892666017483 * x4**2 + + 51.0248662122381 * x4 * y4 + + 0.72892666017483 * y4**2, + ], + -1, + ) + + +__all__ = [ + "rsh_cart_0", + "rsh_cart_1", + "rsh_cart_2", + "rsh_cart_3", + "rsh_cart_4", + "rsh_cart_5", + "rsh_cart_6", + "rsh_cart_7", + "rsh_cart_8", +] + + +from typing import Optional +import torch + + +class SphHarm(torch.nn.Module): + def __init__(self, m, n, dtype=torch.float32) -> None: + super().__init__() + self.dtype = dtype + m = torch.tensor(list(range(-m + 1, m))) + n = torch.tensor(list(range(n))) + self.is_normalized = False + vals = torch.cartesian_prod(m, n).T + vals = vals[:, vals[0] <= vals[1]] + m, n = vals.unbind(0) + + self.register_buffer("m", tensor=m) + self.register_buffer("n", tensor=n) + self.register_buffer("l_max", tensor=torch.max(self.n)) + + f_a, f_b, initial_value, d0_mask_3d, d1_mask_3d = self._init_legendre() + self.register_buffer("f_a", tensor=f_a) + self.register_buffer("f_b", tensor=f_b) + self.register_buffer("d0_mask_3d", tensor=d0_mask_3d) + self.register_buffer("d1_mask_3d", tensor=d1_mask_3d) + self.register_buffer("initial_value", tensor=initial_value) + + @property + def device(self): + return next(self.buffers()).device + + def forward(self, points: torch.Tensor) -> torch.Tensor: + """Computes the spherical harmonics.""" + # Y_l^m = (-1) ^ m c_l^m P_l^m(cos(theta)) exp(i m phi) + B, N, D = points.shape + dtype = points.dtype + theta, phi = points.view(-1, D).to(self.dtype).unbind(-1) + cos_colatitude = torch.cos(phi) + legendre = self._gen_associated_legendre(cos_colatitude) + vals = torch.stack([self.m.abs(), self.n], dim=0) + vals = torch.cat( + [ + vals.repeat(1, theta.shape[0]), + torch.arange(theta.shape[0], device=theta.device) + .unsqueeze(0) + .repeat_interleave(vals.shape[1], dim=1), + ], + dim=0, + ) + legendre_vals = legendre[vals[0], vals[1], vals[2]] + legendre_vals = legendre_vals.reshape(-1, theta.shape[0]) + angle = torch.outer(self.m.abs(), theta) + vandermonde = torch.complex(torch.cos(angle), torch.sin(angle)) + harmonics = torch.complex( + legendre_vals * torch.real(vandermonde), + legendre_vals * torch.imag(vandermonde), + ) + + # Negative order. + m = self.m.unsqueeze(-1) + harmonics = torch.where( + m < 0, (-1.0) ** m.abs() * torch.conj(harmonics), harmonics + ) + harmonics = harmonics.permute(1, 0).reshape(B, N, -1).to(dtype) + return harmonics + + def _gen_recurrence_mask(self) -> tuple[torch.Tensor, torch.Tensor]: + """Generates mask for recurrence relation on the remaining entries. + + The remaining entries are with respect to the diagonal and offdiagonal + entries. + + Args: + l_max: see `gen_normalized_legendre`. + Returns: + torch.Tensors representing the mask used by the recurrence relations. + """ + + # Computes all coefficients. + m_mat, l_mat = torch.meshgrid( + torch.arange(0, self.l_max + 1, device=self.device, dtype=self.dtype), + torch.arange(0, self.l_max + 1, device=self.device, dtype=self.dtype), + indexing="ij", + ) + if self.is_normalized: + c0 = l_mat * l_mat + c1 = m_mat * m_mat + c2 = 2.0 * l_mat + c3 = (l_mat - 1.0) * (l_mat - 1.0) + d0 = torch.sqrt((4.0 * c0 - 1.0) / (c0 - c1)) + d1 = torch.sqrt(((c2 + 1.0) * (c3 - c1)) / ((c2 - 3.0) * (c0 - c1))) + else: + d0 = (2.0 * l_mat - 1.0) / (l_mat - m_mat) + d1 = (l_mat + m_mat - 1.0) / (l_mat - m_mat) + + d0_mask_indices = torch.triu_indices(self.l_max + 1, 1) + d1_mask_indices = torch.triu_indices(self.l_max + 1, 2) + + d_zeros = torch.zeros( + (self.l_max + 1, self.l_max + 1), dtype=self.dtype, device=self.device + ) + d_zeros[d0_mask_indices] = d0[d0_mask_indices] + d0_mask = d_zeros + + d_zeros = torch.zeros( + (self.l_max + 1, self.l_max + 1), dtype=self.dtype, device=self.device + ) + d_zeros[d1_mask_indices] = d1[d1_mask_indices] + d1_mask = d_zeros + + # Creates a 3D mask that contains 1s on the diagonal plane and 0s elsewhere. + i = torch.arange(self.l_max + 1, device=self.device)[:, None, None] + j = torch.arange(self.l_max + 1, device=self.device)[None, :, None] + k = torch.arange(self.l_max + 1, device=self.device)[None, None, :] + mask = (i + j - k == 0).to(self.dtype) + d0_mask_3d = torch.einsum("jk,ijk->ijk", d0_mask, mask) + d1_mask_3d = torch.einsum("jk,ijk->ijk", d1_mask, mask) + return (d0_mask_3d, d1_mask_3d) + + def _recursive(self, i: int, p_val: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + coeff_0 = self.d0_mask_3d[i] + coeff_1 = self.d1_mask_3d[i] + h = torch.einsum( + "ij,ijk->ijk", + coeff_0, + torch.einsum("ijk,k->ijk", torch.roll(p_val, shifts=1, dims=1), x), + ) - torch.einsum("ij,ijk->ijk", coeff_1, torch.roll(p_val, shifts=2, dims=1)) + p_val = p_val + h + return p_val + + def _init_legendre(self): + a_idx = torch.arange(1, self.l_max + 1, dtype=self.dtype, device=self.device) + b_idx = torch.arange(self.l_max, dtype=self.dtype, device=self.device) + if self.is_normalized: + # The initial value p(0,0). + initial_value: torch.Tensor = torch.tensor( + 0.5 / (torch.pi**0.5), device=self.device + ) + f_a = torch.cumprod(-1 * torch.sqrt(1.0 + 0.5 / a_idx), dim=0) + f_b = torch.sqrt(2.0 * b_idx + 3.0) + else: + # The initial value p(0,0). + initial_value = torch.tensor(1.0, device=self.device) + f_a = torch.cumprod(1.0 - 2.0 * a_idx, dim=0) + f_b = 2.0 * b_idx + 1.0 + + d0_mask_3d, d1_mask_3d = self._gen_recurrence_mask() + return f_a, f_b, initial_value, d0_mask_3d, d1_mask_3d + + def _gen_associated_legendre(self, x: torch.Tensor) -> torch.Tensor: + r"""Computes associated Legendre functions (ALFs) of the first kind. + + The ALFs of the first kind are used in spherical harmonics. The spherical + harmonic of degree `l` and order `m` can be written as + `Y_l^m(θ, φ) = N_l^m * P_l^m(cos(θ)) * exp(i m φ)`, where `N_l^m` is the + normalization factor and θ and φ are the colatitude and longitude, + repectively. `N_l^m` is chosen in the way that the spherical harmonics form + a set of orthonormal basis function of L^2(S^2). For the computational + efficiency of spherical harmonics transform, the normalization factor is + used in the computation of the ALFs. In addition, normalizing `P_l^m` + avoids overflow/underflow and achieves better numerical stability. Three + recurrence relations are used in the computation. + + Args: + l_max: The maximum degree of the associated Legendre function. Both the + degrees and orders are `[0, 1, 2, ..., l_max]`. + x: A vector of type `float32`, `float64` containing the sampled points in + spherical coordinates, at which the ALFs are computed; `x` is essentially + `cos(θ)`. For the numerical integration used by the spherical harmonics + transforms, `x` contains the quadrature points in the interval of + `[-1, 1]`. There are several approaches to provide the quadrature points: + Gauss-Legendre method (`scipy.special.roots_legendre`), Gauss-Chebyshev + method (`scipy.special.roots_chebyu`), and Driscoll & Healy + method (Driscoll, James R., and Dennis M. Healy. "Computing Fourier + transforms and convolutions on the 2-sphere." Advances in applied + mathematics 15, no. 2 (1994): 202-250.). The Gauss-Legendre quadrature + points are nearly equal-spaced along θ and provide exact discrete + orthogonality, (P^m)^T W P_m = I, where `T` represents the transpose + operation, `W` is a diagonal matrix containing the quadrature weights, + and `I` is the identity matrix. The Gauss-Chebyshev points are equally + spaced, which only provide approximate discrete orthogonality. The + Driscoll & Healy qudarture points are equally spaced and provide the + exact discrete orthogonality. The number of sampling points is required to + be twice as the number of frequency points (modes) in the Driscoll & Healy + approach, which enables FFT and achieves a fast spherical harmonics + transform. + is_normalized: True if the associated Legendre functions are normalized. + With normalization, `N_l^m` is applied such that the spherical harmonics + form a set of orthonormal basis functions of L^2(S^2). + + Returns: + The 3D array of shape `(l_max + 1, l_max + 1, len(x))` containing the values + of the ALFs at `x`; the dimensions in the sequence of order, degree, and + evalution points. + """ + p = torch.zeros( + (self.l_max + 1, self.l_max + 1, x.shape[0]), dtype=x.dtype, device=x.device + ) + p[0, 0] = self.initial_value + + # Compute the diagonal entries p(l,l) with recurrence. + y = torch.cumprod( + torch.broadcast_to(torch.sqrt(1.0 - x * x), (self.l_max, x.shape[0])), dim=0 + ) + p_diag = self.initial_value * torch.einsum("i,ij->ij", self.f_a, y) + # torch.diag_indices(l_max + 1) + diag_indices = torch.stack( + [torch.arange(0, self.l_max + 1, device=x.device)] * 2, dim=0 + ) + p[(diag_indices[0][1:], diag_indices[1][1:])] = p_diag + + diag_indices = torch.stack( + [torch.arange(0, self.l_max, device=x.device)] * 2, dim=0 + ) + + # Compute the off-diagonal entries with recurrence. + p_offdiag = torch.einsum( + "ij,ij->ij", + torch.einsum("i,j->ij", self.f_b, x), + p[(diag_indices[0], diag_indices[1])], + ) # p[torch.diag_indices(l_max)]) + p[(diag_indices[0][: self.l_max], diag_indices[1][: self.l_max] + 1)] = ( + p_offdiag + ) + + # Compute the remaining entries with recurrence. + if self.l_max > 1: + for i in range(2, self.l_max + 1): + p = self._recursive(i, p, x) + return p diff --git a/src/misc/step_tracker.py b/src/misc/step_tracker.py new file mode 100644 index 0000000000000000000000000000000000000000..7298ffcf5cb028799d67c4dedf7ef1a4bf6fe802 --- /dev/null +++ b/src/misc/step_tracker.py @@ -0,0 +1,23 @@ +from multiprocessing import RLock + +import torch +from jaxtyping import Int64 +from torch import Tensor +from torch.multiprocessing import Manager + + +class StepTracker: + lock: RLock + step: Int64[Tensor, ""] + + def __init__(self): + self.lock = Manager().RLock() + self.step = torch.tensor(0, dtype=torch.int64).share_memory_() + + def set_step(self, step: int) -> None: + with self.lock: + self.step.fill_(step) + + def get_step(self) -> int: + with self.lock: + return self.step.item() diff --git a/src/misc/utils.py b/src/misc/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4e29b8f7e6fa06937b8196988a2a5ef0beb7280b --- /dev/null +++ b/src/misc/utils.py @@ -0,0 +1,73 @@ +import torch + +from src.visualization.color_map import apply_color_map_to_image +import torch.distributed as dist + +def inverse_normalize(tensor, mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)): + mean = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1) + std = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1) + return tensor.mul(std).add(mean) + + +# Color-map the result. +def vis_depth_map(result, near=None, far=None): + if near is None and far is None: + far = result.view(-1)[:16_000_000].quantile(0.99).log() + try: + near = result[result > 0][:16_000_000].quantile(0.01).log() + except: + print("No valid depth values found.") + near = torch.zeros_like(far) + else: + near = near.log() + far = far.log() + + result = result.log() + result = 1 - (result - near) / (far - near) + return apply_color_map_to_image(result, "turbo") + + +def confidence_map(result): + # far = result.view(-1)[:16_000_000].quantile(0.99).log() + # try: + # near = result[result > 0][:16_000_000].quantile(0.01).log() + # except: + # print("No valid depth values found.") + # near = torch.zeros_like(far) + # result = result.log() + # result = 1 - (result - near) / (far - near) + result = result / result.view(-1).max() + return apply_color_map_to_image(result, "magma") + + +def get_overlap_tag(overlap): + if 0.05 <= overlap <= 0.3: + overlap_tag = "small" + elif overlap <= 0.55: + overlap_tag = "medium" + elif overlap <= 0.8: + overlap_tag = "large" + else: + overlap_tag = "ignore" + + return overlap_tag + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() \ No newline at end of file diff --git a/src/misc/wandb_tools.py b/src/misc/wandb_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..a5dc362ad13eef89fcde805695095dbff8ef375c --- /dev/null +++ b/src/misc/wandb_tools.py @@ -0,0 +1,62 @@ +from pathlib import Path + +import wandb + + +def version_to_int(artifact) -> int: + """Convert versions of the form vX to X. For example, v12 to 12.""" + return int(artifact.version[1:]) + + +def download_checkpoint( + run_id: str, + download_dir: Path, + version: str | None, +) -> Path: + api = wandb.Api() + run = api.run(run_id) + + # Find the latest saved model checkpoint. + chosen = None + for artifact in run.logged_artifacts(): + if artifact.type != "model" or artifact.state != "COMMITTED": + continue + + # If no version is specified, use the latest. + if version is None: + if chosen is None or version_to_int(artifact) > version_to_int(chosen): + chosen = artifact + + # If a specific verison is specified, look for it. + elif version == artifact.version: + chosen = artifact + break + + # Download the checkpoint. + download_dir.mkdir(exist_ok=True, parents=True) + root = download_dir / run_id + chosen.download(root=root) + return root / "model.ckpt" + + +def update_checkpoint_path(path: str | None, wandb_cfg: dict) -> Path | None: + if path is None: + return None + + if not str(path).startswith("wandb://"): + return Path(path) + + run_id, *version = path[len("wandb://") :].split(":") + if len(version) == 0: + version = None + elif len(version) == 1: + version = version[0] + else: + raise ValueError("Invalid version specifier!") + + project = wandb_cfg["project"] + return download_checkpoint( + f"{project}/{run_id}", + Path("checkpoints"), + version, + ) diff --git a/src/misc/weight_modify.py b/src/misc/weight_modify.py new file mode 100644 index 0000000000000000000000000000000000000000..b6ed4771c239933615764e91800e513567f42feb --- /dev/null +++ b/src/misc/weight_modify.py @@ -0,0 +1,197 @@ +import logging +from typing import List, Dict + +import math +import torch +from torch import nn as nn +import torch.nn.functional as F + + +_logger = logging.getLogger(__name__) + + +def resample_patch_embed( + patch_embed, + new_size: List[int], + interpolation: str = 'bicubic', + antialias: bool = True, + verbose: bool = False, +): + """Resample the weights of the patch embedding kernel to target resolution. + We resample the patch embedding kernel by approximately inverting the effect + of patch resizing. + + Code based on: + https://github.com/google-research/big_vision/blob/b00544b81f8694488d5f36295aeb7972f3755ffe/big_vision/models/proj/flexi/vit.py + + With this resizing, we can for example load a B/8 filter into a B/16 model + and, on 2x larger input image, the result will match. + + Args: + patch_embed: original parameter to be resized. + new_size (tuple(int, int): target shape (height, width)-only. + interpolation (str): interpolation for resize + antialias (bool): use anti-aliasing filter in resize + verbose (bool): log operation + Returns: + Resized patch embedding kernel. + """ + import numpy as np + try: + import functorch + vmap = functorch.vmap + except ImportError: + if hasattr(torch, 'vmap'): + vmap = torch.vmap + else: + assert False, "functorch or a version of torch with vmap is required for FlexiViT resizing." + + assert len(patch_embed.shape) == 4, "Four dimensions expected" + assert len(new_size) == 2, "New shape should only be hw" + old_size = patch_embed.shape[-2:] + if tuple(old_size) == tuple(new_size): + return patch_embed + + if verbose: + _logger.info(f"Resize patch embedding {patch_embed.shape} to {new_size}, w/ {interpolation} interpolation.") + + def resize(x_np, _new_size): + x_tf = torch.Tensor(x_np)[None, None, ...] + x_upsampled = F.interpolate( + x_tf, size=_new_size, mode=interpolation, antialias=antialias)[0, 0, ...].numpy() + return x_upsampled + + def get_resize_mat(_old_size, _new_size): + mat = [] + for i in range(np.prod(_old_size)): + basis_vec = np.zeros(_old_size) + basis_vec[np.unravel_index(i, _old_size)] = 1. + mat.append(resize(basis_vec, _new_size).reshape(-1)) + return np.stack(mat).T + + resize_mat = get_resize_mat(old_size, new_size) + resize_mat_pinv = torch.tensor(np.linalg.pinv(resize_mat.T), device=patch_embed.device) + + def resample_kernel(kernel): + resampled_kernel = resize_mat_pinv @ kernel.reshape(-1) + return resampled_kernel.reshape(new_size) + + v_resample_kernel = vmap(vmap(resample_kernel, 0, 0), 1, 1) + orig_dtype = patch_embed.dtype + patch_embed = patch_embed.float() + patch_embed = v_resample_kernel(patch_embed) + patch_embed = patch_embed.to(orig_dtype) + return patch_embed + + +def adapt_input_conv(in_chans, conv_weight): + conv_type = conv_weight.dtype + conv_weight = conv_weight.float() # Some weights are in torch.half, ensure it's float for sum on CPU + O, I, J, K = conv_weight.shape + if in_chans == 1: + if I > 3: + assert conv_weight.shape[1] % 3 == 0 + # For models with space2depth stems + conv_weight = conv_weight.reshape(O, I // 3, 3, J, K) + conv_weight = conv_weight.sum(dim=2, keepdim=False) + else: + conv_weight = conv_weight.sum(dim=1, keepdim=True) + elif in_chans != 3: + if I != 3: + raise NotImplementedError('Weight format not supported by conversion.') + else: + # NOTE this strategy should be better than random init, but there could be other combinations of + # the original RGB input layer weights that'd work better for specific cases. + repeat = int(math.ceil(in_chans / 3)) + conv_weight = conv_weight.repeat(1, repeat, 1, 1)[:, :in_chans, :, :] + conv_weight *= (3 / float(in_chans)) + + # instead of assigning the same weight to all channels, we can assign higher weight for original RGB channels + # conv_weight[:, :3, :, :] = conv_weight[:, :3, :, :] * 0.5 + # conv_weight[:, 3:, :, :] = conv_weight[:, 3:, :, :] * 0.5 * (3 / float(in_chans - 3)) + + conv_weight = conv_weight.to(conv_type) + return conv_weight + + +def adapt_head_conv(conv_weight): + conv_type = conv_weight.dtype + conv_weight = conv_weight.float() # Some weights are in torch.half, ensure it's float for sum on CPU + O, I, J, K = conv_weight.shape + + conv_weight_new = torch.chunk(conv_weight, 6, dim=1) + conv_weight_new = [conv_weight_new.mean(dim=1, keepdim=True) for conv_weight_new in conv_weight_new] + conv_weight_new = torch.cat(conv_weight_new, dim=1) * 0.5 + conv_weight = torch.cat([conv_weight, conv_weight_new], dim=1) + conv_weight = conv_weight.to(conv_type) + return conv_weight + + +def adapt_linear(conv_weight): + conv_type = conv_weight.dtype + conv_weight = conv_weight.float() # Some weights are in torch.half, ensure it's float for sum on CPU + O, I = conv_weight.shape + + conv_weight_new = torch.tensor_split(conv_weight, 81, dim=1) + conv_weight_new = [conv_weight_new.mean(dim=1, keepdim=True) for conv_weight_new in conv_weight_new] + conv_weight_new = torch.cat(conv_weight_new, dim=1) + # conv_weight = torch.cat([conv_weight, conv_weight_new], dim=1) + conv_weight = torch.cat([conv_weight * 0.5, conv_weight_new * 0.5], dim=1) + conv_weight = conv_weight.to(conv_type) + return conv_weight + + +def checkpoint_filter_fn( + state_dict: Dict[str, torch.Tensor], + model: nn.Module, + interpolation: str = 'bicubic', + antialias: bool = True, +) -> Dict[str, torch.Tensor]: + """ convert patch embedding weight from manual patchify + linear proj to conv""" + out_dict = {} + # state_dict = state_dict.get('model', state_dict) + # state_dict = state_dict.get('state_dict', state_dict) + prefix = '' + + if prefix: + # filter on & remove prefix string from keys + state_dict = {k[len(prefix):]: v for k, v in state_dict.items() if k.startswith(prefix)} + + for k, v in state_dict.items(): + if 'patch_embed.proj.weight' in k: + O, I, H, W = model.backbone.patch_embed.proj.weight.shape + if len(v.shape) < 4: + # For old models that I trained prior to conv based patchification + O, I, H, W = model.backbone.patch_embed.proj.weight.shape + v = v.reshape(O, -1, H, W) + if v.shape[-1] != W or v.shape[-2] != H: + v = resample_patch_embed( + v, + (H, W), + interpolation=interpolation, + antialias=antialias, + verbose=True, + ) + if v.shape[1] != I: + v = adapt_input_conv(I, v) + # elif 'downstream_head1.dpt.head.0.weight' in k or 'downstream_head2.dpt.head.0.weight' in k: + # v = adapt_head_conv(v) + + elif 'decoder_embed.weight' in k: + O, I = model.backbone.decoder_embed.weight.shape + if v.shape[1] != I: + v = adapt_linear(v) + + out_dict[k] = v + + # add prefix to make our model happy + prefix = 'backbone.' + out_dict = {prefix + k if 'downstream_head' not in k else k: v for k, v in out_dict.items()} + + # # remove the conf head weights + out_dict['downstream_head1.dpt.head.4.weight'] = out_dict['downstream_head1.dpt.head.4.weight'][0:3] + out_dict['downstream_head1.dpt.head.4.bias'] = out_dict['downstream_head1.dpt.head.4.bias'][0:3] + out_dict['downstream_head2.dpt.head.4.weight'] = out_dict['downstream_head2.dpt.head.4.weight'][0:3] + out_dict['downstream_head2.dpt.head.4.bias'] = out_dict['downstream_head2.dpt.head.4.bias'][0:3] + + return out_dict diff --git a/src/model/decoder/__init__.py b/src/model/decoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6d757213bdfaeb3994ab67cbf47852f1a2b5c1d3 --- /dev/null +++ b/src/model/decoder/__init__.py @@ -0,0 +1,12 @@ +from .decoder import Decoder +from .decoder_splatting_cuda import DecoderSplattingCUDA, DecoderSplattingCUDACfg + +DECODERS = { + "splatting_cuda": DecoderSplattingCUDA, +} + +DecoderCfg = DecoderSplattingCUDACfg + + +def get_decoder(decoder_cfg: DecoderCfg) -> Decoder: + return DECODERS[decoder_cfg.name](decoder_cfg) diff --git a/src/model/decoder/cuda_splatting.py b/src/model/decoder/cuda_splatting.py new file mode 100644 index 0000000000000000000000000000000000000000..367c8cc13d998c97b1cdc96c6d076f04fda3be4c --- /dev/null +++ b/src/model/decoder/cuda_splatting.py @@ -0,0 +1,244 @@ +from math import isqrt +from typing import Literal + +import torch +from diff_gaussian_rasterization import ( + GaussianRasterizationSettings, + GaussianRasterizer, +) +from einops import einsum, rearrange, repeat +from jaxtyping import Float, Bool +from torch import Tensor + +from ...geometry.projection import get_fov, homogenize_points + + +def get_projection_matrix( + near: Float[Tensor, " batch"], + far: Float[Tensor, " batch"], + fov_x: Float[Tensor, " batch"], + fov_y: Float[Tensor, " batch"], +) -> Float[Tensor, "batch 4 4"]: + """Maps points in the viewing frustum to (-1, 1) on the X/Y axes and (0, 1) on the Z + axis. Differs from the OpenGL version in that Z doesn't have range (-1, 1) after + transformation and that Z is flipped. + """ + tan_fov_x = (0.5 * fov_x).tan() + tan_fov_y = (0.5 * fov_y).tan() + + top = tan_fov_y * near + bottom = -top + right = tan_fov_x * near + left = -right + + (b,) = near.shape + result = torch.zeros((b, 4, 4), dtype=torch.float32, device=near.device) + result[:, 0, 0] = 2 * near / (right - left) + result[:, 1, 1] = 2 * near / (top - bottom) + result[:, 0, 2] = (right + left) / (right - left) + result[:, 1, 2] = (top + bottom) / (top - bottom) + result[:, 3, 2] = 1 + result[:, 2, 2] = far / (far - near) + result[:, 2, 3] = -(far * near) / (far - near) + return result + + +def render_cuda( + extrinsics: Float[Tensor, "batch 4 4"], + intrinsics: Float[Tensor, "batch 3 3"], + near: Float[Tensor, " batch"], + far: Float[Tensor, " batch"], + image_shape: tuple[int, int], + background_color: Float[Tensor, "batch 3"], + gaussian_means: Float[Tensor, "batch gaussian 3"], + gaussian_covariances: Float[Tensor, "batch gaussian 3 3"], + gaussian_sh_coefficients: Float[Tensor, "batch gaussian 3 d_sh"], + gaussian_opacities: Float[Tensor, "batch gaussian"], + scale_invariant: bool = True, + use_sh: bool = True, + cam_rot_delta: Float[Tensor, "batch 3"] | None = None, + cam_trans_delta: Float[Tensor, "batch 3"] | None = None, + voxel_masks: Bool[Tensor, "batch gaussian"] | None = None, +) -> tuple[Float[Tensor, "batch 3 height width"], Float[Tensor, "batch height width"]]: + assert use_sh or gaussian_sh_coefficients.shape[-1] == 1 + + # Make sure everything is in a range where numerical issues don't appear. + if scale_invariant: + scale = 1 / near + extrinsics = extrinsics.clone() + extrinsics[..., :3, 3] = extrinsics[..., :3, 3] * scale[:, None] + gaussian_covariances = gaussian_covariances * (scale[:, None, None, None] ** 2) + gaussian_means = gaussian_means * scale[:, None, None] + near = near * scale + far = far * scale + + _, _, _, n = gaussian_sh_coefficients.shape + degree = isqrt(n) - 1 + shs = rearrange(gaussian_sh_coefficients, "b g xyz n -> b g n xyz").contiguous() + + b, _, _ = extrinsics.shape + h, w = image_shape + + fov_x, fov_y = get_fov(intrinsics).unbind(dim=-1) + tan_fov_x = (0.5 * fov_x).tan() + tan_fov_y = (0.5 * fov_y).tan() + + projection_matrix = get_projection_matrix(near, far, fov_x, fov_y) + projection_matrix = rearrange(projection_matrix, "b i j -> b j i") + view_matrix = rearrange(extrinsics.inverse(), "b i j -> b j i") + full_projection = view_matrix @ projection_matrix + + all_images = [] + all_radii = [] + all_depths = [] + for i in range(b): + # Set up a tensor for the gradients of the screen-space means. + mean_gradients = torch.zeros_like(gaussian_means[i], requires_grad=True) + try: + mean_gradients.retain_grad() + except Exception: + pass + + settings = GaussianRasterizationSettings( + image_height=h, + image_width=w, + tanfovx=tan_fov_x[i].item(), + tanfovy=tan_fov_y[i].item(), + bg=background_color[i], + scale_modifier=1.0, + viewmatrix=view_matrix[i], + projmatrix=full_projection[i], + projmatrix_raw=projection_matrix[i], + sh_degree=degree, + campos=extrinsics[i, :3, 3], + prefiltered=False, # This matches the original usage. + debug=False, + ) + rasterizer = GaussianRasterizer(settings) + + row, col = torch.triu_indices(3, 3) + + if voxel_masks is not None: + voxel_mask = voxel_masks[i] + image, radii, depth, opacity, n_touched = rasterizer( + means3D=gaussian_means[i][voxel_mask], + means2D=mean_gradients[voxel_mask], + shs=shs[i][voxel_mask] if use_sh else None, + colors_precomp=None if use_sh else shs[i, :, 0, :][voxel_mask], + opacities=gaussian_opacities[i][voxel_mask, ..., None], + cov3D_precomp=gaussian_covariances[i, :, row, col][voxel_mask], + theta=cam_rot_delta[i] if cam_rot_delta is not None else None, + rho=cam_trans_delta[i] if cam_trans_delta is not None else None, + ) + else: + image, radii, depth, opacity, n_touched = rasterizer( + means3D=gaussian_means[i], + means2D=mean_gradients, + shs=shs[i] if use_sh else None, + colors_precomp=None if use_sh else shs[i, :, 0, :], + opacities=gaussian_opacities[i, ..., None], + cov3D_precomp=gaussian_covariances[i, :, row, col], + theta=cam_rot_delta[i] if cam_rot_delta is not None else None, + rho=cam_trans_delta[i] if cam_trans_delta is not None else None, + ) + all_images.append(image) + all_radii.append(radii) + all_depths.append(depth.squeeze(0)) + return torch.stack(all_images), torch.stack(all_depths) + + +def render_cuda_orthographic( + extrinsics: Float[Tensor, "batch 4 4"], + width: Float[Tensor, " batch"], + height: Float[Tensor, " batch"], + near: Float[Tensor, " batch"], + far: Float[Tensor, " batch"], + image_shape: tuple[int, int], + background_color: Float[Tensor, "batch 3"], + gaussian_means: Float[Tensor, "batch gaussian 3"], + gaussian_covariances: Float[Tensor, "batch gaussian 3 3"], + gaussian_sh_coefficients: Float[Tensor, "batch gaussian 3 d_sh"], + gaussian_opacities: Float[Tensor, "batch gaussian"], + fov_degrees: float = 0.1, + use_sh: bool = True, + dump: dict | None = None, +) -> Float[Tensor, "batch 3 height width"]: + b, _, _ = extrinsics.shape + h, w = image_shape + assert use_sh or gaussian_sh_coefficients.shape[-1] == 1 + + _, _, _, n = gaussian_sh_coefficients.shape + degree = isqrt(n) - 1 + shs = rearrange(gaussian_sh_coefficients, "b g xyz n -> b g n xyz").contiguous() + + # Create fake "orthographic" projection by moving the camera back and picking a + # small field of view. + fov_x = torch.tensor(fov_degrees, device=extrinsics.device).deg2rad() + tan_fov_x = (0.5 * fov_x).tan() + distance_to_near = (0.5 * width) / tan_fov_x + tan_fov_y = 0.5 * height / distance_to_near + fov_y = (2 * tan_fov_y).atan() + near = near + distance_to_near + far = far + distance_to_near + move_back = torch.eye(4, dtype=torch.float32, device=extrinsics.device) + move_back[2, 3] = -distance_to_near + extrinsics = extrinsics @ move_back + + # Escape hatch for visualization/figures. + if dump is not None: + dump["extrinsics"] = extrinsics + dump["fov_x"] = fov_x + dump["fov_y"] = fov_y + dump["near"] = near + dump["far"] = far + + projection_matrix = get_projection_matrix( + near, far, repeat(fov_x, "-> b", b=b), fov_y + ) + projection_matrix = rearrange(projection_matrix, "b i j -> b j i") + view_matrix = rearrange(extrinsics.inverse(), "b i j -> b j i") + full_projection = view_matrix @ projection_matrix + + all_images = [] + all_radii = [] + for i in range(b): + # Set up a tensor for the gradients of the screen-space means. + mean_gradients = torch.zeros_like(gaussian_means[i], requires_grad=True) + try: + mean_gradients.retain_grad() + except Exception: + pass + + settings = GaussianRasterizationSettings( + image_height=h, + image_width=w, + tanfovx=tan_fov_x, + tanfovy=tan_fov_y, + bg=background_color[i], + scale_modifier=1.0, + viewmatrix=view_matrix[i], + projmatrix=full_projection[i], + projmatrix_raw=projection_matrix[i], + sh_degree=degree, + campos=extrinsics[i, :3, 3], + prefiltered=False, # This matches the original usage. + debug=False, + ) + rasterizer = GaussianRasterizer(settings) + + row, col = torch.triu_indices(3, 3) + + image, radii, depth, opacity, n_touched = rasterizer( + means3D=gaussian_means[i], + means2D=mean_gradients, + shs=shs[i] if use_sh else None, + colors_precomp=None if use_sh else shs[i, :, 0, :], + opacities=gaussian_opacities[i, ..., None], + cov3D_precomp=gaussian_covariances[i, :, row, col], + ) + all_images.append(image) + all_radii.append(radii) + return torch.stack(all_images) + + +DepthRenderingMode = Literal["depth", "disparity", "relative_disparity", "log"] diff --git a/src/model/decoder/decoder.py b/src/model/decoder/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..296d3301b7b11c09da1d674a62050ed9f672fcb6 --- /dev/null +++ b/src/model/decoder/decoder.py @@ -0,0 +1,46 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Generic, Literal, TypeVar + +from jaxtyping import Float +from torch import Tensor, nn + +from ..types import Gaussians + +DepthRenderingMode = Literal[ + "depth", + "log", + "disparity", + "relative_disparity", +] + + +@dataclass +class DecoderOutput: + color: Float[Tensor, "batch view 3 height width"] + depth: Float[Tensor, "batch view height width"] | None + alpha: Float[Tensor, "batch view height width"] | None + lod_rendering: dict | None + +T = TypeVar("T") + + +class Decoder(nn.Module, ABC, Generic[T]): + cfg: T + + def __init__(self, cfg: T) -> None: + super().__init__() + self.cfg = cfg + + @abstractmethod + def forward( + self, + gaussians: Gaussians, + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + near: Float[Tensor, "batch view"], + far: Float[Tensor, "batch view"], + image_shape: tuple[int, int], + depth_mode: DepthRenderingMode | None = None, + ) -> DecoderOutput: + pass diff --git a/src/model/decoder/decoder_splatting_cuda.py b/src/model/decoder/decoder_splatting_cuda.py new file mode 100644 index 0000000000000000000000000000000000000000..59681e6bb95b0a129437f913edd9c60331d49f88 --- /dev/null +++ b/src/model/decoder/decoder_splatting_cuda.py @@ -0,0 +1,111 @@ +from dataclasses import dataclass +from typing import Literal + +import torch +from einops import rearrange, repeat +from jaxtyping import Float +from torch import Tensor +import torchvision + +from ..types import Gaussians +# from .cuda_splatting import DepthRenderingMode, render_cuda +from .decoder import Decoder, DecoderOutput +from math import sqrt +from gsplat import rasterization + +from ...misc.utils import vis_depth_map + +DepthRenderingMode = Literal["depth", "disparity", "relative_disparity", "log"] + +@dataclass +class DecoderSplattingCUDACfg: + name: Literal["splatting_cuda"] + background_color: list[float] + make_scale_invariant: bool + + +class DecoderSplattingCUDA(Decoder[DecoderSplattingCUDACfg]): + background_color: Float[Tensor, "3"] + + def __init__( + self, + cfg: DecoderSplattingCUDACfg, + ) -> None: + super().__init__(cfg) + self.make_scale_invariant = cfg.make_scale_invariant + self.register_buffer( + "background_color", + torch.tensor(cfg.background_color, dtype=torch.float32), + persistent=False, + ) + + def rendering_fn( + self, + gaussians: Gaussians, + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + near: Float[Tensor, "batch view"], + far: Float[Tensor, "batch view"], + image_shape: tuple[int, int], + depth_mode: DepthRenderingMode | None = None, + cam_rot_delta: Float[Tensor, "batch view 3"] | None = None, + cam_trans_delta: Float[Tensor, "batch view 3"] | None = None, + ) -> DecoderOutput: + B, V, _, _ = intrinsics.shape + H, W = image_shape + rendered_imgs, rendered_depths, rendered_alphas = [], [], [] + xyzs, opacitys, rotations, scales, features = gaussians.means, gaussians.opacities, gaussians.rotations, gaussians.scales, gaussians.harmonics.permute(0, 1, 3, 2).contiguous() + covariances = gaussians.covariances + for i in range(B): + xyz_i = xyzs[i].float() + feature_i = features[i].float() + covar_i = covariances[i].float() + scale_i = scales[i].float() + rotation_i = rotations[i].float() + opacity_i = opacitys[i].squeeze().float() + test_w2c_i = extrinsics[i].float().inverse() # (V, 4, 4) + test_intr_i_normalized = intrinsics[i].float() + # Denormalize the intrinsics into standred format + test_intr_i = test_intr_i_normalized.clone() + test_intr_i[:, 0] = test_intr_i_normalized[:, 0] * W + test_intr_i[:, 1] = test_intr_i_normalized[:, 1] * H + sh_degree = (int(sqrt(feature_i.shape[-2])) - 1) + + rendering_list = [] + rendering_depth_list = [] + rendering_alpha_list = [] + for j in range(V): + rendering, alpha, _ = rasterization(xyz_i, rotation_i, scale_i, opacity_i, feature_i, + test_w2c_i[j:j+1], test_intr_i[j:j+1], W, H, sh_degree=sh_degree, + # near_plane=near[i].mean(), far_plane=far[i].mean(), + render_mode="RGB+D", packed=False, + near_plane=1e-10, + backgrounds=self.background_color.unsqueeze(0).repeat(1, 1), + radius_clip=0.1, + covars=covar_i, + rasterize_mode='classic') # (V, H, W, 3) + rendering_img, rendering_depth = torch.split(rendering, [3, 1], dim=-1) + rendering_img = rendering_img.clamp(0.0, 1.0) + rendering_list.append(rendering_img.permute(0, 3, 1, 2)) + rendering_depth_list.append(rendering_depth) + rendering_alpha_list.append(alpha) + rendered_depths.append(torch.cat(rendering_depth_list, dim=0).squeeze()) + rendered_imgs.append(torch.cat(rendering_list, dim=0)) + rendered_alphas.append(torch.cat(rendering_alpha_list, dim=0).squeeze()) + return DecoderOutput(torch.stack(rendered_imgs), torch.stack(rendered_depths), torch.stack(rendered_alphas), lod_rendering=None) + + def forward( + self, + gaussians: Gaussians, + extrinsics: Float[Tensor, "batch view 4 4"], + intrinsics: Float[Tensor, "batch view 3 3"], + near: Float[Tensor, "batch view"], + far: Float[Tensor, "batch view"], + image_shape: tuple[int, int], + depth_mode: DepthRenderingMode | None = None, + cam_rot_delta: Float[Tensor, "batch view 3"] | None = None, + cam_trans_delta: Float[Tensor, "batch view 3"] | None = None, + ) -> DecoderOutput: + + return self.rendering_fn(gaussians, extrinsics, intrinsics, near, far, image_shape, depth_mode, cam_rot_delta, cam_trans_delta) + diff --git a/src/model/encoder/__init__.py b/src/model/encoder/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4da77b7a2bd380ead96672170135b88ca2ff66c9 --- /dev/null +++ b/src/model/encoder/__init__.py @@ -0,0 +1,19 @@ +from typing import Optional, Union + +from .encoder import Encoder +from .visualization.encoder_visualizer import EncoderVisualizer +from .anysplat import EncoderAnySplat, EncoderAnySplatCfg + +ENCODERS = { + "anysplat": (EncoderAnySplat, None), +} + +EncoderCfg = Union[EncoderAnySplatCfg] + + +def get_encoder(cfg: EncoderCfg) -> tuple[Encoder, Optional[EncoderVisualizer]]: + encoder, visualizer = ENCODERS[cfg.name] + encoder = encoder(cfg) + if visualizer is not None: + visualizer = visualizer(cfg.visualizer, encoder) + return encoder, visualizer diff --git a/src/model/encoder/anysplat.py b/src/model/encoder/anysplat.py new file mode 100644 index 0000000000000000000000000000000000000000..8e6c73497b696ac4689a9120fc2f3c7bbd20c9d9 --- /dev/null +++ b/src/model/encoder/anysplat.py @@ -0,0 +1,593 @@ +import copy + +# VGGT parts +import os +import sys +from copy import deepcopy +from dataclasses import dataclass +from typing import List, Literal, Optional + +import torch +import torch.nn.functional as F +import torchvision +from einops import rearrange +from huggingface_hub import PyTorchModelHubMixin +from jaxtyping import Float +from src.dataset.shims.bounds_shim import apply_bounds_shim +from src.dataset.shims.normalize_shim import apply_normalize_shim +from src.dataset.shims.patch_shim import apply_patch_shim +from src.dataset.types import BatchedExample, DataShim +from src.geometry.projection import sample_image_grid + +from src.model.encoder.heads.vggt_dpt_gs_head import VGGT_DPT_GS_Head +from src.model.encoder.vggt.utils.geometry import ( + batchify_unproject_depth_map_to_point_map, + unproject_depth_map_to_point_map, +) +from src.model.encoder.vggt.utils.pose_enc import pose_encoding_to_extri_intri +from src.utils.geometry import get_rel_pos # used for model hub +from torch import nn, Tensor +from torch_scatter import scatter_add, scatter_max + +from ..types import Gaussians +from .backbone import Backbone, BackboneCfg, get_backbone + +from .backbone.croco.misc import transpose_to_landscape +from .common.gaussian_adapter import ( + GaussianAdapter, + GaussianAdapterCfg, + UnifiedGaussianAdapter, +) +from .encoder import Encoder, EncoderOutput +from .heads import head_factory +from .visualization.encoder_visualizer_epipolar_cfg import EncoderVisualizerEpipolarCfg + +root_path = os.path.abspath(".") +sys.path.append(root_path) +from src.model.encoder.heads.head_modules import TransformerBlockSelfAttn +from src.model.encoder.vggt.heads.dpt_head import DPTHead +from src.model.encoder.vggt.layers.mlp import Mlp +from src.model.encoder.vggt.models.vggt import VGGT + +inf = float("inf") + + +@dataclass +class OpacityMappingCfg: + initial: float + final: float + warm_up: int + + +@dataclass +class GSHeadParams: + dec_depth: int = 23 + patch_size: tuple[int, int] = (14, 14) + enc_embed_dim: int = 2048 + dec_embed_dim: int = 2048 + feature_dim: int = 256 + depth_mode = ("exp", -inf, inf) + conf_mode = True + + +@dataclass +class EncoderAnySplatCfg: + name: Literal["anysplat"] + anchor_feat_dim: int + voxel_size: float + n_offsets: int + d_feature: int + add_view: bool + num_monocular_samples: int + backbone: BackboneCfg + visualizer: EncoderVisualizerEpipolarCfg + gaussian_adapter: GaussianAdapterCfg + apply_bounds_shim: bool + opacity_mapping: OpacityMappingCfg + gaussians_per_pixel: int + num_surfaces: int + gs_params_head_type: str + input_mean: tuple[float, float, float] = (0.5, 0.5, 0.5) + input_std: tuple[float, float, float] = (0.5, 0.5, 0.5) + pretrained_weights: str = "" + pose_free: bool = True + pred_pose: bool = True + gt_pose_to_pts: bool = False + gs_prune: bool = False + opacity_threshold: float = 0.001 + gs_keep_ratio: float = 1.0 + pred_head_type: Literal["depth", "point"] = "point" + freeze_backbone: bool = False + freeze_module: Literal[ + "all", + "global", + "frame", + "patch_embed", + "patch_embed+frame", + "patch_embed+global", + "global+frame", + "None", + ] = "None" + distill: bool = False + render_conf: bool = False + opacity_conf: bool = False + conf_threshold: float = 0.1 + intermediate_layer_idx: Optional[List[int]] = None + voxelize: bool = False + + +def rearrange_head(feat, patch_size, H, W): + B = feat.shape[0] + feat = feat.transpose(-1, -2).view(B, -1, H // patch_size, W // patch_size) + feat = F.pixel_shuffle(feat, patch_size) # B,D,H,W + feat = rearrange(feat, "b d h w -> b (h w) d") + return feat + + +class EncoderAnySplat(Encoder[EncoderAnySplatCfg]): + backbone: nn.Module + gaussian_adapter: GaussianAdapter + + def __init__(self, cfg: EncoderAnySplatCfg) -> None: + super().__init__(cfg) + model_full = VGGT.from_pretrained("facebook/VGGT-1B") + # model_full = VGGT() + self.aggregator = model_full.aggregator.to(torch.bfloat16) + self.freeze_backbone = cfg.freeze_backbone + self.distill = cfg.distill + self.pred_pose = cfg.pred_pose + + self.camera_head = model_full.camera_head + if self.cfg.pred_head_type == "depth": + self.depth_head = model_full.depth_head + else: + self.point_head = model_full.point_head + + if self.distill: + self.distill_aggregator = copy.deepcopy(self.aggregator) + self.distill_camera_head = copy.deepcopy(self.camera_head) + self.distill_depth_head = copy.deepcopy(self.depth_head) + for module in [ + self.distill_aggregator, + self.distill_camera_head, + self.distill_depth_head, + ]: + for param in module.parameters(): + param.requires_grad = False + param.data = param.data.cpu() + + if self.freeze_backbone: + # Freeze backbone components + if self.cfg.pred_head_type == "depth": + for module in [self.aggregator, self.camera_head, self.depth_head]: + for param in module.parameters(): + param.requires_grad = False + else: + for module in [self.aggregator, self.camera_head, self.point_head]: + for param in module.parameters(): + param.requires_grad = False + else: + # aggregator freeze + freeze_module = self.cfg.freeze_module + if freeze_module == "None": + pass + + elif freeze_module == "all": + for param in self.aggregator.parameters(): + param.requires_grad = False + + else: + module_pairs = { + "patch_embed+frame": ["patch_embed", "frame"], + "patch_embed+global": ["patch_embed", "global"], + "global+frame": ["global", "frame"], + } + + if freeze_module in module_pairs: + for name, param in self.aggregator.named_parameters(): + if any(m in name for m in module_pairs[freeze_module]): + param.requires_grad = False + else: + for name, param in self.named_parameters(): + param.requires_grad = ( + freeze_module not in name and "distill" not in name + ) + + self.pose_free = cfg.pose_free + if self.pose_free: + self.gaussian_adapter = UnifiedGaussianAdapter(cfg.gaussian_adapter) + else: + self.gaussian_adapter = GaussianAdapter(cfg.gaussian_adapter) + + self.raw_gs_dim = 1 + self.gaussian_adapter.d_in # 1 for opacity + self.voxel_size = cfg.voxel_size + self.gs_params_head_type = cfg.gs_params_head_type + # fake backbone for head parameters + head_params = GSHeadParams() + self.gaussian_param_head = VGGT_DPT_GS_Head( + dim_in=2048, + patch_size=head_params.patch_size, + output_dim=self.raw_gs_dim + 1, + activation="norm_exp", + conf_activation="expp1", + features=head_params.feature_dim, + ) + + def map_pdf_to_opacity( + self, + pdf: Float[Tensor, " *batch"], + global_step: int, + ) -> Float[Tensor, " *batch"]: + # https://www.desmos.com/calculator/opvwti3ba9 + + # Figure out the exponent. + cfg = self.cfg.opacity_mapping + x = cfg.initial + min(global_step / cfg.warm_up, 1) * (cfg.final - cfg.initial) + exponent = 2**x + + # Map the probability density to an opacity. + return 0.5 * (1 - (1 - pdf) ** exponent + pdf ** (1 / exponent)) + + def normalize_pts3d(self, pts3ds, valid_masks, original_extrinsics=None): + # normalize pts_all + B = pts3ds.shape[0] + pts3d_norms = [] + scale_factors = [] + for bs in range(B): + pts3d, valid_mask = pts3ds[bs], valid_masks[bs] + if original_extrinsics is not None: + camera_c2w = original_extrinsics[bs] + first_camera_w2c = ( + camera_c2w[0].inverse().unsqueeze(0).repeat(pts3d.shape[0], 1, 1) + ) + + pts3d_homo = torch.cat( + [pts3d, torch.ones_like(pts3d[:, :, :, :1])], dim=-1 + ) + transformed_pts3d = torch.bmm( + first_camera_w2c, pts3d_homo.flatten(1, 2).transpose(1, 2) + ).transpose(1, 2)[..., :3] + scene_scale = torch.norm( + transformed_pts3d.flatten(0, 1)[valid_mask.flatten(0, 2).bool()], + dim=-1, + ).mean() + else: + transformed_pts3d = pts3d[valid_mask] + dis = transformed_pts3d.norm(dim=-1) + scene_scale = dis.mean().clip(min=1e-8) + # pts3d_norm[bs] = pts3d[bs] / scene_scale + pts3d_norms.append(pts3d / scene_scale) + scale_factors.append(scene_scale) + return torch.stack(pts3d_norms, dim=0), torch.stack(scale_factors, dim=0) + + def align_pts_all_with_pts3d( + self, pts_all, pts3d, valid_mask, original_extrinsics=None + ): + # align pts_all with pts3d + B = pts_all.shape[0] + + # follow vggt's normalization implementation + pts3d_norm, scale_factor = self.normalize_pts3d( + pts3d, valid_mask, original_extrinsics + ) # check if this is correct + pts_all = pts_all * scale_factor.view(B, 1, 1, 1, 1) + + return pts_all + + def pad_tensor_list(self, tensor_list, pad_shape, value=0.0): + padded = [] + for t in tensor_list: + pad_len = pad_shape[0] - t.shape[0] + if pad_len > 0: + padding = torch.full( + (pad_len, *t.shape[1:]), value, device=t.device, dtype=t.dtype + ) + t = torch.cat([t, padding], dim=0) + padded.append(t) + return torch.stack(padded) + + def voxelizaton_with_fusion(self, img_feat, pts3d, voxel_size, conf=None): + # img_feat: B*V, C, H, W + # pts3d: B*V, 3, H, W + V, C, H, W = img_feat.shape + pts3d_flatten = pts3d.permute(0, 2, 3, 1).flatten(0, 2) + + voxel_indices = (pts3d_flatten / voxel_size).round().int() # [B*V*N, 3] + unique_voxels, inverse_indices, counts = torch.unique( + voxel_indices, dim=0, return_inverse=True, return_counts=True + ) + + # Flatten confidence scores and features + conf_flat = conf.flatten() # [B*V*N] + anchor_feats_flat = img_feat.permute(0, 2, 3, 1).flatten(0, 2) # [B*V*N, ...] + + # Compute softmax weights per voxel + conf_voxel_max, _ = scatter_max(conf_flat, inverse_indices, dim=0) + conf_exp = torch.exp(conf_flat - conf_voxel_max[inverse_indices]) + voxel_weights = scatter_add( + conf_exp, inverse_indices, dim=0 + ) # [num_unique_voxels] + weights = (conf_exp / (voxel_weights[inverse_indices] + 1e-6)).unsqueeze( + -1 + ) # [B*V*N, 1] + + # Compute weighted average of positions and features + weighted_pts = pts3d_flatten * weights + weighted_feats = anchor_feats_flat.squeeze(1) * weights + + # Aggregate per voxel + voxel_pts = scatter_add( + weighted_pts, inverse_indices, dim=0 + ) # [num_unique_voxels, 3] + voxel_feats = scatter_add( + weighted_feats, inverse_indices, dim=0 + ) # [num_unique_voxels, feat_dim] + + return voxel_pts, voxel_feats + + def forward( + self, + image: torch.Tensor, + global_step: int = 0, + visualization_dump: Optional[dict] = None, + ) -> Gaussians: + device = image.device + b, v, _, h, w = image.shape + distill_infos = {} + if self.distill: + distill_image = image.clone().detach() + for module in [ + self.distill_aggregator, + self.distill_camera_head, + self.distill_depth_head, + ]: + for param in module.parameters(): + param.data = param.data.to(device, non_blocking=True) + + with torch.no_grad(): + # Process with bfloat16 precision + with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16): + distill_aggregated_tokens_list, distill_patch_start_idx = ( + self.distill_aggregator( + distill_image.to(torch.bfloat16), + intermediate_layer_idx=self.cfg.intermediate_layer_idx, + ) + ) + + # Process with default precision + with torch.amp.autocast("cuda", enabled=False): + # Get camera pose information + distill_pred_pose_enc_list = self.distill_camera_head( + distill_aggregated_tokens_list + ) + last_distill_pred_pose_enc = distill_pred_pose_enc_list[-1] + distill_extrinsic, distill_intrinsic = pose_encoding_to_extri_intri( + last_distill_pred_pose_enc, image.shape[-2:] + ) + + # Get depth information + distill_depth_map, distill_depth_conf = self.distill_depth_head( + distill_aggregated_tokens_list, + images=distill_image, + patch_start_idx=distill_patch_start_idx, + ) + + # Convert depth to 3D points + distill_pts_all = batchify_unproject_depth_map_to_point_map( + distill_depth_map, distill_extrinsic, distill_intrinsic + ) + # Store results + distill_infos["pred_pose_enc_list"] = distill_pred_pose_enc_list + distill_infos["pts_all"] = distill_pts_all + distill_infos["depth_map"] = distill_depth_map + + conf_threshold = torch.quantile( + distill_depth_conf.flatten(2, 3), 0.3, dim=-1, keepdim=True + ) # Get threshold for each view + conf_mask = distill_depth_conf > conf_threshold.unsqueeze(-1) + distill_infos["conf_mask"] = conf_mask + + for module in [ + self.distill_aggregator, + self.distill_camera_head, + self.distill_depth_head, + ]: + for param in module.parameters(): + param.data = param.data.cpu() + # Clean up to save memory + del distill_aggregated_tokens_list, distill_patch_start_idx + del distill_pred_pose_enc_list, last_distill_pred_pose_enc + del distill_extrinsic, distill_intrinsic + del distill_depth_map, distill_depth_conf + torch.cuda.empty_cache() + + with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16): + aggregated_tokens_list, patch_start_idx = self.aggregator( + image.to(torch.bfloat16), + intermediate_layer_idx=self.cfg.intermediate_layer_idx, + ) + + with torch.amp.autocast("cuda", enabled=False): + pred_pose_enc_list = self.camera_head(aggregated_tokens_list) + last_pred_pose_enc = pred_pose_enc_list[-1] + extrinsic, intrinsic = pose_encoding_to_extri_intri( + last_pred_pose_enc, image.shape[-2:] + ) # only for debug + + if self.cfg.pred_head_type == "point": + pts_all, pts_conf = self.point_head( + aggregated_tokens_list, + images=image, + patch_start_idx=patch_start_idx, + ) + elif self.cfg.pred_head_type == "depth": + depth_map, depth_conf = self.depth_head( + aggregated_tokens_list, + images=image, + patch_start_idx=patch_start_idx, + ) + pts_all = batchify_unproject_depth_map_to_point_map( + depth_map, extrinsic, intrinsic + ) + else: + raise ValueError(f"Invalid pred_head_type: {self.cfg.pred_head_type}") + + if self.cfg.render_conf: + conf_valid = torch.quantile( + depth_conf.flatten(0, 1), self.cfg.conf_threshold + ) + conf_valid_mask = depth_conf > conf_valid + else: + conf_valid_mask = torch.ones_like(depth_conf, dtype=torch.bool) + + # dpt style gs_head input format + out = self.gaussian_param_head( + aggregated_tokens_list, + pts_all.flatten(0, 1).permute(0, 3, 1, 2), + image, + patch_start_idx=patch_start_idx, + image_size=(h, w), + ) + + del aggregated_tokens_list, patch_start_idx + torch.cuda.empty_cache() + + pts_flat = pts_all.flatten(2, 3) + scene_scale = pts_flat.norm(dim=-1).mean().clip(min=1e-8) + + anchor_feats, conf = out[:, :, : self.raw_gs_dim], out[:, :, self.raw_gs_dim] + + neural_feats_list, neural_pts_list = [], [] + if self.cfg.voxelize: + for b_i in range(b): + neural_pts, neural_feats = self.voxelizaton_with_fusion( + anchor_feats[b_i], + pts_all[b_i].permute(0, 3, 1, 2).contiguous(), + self.voxel_size, + conf=conf[b_i], + ) + neural_feats_list.append(neural_feats) + neural_pts_list.append(neural_pts) + else: + for b_i in range(b): + neural_feats_list.append( + anchor_feats[b_i].permute(0, 2, 3, 1)[conf_valid_mask[b_i]] + ) + neural_pts_list.append(pts_all[b_i][conf_valid_mask[b_i]]) + + max_voxels = max(f.shape[0] for f in neural_feats_list) + neural_feats = self.pad_tensor_list( + neural_feats_list, (max_voxels,), value=-1e10 + ) + + neural_pts = self.pad_tensor_list( + neural_pts_list, (max_voxels,), -1e4 + ) # -1 == invalid voxel + + depths = neural_pts[..., -1].unsqueeze(-1) + densities = neural_feats[..., 0].sigmoid() + + assert len(densities.shape) == 2, "the shape of densities should be (B, N)" + assert neural_pts.shape[1] > 1, "the number of voxels should be greater than 1" + + opacity = self.map_pdf_to_opacity(densities, global_step).squeeze(-1) + if self.cfg.opacity_conf: + shift = torch.quantile(depth_conf, self.cfg.conf_threshold) + opacity = opacity * torch.sigmoid(depth_conf - shift)[ + conf_valid_mask + ].unsqueeze( + 0 + ) # little bit hacky + + # GS Prune, but only works when bs = 1 + # if want to support bs > 1, need to random prune gaussians based on the rank of opacity like LongLRM + # Note: we not prune gaussians here, but we will try it in the future + if self.cfg.gs_prune and b == 1: + opacity_threshold = self.cfg.opacity_threshold + gaussian_usage = opacity > opacity_threshold # (B, N) + + print( + f"based on opacity threshold {opacity_threshold}, pruned {gaussian_usage.shape[1] - neural_pts.shape[1]} gaussians out of {gaussian_usage.shape[1]}" + ) + + if (gaussian_usage.sum() / gaussian_usage.numel()) > self.cfg.gs_keep_ratio: + # rank by opacity + num_keep = int(gaussian_usage.shape[1] * self.cfg.gs_keep_ratio) + idx_sort = opacity.argsort(dim=1, descending=True) + keep_idx = idx_sort[:, :num_keep] + gaussian_usage = torch.zeros_like(gaussian_usage, dtype=torch.bool) + gaussian_usage.scatter_(1, keep_idx, True) + + neural_pts = neural_pts[gaussian_usage].view(b, -1, 3).contiguous() + depths = depths[gaussian_usage].view(b, -1, 1).contiguous() + neural_feats = ( + neural_feats[gaussian_usage].view(b, -1, self.raw_gs_dim).contiguous() + ) + opacity = opacity[gaussian_usage].view(b, -1).contiguous() + + print( + f"finally pruned {gaussian_usage.shape[1] - neural_pts.shape[1]} gaussians out of {gaussian_usage.shape[1]}" + ) + + gaussians = self.gaussian_adapter.forward( + neural_pts, + depths, + opacity, + neural_feats[..., 1:].squeeze(2), + ) + + if visualization_dump is not None: + visualization_dump["depth"] = rearrange( + pts_all[..., -1].flatten(2, 3).unsqueeze(-1).unsqueeze(-1), + "b v (h w) srf s -> b v h w srf s", + h=h, + w=w, + ) + + infos = {} + infos["scene_scale"] = scene_scale + infos["voxelize_ratio"] = densities.shape[1] / (h * w * v) + + print( + f"scene scale: {scene_scale:.3f}, pixel-wise num: {h*w*v}, after voxelize: {neural_pts.shape[1]}, voxelize ratio: {infos['voxelize_ratio']:.3f}" + ) + print( + f"Gaussians attributes: \n" + f"opacities: mean: {gaussians.opacities.mean()}, min: {gaussians.opacities.min()}, max: {gaussians.opacities.max()} \n" + f"scales: mean: {gaussians.scales.mean()}, min: {gaussians.scales.min()}, max: {gaussians.scales.max()}" + ) + + print("B:", b, "V:", v, "H:", h, "W:", w) + extrinsic_padding = ( + torch.tensor([0, 0, 0, 1], device=device, dtype=extrinsic.dtype) + .view(1, 1, 1, 4) + .repeat(b, v, 1, 1) + ) + intrinsic = intrinsic.clone() # Create a new tensor + intrinsic = torch.stack( + [intrinsic[:, :, 0] / w, intrinsic[:, :, 1] / h, intrinsic[:, :, 2]], dim=2 + ) + + return EncoderOutput( + gaussians=gaussians, + pred_pose_enc_list=pred_pose_enc_list, + pred_context_pose=dict( + extrinsic=torch.cat([extrinsic, extrinsic_padding], dim=2).inverse(), + intrinsic=intrinsic, + ), + depth_dict=dict(depth=depth_map, conf_valid_mask=conf_valid_mask), + infos=infos, + distill_infos=distill_infos, + ) + + def get_data_shim(self) -> DataShim: + def data_shim(batch: BatchedExample) -> BatchedExample: + batch = apply_normalize_shim( + batch, + self.cfg.input_mean, + self.cfg.input_std, + ) + + return batch + + return data_shim diff --git a/src/model/encoder/backbone/__init__.py b/src/model/encoder/backbone/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f6a0d24132a5fa7ea3d42ed06141d55b93d21036 --- /dev/null +++ b/src/model/encoder/backbone/__init__.py @@ -0,0 +1,21 @@ +from typing import Any +import torch.nn as nn + +from .backbone import Backbone +from .backbone_croco_multiview import AsymmetricCroCoMulti +from .backbone_dino import BackboneDino, BackboneDinoCfg +from .backbone_resnet import BackboneResnet, BackboneResnetCfg +from .backbone_croco import AsymmetricCroCo, BackboneCrocoCfg + +BACKBONES: dict[str, Backbone[Any]] = { + "resnet": BackboneResnet, + "dino": BackboneDino, + "croco": AsymmetricCroCo, + "croco_multi": AsymmetricCroCoMulti, +} + +BackboneCfg = BackboneResnetCfg | BackboneDinoCfg | BackboneCrocoCfg + + +def get_backbone(cfg: BackboneCfg, d_in: int = 3) -> nn.Module: + return BACKBONES[cfg.name](cfg, d_in) diff --git a/src/model/encoder/backbone/backbone.py b/src/model/encoder/backbone/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..2644e44eb38279165314fcb47bd3a7757abb76cc --- /dev/null +++ b/src/model/encoder/backbone/backbone.py @@ -0,0 +1,29 @@ +from abc import ABC, abstractmethod +from typing import Generic, TypeVar + +from jaxtyping import Float +from torch import Tensor, nn + +from src.dataset.types import BatchedViews + +T = TypeVar("T") + + +class Backbone(nn.Module, ABC, Generic[T]): + cfg: T + + def __init__(self, cfg: T) -> None: + super().__init__() + self.cfg = cfg + + @abstractmethod + def forward( + self, + context: BatchedViews, + ) -> Float[Tensor, "batch view d_out height width"]: + pass + + @property + @abstractmethod + def d_out(self) -> int: + pass diff --git a/src/model/encoder/backbone/backbone_croco.py b/src/model/encoder/backbone/backbone_croco.py new file mode 100644 index 0000000000000000000000000000000000000000..8eb26b5cfff2ab4fa5961d18d7166d37cde54805 --- /dev/null +++ b/src/model/encoder/backbone/backbone_croco.py @@ -0,0 +1,281 @@ +from copy import deepcopy +from dataclasses import dataclass +from typing import Literal + +import torch +from torch import nn + +from .croco.blocks import DecoderBlock +from .croco.croco import CroCoNet +from .croco.misc import fill_default_args, freeze_all_params, transpose_to_landscape, is_symmetrized, interleave, \ + make_batch_symmetric +from .croco.patch_embed import get_patch_embed +from .backbone import Backbone +from src.geometry.camera_emb import get_intrinsic_embedding + + +inf = float('inf') + + +croco_params = { + 'ViTLarge_BaseDecoder': { + 'enc_depth': 24, + 'dec_depth': 12, + 'enc_embed_dim': 1024, + 'dec_embed_dim': 768, + 'enc_num_heads': 16, + 'dec_num_heads': 12, + 'pos_embed': 'RoPE100', + 'img_size': (512, 512), + }, +} + +default_dust3r_params = { + 'enc_depth': 24, + 'dec_depth': 12, + 'enc_embed_dim': 1024, + 'dec_embed_dim': 768, + 'enc_num_heads': 16, + 'dec_num_heads': 12, + 'pos_embed': 'RoPE100', + 'patch_embed_cls': 'PatchEmbedDust3R', + 'img_size': (512, 512), + 'head_type': 'dpt', + 'output_mode': 'pts3d', + 'depth_mode': ('exp', -inf, inf), + 'conf_mode': ('exp', 1, inf) +} + + +@dataclass +class BackboneCrocoCfg: + name: Literal["croco", "croco_multi"] + model: Literal["ViTLarge_BaseDecoder", "ViTBase_SmallDecoder", "ViTBase_BaseDecoder"] # keep interface for the last two models, but they are not supported + patch_embed_cls: str = 'PatchEmbedDust3R' # PatchEmbedDust3R or ManyAR_PatchEmbed + asymmetry_decoder: bool = True + intrinsics_embed_loc: Literal["encoder", "decoder", "none"] = 'none' + intrinsics_embed_degree: int = 0 + intrinsics_embed_type: Literal["pixelwise", "linear", "token"] = 'token' # linear or dpt + + +class AsymmetricCroCo(CroCoNet): + """ Two siamese encoders, followed by two decoders. + The goal is to output 3d points directly, both images in view1's frame + (hence the asymmetry). + """ + + def __init__(self, cfg: BackboneCrocoCfg, d_in: int) -> None: + + self.intrinsics_embed_loc = cfg.intrinsics_embed_loc + self.intrinsics_embed_degree = cfg.intrinsics_embed_degree + self.intrinsics_embed_type = cfg.intrinsics_embed_type + self.intrinsics_embed_encoder_dim = 0 + self.intrinsics_embed_decoder_dim = 0 + if self.intrinsics_embed_loc == 'encoder' and self.intrinsics_embed_type == 'pixelwise': + self.intrinsics_embed_encoder_dim = (self.intrinsics_embed_degree + 1) ** 2 if self.intrinsics_embed_degree > 0 else 3 + elif self.intrinsics_embed_loc == 'decoder' and self.intrinsics_embed_type == 'pixelwise': + self.intrinsics_embed_decoder_dim = (self.intrinsics_embed_degree + 1) ** 2 if self.intrinsics_embed_degree > 0 else 3 + + self.patch_embed_cls = cfg.patch_embed_cls + self.croco_args = fill_default_args(croco_params[cfg.model], CroCoNet.__init__) + + super().__init__(**croco_params[cfg.model]) + + if cfg.asymmetry_decoder: + self.dec_blocks2 = deepcopy(self.dec_blocks) # This is used in DUSt3R and MASt3R + + if self.intrinsics_embed_type == 'linear' or self.intrinsics_embed_type == 'token': + self.intrinsic_encoder = nn.Linear(9, 1024) + + # self.set_freeze(freeze) + + def _set_patch_embed(self, img_size=224, patch_size=16, enc_embed_dim=768, in_chans=3): + in_chans = in_chans + self.intrinsics_embed_encoder_dim + self.patch_embed = get_patch_embed(self.patch_embed_cls, img_size, patch_size, enc_embed_dim, in_chans) + + def _set_decoder(self, enc_embed_dim, dec_embed_dim, dec_num_heads, dec_depth, mlp_ratio, norm_layer, norm_im2_in_dec): + self.dec_depth = dec_depth + self.dec_embed_dim = dec_embed_dim + # transfer from encoder to decoder + enc_embed_dim = enc_embed_dim + self.intrinsics_embed_decoder_dim + self.decoder_embed = nn.Linear(enc_embed_dim, dec_embed_dim, bias=True) + # transformer for the decoder + self.dec_blocks = nn.ModuleList([ + DecoderBlock(dec_embed_dim, dec_num_heads, mlp_ratio=mlp_ratio, qkv_bias=True, norm_layer=norm_layer, norm_mem=norm_im2_in_dec, rope=self.rope) + for i in range(dec_depth)]) + # final norm layer + self.dec_norm = norm_layer(dec_embed_dim) + + def load_state_dict(self, ckpt, **kw): + # duplicate all weights for the second decoder if not present + new_ckpt = dict(ckpt) + if not any(k.startswith('dec_blocks2') for k in ckpt): + for key, value in ckpt.items(): + if key.startswith('dec_blocks'): + new_ckpt[key.replace('dec_blocks', 'dec_blocks2')] = value + return super().load_state_dict(new_ckpt, **kw) + + def set_freeze(self, freeze): # this is for use by downstream models + assert freeze in ['none', 'mask', 'encoder'], f"unexpected freeze={freeze}" + to_be_frozen = { + 'none': [], + 'mask': [self.mask_token], + 'encoder': [self.mask_token, self.patch_embed, self.enc_blocks], + 'encoder_decoder': [self.mask_token, self.patch_embed, self.enc_blocks, self.enc_norm, self.decoder_embed, self.dec_blocks, self.dec_blocks2, self.dec_norm], + } + freeze_all_params(to_be_frozen[freeze]) + + def _set_prediction_head(self, *args, **kwargs): + """ No prediction head """ + return + + def _encode_image(self, image, true_shape, intrinsics_embed=None): + # embed the image into patches (x has size B x Npatches x C) + x, pos = self.patch_embed(image, true_shape=true_shape) + + if intrinsics_embed is not None: + + if self.intrinsics_embed_type == 'linear': + x = x + intrinsics_embed + elif self.intrinsics_embed_type == 'token': + x = torch.cat((x, intrinsics_embed), dim=1) + add_pose = pos[:, 0:1, :].clone() + add_pose[:, :, 0] += (pos[:, -1, 0].unsqueeze(-1) + 1) + pos = torch.cat((pos, add_pose), dim=1) + + # add positional embedding without cls token + assert self.enc_pos_embed is None + + # now apply the transformer encoder and normalization + for blk in self.enc_blocks: + x = blk(x, pos) + + x = self.enc_norm(x) + return x, pos, None + + def _encode_image_pairs(self, img1, img2, true_shape1, true_shape2, intrinsics_embed1=None, intrinsics_embed2=None): + if img1.shape[-2:] == img2.shape[-2:]: + out, pos, _ = self._encode_image(torch.cat((img1, img2), dim=0), + torch.cat((true_shape1, true_shape2), dim=0), + torch.cat((intrinsics_embed1, intrinsics_embed2), dim=0) if intrinsics_embed1 is not None else None) + out, out2 = out.chunk(2, dim=0) + pos, pos2 = pos.chunk(2, dim=0) + else: + out, pos, _ = self._encode_image(img1, true_shape1, intrinsics_embed1) + out2, pos2, _ = self._encode_image(img2, true_shape2, intrinsics_embed2) + return out, out2, pos, pos2 + + def _encode_symmetrized(self, view1, view2, force_asym=False): + img1 = view1['img'] + img2 = view2['img'] + B = img1.shape[0] + # Recover true_shape when available, otherwise assume that the img shape is the true one + shape1 = view1.get('true_shape', torch.tensor(img1.shape[-2:])[None].repeat(B, 1)) + shape2 = view2.get('true_shape', torch.tensor(img2.shape[-2:])[None].repeat(B, 1)) + # warning! maybe the images have different portrait/landscape orientations + + intrinsics_embed1 = view1.get('intrinsics_embed', None) + intrinsics_embed2 = view2.get('intrinsics_embed', None) + + if force_asym or not is_symmetrized(view1, view2): + feat1, feat2, pos1, pos2 = self._encode_image_pairs(img1, img2, shape1, shape2, intrinsics_embed1, intrinsics_embed2) + else: + # computing half of forward pass!' + feat1, feat2, pos1, pos2 = self._encode_image_pairs(img1[::2], img2[::2], shape1[::2], shape2[::2]) + feat1, feat2 = interleave(feat1, feat2) + pos1, pos2 = interleave(pos1, pos2) + + return (shape1, shape2), (feat1, feat2), (pos1, pos2) + + def _decoder(self, f1, pos1, f2, pos2, extra_embed1=None, extra_embed2=None): + final_output = [(f1, f2)] # before projection + + if extra_embed1 is not None: + f1 = torch.cat((f1, extra_embed1), dim=-1) + if extra_embed2 is not None: + f2 = torch.cat((f2, extra_embed2), dim=-1) + + # project to decoder dim + f1 = self.decoder_embed(f1) + f2 = self.decoder_embed(f2) + + final_output.append((f1, f2)) + for blk1, blk2 in zip(self.dec_blocks, self.dec_blocks2): + # img1 side + f1, _ = blk1(*final_output[-1][::+1], pos1, pos2) + # img2 side + f2, _ = blk2(*final_output[-1][::-1], pos2, pos1) + # store the result + final_output.append((f1, f2)) + + # normalize last output + del final_output[1] # duplicate with final_output[0] + final_output[-1] = tuple(map(self.dec_norm, final_output[-1])) + return zip(*final_output) + + def _downstream_head(self, head_num, decout, img_shape): + B, S, D = decout[-1].shape + # img_shape = tuple(map(int, img_shape)) + head = getattr(self, f'head{head_num}') + return head(decout, img_shape) + + def forward(self, + context: dict, + symmetrize_batch=False, + return_views=False, + ): + b, v, _, h, w = context["image"].shape + device = context["image"].device + + view1, view2 = ({'img': context["image"][:, 0]}, + {'img': context["image"][:, 1]}) + + # camera embedding in the encoder + if self.intrinsics_embed_loc == 'encoder' and self.intrinsics_embed_type == 'pixelwise': + intrinsic_emb = get_intrinsic_embedding(context, degree=self.intrinsics_embed_degree) + view1['img'] = torch.cat((view1['img'], intrinsic_emb[:, 0]), dim=1) + view2['img'] = torch.cat((view2['img'], intrinsic_emb[:, 1]), dim=1) + + if self.intrinsics_embed_loc == 'encoder' and (self.intrinsics_embed_type == 'token' or self.intrinsics_embed_type == 'linear'): + intrinsic_embedding = self.intrinsic_encoder(context["intrinsics"].flatten(2)) + view1['intrinsics_embed'] = intrinsic_embedding[:, 0].unsqueeze(1) + view2['intrinsics_embed'] = intrinsic_embedding[:, 1].unsqueeze(1) + + if symmetrize_batch: + instance_list_view1, instance_list_view2 = [0 for _ in range(b)], [1 for _ in range(b)] + view1['instance'] = instance_list_view1 + view2['instance'] = instance_list_view2 + view1['idx'] = instance_list_view1 + view2['idx'] = instance_list_view2 + view1, view2 = make_batch_symmetric(view1, view2) + + # encode the two images --> B,S,D + (shape1, shape2), (feat1, feat2), (pos1, pos2) = self._encode_symmetrized(view1, view2, force_asym=False) + else: + # encode the two images --> B,S,D + (shape1, shape2), (feat1, feat2), (pos1, pos2) = self._encode_symmetrized(view1, view2, force_asym=True) + + if self.intrinsics_embed_loc == 'decoder': + # FIXME: downsample is hardcoded to 16 + intrinsic_emb = get_intrinsic_embedding(context, degree=self.intrinsics_embed_degree, downsample=16, merge_hw=True) + dec1, dec2 = self._decoder(feat1, pos1, feat2, pos2, intrinsic_emb[:, 0], intrinsic_emb[:, 1]) + else: + dec1, dec2 = self._decoder(feat1, pos1, feat2, pos2) + + if self.intrinsics_embed_loc == 'encoder' and self.intrinsics_embed_type == 'token': + dec1, dec2 = list(dec1), list(dec2) + for i in range(len(dec1)): + dec1[i] = dec1[i][:, :-1] + dec2[i] = dec2[i][:, :-1] + + if return_views: + return dec1, dec2, shape1, shape2, view1, view2 + return dec1, dec2, shape1, shape2 + + @property + def patch_size(self) -> int: + return 16 + + @property + def d_out(self) -> int: + return 1024 diff --git a/src/model/encoder/backbone/backbone_croco_multiview.py b/src/model/encoder/backbone/backbone_croco_multiview.py new file mode 100644 index 0000000000000000000000000000000000000000..9d5689e0899c3b77dd49a9bed97bfa6760ecea6e --- /dev/null +++ b/src/model/encoder/backbone/backbone_croco_multiview.py @@ -0,0 +1,245 @@ +from copy import deepcopy +from dataclasses import dataclass +from typing import Literal + +import torch +from einops import rearrange +from torch import nn + +from .croco.blocks import DecoderBlock +from .croco.croco import CroCoNet +from .croco.misc import fill_default_args, freeze_all_params, transpose_to_landscape, is_symmetrized, interleave, \ + make_batch_symmetric +from .croco.patch_embed import get_patch_embed +from .backbone import Backbone +from src.geometry.camera_emb import get_intrinsic_embedding + +inf = float('inf') + + +croco_params = { + 'ViTLarge_BaseDecoder': { + 'enc_depth': 24, + 'dec_depth': 12, + 'enc_embed_dim': 1024, + 'dec_embed_dim': 768, + 'enc_num_heads': 16, + 'dec_num_heads': 12, + 'pos_embed': 'RoPE100', + 'img_size': (512, 512), + }, +} + +default_dust3r_params = { + 'enc_depth': 24, + 'dec_depth': 12, + 'enc_embed_dim': 1024, + 'dec_embed_dim': 768, + 'enc_num_heads': 16, + 'dec_num_heads': 12, + 'pos_embed': 'RoPE100', + 'patch_embed_cls': 'PatchEmbedDust3R', + 'img_size': (512, 512), + 'head_type': 'dpt', + 'output_mode': 'pts3d', + 'depth_mode': ('exp', -inf, inf), + 'conf_mode': ('exp', 1, inf) +} + + +@dataclass +class BackboneCrocoCfg: + name: Literal["croco"] + model: Literal["ViTLarge_BaseDecoder", "ViTBase_SmallDecoder", "ViTBase_BaseDecoder"] # keep interface for the last two models, but they are not supported + patch_embed_cls: str = 'PatchEmbedDust3R' # PatchEmbedDust3R or ManyAR_PatchEmbed + asymmetry_decoder: bool = True + intrinsics_embed_loc: Literal["encoder", "decoder", "none"] = 'none' + intrinsics_embed_degree: int = 0 + intrinsics_embed_type: Literal["pixelwise", "linear", "token"] = 'token' # linear or dpt + + +class AsymmetricCroCoMulti(CroCoNet): + """ Two siamese encoders, followed by two decoders. + The goal is to output 3d points directly, both images in view1's frame + (hence the asymmetry). + """ + + def __init__(self, cfg: BackboneCrocoCfg, d_in: int) -> None: + + self.intrinsics_embed_loc = cfg.intrinsics_embed_loc + self.intrinsics_embed_degree = cfg.intrinsics_embed_degree + self.intrinsics_embed_type = cfg.intrinsics_embed_type + self.intrinsics_embed_encoder_dim = 0 + self.intrinsics_embed_decoder_dim = 0 + if self.intrinsics_embed_loc == 'encoder' and self.intrinsics_embed_type == 'pixelwise': + self.intrinsics_embed_encoder_dim = (self.intrinsics_embed_degree + 1) ** 2 if self.intrinsics_embed_degree > 0 else 3 + elif self.intrinsics_embed_loc == 'decoder' and self.intrinsics_embed_type == 'pixelwise': + self.intrinsics_embed_decoder_dim = (self.intrinsics_embed_degree + 1) ** 2 if self.intrinsics_embed_degree > 0 else 3 + + self.patch_embed_cls = cfg.patch_embed_cls + self.croco_args = fill_default_args(croco_params[cfg.model], CroCoNet.__init__) + + super().__init__(**croco_params[cfg.model]) + + if cfg.asymmetry_decoder: + self.dec_blocks2 = deepcopy(self.dec_blocks) # This is used in DUSt3R and MASt3R + + if self.intrinsics_embed_type == 'linear' or self.intrinsics_embed_type == 'token': + self.intrinsic_encoder = nn.Linear(9, 1024) + + # self.set_downstream_head(output_mode, head_type, landscape_only, depth_mode, conf_mode, **croco_kwargs) + # self.set_freeze(freeze) + + def _set_patch_embed(self, img_size=224, patch_size=16, enc_embed_dim=768, in_chans=3): + in_chans = in_chans + self.intrinsics_embed_encoder_dim + self.patch_embed = get_patch_embed(self.patch_embed_cls, img_size, patch_size, enc_embed_dim, in_chans) + + def _set_decoder(self, enc_embed_dim, dec_embed_dim, dec_num_heads, dec_depth, mlp_ratio, norm_layer, norm_im2_in_dec): + self.dec_depth = dec_depth + self.dec_embed_dim = dec_embed_dim + # transfer from encoder to decoder + enc_embed_dim = enc_embed_dim + self.intrinsics_embed_decoder_dim + self.decoder_embed = nn.Linear(enc_embed_dim, dec_embed_dim, bias=True) + # transformer for the decoder + self.dec_blocks = nn.ModuleList([ + DecoderBlock(dec_embed_dim, dec_num_heads, mlp_ratio=mlp_ratio, qkv_bias=True, norm_layer=norm_layer, norm_mem=norm_im2_in_dec, rope=self.rope) + for i in range(dec_depth)]) + # final norm layer + self.dec_norm = norm_layer(dec_embed_dim) + + def load_state_dict(self, ckpt, **kw): + # duplicate all weights for the second decoder if not present + new_ckpt = dict(ckpt) + if not any(k.startswith('dec_blocks2') for k in ckpt): + for key, value in ckpt.items(): + if key.startswith('dec_blocks'): + new_ckpt[key.replace('dec_blocks', 'dec_blocks2')] = value + return super().load_state_dict(new_ckpt, **kw) + + def set_freeze(self, freeze): # this is for use by downstream models + assert freeze in ['none', 'mask', 'encoder'], f"unexpected freeze={freeze}" + to_be_frozen = { + 'none': [], + 'mask': [self.mask_token], + 'encoder': [self.mask_token, self.patch_embed, self.enc_blocks], + 'encoder_decoder': [self.mask_token, self.patch_embed, self.enc_blocks, self.enc_norm, self.decoder_embed, self.dec_blocks, self.dec_blocks2, self.dec_norm], + } + freeze_all_params(to_be_frozen[freeze]) + + def _set_prediction_head(self, *args, **kwargs): + """ No prediction head """ + return + + def _encode_image(self, image, true_shape, intrinsics_embed=None): + # embed the image into patches (x has size B x Npatches x C) + x, pos = self.patch_embed(image, true_shape=true_shape) + + if intrinsics_embed is not None: + + if self.intrinsics_embed_type == 'linear': + x = x + intrinsics_embed + elif self.intrinsics_embed_type == 'token': + x = torch.cat((x, intrinsics_embed), dim=1) + add_pose = pos[:, 0:1, :].clone() + add_pose[:, :, 0] += (pos[:, -1, 0].unsqueeze(-1) + 1) + pos = torch.cat((pos, add_pose), dim=1) + + # add positional embedding without cls token + assert self.enc_pos_embed is None + + # now apply the transformer encoder and normalization + for blk in self.enc_blocks: + x = blk(x, pos) + + x = self.enc_norm(x) + return x, pos, None + + def _decoder(self, feat, pose, extra_embed=None): + b, v, l, c = feat.shape + final_output = [feat] # before projection + if extra_embed is not None: + feat = torch.cat((feat, extra_embed), dim=-1) + + # project to decoder dim + f = rearrange(feat, "b v l c -> (b v) l c") + f = self.decoder_embed(f) + f = rearrange(f, "(b v) l c -> b v l c", b=b, v=v) + final_output.append(f) + + def generate_ctx_views(x): + b, v, l, c = x.shape + ctx_views = x.unsqueeze(1).expand(b, v, v, l, c) + mask = torch.arange(v).unsqueeze(0) != torch.arange(v).unsqueeze(1) + ctx_views = ctx_views[:, mask].reshape(b, v, v - 1, l, c) # B, V, V-1, L, C + ctx_views = ctx_views.flatten(2, 3) # B, V, (V-1)*L, C + return ctx_views.contiguous() + + pos_ctx = generate_ctx_views(pose) + for blk1, blk2 in zip(self.dec_blocks, self.dec_blocks2): + feat_current = final_output[-1] + feat_current_ctx = generate_ctx_views(feat_current) + # img1 side + f1, _ = blk1(feat_current[:, 0].contiguous(), feat_current_ctx[:, 0].contiguous(), pose[:, 0].contiguous(), pos_ctx[:, 0].contiguous()) + f1 = f1.unsqueeze(1) + # img2 side + f2, _ = blk2(rearrange(feat_current[:, 1:], "b v l c -> (b v) l c"), + rearrange(feat_current_ctx[:, 1:], "b v l c -> (b v) l c"), + rearrange(pose[:, 1:], "b v l c -> (b v) l c"), + rearrange(pos_ctx[:, 1:], "b v l c -> (b v) l c")) + f2 = rearrange(f2, "(b v) l c -> b v l c", b=b, v=v-1) + # store the result + final_output.append(torch.cat((f1, f2), dim=1)) + + # normalize last output + del final_output[1] # duplicate with final_output[0] + last_feat = rearrange(final_output[-1], "b v l c -> (b v) l c") + last_feat = self.dec_norm(last_feat) + final_output[-1] = rearrange(last_feat, "(b v) l c -> b v l c", b=b, v=v) + return final_output + + def forward(self, + context: dict, + symmetrize_batch=False, + return_views=False, + ): + b, v, _, h, w = context["image"].shape + images_all = context["image"] + + # camera embedding in the encoder + if self.intrinsics_embed_loc == 'encoder' and self.intrinsics_embed_type == 'pixelwise': + intrinsic_embedding = get_intrinsic_embedding(context, degree=self.intrinsics_embed_degree) + images_all = torch.cat((images_all, intrinsic_embedding), dim=2) + + intrinsic_embedding_all = None + if self.intrinsics_embed_loc == 'encoder' and (self.intrinsics_embed_type == 'token' or self.intrinsics_embed_type == 'linear'): + intrinsic_embedding = self.intrinsic_encoder(context["intrinsics"].flatten(2)) + intrinsic_embedding_all = rearrange(intrinsic_embedding, "b v c -> (b v) c").unsqueeze(1) + + # step 1: encoder input images + images_all = rearrange(images_all, "b v c h w -> (b v) c h w") + shape_all = torch.tensor(images_all.shape[-2:])[None].repeat(b*v, 1) + + feat, pose, _ = self._encode_image(images_all, shape_all, intrinsic_embedding_all) + + feat = rearrange(feat, "(b v) l c -> b v l c", b=b, v=v) + pose = rearrange(pose, "(b v) l c -> b v l c", b=b, v=v) + + # step 2: decoder + dec_feat = self._decoder(feat, pose) + shape = rearrange(shape_all, "(b v) c -> b v c", b=b, v=v) + images = rearrange(images_all, "(b v) c h w -> b v c h w", b=b, v=v) + + if self.intrinsics_embed_loc == 'encoder' and self.intrinsics_embed_type == 'token': + dec_feat = list(dec_feat) + for i in range(len(dec_feat)): + dec_feat[i] = dec_feat[i][:, :, :-1] + + return dec_feat, shape, images + + @property + def patch_size(self) -> int: + return 16 + + @property + def d_out(self) -> int: + return 1024 diff --git a/src/model/encoder/backbone/backbone_dino.py b/src/model/encoder/backbone/backbone_dino.py new file mode 100644 index 0000000000000000000000000000000000000000..a69fce1c633e25b7e93cba45f0a5895587fd293b --- /dev/null +++ b/src/model/encoder/backbone/backbone_dino.py @@ -0,0 +1,79 @@ +from dataclasses import dataclass +from typing import Literal + +import torch +from einops import rearrange, repeat +from jaxtyping import Float +from torch import Tensor, nn + +from src.dataset.types import BatchedViews +from .backbone import Backbone +from .backbone_resnet import BackboneResnet, BackboneResnetCfg + + +@dataclass +class BackboneDinoCfg: + name: Literal["dino"] + model: Literal["dino_vits16", "dino_vits8", "dino_vitb16", "dino_vitb8"] + d_out: int + + +class BackboneDino(Backbone[BackboneDinoCfg]): + def __init__(self, cfg: BackboneDinoCfg, d_in: int) -> None: + super().__init__(cfg) + assert d_in == 3 + self.dino = torch.hub.load("facebookresearch/dino:main", cfg.model) + self.resnet_backbone = BackboneResnet( + BackboneResnetCfg("resnet", "dino_resnet50", 4, False, cfg.d_out), + d_in, + ) + self.global_token_mlp = nn.Sequential( + nn.Linear(768, 768), + nn.ReLU(), + nn.Linear(768, cfg.d_out), + ) + self.local_token_mlp = nn.Sequential( + nn.Linear(768, 768), + nn.ReLU(), + nn.Linear(768, cfg.d_out), + ) + + def forward( + self, + context: BatchedViews, + ) -> Float[Tensor, "batch view d_out height width"]: + # Compute features from the DINO-pretrained resnet50. + resnet_features = self.resnet_backbone(context) + + # Compute features from the DINO-pretrained ViT. + b, v, _, h, w = context["image"].shape + assert h % self.patch_size == 0 and w % self.patch_size == 0 + tokens = rearrange(context["image"], "b v c h w -> (b v) c h w") + tokens = self.dino.get_intermediate_layers(tokens)[0] + global_token = self.global_token_mlp(tokens[:, 0]) + local_tokens = self.local_token_mlp(tokens[:, 1:]) + + # Repeat the global token to match the image shape. + global_token = repeat(global_token, "(b v) c -> b v c h w", b=b, v=v, h=h, w=w) + + # Repeat the local tokens to match the image shape. + local_tokens = repeat( + local_tokens, + "(b v) (h w) c -> b v c (h hps) (w wps)", + b=b, + v=v, + h=h // self.patch_size, + hps=self.patch_size, + w=w // self.patch_size, + wps=self.patch_size, + ) + + return resnet_features + local_tokens + global_token + + @property + def patch_size(self) -> int: + return int("".join(filter(str.isdigit, self.cfg.model))) + + @property + def d_out(self) -> int: + return self.cfg.d_out diff --git a/src/model/encoder/backbone/backbone_resnet.py b/src/model/encoder/backbone/backbone_resnet.py new file mode 100644 index 0000000000000000000000000000000000000000..97c0c320ad500acae43fb3483965876007f16e7b --- /dev/null +++ b/src/model/encoder/backbone/backbone_resnet.py @@ -0,0 +1,100 @@ +import functools +from dataclasses import dataclass +from typing import Literal + +import torch +import torch.nn.functional as F +import torchvision +from einops import rearrange +from jaxtyping import Float +from torch import Tensor, nn +from torchvision.models import ResNet + +from src.dataset.types import BatchedViews +from .backbone import Backbone + + +@dataclass +class BackboneResnetCfg: + name: Literal["resnet"] + model: Literal[ + "resnet18", "resnet34", "resnet50", "resnet101", "resnet152", "dino_resnet50" + ] + num_layers: int + use_first_pool: bool + d_out: int + + +class BackboneResnet(Backbone[BackboneResnetCfg]): + model: ResNet + + def __init__(self, cfg: BackboneResnetCfg, d_in: int) -> None: + super().__init__(cfg) + + assert d_in == 3 + + norm_layer = functools.partial( + nn.InstanceNorm2d, + affine=False, + track_running_stats=False, + ) + + if cfg.model == "dino_resnet50": + self.model = torch.hub.load("facebookresearch/dino:main", "dino_resnet50") + else: + self.model = getattr(torchvision.models, cfg.model)(norm_layer=norm_layer) + + # Set up projections + self.projections = nn.ModuleDict({}) + for index in range(1, cfg.num_layers): + key = f"layer{index}" + block = getattr(self.model, key) + conv_index = 1 + try: + while True: + d_layer_out = getattr(block[-1], f"conv{conv_index}").out_channels + conv_index += 1 + except AttributeError: + pass + self.projections[key] = nn.Conv2d(d_layer_out, cfg.d_out, 1) + + # Add a projection for the first layer. + self.projections["layer0"] = nn.Conv2d( + self.model.conv1.out_channels, cfg.d_out, 1 + ) + + def forward( + self, + context: BatchedViews, + ) -> Float[Tensor, "batch view d_out height width"]: + # Merge the batch dimensions. + b, v, _, h, w = context["image"].shape + x = rearrange(context["image"], "b v c h w -> (b v) c h w") + + # Run the images through the resnet. + x = self.model.conv1(x) + x = self.model.bn1(x) + x = self.model.relu(x) + features = [self.projections["layer0"](x)] + + # Propagate the input through the resnet's layers. + for index in range(1, self.cfg.num_layers): + key = f"layer{index}" + if index == 0 and self.cfg.use_first_pool: + x = self.model.maxpool(x) + x = getattr(self.model, key)(x) + features.append(self.projections[key](x)) + + # Upscale the features. + features = [ + F.interpolate(f, (h, w), mode="bilinear", align_corners=True) + for f in features + ] + features = torch.stack(features).sum(dim=0) + + # Separate batch dimensions. + return rearrange(features, "(b v) c h w -> b v c h w", b=b, v=v) + + @property + def d_out(self) -> int: + return self.cfg.d_out diff --git a/src/model/encoder/backbone/croco/README.md b/src/model/encoder/backbone/croco/README.md new file mode 100644 index 0000000000000000000000000000000000000000..af8a9f0dfce215af0e9a4a17fbdabe3a0c18f0ed --- /dev/null +++ b/src/model/encoder/backbone/croco/README.md @@ -0,0 +1,7 @@ +Most of the code under src/model/encoder/backbone/croco/ is from the original CROCO implementation. +The code is not modified in any way except the relative module path. +The original code can be found at [croco Github Repo](https://github.com/naver/croco/tree/743ee71a2a9bf57cea6832a9064a70a0597fcfcb/models). + + +Except: +- 'misc.py', 'patch_embed.py' is from DUSt3R. \ No newline at end of file diff --git a/src/model/encoder/backbone/croco/__init__.py b/src/model/encoder/backbone/croco/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/model/encoder/backbone/croco/blocks.py b/src/model/encoder/backbone/croco/blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..18133524f0ae265b0bd8d062d7c9eeaa63858a9b --- /dev/null +++ b/src/model/encoder/backbone/croco/blocks.py @@ -0,0 +1,241 @@ +# Copyright (C) 2022-present Naver Corporation. All rights reserved. +# Licensed under CC BY-NC-SA 4.0 (non-commercial use only). + + +# -------------------------------------------------------- +# Main encoder/decoder blocks +# -------------------------------------------------------- +# References: +# timm +# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py +# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/helpers.py +# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py +# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/mlp.py +# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/patch_embed.py + + +import torch +import torch.nn as nn + +from itertools import repeat +import collections.abc + + +def _ntuple(n): + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return x + return tuple(repeat(x, n)) + return parse +to_2tuple = _ntuple(2) + +def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + """ + if drop_prob == 0. or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + """ + def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) + + def extra_repr(self): + return f'drop_prob={round(self.drop_prob,3):0.3f}' + +class Mlp(nn.Module): + """ MLP as used in Vision Transformer, MLP-Mixer and related networks""" + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias[0]) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias[1]) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.fc2(x) + x = self.drop2(x) + return x + +class Attention(nn.Module): + + def __init__(self, dim, rope=None, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim ** -0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + self.rope = rope + + def forward(self, x, xpos): + B, N, C = x.shape + + qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).transpose(1,3) + q, k, v = [qkv[:,:,i] for i in range(3)] + # q,k,v = qkv.unbind(2) # make torchscript happy (cannot use tensor as tuple) + + if self.rope is not None: + q = self.rope(q, xpos) + k = self.rope(k, xpos) + + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + +class Block(nn.Module): + + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, rope=None): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + def forward(self, x, xpos): + x = x + self.drop_path(self.attn(self.norm1(x), xpos)) + x = x + self.drop_path(self.mlp(self.norm2(x))) + return x + +class CrossAttention(nn.Module): + + def __init__(self, dim, rope=None, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim ** -0.5 + + self.projq = nn.Linear(dim, dim, bias=qkv_bias) + self.projk = nn.Linear(dim, dim, bias=qkv_bias) + self.projv = nn.Linear(dim, dim, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.rope = rope + + def forward(self, query, key, value, qpos, kpos): + B, Nq, C = query.shape + Nk = key.shape[1] + Nv = value.shape[1] + + q = self.projq(query).reshape(B,Nq,self.num_heads, C// self.num_heads).permute(0, 2, 1, 3) + k = self.projk(key).reshape(B,Nk,self.num_heads, C// self.num_heads).permute(0, 2, 1, 3) + v = self.projv(value).reshape(B,Nv,self.num_heads, C// self.num_heads).permute(0, 2, 1, 3) + + if self.rope is not None: + q = self.rope(q, qpos) + k = self.rope(k, kpos) + + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, Nq, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + +class DecoderBlock(nn.Module): + + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, norm_mem=True, rope=None): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) + self.cross_attn = CrossAttention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + self.norm3 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + self.norm_y = norm_layer(dim) if norm_mem else nn.Identity() + + def forward(self, x, y, xpos, ypos): + x = x + self.drop_path(self.attn(self.norm1(x), xpos)) + y_ = self.norm_y(y) + x = x + self.drop_path(self.cross_attn(self.norm2(x), y_, y_, xpos, ypos)) + x = x + self.drop_path(self.mlp(self.norm3(x))) + return x, y + + +# patch embedding +class PositionGetter(object): + """ return positions of patches """ + + def __init__(self): + self.cache_positions = {} + + def __call__(self, b, h, w, device): + if not (h,w) in self.cache_positions: + x = torch.arange(w, device=device) + y = torch.arange(h, device=device) + self.cache_positions[h,w] = torch.cartesian_prod(y, x) # (h, w, 2) + pos = self.cache_positions[h,w].view(1, h*w, 2).expand(b, -1, 2).clone() + return pos + +class PatchEmbed(nn.Module): + """ just adding _init_weights + position getter compared to timm.models.layers.patch_embed.PatchEmbed""" + + def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + self.img_size = img_size + self.patch_size = patch_size + self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) + self.num_patches = self.grid_size[0] * self.grid_size[1] + self.flatten = flatten + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + self.position_getter = PositionGetter() + + def forward(self, x): + B, C, H, W = x.shape + torch._assert(H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).") + torch._assert(W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).") + x = self.proj(x) + pos = self.position_getter(B, x.size(2), x.size(3), x.device) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # BCHW -> BNC + x = self.norm(x) + return x, pos + + def _init_weights(self): + w = self.proj.weight.data + torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1])) + diff --git a/src/model/encoder/backbone/croco/croco.py b/src/model/encoder/backbone/croco/croco.py new file mode 100644 index 0000000000000000000000000000000000000000..74bcb8521a1ef0f198f911ab3e7dac166b281f5c --- /dev/null +++ b/src/model/encoder/backbone/croco/croco.py @@ -0,0 +1,249 @@ +# Copyright (C) 2022-present Naver Corporation. All rights reserved. +# Licensed under CC BY-NC-SA 4.0 (non-commercial use only). + + +# -------------------------------------------------------- +# CroCo model during pretraining +# -------------------------------------------------------- + + + +import torch +import torch.nn as nn +torch.backends.cuda.matmul.allow_tf32 = True # for gpu >= Ampere and pytorch >= 1.12 +from functools import partial + +from .blocks import Block, DecoderBlock, PatchEmbed +from .pos_embed import get_2d_sincos_pos_embed, RoPE2D +from .masking import RandomMask + + +class CroCoNet(nn.Module): + + def __init__(self, + img_size=224, # input image size + patch_size=16, # patch_size + mask_ratio=0.9, # ratios of masked tokens + enc_embed_dim=768, # encoder feature dimension + enc_depth=12, # encoder depth + enc_num_heads=12, # encoder number of heads in the transformer block + dec_embed_dim=512, # decoder feature dimension + dec_depth=8, # decoder depth + dec_num_heads=16, # decoder number of heads in the transformer block + mlp_ratio=4, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + norm_im2_in_dec=True, # whether to apply normalization of the 'memory' = (second image) in the decoder + pos_embed='cosine', # positional embedding (either cosine or RoPE100) + ): + + super(CroCoNet, self).__init__() + + # patch embeddings (with initialization done as in MAE) + self._set_patch_embed(img_size, patch_size, enc_embed_dim) + + # mask generations + self._set_mask_generator(self.patch_embed.num_patches, mask_ratio) + + self.pos_embed = pos_embed + if pos_embed=='cosine': + # positional embedding of the encoder + enc_pos_embed = get_2d_sincos_pos_embed(enc_embed_dim, int(self.patch_embed.num_patches**.5), n_cls_token=0) + self.register_buffer('enc_pos_embed', torch.from_numpy(enc_pos_embed).float()) + # positional embedding of the decoder + dec_pos_embed = get_2d_sincos_pos_embed(dec_embed_dim, int(self.patch_embed.num_patches**.5), n_cls_token=0) + self.register_buffer('dec_pos_embed', torch.from_numpy(dec_pos_embed).float()) + # pos embedding in each block + self.rope = None # nothing for cosine + elif pos_embed.startswith('RoPE'): # eg RoPE100 + self.enc_pos_embed = None # nothing to add in the encoder with RoPE + self.dec_pos_embed = None # nothing to add in the decoder with RoPE + if RoPE2D is None: raise ImportError("Cannot find cuRoPE2D, please install it following the README instructions") + freq = float(pos_embed[len('RoPE'):]) + self.rope = RoPE2D(freq=freq) + else: + raise NotImplementedError('Unknown pos_embed '+pos_embed) + + # transformer for the encoder + self.enc_depth = enc_depth + self.enc_embed_dim = enc_embed_dim + self.enc_blocks = nn.ModuleList([ + Block(enc_embed_dim, enc_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer, rope=self.rope) + for i in range(enc_depth)]) + self.enc_norm = norm_layer(enc_embed_dim) + + # masked tokens + self._set_mask_token(dec_embed_dim) + + # decoder + self._set_decoder(enc_embed_dim, dec_embed_dim, dec_num_heads, dec_depth, mlp_ratio, norm_layer, norm_im2_in_dec) + + # prediction head + self._set_prediction_head(dec_embed_dim, patch_size) + + # initializer weights + self.initialize_weights() + + def _set_patch_embed(self, img_size=224, patch_size=16, enc_embed_dim=768): + self.patch_embed = PatchEmbed(img_size, patch_size, 3, enc_embed_dim) + + def _set_mask_generator(self, num_patches, mask_ratio): + self.mask_generator = RandomMask(num_patches, mask_ratio) + + def _set_mask_token(self, dec_embed_dim): + self.mask_token = nn.Parameter(torch.zeros(1, 1, dec_embed_dim)) + + def _set_decoder(self, enc_embed_dim, dec_embed_dim, dec_num_heads, dec_depth, mlp_ratio, norm_layer, norm_im2_in_dec): + self.dec_depth = dec_depth + self.dec_embed_dim = dec_embed_dim + # transfer from encoder to decoder + self.decoder_embed = nn.Linear(enc_embed_dim+0, dec_embed_dim, bias=True) + # transformer for the decoder + self.dec_blocks = nn.ModuleList([ + DecoderBlock(dec_embed_dim, dec_num_heads, mlp_ratio=mlp_ratio, qkv_bias=True, norm_layer=norm_layer, norm_mem=norm_im2_in_dec, rope=self.rope) + for i in range(dec_depth)]) + # final norm layer + self.dec_norm = norm_layer(dec_embed_dim) + + def _set_prediction_head(self, dec_embed_dim, patch_size): + self.prediction_head = nn.Linear(dec_embed_dim, patch_size**2 * 3, bias=True) + + + def initialize_weights(self): + # patch embed + self.patch_embed._init_weights() + # mask tokens + if self.mask_token is not None: torch.nn.init.normal_(self.mask_token, std=.02) + # linears and layer norms + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + # we use xavier_uniform following official JAX ViT: + torch.nn.init.xavier_uniform_(m.weight) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def _encode_image(self, image, do_mask=False, return_all_blocks=False): + """ + image has B x 3 x img_size x img_size + do_mask: whether to perform masking or not + return_all_blocks: if True, return the features at the end of every block + instead of just the features from the last block (eg for some prediction heads) + """ + # embed the image into patches (x has size B x Npatches x C) + # and get position if each return patch (pos has size B x Npatches x 2) + x, pos = self.patch_embed(image) + # add positional embedding without cls token + if self.enc_pos_embed is not None: + x = x + self.enc_pos_embed[None,...] + # apply masking + B,N,C = x.size() + if do_mask: + masks = self.mask_generator(x) + x = x[~masks].view(B, -1, C) + posvis = pos[~masks].view(B, -1, 2) + else: + B,N,C = x.size() + masks = torch.zeros((B,N), dtype=bool) + posvis = pos + # now apply the transformer encoder and normalization + if return_all_blocks: + out = [] + for blk in self.enc_blocks: + x = blk(x, posvis) + out.append(x) + out[-1] = self.enc_norm(out[-1]) + return out, pos, masks + else: + for blk in self.enc_blocks: + x = blk(x, posvis) + x = self.enc_norm(x) + return x, pos, masks + + def _decoder(self, feat1, pos1, masks1, feat2, pos2, return_all_blocks=False): + """ + return_all_blocks: if True, return the features at the end of every block + instead of just the features from the last block (eg for some prediction heads) + + masks1 can be None => assume image1 fully visible + """ + # encoder to decoder layer + visf1 = self.decoder_embed(feat1) + f2 = self.decoder_embed(feat2) + # append masked tokens to the sequence + B,Nenc,C = visf1.size() + if masks1 is None: # downstreams + f1_ = visf1 + else: # pretraining + Ntotal = masks1.size(1) + f1_ = self.mask_token.repeat(B, Ntotal, 1).to(dtype=visf1.dtype) + f1_[~masks1] = visf1.view(B * Nenc, C) + # add positional embedding + if self.dec_pos_embed is not None: + f1_ = f1_ + self.dec_pos_embed + f2 = f2 + self.dec_pos_embed + # apply Transformer blocks + out = f1_ + out2 = f2 + if return_all_blocks: + _out, out = out, [] + for blk in self.dec_blocks: + _out, out2 = blk(_out, out2, pos1, pos2) + out.append(_out) + out[-1] = self.dec_norm(out[-1]) + else: + for blk in self.dec_blocks: + out, out2 = blk(out, out2, pos1, pos2) + out = self.dec_norm(out) + return out + + def patchify(self, imgs): + """ + imgs: (B, 3, H, W) + x: (B, L, patch_size**2 *3) + """ + p = self.patch_embed.patch_size[0] + assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0 + + h = w = imgs.shape[2] // p + x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p)) + x = torch.einsum('nchpwq->nhwpqc', x) + x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3)) + + return x + + def unpatchify(self, x, channels=3): + """ + x: (N, L, patch_size**2 *channels) + imgs: (N, 3, H, W) + """ + patch_size = self.patch_embed.patch_size[0] + h = w = int(x.shape[1]**.5) + assert h * w == x.shape[1] + x = x.reshape(shape=(x.shape[0], h, w, patch_size, patch_size, channels)) + x = torch.einsum('nhwpqc->nchpwq', x) + imgs = x.reshape(shape=(x.shape[0], channels, h * patch_size, h * patch_size)) + return imgs + + def forward(self, img1, img2): + """ + img1: tensor of size B x 3 x img_size x img_size + img2: tensor of size B x 3 x img_size x img_size + + out will be B x N x (3*patch_size*patch_size) + masks are also returned as B x N just in case + """ + # encoder of the masked first image + feat1, pos1, mask1 = self._encode_image(img1, do_mask=True) + # encoder of the second image + feat2, pos2, _ = self._encode_image(img2, do_mask=False) + # decoder + decfeat = self._decoder(feat1, pos1, mask1, feat2, pos2) + # prediction head + out = self.prediction_head(decfeat) + # get target + target = self.patchify(img1) + return out, mask1, target diff --git a/src/model/encoder/backbone/croco/curope/__init__.py b/src/model/encoder/backbone/croco/curope/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..25e3d48a162760260826080f6366838e83e26878 --- /dev/null +++ b/src/model/encoder/backbone/croco/curope/__init__.py @@ -0,0 +1,4 @@ +# Copyright (C) 2022-present Naver Corporation. All rights reserved. +# Licensed under CC BY-NC-SA 4.0 (non-commercial use only). + +from .curope2d import cuRoPE2D diff --git a/src/model/encoder/backbone/croco/curope/curope.cpp b/src/model/encoder/backbone/croco/curope/curope.cpp new file mode 100644 index 0000000000000000000000000000000000000000..8fe9058e05aa1bf3f37b0d970edc7312bc68455b --- /dev/null +++ b/src/model/encoder/backbone/croco/curope/curope.cpp @@ -0,0 +1,69 @@ +/* + Copyright (C) 2022-present Naver Corporation. All rights reserved. + Licensed under CC BY-NC-SA 4.0 (non-commercial use only). +*/ + +#include
+
+- Low-rank 4D bilateral grid for MR enhancement:
+
+
+
+
+Below is the API reference.
+
+"""
+
+import tensorly as tl
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+tl.set_backend("pytorch")
+
+
+def color_correct(
+ img: torch.Tensor, ref: torch.Tensor, num_iters: int = 5, eps: float = 0.5 / 255
+) -> torch.Tensor:
+ """
+ Warp `img` to match the colors in `ref_img` using iterative color matching.
+
+ This function performs color correction by warping the colors of the input image
+ to match those of a reference image. It uses a least squares method to find a
+ transformation that maps the input image's colors to the reference image's colors.
+
+ The algorithm iteratively solves a system of linear equations, updating the set of
+ unsaturated pixels in each iteration. This approach helps handle non-linear color
+ transformations and reduces the impact of clipping.
+
+ Args:
+ img (torch.Tensor): Input image to be color corrected. Shape: [..., num_channels]
+ ref (torch.Tensor): Reference image to match colors. Shape: [..., num_channels]
+ num_iters (int, optional): Number of iterations for the color matching process.
+ Default is 5.
+ eps (float, optional): Small value to determine the range of unclipped pixels.
+ Default is 0.5 / 255.
+
+ Returns:
+ torch.Tensor: Color corrected image with the same shape as the input image.
+
+ Note:
+ - Both input and reference images should be in the range [0, 1].
+ - The function works with any number of channels, but typically used with 3 (RGB).
+ """
+ if img.shape[-1] != ref.shape[-1]:
+ raise ValueError(
+ f"img's {img.shape[-1]} and ref's {ref.shape[-1]} channels must match"
+ )
+ num_channels = img.shape[-1]
+ img_mat = img.reshape([-1, num_channels])
+ ref_mat = ref.reshape([-1, num_channels])
+
+ def is_unclipped(z):
+ return (z >= eps) & (z <= 1 - eps) # z \in [eps, 1-eps].
+
+ mask0 = is_unclipped(img_mat)
+ # Because the set of saturated pixels may change after solving for a
+ # transformation, we repeatedly solve a system `num_iters` times and update
+ # our estimate of which pixels are saturated.
+ for _ in range(num_iters):
+ # Construct the left hand side of a linear system that contains a quadratic
+ # expansion of each pixel of `img`.
+ a_mat = []
+ for c in range(num_channels):
+ a_mat.append(img_mat[:, c : (c + 1)] * img_mat[:, c:]) # Quadratic term.
+ a_mat.append(img_mat) # Linear term.
+ a_mat.append(torch.ones_like(img_mat[:, :1])) # Bias term.
+ a_mat = torch.cat(a_mat, dim=-1)
+ warp = []
+ for c in range(num_channels):
+ # Construct the right hand side of a linear system containing each color
+ # of `ref`.
+ b = ref_mat[:, c]
+ # Ignore rows of the linear system that were saturated in the input or are
+ # saturated in the current corrected color estimate.
+ mask = mask0[:, c] & is_unclipped(img_mat[:, c]) & is_unclipped(b)
+ ma_mat = torch.where(mask[:, None], a_mat, torch.zeros_like(a_mat))
+ mb = torch.where(mask, b, torch.zeros_like(b))
+ w = torch.linalg.lstsq(ma_mat, mb, rcond=-1)[0]
+ assert torch.all(torch.isfinite(w))
+ warp.append(w)
+ warp = torch.stack(warp, dim=-1)
+ # Apply the warp to update img_mat.
+ img_mat = torch.clip(torch.matmul(a_mat, warp), 0, 1)
+ corrected_img = torch.reshape(img_mat, img.shape)
+ return corrected_img
+
+
+def bilateral_grid_tv_loss(model, config):
+ """Computes total variations of bilateral grids."""
+ total_loss = 0.0
+
+ for bil_grids in model.bil_grids:
+ total_loss += config.bilgrid_tv_loss_mult * total_variation_loss(
+ bil_grids.grids
+ )
+
+ return total_loss
+
+
+def color_affine_transform(affine_mats, rgb):
+ """Applies color affine transformations.
+
+ Args:
+ affine_mats (torch.Tensor): Affine transformation matrices. Supported shape: $(..., 3, 4)$.
+ rgb (torch.Tensor): Input RGB values. Supported shape: $(..., 3)$.
+
+ Returns:
+ Output transformed colors of shape $(..., 3)$.
+ """
+ return (
+ torch.matmul(affine_mats[..., :3], rgb.unsqueeze(-1)).squeeze(-1)
+ + affine_mats[..., 3]
+ )
+
+
+def _num_tensor_elems(t):
+ return max(torch.prod(torch.tensor(t.size()[1:]).float()).item(), 1.0)
+
+
+def total_variation_loss(x): # noqa: F811
+ """Returns total variation on multi-dimensional tensors.
+
+ Args:
+ x (torch.Tensor): The input tensor with shape $(B, C, ...)$, where $B$ is the batch size and $C$ is the channel size.
+ """
+ batch_size = x.shape[0]
+ tv = 0
+ for i in range(2, len(x.shape)):
+ n_res = x.shape[i]
+ idx1 = torch.arange(1, n_res, device=x.device)
+ idx2 = torch.arange(0, n_res - 1, device=x.device)
+ x1 = x.index_select(i, idx1)
+ x2 = x.index_select(i, idx2)
+ count = _num_tensor_elems(x1)
+ tv += torch.pow((x1 - x2), 2).sum() / count
+ return tv / batch_size
+
+
+def slice(bil_grids, xy, rgb, grid_idx):
+ """Slices a batch of 3D bilateral grids by pixel coordinates `xy` and gray-scale guidances of pixel colors `rgb`.
+
+ Supports 2-D, 3-D, and 4-D input shapes. The first dimension of the input is the batch size
+ and the last dimension is 2 for `xy`, 3 for `rgb`, and 1 for `grid_idx`.
+
+ The return value is a dictionary containing the affine transformations `affine_mats` sliced from bilateral grids and
+ the output color `rgb_out` after applying the afffine transformations.
+
+ In the 2-D input case, `xy` is a $(N, 2)$ tensor, `rgb` is a $(N, 3)$ tensor, and `grid_idx` is a $(N, 1)$ tensor.
+ Then `affine_mats[i]` can be obtained via slicing the bilateral grid indexed at `grid_idx[i]` by `xy[i, :]` and `rgb2gray(rgb[i, :])`.
+ For 3-D and 4-D input cases, the behavior of indexing bilateral grids and coordinates is the same with the 2-D case.
+
+ .. note::
+ This function can be regarded as a wrapper of `color_affine_transform` and `BilateralGrid` with a slight performance improvement.
+ When `grid_idx` contains a unique index, only a single bilateral grid will used during the slicing. In this case, this function will not
+ perform tensor indexing to avoid data copy and extra memory
+ (see [this](https://discuss.pytorch.org/t/does-indexing-a-tensor-return-a-copy-of-it/164905)).
+
+ Args:
+ bil_grids (`BilateralGrid`): An instance of $N$ bilateral grids.
+ xy (torch.Tensor): The x-y coordinates of shape $(..., 2)$ in the range of $[0,1]$.
+ rgb (torch.Tensor): The RGB values of shape $(..., 3)$ for computing the guidance coordinates, ranging in $[0,1]$.
+ grid_idx (torch.Tensor): The indices of bilateral grids for each slicing. Shape: $(..., 1)$.
+
+ Returns:
+ A dictionary with keys and values as follows:
+ ```
+ {
+ "rgb": Transformed RGB colors. Shape: (..., 3),
+ "rgb_affine_mats": The sliced affine transformation matrices from bilateral grids. Shape: (..., 3, 4)
+ }
+ ```
+ """
+
+ sh_ = rgb.shape
+
+ grid_idx_unique = torch.unique(grid_idx)
+ if len(grid_idx_unique) == 1:
+ # All pixels are from a single view.
+ grid_idx = grid_idx_unique # (1,)
+ xy = xy.unsqueeze(0) # (1, ..., 2)
+ rgb = rgb.unsqueeze(0) # (1, ..., 3)
+ else:
+ # Pixels are randomly sampled from different views.
+ if len(grid_idx.shape) == 4:
+ grid_idx = grid_idx[:, 0, 0, 0] # (chunk_size,)
+ elif len(grid_idx.shape) == 3:
+ grid_idx = grid_idx[:, 0, 0] # (chunk_size,)
+ elif len(grid_idx.shape) == 2:
+ grid_idx = grid_idx[:, 0] # (chunk_size,)
+ else:
+ raise ValueError(
+ "The input to bilateral grid slicing is not supported yet."
+ )
+
+ affine_mats = bil_grids(xy, rgb, grid_idx)
+ rgb = color_affine_transform(affine_mats, rgb)
+
+ return {
+ "rgb": rgb.reshape(*sh_),
+ "rgb_affine_mats": affine_mats.reshape(
+ *sh_[:-1], affine_mats.shape[-2], affine_mats.shape[-1]
+ ),
+ }
+
+
+class BilateralGrid(nn.Module):
+ """Class for 3D bilateral grids.
+
+ Holds one or more than one bilateral grids.
+ """
+
+ def __init__(self, num, grid_X=16, grid_Y=16, grid_W=8):
+ """
+ Args:
+ num (int): The number of bilateral grids (i.e., the number of views).
+ grid_X (int): Defines grid width $W$.
+ grid_Y (int): Defines grid height $H$.
+ grid_W (int): Defines grid guidance dimension $L$.
+ """
+ super(BilateralGrid, self).__init__()
+
+ self.grid_width = grid_X
+ """Grid width. Type: int."""
+ self.grid_height = grid_Y
+ """Grid height. Type: int."""
+ self.grid_guidance = grid_W
+ """Grid guidance dimension. Type: int."""
+
+ # Initialize grids.
+ grid = self._init_identity_grid()
+ self.grids = nn.Parameter(grid.tile(num, 1, 1, 1, 1)) # (N, 12, L, H, W)
+ """ A 5-D tensor of shape $(N, 12, L, H, W)$."""
+
+ # Weights of BT601 RGB-to-gray.
+ self.register_buffer("rgb2gray_weight", torch.Tensor([[0.299, 0.587, 0.114]]))
+ self.rgb2gray = lambda rgb: (rgb @ self.rgb2gray_weight.T) * 2.0 - 1.0
+ """ A function that converts RGB to gray-scale guidance in $[-1, 1]$."""
+
+ def _init_identity_grid(self):
+ grid = torch.tensor(
+ [
+ 1.0,
+ 0,
+ 0,
+ 0,
+ 0,
+ 1.0,
+ 0,
+ 0,
+ 0,
+ 0,
+ 1.0,
+ 0,
+ ]
+ ).float()
+ grid = grid.repeat(
+ [self.grid_guidance * self.grid_height * self.grid_width, 1]
+ ) # (L * H * W, 12)
+ grid = grid.reshape(
+ 1, self.grid_guidance, self.grid_height, self.grid_width, -1
+ ) # (1, L, H, W, 12)
+ grid = grid.permute(0, 4, 1, 2, 3) # (1, 12, L, H, W)
+ return grid
+
+ def tv_loss(self):
+ """Computes and returns total variation loss on the bilateral grids."""
+ return total_variation_loss(self.grids)
+
+ def forward(self, grid_xy, rgb, idx=None):
+ """Bilateral grid slicing. Supports 2-D, 3-D, 4-D, and 5-D input.
+ For the 2-D, 3-D, and 4-D cases, please refer to `slice`.
+ For the 5-D cases, `idx` will be unused and the first dimension of `xy` should be
+ equal to the number of bilateral grids. Then this function becomes PyTorch's
+ [`F.grid_sample`](https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html).
+
+ Args:
+ grid_xy (torch.Tensor): The x-y coordinates in the range of $[0,1]$.
+ rgb (torch.Tensor): The RGB values in the range of $[0,1]$.
+ idx (torch.Tensor): The bilateral grid indices.
+
+ Returns:
+ Sliced affine matrices of shape $(..., 3, 4)$.
+ """
+
+ grids = self.grids
+ input_ndims = len(grid_xy.shape)
+ assert len(rgb.shape) == input_ndims
+
+ if input_ndims > 1 and input_ndims < 5:
+ # Convert input into 5D
+ for i in range(5 - input_ndims):
+ grid_xy = grid_xy.unsqueeze(1)
+ rgb = rgb.unsqueeze(1)
+ assert idx is not None
+ elif input_ndims != 5:
+ raise ValueError(
+ "Bilateral grid slicing only takes either 2D, 3D, 4D and 5D inputs"
+ )
+
+ grids = self.grids
+ if idx is not None:
+ grids = grids[idx]
+ assert grids.shape[0] == grid_xy.shape[0]
+
+ # Generate slicing coordinates.
+ grid_xy = (grid_xy - 0.5) * 2 # Rescale to [-1, 1].
+ grid_z = self.rgb2gray(rgb)
+
+ # print(grid_xy.shape, grid_z.shape)
+ # exit()
+ grid_xyz = torch.cat([grid_xy, grid_z], dim=-1) # (N, m, h, w, 3)
+
+ affine_mats = F.grid_sample(
+ grids, grid_xyz, mode="bilinear", align_corners=True, padding_mode="border"
+ ) # (N, 12, m, h, w)
+ affine_mats = affine_mats.permute(0, 2, 3, 4, 1) # (N, m, h, w, 12)
+ affine_mats = affine_mats.reshape(
+ *affine_mats.shape[:-1], 3, 4
+ ) # (N, m, h, w, 3, 4)
+
+ for _ in range(5 - input_ndims):
+ affine_mats = affine_mats.squeeze(1)
+
+ return affine_mats
+
+
+def slice4d(bil_grid4d, xyz, rgb):
+ """Slices a 4D bilateral grid by point coordinates `xyz` and gray-scale guidances of radiance colors `rgb`.
+
+ Args:
+ bil_grid4d (`BilateralGridCP4D`): The input 4D bilateral grid.
+ xyz (torch.Tensor): The xyz coordinates with shape $(..., 3)$.
+ rgb (torch.Tensor): The RGB values with shape $(..., 3)$.
+
+ Returns:
+ A dictionary with keys and values as follows:
+ ```
+ {
+ "rgb": Transformed radiance RGB colors. Shape: (..., 3),
+ "rgb_affine_mats": The sliced affine transformation matrices from the 4D bilateral grid. Shape: (..., 3, 4)
+ }
+ ```
+ """
+
+ affine_mats = bil_grid4d(xyz, rgb)
+ rgb = color_affine_transform(affine_mats, rgb)
+
+ return {"rgb": rgb, "rgb_affine_mats": affine_mats}
+
+
+class _ScaledTanh(nn.Module):
+ def __init__(self, s=2.0):
+ super().__init__()
+ self.scaler = s
+
+ def forward(self, x):
+ return torch.tanh(self.scaler * x)
+
+
+class BilateralGridCP4D(nn.Module):
+ """Class for low-rank 4D bilateral grids."""
+
+ def __init__(
+ self,
+ grid_X=16,
+ grid_Y=16,
+ grid_Z=16,
+ grid_W=8,
+ rank=5,
+ learn_gray=True,
+ gray_mlp_width=8,
+ gray_mlp_depth=2,
+ init_noise_scale=1e-6,
+ bound=2.0,
+ ):
+ """
+ Args:
+ grid_X (int): Defines grid width.
+ grid_Y (int): Defines grid height.
+ grid_Z (int): Defines grid depth.
+ grid_W (int): Defines grid guidance dimension.
+ rank (int): Rank of the 4D bilateral grid.
+ learn_gray (bool): If True, an MLP will be learned to convert RGB colors to gray-scale guidances.
+ gray_mlp_width (int): The MLP width for learnable guidance.
+ gray_mlp_depth (int): The number of MLP layers for learnable guidance.
+ init_noise_scale (float): The noise scale of the initialized factors.
+ bound (float): The bound of the xyz coordinates.
+ """
+ super(BilateralGridCP4D, self).__init__()
+
+ self.grid_X = grid_X
+ """Grid width. Type: int."""
+ self.grid_Y = grid_Y
+ """Grid height. Type: int."""
+ self.grid_Z = grid_Z
+ """Grid depth. Type: int."""
+ self.grid_W = grid_W
+ """Grid guidance dimension. Type: int."""
+ self.rank = rank
+ """Rank of the 4D bilateral grid. Type: int."""
+ self.learn_gray = learn_gray
+ """Flags of learnable guidance is used. Type: bool."""
+ self.gray_mlp_width = gray_mlp_width
+ """The MLP width for learnable guidance. Type: int."""
+ self.gray_mlp_depth = gray_mlp_depth
+ """The MLP depth for learnable guidance. Type: int."""
+ self.init_noise_scale = init_noise_scale
+ """The noise scale of the initialized factors. Type: float."""
+ self.bound = bound
+ """The bound of the xyz coordinates. Type: float."""
+
+ self._init_cp_factors_parafac()
+
+ self.rgb2gray = None
+ """ A function that converts RGB to gray-scale guidances in $[-1, 1]$.
+ If `learn_gray` is True, this will be an MLP network."""
+
+ if self.learn_gray:
+
+ def rgb2gray_mlp_linear(layer):
+ return nn.Linear(
+ self.gray_mlp_width,
+ self.gray_mlp_width if layer < self.gray_mlp_depth - 1 else 1,
+ )
+
+ def rgb2gray_mlp_actfn(_):
+ return nn.ReLU(inplace=True)
+
+ self.rgb2gray = nn.Sequential(
+ *(
+ [nn.Linear(3, self.gray_mlp_width)]
+ + [
+ nn_module(layer)
+ for layer in range(1, self.gray_mlp_depth)
+ for nn_module in [rgb2gray_mlp_actfn, rgb2gray_mlp_linear]
+ ]
+ + [_ScaledTanh(2.0)]
+ )
+ )
+ else:
+ # Weights of BT601/BT470 RGB-to-gray.
+ self.register_buffer(
+ "rgb2gray_weight", torch.Tensor([[0.299, 0.587, 0.114]])
+ )
+ self.rgb2gray = lambda rgb: (rgb @ self.rgb2gray_weight.T) * 2.0 - 1.0
+
+ def _init_identity_grid(self):
+ grid = torch.tensor(
+ [
+ 1.0,
+ 0,
+ 0,
+ 0,
+ 0,
+ 1.0,
+ 0,
+ 0,
+ 0,
+ 0,
+ 1.0,
+ 0,
+ ]
+ ).float()
+ grid = grid.repeat([self.grid_W * self.grid_Z * self.grid_Y * self.grid_X, 1])
+ grid = grid.reshape(self.grid_W, self.grid_Z, self.grid_Y, self.grid_X, -1)
+ grid = grid.permute(4, 0, 1, 2, 3) # (12, grid_W, grid_Z, grid_Y, grid_X)
+ return grid
+
+ def _init_cp_factors_parafac(self):
+ # Initialize identity grids.
+ init_grids = self._init_identity_grid()
+ # Random noises are added to avoid singularity.
+ init_grids = torch.randn_like(init_grids) * self.init_noise_scale + init_grids
+ from tensorly.decomposition import parafac
+
+ # Initialize grid CP factors
+ _, facs = parafac(init_grids.clone().detach(), rank=self.rank)
+
+ self.num_facs = len(facs)
+
+ self.fac_0 = nn.Linear(facs[0].shape[0], facs[0].shape[1], bias=False)
+ self.fac_0.weight = nn.Parameter(facs[0]) # (12, rank)
+
+ for i in range(1, self.num_facs):
+ fac = facs[i].T # (rank, grid_size)
+ fac = fac.view(1, fac.shape[0], fac.shape[1], 1) # (1, rank, grid_size, 1)
+ self.register_buffer(f"fac_{i}_init", fac)
+
+ fac_resid = torch.zeros_like(fac)
+ self.register_parameter(f"fac_{i}", nn.Parameter(fac_resid))
+
+ def tv_loss(self):
+ """Computes and returns total variation loss on the factors of the low-rank 4D bilateral grids."""
+
+ total_loss = 0
+ for i in range(1, self.num_facs):
+ fac = self.get_parameter(f"fac_{i}")
+ total_loss += total_variation_loss(fac)
+
+ return total_loss
+
+ def forward(self, xyz, rgb):
+ """Low-rank 4D bilateral grid slicing.
+
+ Args:
+ xyz (torch.Tensor): The xyz coordinates with shape $(..., 3)$.
+ rgb (torch.Tensor): The corresponding RGB values with shape $(..., 3)$.
+
+ Returns:
+ Sliced affine matrices with shape $(..., 3, 4)$.
+ """
+ sh_ = xyz.shape
+ xyz = xyz.reshape(-1, 3) # flatten (N, 3)
+ rgb = rgb.reshape(-1, 3) # flatten (N, 3)
+
+ xyz = xyz / self.bound
+ assert self.rgb2gray is not None
+ gray = self.rgb2gray(rgb)
+ xyzw = torch.cat([xyz, gray], dim=-1) # (N, 4)
+ xyzw = xyzw.transpose(0, 1) # (4, N)
+ coords = torch.stack([torch.zeros_like(xyzw), xyzw], dim=-1) # (4, N, 2)
+ coords = coords.unsqueeze(1) # (4, 1, N, 2)
+
+ coef = 1.0
+ for i in range(1, self.num_facs):
+ fac = self.get_parameter(f"fac_{i}") + self.get_buffer(f"fac_{i}_init")
+ coef = coef * F.grid_sample(
+ fac, coords[[i - 1]], align_corners=True, padding_mode="border"
+ ) # [1, rank, 1, N]
+ coef = coef.squeeze([0, 2]).transpose(0, 1) # (N, rank) #type: ignore
+ mat = self.fac_0(coef)
+ return mat.reshape(*sh_[:-1], 3, 4)
diff --git a/src/post_opt/simple_trainer.py b/src/post_opt/simple_trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..207f2d63266246140d161a2dda449ac0a2f281dc
--- /dev/null
+++ b/src/post_opt/simple_trainer.py
@@ -0,0 +1,1433 @@
+import json
+import math
+import os
+import time
+from collections import defaultdict
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+import imageio
+import matplotlib
+import torchvision
+import numpy as np
+import torch
+import torch.nn.functional as F
+import tqdm
+import tyro
+import viser
+import yaml
+import torchvision
+import sys
+from plyfile import PlyData, PlyElement
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
+from src.model.encoder.vggt.utils.pose_enc import pose_encoding_to_extri_intri
+from src.model.types import Gaussians
+from src.post_opt.datasets.colmap import Dataset, Parser
+from src.post_opt.datasets.traj import (
+ generate_ellipse_path_z,
+ generate_interpolated_path,
+ generate_spiral_path,
+)
+from fused_ssim import fused_ssim
+
+from src.utils.image import process_image
+from src.post_opt.exporter import export_splats
+from src.post_opt.lib_bilagrid import BilateralGrid, color_correct, slice, total_variation_loss
+from torch import Tensor
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from typing_extensions import Literal, assert_never
+from src.post_opt.utils import AppearanceOptModule, CameraOptModule, knn, rgb_to_sh, set_random_seed
+
+# from gsplat import export_splats
+from gsplat.compression import PngCompression
+from gsplat.distributed import cli
+# from gsplat.optimizers import SelectiveAdam
+# from gsplat.rendering import rasterization
+from gsplat import rasterization
+from gsplat.strategy import DefaultStrategy, MCMCStrategy
+from src.post_opt.gsplat_viewer import GsplatViewer, GsplatRenderTabState
+from nerfview import CameraState, RenderTabState, apply_float_colormap
+
+import torch
+from einops import rearrange
+from jaxtyping import Float
+from torch import Tensor
+from scipy.spatial.transform import Rotation as R
+
+from src.model.model.anysplat import AnySplat
+
+
+# pytorch3d/pytorch3d/transforms/rotation_conversions.py at main · facebookresearch/pytorch3d
+def quaternion_to_matrix(
+ quaternions: Float[Tensor, "*batch 4"],
+ eps: float = 1e-8,
+) -> Float[Tensor, "*batch 3 3"]:
+ # Order changed to match scipy format!
+ i, j, k, r = torch.unbind(quaternions, dim=-1)
+ two_s = 2 / ((quaternions * quaternions).sum(dim=-1) + eps)
+
+ o = torch.stack(
+ (
+ 1 - two_s * (j * j + k * k),
+ two_s * (i * j - k * r),
+ two_s * (i * k + j * r),
+ two_s * (i * j + k * r),
+ 1 - two_s * (i * i + k * k),
+ two_s * (j * k - i * r),
+ two_s * (i * k - j * r),
+ two_s * (j * k + i * r),
+ 1 - two_s * (i * i + j * j),
+ ),
+ -1,
+ )
+ return rearrange(o, "... (i j) -> ... i j", i=3, j=3)
+
+def construct_list_of_attributes(num_rest: int) -> list[str]:
+ attributes = ["x", "y", "z", "nx", "ny", "nz"]
+ for i in range(3):
+ attributes.append(f"f_dc_{i}")
+ for i in range(num_rest):
+ attributes.append(f"f_rest_{i}")
+ attributes.append("opacity")
+ for i in range(3):
+ attributes.append(f"scale_{i}")
+ for i in range(4):
+ attributes.append(f"rot_{i}")
+ return attributes
+
+def export_ply(
+ means: Float[Tensor, "gaussian 3"],
+ scales: Float[Tensor, "gaussian 3"],
+ rotations: Float[Tensor, "gaussian 4"],
+ harmonics: Float[Tensor, "gaussian 3 d_sh"],
+ opacities: Float[Tensor, " gaussian"],
+ path: Path,
+ shift_and_scale: bool = False,
+ save_sh_dc_only: bool = True,
+):
+ if shift_and_scale:
+ # Shift the scene so that the median Gaussian is at the origin.
+ means = means - means.median(dim=0).values
+
+ # Rescale the scene so that most Gaussians are within range [-1, 1].
+ scale_factor = means.abs().quantile(0.95, dim=0).max()
+ means = means / scale_factor
+ scales = scales / scale_factor
+
+ # Apply the rotation to the Gaussian rotations.
+ rotations = R.from_quat(rotations.detach().cpu().numpy()).as_matrix()
+ rotations = R.from_matrix(rotations).as_quat()
+ x, y, z, w = rearrange(rotations, "g xyzw -> xyzw g")
+ rotations = np.stack((w, x, y, z), axis=-1)
+
+ # Since current model use SH_degree = 4,
+ # which require large memory to store, we can only save the DC band to save memory.
+ f_dc = harmonics[..., 0]
+ f_rest = harmonics[..., 1:].flatten(start_dim=1)
+
+ dtype_full = [(attribute, "f4") for attribute in construct_list_of_attributes(0 if save_sh_dc_only else f_rest.shape[1])]
+ elements = np.empty(means.shape[0], dtype=dtype_full)
+ attributes = [
+ means.detach().cpu().numpy(),
+ torch.zeros_like(means).detach().cpu().numpy(),
+ f_dc.detach().cpu().contiguous().numpy(),
+ f_rest.detach().cpu().contiguous().numpy(),
+ opacities[..., None].detach().cpu().numpy(),
+ scales.detach().cpu().numpy(),
+ rotations,
+ ]
+ if save_sh_dc_only:
+ # remove f_rest from attributes
+ attributes.pop(3)
+
+ attributes = np.concatenate(attributes, axis=1)
+ elements[:] = list(map(tuple, attributes))
+ path.parent.mkdir(exist_ok=True, parents=True)
+ PlyData([PlyElement.describe(elements, "vertex")]).write(path)
+
+def colorize_depth_maps(depth_map, min_depth=0.0, max_depth=1.0, cmap="Spectral", valid_mask=None):
+ """
+ Colorize depth maps.
+ """
+ assert len(depth_map.shape) >= 2, "Invalid dimension"
+
+ if isinstance(depth_map, torch.Tensor):
+ depth = depth_map.detach().clone().squeeze().numpy()
+ elif isinstance(depth_map, np.ndarray):
+ depth = depth_map.copy().squeeze()
+ # reshape to [ (B,) H, W ]
+ if depth.ndim < 3:
+ depth = depth[np.newaxis, :, :]
+
+ # colorize
+ cm = matplotlib.colormaps[cmap]
+ # depth = ((depth - min_depth) / (max_depth - min_depth)).clip(0, 1)
+ depth = ((depth - depth.min()) / (depth.max() - depth.min())).clip(0, 1)
+ img_colored_np = cm(depth, bytes=False)[:, :, :, 0:3] # value from 0 to 1
+ img_colored_np = np.rollaxis(img_colored_np, 3, 1)
+
+ if valid_mask is not None:
+ if isinstance(depth_map, torch.Tensor):
+ valid_mask = valid_mask.detach().numpy()
+ valid_mask = valid_mask.squeeze() # [H, W] or [B, H, W]
+ if valid_mask.ndim < 3:
+ valid_mask = valid_mask[np.newaxis, np.newaxis, :, :]
+ else:
+ valid_mask = valid_mask[:, np.newaxis, :, :]
+ valid_mask = np.repeat(valid_mask, 3, axis=1)
+ img_colored_np[~valid_mask] = 0
+
+ if isinstance(depth_map, torch.Tensor):
+ img_colored = torch.from_numpy(img_colored_np).float()
+ elif isinstance(depth_map, np.ndarray):
+ img_colored = img_colored_np
+
+ return img_colored
+
+def build_covariance(
+ scale: Float[Tensor, "*#batch 3"],
+ rotation_xyzw: Float[Tensor, "*#batch 4"],
+) -> Float[Tensor, "*batch 3 3"]:
+ scale = scale.diag_embed()
+ rotation = quaternion_to_matrix(rotation_xyzw)
+ return (
+ rotation
+ @ scale
+ @ rearrange(scale, "... i j -> ... j i")
+ @ rearrange(rotation, "... i j -> ... j i")
+ )
+
+
+@dataclass
+class Config:
+ # Disable viewer
+ disable_viewer: bool = True
+ # Path to the .pt files. If provide, it will skip training and run evaluation only.
+ ckpt: Optional[List[str]] = None
+ # Name of compression strategy to use
+ compression: Optional[Literal["png"]] = None
+ # Render trajectory path
+ render_traj_path: str = "interp"
+
+ data_dir: str = "data/360_v2/garden"
+ # Downsample factor for the dataset
+ data_factor: int = 4
+ # Directory to save results
+ result_dir: str = "results/garden"
+ # Every N images there is a test image
+ test_every: int = 8
+ # Random crop size for training (experimental)
+ patch_size: Optional[int] = None
+ # A global scaler that applies to the scene size related parameters
+ global_scale: float = 1.0
+ # Normalize the world space
+ normalize_world_space: bool = True
+ # Camera model
+ camera_model: Literal["pinhole", "ortho", "fisheye"] = "pinhole"
+
+ # Port for the viewer server
+ port: int = 8080
+
+ # Batch size for training. Learning rates are scaled automatically
+ batch_size: int = 1
+ # A global factor to scale the number of training steps
+ steps_scaler: float = 1.0
+
+ # Number of training steps
+ max_steps: int = 3_000
+ # Steps to evaluate the model
+ eval_steps: List[int] = field(default_factory=lambda: [1, 1_000, 3_000, 7_000, 10_000])
+ # Steps to save the model
+ save_steps: List[int] = field(default_factory=lambda: [1, 1_000, 3_000, 7_000, 10_000])
+ # Whether to save ply file (storage size can be large)
+ save_ply: bool = False
+ # Steps to save the model as ply
+ ply_steps: List[int] = field(default_factory=lambda: [1, 1_000, 3_000, 7_000, 10_000])
+ # Whether to disable video generation during training and evaluation
+ disable_video: bool = False
+
+ # Initialization strategy
+ init_type: str = "sfm"
+ # Initial number of GSs. Ignored if using sfm
+ init_num_pts: int = 100_000
+ # Initial extent of GSs as a multiple of the camera extent. Ignored if using sfm
+ init_extent: float = 3.0
+ # Degree of spherical harmonics
+ sh_degree: int = 4
+ # Turn on another SH degree every this steps
+ sh_degree_interval: int = 1000
+ # Initial opacity of GS
+ init_opa: float = 0.1
+ # Initial scale of GS
+ init_scale: float = 1.0
+ # Weight for SSIM loss
+ ssim_lambda: float = 0.2
+
+ # Near plane clipping distance
+ near_plane: float = 1e-10
+ # Far plane clipping distance
+ far_plane: float = 1e10
+
+ # Strategy for GS densification
+ strategy: Union[DefaultStrategy, MCMCStrategy] = field(
+ default_factory=DefaultStrategy
+ )
+ # Use packed mode for rasterization, this leads to less memory usage but slightly slower.
+ packed: bool = False
+ # Use sparse gradients for optimization. (experimental)
+ sparse_grad: bool = False
+ # Use visible adam from Taming 3DGS. (experimental)
+ visible_adam: bool = False
+ # Anti-aliasing in rasterization. Might slightly hurt quantitative metrics.
+ antialiased: bool = False
+
+ # Use random background for training to discourage transparency
+ random_bkgd: bool = False
+
+ # Opacity regularization
+ opacity_reg: float = 0.0
+ # Scale regularization
+ scale_reg: float = 0.0
+
+ # Enable camera optimization.
+ pose_opt: bool = True
+ # Learning rate for camera optimization
+ pose_opt_lr: float = 1e-5
+ # Regularization for camera optimization as weight decay
+ pose_opt_reg: float = 1e-6
+ # Add noise to camera extrinsics. This is only to test the camera pose optimization.
+ pose_noise: float = 0.0
+
+ # Enable appearance optimization. (experimental)
+ app_opt: bool = False
+ # Appearance embedding dimension
+ app_embed_dim: int = 16
+ # Learning rate for appearance optimization
+ app_opt_lr: float = 1e-3
+ # Regularization for appearance optimization as weight decay
+ app_opt_reg: float = 1e-6
+
+ # Enable bilateral grid. (experimental)
+ use_bilateral_grid: bool = False
+ # Shape of the bilateral grid (X, Y, W)
+ bilateral_grid_shape: Tuple[int, int, int] = (16, 16, 8)
+
+ # Enable depth loss. (experimental)
+ depth_loss: bool = False
+ # Weight for depth loss
+ depth_lambda: float = 1e-2
+
+ # Dump information to tensorboard every this steps
+ tb_every: int = 100
+ # Save training images to tensorboard
+ tb_save_image: bool = False
+
+ lpips_net: Literal["vgg", "alex"] = "vgg"
+
+ lr_means: float = 1.6e-4
+ lr_scales: float = 5e-3
+ lr_quats: float = 1e-3
+ lr_opacities: float = 5e-2
+ lr_sh: float = 2.5e-3
+
+ def adjust_steps(self, factor: float):
+ self.eval_steps = [int(i * factor) for i in self.eval_steps]
+ self.save_steps = [int(i * factor) for i in self.save_steps]
+ self.ply_steps = [int(i * factor) for i in self.ply_steps]
+ self.max_steps = int(self.max_steps * factor)
+ self.sh_degree_interval = int(self.sh_degree_interval * factor)
+
+ strategy = self.strategy
+ if isinstance(strategy, DefaultStrategy):
+ # strategy.refine_start_iter = int(strategy.refine_start_iter * factor)
+ # strategy.refine_stop_iter = int(strategy.refine_stop_iter * factor)
+ # strategy.reset_every = int(strategy.reset_every * factor)
+ # strategy.refine_every = int(strategy.refine_every * factor)
+
+ strategy.refine_start_iter = 30000
+ strategy.refine_stop_iter = 0
+ strategy.reset_every = 30000
+ strategy.refine_every = 30000
+
+ elif isinstance(strategy, MCMCStrategy):
+ strategy.refine_start_iter = int(strategy.refine_start_iter * factor)
+ strategy.refine_stop_iter = int(strategy.refine_stop_iter * factor)
+ strategy.refine_every = int(strategy.refine_every * factor)
+ else:
+ assert_never(strategy)
+
+
+def create_splats_with_optimizers(
+ gaussians: Gaussians,
+ init_num_pts: int = 100_000,
+ init_extent: float = 3.0,
+ init_opacity: float = 0.1,
+ init_scale: float = 1.0,
+ sh_degree: int = 3,
+ sparse_grad: bool = False,
+ visible_adam: bool = False,
+ batch_size: int = 1,
+ feature_dim: Optional[int] = None,
+ device: str = "cuda",
+ world_rank: int = 0,
+ world_size: int = 1,
+ cfg: Config = None,
+) -> Tuple[torch.nn.ParameterDict, Dict[str, torch.optim.Optimizer]]:
+
+ points = gaussians.means[0].detach().float()
+ scales = torch.log(gaussians.scales[0].detach().float())
+ quats = gaussians.rotations[0].detach().float()
+ opacities = torch.logit(gaussians.opacities[0].detach().float())
+ harmonics = gaussians.harmonics[0].detach().float().permute(0, 2, 1).contiguous()
+
+ N = points.shape[0]
+
+ scene_scale = 1.0
+ masks = opacities.sigmoid() > 0.01
+ harmonics = harmonics[masks]
+ params = [
+ # name, value, lr
+ ("means", torch.nn.Parameter(points[masks]), cfg.lr_means * scene_scale),
+ ("scales", torch.nn.Parameter(scales[masks]), cfg.lr_scales),
+ ("quats", torch.nn.Parameter(quats[masks]), cfg.lr_quats),
+ ("opacities", torch.nn.Parameter(opacities[masks]), cfg.lr_opacities),
+ ]
+
+ params.append(("sh0", torch.nn.Parameter(harmonics[:, :1, :]), cfg.lr_sh))
+ params.append(("shN", torch.nn.Parameter(harmonics[:, 1:, :]), cfg.lr_sh/20))
+
+ splats = torch.nn.ParameterDict({n: v for n, v, _ in params}).to(device)
+ # Scale learning rate based on batch size, reference:
+ # https://www.cs.princeton.edu/~smalladi/blog/2024/01/22/SDEs-ScalingRules/
+ # Note that this would not make the training exactly equivalent, see
+ # https://arxiv.org/pdf/2402.18824v1
+ BS = batch_size * world_size
+ optimizer_class = None
+ if sparse_grad:
+ optimizer_class = torch.optim.SparseAdam
+ elif visible_adam:
+ optimizer_class = SelectiveAdam
+ else:
+ optimizer_class = torch.optim.Adam
+ optimizers = {
+ name: optimizer_class(
+ [{"params": splats[name], "lr": lr * math.sqrt(BS), "name": name}],
+ eps=1e-15 / math.sqrt(BS),
+ # TODO: check betas logic when BS is larger than 10 betas[0] will be zero.
+ betas=(1 - BS * (1 - 0.9), 1 - BS * (1 - 0.999)),
+ )
+ for name, _, lr in params
+ }
+ return splats, optimizers
+
+
+class Runner:
+ """Engine for training and testing."""
+
+ def __init__(
+ self, local_rank: int, world_rank, world_size: int, cfg: Config
+ ) -> None:
+ set_random_seed(42 + local_rank)
+
+ self.cfg = cfg
+ self.world_rank = world_rank
+ self.local_rank = local_rank
+ self.world_size = world_size
+ self.device = f"cuda:{local_rank}"
+
+ # Where to dump results.
+ os.makedirs(cfg.result_dir, exist_ok=True)
+
+ # Setup output directories.
+ self.ckpt_dir = f"{cfg.result_dir}/ckpts"
+ os.makedirs(self.ckpt_dir, exist_ok=True)
+ self.stats_dir = f"{cfg.result_dir}/stats"
+ os.makedirs(self.stats_dir, exist_ok=True)
+ self.render_dir = f"{cfg.result_dir}/renders"
+ os.makedirs(self.render_dir, exist_ok=True)
+ self.ply_dir = f"{cfg.result_dir}/ply"
+ os.makedirs(self.ply_dir, exist_ok=True)
+
+ # Tensorboard
+ self.writer = SummaryWriter(log_dir=f"{cfg.result_dir}/tb")
+
+ # first get the initial 3DGS and camera poses
+ model = AnySplat.from_pretrained("lhjiang/anysplat")
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ model.to(device)
+ model.eval()
+ for param in model.parameters():
+ param.requires_grad = False
+
+ image_folder = cfg.data_dir
+ image_names = sorted([os.path.join(image_folder, f) for f in os.listdir(image_folder) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
+ images = [process_image(img_path) for img_path in image_names]
+ ctx_indices = [idx for idx, name in enumerate(image_names) if idx % cfg.test_every != 0]
+ tgt_indices = [idx for idx, name in enumerate(image_names) if idx % cfg.test_every == 0]
+
+ ctx_images = torch.stack([images[i] for i in ctx_indices], dim=0).unsqueeze(0).to(device)
+ tgt_images = torch.stack([images[i] for i in tgt_indices], dim=0).unsqueeze(0).to(device)
+ ctx_images = (ctx_images+1)*0.5
+ tgt_images = (tgt_images+1)*0.5
+ b, v, _, h, w = tgt_images.shape
+
+ # run inference
+ encoder_output = model.encoder(
+ ctx_images,
+ global_step=0,
+ visualization_dump={},
+ )
+ gaussians, pred_context_pose = encoder_output.gaussians, encoder_output.pred_context_pose
+
+ num_context_view = ctx_images.shape[1]
+ vggt_input_image = torch.cat((ctx_images, tgt_images), dim=1).to(torch.bfloat16)
+ with torch.no_grad(), torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+ aggregated_tokens_list, patch_start_idx = model.encoder.aggregator(vggt_input_image, intermediate_layer_idx=model.encoder.cfg.intermediate_layer_idx)
+ with torch.cuda.amp.autocast(enabled=False):
+ fp32_tokens = [token.float() for token in aggregated_tokens_list]
+ pred_all_pose_enc = model.encoder.camera_head(fp32_tokens)[-1]
+ pred_all_extrinsic, pred_all_intrinsic = pose_encoding_to_extri_intri(pred_all_pose_enc, vggt_input_image.shape[-2:])
+
+ extrinsic_padding = torch.tensor([0, 0, 0, 1], device=pred_all_extrinsic.device, dtype=pred_all_extrinsic.dtype).view(1, 1, 1, 4).repeat(b, vggt_input_image.shape[1], 1, 1)
+ pred_all_extrinsic = torch.cat([pred_all_extrinsic, extrinsic_padding], dim=2).inverse()
+
+ pred_all_intrinsic[:, :, 0] = pred_all_intrinsic[:, :, 0] / w
+ pred_all_intrinsic[:, :, 1] = pred_all_intrinsic[:, :, 1] / h
+ pred_all_context_extrinsic, pred_all_target_extrinsic = pred_all_extrinsic[:, :num_context_view], pred_all_extrinsic[:, num_context_view:]
+ pred_all_context_intrinsic, pred_all_target_intrinsic = pred_all_intrinsic[:, :num_context_view], pred_all_intrinsic[:, num_context_view:]
+
+ scale_factor = pred_context_pose['extrinsic'][:, :, :3, 3].mean() / pred_all_context_extrinsic[:, :, :3, 3].mean()
+ pred_all_target_extrinsic[..., :3, 3] = pred_all_target_extrinsic[..., :3, 3] * scale_factor
+ pred_all_context_extrinsic[..., :3, 3] = pred_all_context_extrinsic[..., :3, 3] * scale_factor
+ print("scale_factor:", scale_factor)
+
+ # Load data: Training data should contain initial points and colors.
+ # self.parser = Parser(
+ # data_dir=cfg.data_dir,
+ # factor=cfg.data_factor,
+ # normalize=cfg.normalize_world_space,
+ # test_every=cfg.test_every,
+ # )
+ self.trainset = Dataset(
+ split="train",
+ images=ctx_images[0].detach().cpu().numpy(),
+ camtoworlds=pred_all_context_extrinsic[0].detach().cpu().numpy(),
+ Ks=pred_all_context_intrinsic[0].detach().cpu().numpy(),
+ patch_size=cfg.patch_size,
+ load_depths=cfg.depth_loss,
+ )
+ self.valset = Dataset(
+ images=tgt_images[0].detach().cpu().numpy(),
+ camtoworlds=pred_all_target_extrinsic[0].detach().cpu().numpy(),
+ Ks=pred_all_target_intrinsic[0].detach().cpu().numpy(),
+ split="val"
+ )
+
+ # Model
+ feature_dim = 32 if cfg.app_opt else None
+ self.splats, self.optimizers = create_splats_with_optimizers(
+ gaussians=gaussians,
+ init_num_pts=cfg.init_num_pts,
+ init_extent=cfg.init_extent,
+ init_opacity=cfg.init_opa,
+ init_scale=cfg.init_scale,
+ sh_degree=cfg.sh_degree,
+ sparse_grad=cfg.sparse_grad,
+ visible_adam=cfg.visible_adam,
+ batch_size=cfg.batch_size,
+ feature_dim=feature_dim,
+ device=self.device,
+ world_rank=world_rank,
+ world_size=world_size,
+ cfg=cfg,
+ )
+ print("Model initialized. Number of GS:", len(self.splats["means"]))
+
+ # Densification Strategy
+ self.cfg.strategy.check_sanity(self.splats, self.optimizers)
+
+ if isinstance(self.cfg.strategy, DefaultStrategy):
+ self.strategy_state = self.cfg.strategy.initialize_state(
+ scene_scale=1.0
+ )
+ elif isinstance(self.cfg.strategy, MCMCStrategy):
+ self.strategy_state = self.cfg.strategy.initialize_state()
+ else:
+ assert_never(self.cfg.strategy)
+
+ # Compression Strategy
+ self.compression_method = None
+ if cfg.compression is not None:
+ if cfg.compression == "png":
+ self.compression_method = PngCompression()
+ else:
+ raise ValueError(f"Unknown compression strategy: {cfg.compression}")
+
+ self.pose_optimizers = []
+ if cfg.pose_opt:
+ self.pose_adjust = CameraOptModule(len(self.trainset)).to(self.device)
+ self.pose_adjust.zero_init()
+ self.pose_optimizers = [
+ torch.optim.Adam(
+ self.pose_adjust.parameters(),
+ lr=cfg.pose_opt_lr * math.sqrt(cfg.batch_size),
+ weight_decay=cfg.pose_opt_reg,
+ )
+ ]
+ if world_size > 1:
+ self.pose_adjust = DDP(self.pose_adjust)
+
+ if cfg.pose_noise > 0.0:
+ self.pose_perturb = CameraOptModule(len(self.trainset)).to(self.device)
+ self.pose_perturb.random_init(cfg.pose_noise)
+ if world_size > 1:
+ self.pose_perturb = DDP(self.pose_perturb)
+
+ self.app_optimizers = []
+ if cfg.app_opt:
+ assert feature_dim is not None
+ self.app_module = AppearanceOptModule(
+ len(self.trainset), feature_dim, cfg.app_embed_dim, cfg.sh_degree
+ ).to(self.device)
+ # initialize the last layer to be zero so that the initial output is zero.
+ torch.nn.init.zeros_(self.app_module.color_head[-1].weight)
+ torch.nn.init.zeros_(self.app_module.color_head[-1].bias)
+ self.app_optimizers = [
+ torch.optim.Adam(
+ self.app_module.embeds.parameters(),
+ lr=cfg.app_opt_lr * math.sqrt(cfg.batch_size) * 10.0,
+ weight_decay=cfg.app_opt_reg,
+ ),
+ torch.optim.Adam(
+ self.app_module.color_head.parameters(),
+ lr=cfg.app_opt_lr * math.sqrt(cfg.batch_size),
+ ),
+ ]
+ if world_size > 1:
+ self.app_module = DDP(self.app_module)
+
+ self.bil_grid_optimizers = []
+ if cfg.use_bilateral_grid:
+ self.bil_grids = BilateralGrid(
+ len(self.trainset),
+ grid_X=cfg.bilateral_grid_shape[0],
+ grid_Y=cfg.bilateral_grid_shape[1],
+ grid_W=cfg.bilateral_grid_shape[2],
+ ).to(self.device)
+ self.bil_grid_optimizers = [
+ torch.optim.Adam(
+ self.bil_grids.parameters(),
+ lr=2e-3 * math.sqrt(cfg.batch_size),
+ eps=1e-15,
+ ),
+ ]
+
+ # Losses & Metrics.
+ self.ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(self.device)
+ self.psnr = PeakSignalNoiseRatio(data_range=1.0).to(self.device)
+
+ if cfg.lpips_net == "alex":
+ self.lpips = LearnedPerceptualImagePatchSimilarity(
+ net_type="alex", normalize=True
+ ).to(self.device)
+ elif cfg.lpips_net == "vgg":
+ # The 3DGS official repo uses lpips vgg, which is equivalent with the following:
+ self.lpips = LearnedPerceptualImagePatchSimilarity(
+ net_type="vgg", normalize=False
+ ).to(self.device)
+ else:
+ raise ValueError(f"Unknown LPIPS network: {cfg.lpips_net}")
+
+ # Viewer
+ if not self.cfg.disable_viewer:
+ self.server = viser.ViserServer(port=cfg.port, verbose=False)
+ self.viewer = GsplatViewer(
+ server=self.server,
+ render_fn=self._viewer_render_fn,
+ output_dir=Path(cfg.result_dir),
+ mode="training",
+ )
+
+ def rasterize_splats(
+ self,
+ camtoworlds: Tensor,
+ Ks: Tensor,
+ width: int,
+ height: int,
+ masks: Optional[Tensor] = None,
+ rasterize_mode: Optional[Literal["classic", "antialiased"]] = None,
+ camera_model: Optional[Literal["pinhole", "ortho", "fisheye"]] = None,
+ **kwargs,
+ ) -> Tuple[Tensor, Tensor, Dict]:
+ means = self.splats["means"] # [N, 3]
+ # quats = F.normalize(self.splats["quats"], dim=-1) # [N, 4]
+ # rasterization does normalization internally
+ quats = self.splats["quats"] # [N, 4]
+ scales = torch.exp(self.splats["scales"]) # [N, 3]
+ opacities = torch.sigmoid(self.splats["opacities"]) # [N,]
+
+ image_ids = kwargs.pop("image_ids", None)
+ if self.cfg.app_opt:
+ colors = self.app_module(
+ features=self.splats["features"],
+ embed_ids=image_ids,
+ dirs=means[None, :, :] - camtoworlds[:, None, :3, 3],
+ sh_degree=kwargs.pop("sh_degree", self.cfg.sh_degree),
+ )
+ colors = colors + self.splats["colors"]
+ colors = torch.sigmoid(colors)
+ else:
+ colors = torch.cat([self.splats["sh0"], self.splats["shN"]], 1) # [N, K, 3]
+
+ if rasterize_mode is None:
+ rasterize_mode = "antialiased" if self.cfg.antialiased else "classic"
+ if camera_model is None:
+ camera_model = self.cfg.camera_model
+
+ # covariance = build_covariance(scales[None], quats[None]).squeeze(0)
+ render_colors, render_alphas, info = rasterization(
+ means=means,
+ quats=quats,
+ scales=scales,
+ opacities=opacities,
+ colors=colors,
+ # covars=covariance,
+ viewmats=torch.linalg.inv(camtoworlds), # [C, 4, 4]
+ Ks=Ks, # [C, 3, 3]
+ width=width,
+ height=height,
+ packed=self.cfg.packed,
+ absgrad=(
+ self.cfg.strategy.absgrad
+ if isinstance(self.cfg.strategy, DefaultStrategy)
+ else False
+ ),
+ sparse_grad=self.cfg.sparse_grad,
+ rasterize_mode=rasterize_mode,
+ distributed=self.world_size > 1,
+ camera_model=self.cfg.camera_model,
+ radius_clip=0.1,
+ backgrounds=torch.tensor([0.0, 0.0, 0.0]).cuda().unsqueeze(0).repeat(1, 1),
+ **kwargs,
+ )
+ if masks is not None:
+ render_colors[~masks] = 0
+ return render_colors, render_alphas, info
+
+ def train(self):
+ cfg = self.cfg
+ device = self.device
+ world_rank = self.world_rank
+ world_size = self.world_size
+
+ # Dump cfg.
+ if world_rank == 0:
+ with open(f"{cfg.result_dir}/cfg.yml", "w") as f:
+ yaml.dump(vars(cfg), f)
+
+ max_steps = cfg.max_steps
+ init_step = 0
+
+ schedulers = [
+ # means has a learning rate schedule, that end at 0.01 of the initial value
+ torch.optim.lr_scheduler.ExponentialLR(
+ self.optimizers["means"], gamma=0.01 ** (1.0 / max_steps)
+ ),
+ ]
+ if cfg.pose_opt:
+ # pose optimization has a learning rate schedule
+ schedulers.append(
+ torch.optim.lr_scheduler.ExponentialLR(
+ self.pose_optimizers[0], gamma=0.01 ** (1.0 / max_steps)
+ )
+ )
+ if cfg.use_bilateral_grid:
+ # bilateral grid has a learning rate schedule. Linear warmup for 1000 steps.
+ schedulers.append(
+ torch.optim.lr_scheduler.ChainedScheduler(
+ [
+ torch.optim.lr_scheduler.LinearLR(
+ self.bil_grid_optimizers[0],
+ start_factor=0.01,
+ total_iters=1000,
+ ),
+ torch.optim.lr_scheduler.ExponentialLR(
+ self.bil_grid_optimizers[0], gamma=0.01 ** (1.0 / max_steps)
+ ),
+ ]
+ )
+ )
+
+ trainloader = torch.utils.data.DataLoader(
+ self.trainset,
+ batch_size=cfg.batch_size,
+ shuffle=True,
+ num_workers=4,
+ persistent_workers=True,
+ pin_memory=True,
+ )
+ trainloader_iter = iter(trainloader)
+
+ # Training loop.
+ global_tic = time.time()
+ pbar = tqdm.tqdm(range(init_step, max_steps))
+ for step in pbar:
+ if not cfg.disable_viewer:
+ while self.viewer.state == "paused":
+ time.sleep(0.01)
+ self.viewer.lock.acquire()
+ tic = time.time()
+
+ try:
+ data = next(trainloader_iter)
+ except StopIteration:
+ trainloader_iter = iter(trainloader)
+ data = next(trainloader_iter)
+
+ camtoworlds = camtoworlds_gt = data["camtoworld"].to(device) # [1, 4, 4]
+ Ks = data["K"].to(device) # [1, 3, 3]
+ pixels = data["image"].to(device) / 255.0 # [1, H, W, 3]
+ num_train_rays_per_step = (
+ pixels.shape[0] * pixels.shape[1] * pixels.shape[2]
+ )
+ image_ids = data["image_id"].to(device)
+ masks = data["mask"].to(device) if "mask" in data else None # [1, H, W]
+ if cfg.depth_loss:
+ points = data["points"].to(device) # [1, M, 2]
+ depths_gt = data["depths"].to(device) # [1, M]
+
+ height, width = pixels.shape[1:3]
+
+ if cfg.pose_noise:
+ camtoworlds = self.pose_perturb(camtoworlds, image_ids)
+
+ if cfg.pose_opt:
+ camtoworlds = self.pose_adjust(camtoworlds, image_ids)
+
+ # sh schedule
+ # sh_degree_to_use = min(step // cfg.sh_degree_interval, cfg.sh_degree)
+ sh_degree_to_use = cfg.sh_degree
+
+ # forward
+ renders, alphas, info = self.rasterize_splats(
+ camtoworlds=camtoworlds,
+ Ks=Ks,
+ width=width,
+ height=height,
+ sh_degree=sh_degree_to_use,
+ near_plane=cfg.near_plane,
+ far_plane=cfg.far_plane,
+ image_ids=image_ids,
+ render_mode="RGB+ED" if cfg.depth_loss else "RGB",
+ masks=masks,
+ )
+ if renders.shape[-1] == 4:
+ colors, depths = renders[..., 0:3], renders[..., 3:4]
+ else:
+ colors, depths = renders, None
+
+ if cfg.use_bilateral_grid:
+ grid_y, grid_x = torch.meshgrid(
+ (torch.arange(height, device=self.device) + 0.5) / height,
+ (torch.arange(width, device=self.device) + 0.5) / width,
+ indexing="ij",
+ )
+ grid_xy = torch.stack([grid_x, grid_y], dim=-1).unsqueeze(0)
+ colors = slice(self.bil_grids, grid_xy, colors, image_ids)["rgb"]
+
+ if cfg.random_bkgd:
+ bkgd = torch.rand(1, 3, device=device)
+ colors = colors + bkgd * (1.0 - alphas)
+
+ self.cfg.strategy.step_pre_backward(
+ params=self.splats,
+ optimizers=self.optimizers,
+ state=self.strategy_state,
+ step=step,
+ info=info,
+ )
+
+ # loss
+ l1loss = F.l1_loss(colors, pixels)
+ ssimloss = 1.0 - fused_ssim(
+ colors.permute(0, 3, 1, 2), pixels.permute(0, 3, 1, 2), padding="valid"
+ )
+ loss = l1loss * (1.0 - cfg.ssim_lambda) + ssimloss * cfg.ssim_lambda
+ if cfg.depth_loss:
+ # query depths from depth map
+ points = torch.stack(
+ [
+ points[:, :, 0] / (width - 1) * 2 - 1,
+ points[:, :, 1] / (height - 1) * 2 - 1,
+ ],
+ dim=-1,
+ ) # normalize to [-1, 1]
+ grid = points.unsqueeze(2) # [1, M, 1, 2]
+ depths = F.grid_sample(
+ depths.permute(0, 3, 1, 2), grid, align_corners=True
+ ) # [1, 1, M, 1]
+ depths = depths.squeeze(3).squeeze(1) # [1, M]
+ # calculate loss in disparity space
+ disp = torch.where(depths > 0.0, 1.0 / depths, torch.zeros_like(depths))
+ disp_gt = 1.0 / depths_gt # [1, M]
+ depthloss = F.l1_loss(disp, disp_gt) * self.scene_scale
+ loss += depthloss * cfg.depth_lambda
+ if cfg.use_bilateral_grid:
+ tvloss = 10 * total_variation_loss(self.bil_grids.grids)
+ loss += tvloss
+
+ # regularizations
+ if cfg.opacity_reg > 0.0:
+ loss = (
+ loss
+ + cfg.opacity_reg
+ * torch.abs(torch.sigmoid(self.splats["opacities"])).mean()
+ )
+ if cfg.scale_reg > 0.0:
+ loss = (
+ loss
+ + cfg.scale_reg * torch.abs(torch.exp(self.splats["scales"])).mean()
+ )
+
+ loss.backward()
+
+ desc = f"loss={loss.item():.3f}| " f"sh degree={sh_degree_to_use}| "
+ if cfg.depth_loss:
+ desc += f"depth loss={depthloss.item():.6f}| "
+ if cfg.pose_opt and cfg.pose_noise:
+ # monitor the pose error if we inject noise
+ pose_err = F.l1_loss(camtoworlds_gt, camtoworlds)
+ desc += f"pose err={pose_err.item():.6f}| "
+ pbar.set_description(desc)
+
+ # write images (gt and render)
+ # if world_rank == 0 and step % 800 == 0:
+ # canvas = torch.cat([pixels, colors], dim=2).detach().cpu().numpy()
+ # canvas = canvas.reshape(-1, *canvas.shape[2:])
+ # imageio.imwrite(
+ # f"{self.render_dir}/train_rank{self.world_rank}.png",
+ # (canvas * 255).astype(np.uint8),
+ # )
+
+ if world_rank == 0 and cfg.tb_every > 0 and step % cfg.tb_every == 0:
+ mem = torch.cuda.max_memory_allocated() / 1024**3
+ self.writer.add_scalar("train/loss", loss.item(), step)
+ self.writer.add_scalar("train/l1loss", l1loss.item(), step)
+ self.writer.add_scalar("train/ssimloss", ssimloss.item(), step)
+ self.writer.add_scalar("train/num_GS", len(self.splats["means"]), step)
+ self.writer.add_scalar("train/mem", mem, step)
+ if cfg.depth_loss:
+ self.writer.add_scalar("train/depthloss", depthloss.item(), step)
+ if cfg.use_bilateral_grid:
+ self.writer.add_scalar("train/tvloss", tvloss.item(), step)
+ if cfg.tb_save_image:
+ canvas = torch.cat([pixels, colors], dim=2).detach().cpu().numpy()
+ canvas = canvas.reshape(-1, *canvas.shape[2:])
+ self.writer.add_image("train/render", canvas, step)
+ self.writer.flush()
+
+ # save checkpoint before updating the model
+ if step in [i - 1 for i in cfg.save_steps] or step == max_steps - 1:
+ mem = torch.cuda.max_memory_allocated() / 1024**3
+ stats = {
+ "mem": mem,
+ "ellipse_time": time.time() - global_tic,
+ "num_GS": len(self.splats["means"]),
+ }
+ print("Step: ", step, stats)
+ with open(
+ f"{self.stats_dir}/train_step{step:04d}_rank{self.world_rank}.json",
+ "w",
+ ) as f:
+ json.dump(stats, f)
+ data = {"step": step, "splats": self.splats.state_dict()}
+ if cfg.pose_opt:
+ if world_size > 1:
+ data["pose_adjust"] = self.pose_adjust.module.state_dict()
+ else:
+ data["pose_adjust"] = self.pose_adjust.state_dict()
+ if cfg.app_opt:
+ if world_size > 1:
+ data["app_module"] = self.app_module.module.state_dict()
+ else:
+ data["app_module"] = self.app_module.state_dict()
+ torch.save(
+ data, f"{self.ckpt_dir}/ckpt_{step}_rank{self.world_rank}.pt"
+ )
+ if (
+ step in [i - 1 for i in cfg.ply_steps] or step == max_steps - 1
+ ) and cfg.save_ply:
+
+ if self.cfg.app_opt:
+ # eval at origin to bake the appeareance into the colors
+ rgb = self.app_module(
+ features=self.splats["features"],
+ embed_ids=None,
+ dirs=torch.zeros_like(self.splats["means"][None, :, :]),
+ sh_degree=sh_degree_to_use,
+ )
+ rgb = rgb + self.splats["colors"]
+ rgb = torch.sigmoid(rgb).squeeze(0).unsqueeze(1)
+ sh0 = rgb_to_sh(rgb)
+ shN = torch.empty([sh0.shape[0], 0, 3], device=sh0.device)
+ else:
+ sh0 = self.splats["sh0"]
+ shN = self.splats["shN"]
+ # shN = torch.empty([sh0.shape[0], 0, 3], device=sh0.device)
+
+ means = self.splats["means"]
+ scales = self.splats["scales"]
+ quats = self.splats["quats"]
+ opacities = self.splats["opacities"]
+
+ # export_splats(
+ # means=means,
+ # scales=scales,
+ # quats=quats,
+ # opacities=opacities,
+ # sh0=sh0,
+ # shN=shN,
+ # format="ply",
+ # save_to=f"{self.ply_dir}/point_cloud_{step}.ply",
+ # )
+ export_ply(
+ means=means,
+ scales=scales,
+ rotations=quats,
+ harmonics=torch.cat([sh0, shN], dim=1).permute(0, 2, 1),
+ opacities=opacities.sigmoid(),
+ path=Path(f"{self.ply_dir}/point_cloud_{step}.ply"),
+ )
+
+ # Turn Gradients into Sparse Tensor before running optimizer
+ if cfg.sparse_grad:
+ assert cfg.packed, "Sparse gradients only work with packed mode."
+ gaussian_ids = info["gaussian_ids"]
+ for k in self.splats.keys():
+ grad = self.splats[k].grad
+ if grad is None or grad.is_sparse:
+ continue
+ self.splats[k].grad = torch.sparse_coo_tensor(
+ indices=gaussian_ids[None], # [1, nnz]
+ values=grad[gaussian_ids], # [nnz, ...]
+ size=self.splats[k].size(), # [N, ...]
+ is_coalesced=len(Ks) == 1,
+ )
+
+ if cfg.visible_adam:
+ gaussian_cnt = self.splats.means.shape[0]
+ if cfg.packed:
+ visibility_mask = torch.zeros_like(
+ self.splats["opacities"], dtype=bool
+ )
+ visibility_mask.scatter_(0, info["gaussian_ids"], 1)
+ else:
+ visibility_mask = (info["radii"] > 0).all(-1).any(0)
+
+ # optimize
+ for optimizer in self.optimizers.values():
+ if cfg.visible_adam:
+ optimizer.step(visibility_mask)
+ else:
+ optimizer.step()
+ optimizer.zero_grad(set_to_none=True)
+ for optimizer in self.pose_optimizers:
+ optimizer.step()
+ optimizer.zero_grad(set_to_none=True)
+ for optimizer in self.app_optimizers:
+ optimizer.step()
+ optimizer.zero_grad(set_to_none=True)
+ for optimizer in self.bil_grid_optimizers:
+ optimizer.step()
+ optimizer.zero_grad(set_to_none=True)
+ for scheduler in schedulers:
+ scheduler.step()
+
+ # Run post-backward steps after backward and optimizer
+ if isinstance(self.cfg.strategy, DefaultStrategy):
+ self.cfg.strategy.step_post_backward(
+ params=self.splats,
+ optimizers=self.optimizers,
+ state=self.strategy_state,
+ step=step,
+ info=info,
+ packed=cfg.packed,
+ )
+ elif isinstance(self.cfg.strategy, MCMCStrategy):
+ self.cfg.strategy.step_post_backward(
+ params=self.splats,
+ optimizers=self.optimizers,
+ state=self.strategy_state,
+ step=step,
+ info=info,
+ lr=schedulers[0].get_last_lr()[0],
+ )
+ else:
+ assert_never(self.cfg.strategy)
+
+ # eval the full set
+ if step in [i - 1 for i in cfg.eval_steps]:
+ self.eval(step)
+ # self.render_traj(step)
+
+ # run compression
+ if cfg.compression is not None and step in [i - 1 for i in cfg.eval_steps]:
+ self.run_compression(step=step)
+
+ if not cfg.disable_viewer:
+ self.viewer.lock.release()
+ num_train_steps_per_sec = 1.0 / (time.time() - tic)
+ num_train_rays_per_sec = (
+ num_train_rays_per_step * num_train_steps_per_sec
+ )
+ # Update the viewer state.
+ self.viewer.render_tab_state.num_train_rays_per_sec = (
+ num_train_rays_per_sec
+ )
+ # Update the scene.
+ self.viewer.update(step, num_train_rays_per_step)
+
+ @torch.no_grad()
+ def eval(self, step: int, stage: str = "val"):
+ """Entry for evaluation."""
+ print("Running evaluation...")
+ cfg = self.cfg
+ device = self.device
+ world_rank = self.world_rank
+ world_size = self.world_size
+
+ valloader = torch.utils.data.DataLoader(
+ self.valset, batch_size=1, shuffle=False, num_workers=1
+ )
+ ellipse_time = 0
+ metrics = defaultdict(list)
+ for i, data in enumerate(valloader):
+ camtoworlds = data["camtoworld"].to(device)
+ Ks = data["K"].to(device)
+ pixels = data["image"].to(device) / 255.0
+ masks = data["mask"].to(device) if "mask" in data else None
+ height, width = pixels.shape[1:3]
+
+ torch.cuda.synchronize()
+ tic = time.time()
+ render_colors, _, _ = self.rasterize_splats(
+ camtoworlds=camtoworlds,
+ Ks=Ks,
+ width=width,
+ height=height,
+ sh_degree=cfg.sh_degree,
+ near_plane=cfg.near_plane,
+ far_plane=cfg.far_plane,
+ # radius_clip=0.1,
+ render_mode="RGB+ED",
+ masks=masks,
+ ) # [1, H, W, 3]
+ torch.cuda.synchronize()
+ ellipse_time += time.time() - tic
+
+ colors = render_colors[..., :3]
+ depths = render_colors[..., 3]
+
+ colors = torch.clamp(colors, 0.0, 1.0)
+ canvas_list = [pixels, colors]
+
+ if world_rank == 0:
+ # write images
+ canvas = torch.cat(canvas_list, dim=2).squeeze(0).cpu().numpy()
+ canvas = (canvas * 255).astype(np.uint8)
+ imageio.imwrite(
+ f"{self.render_dir}/{stage}_step{step}_{i:04d}.png",
+ canvas,
+ )
+ torchvision.utils.save_image(pixels.permute(0, 3, 1, 2), f"{self.render_dir}/gt_rgb_{stage}_step{step}_{i:04d}.png")
+ torchvision.utils.save_image(colors.permute(0, 3, 1, 2), f"{self.render_dir}/render_rgb_{stage}_step{step}_{i:04d}.png")
+ # save depth & normal map
+
+
+ pixels_p = pixels.permute(0, 3, 1, 2) # [1, 3, H, W]
+ colors_p = colors.permute(0, 3, 1, 2) # [1, 3, H, W]
+
+ metrics["psnr"].append(self.psnr(colors_p, pixels_p))
+ metrics["ssim"].append(self.ssim(colors_p, pixels_p))
+ metrics["lpips"].append(self.lpips(colors_p, pixels_p))
+ if cfg.use_bilateral_grid:
+ cc_colors = color_correct(colors, pixels)
+ cc_colors_p = cc_colors.permute(0, 3, 1, 2) # [1, 3, H, W]
+ metrics["cc_psnr"].append(self.psnr(cc_colors_p, pixels_p))
+
+ if world_rank == 0:
+ ellipse_time /= len(valloader)
+
+ stats = {k: torch.stack(v).mean().item() for k, v in metrics.items()}
+ stats.update(
+ {
+ "ellipse_time": ellipse_time,
+ "num_GS": len(self.splats["means"]),
+ }
+ )
+ print(
+ f"PSNR: {stats['psnr']:.3f}, SSIM: {stats['ssim']:.4f}, LPIPS: {stats['lpips']:.3f} "
+ f"Time: {stats['ellipse_time']:.3f}s/image "
+ f"Number of GS: {stats['num_GS']}"
+ )
+ # save stats as json
+ with open(f"{self.stats_dir}/{stage}_step{step:04d}.json", "w") as f:
+ json.dump(stats, f)
+ # save stats to tensorboard
+ for k, v in stats.items():
+ self.writer.add_scalar(f"{stage}/{k}", v, step)
+ self.writer.flush()
+
+ @torch.no_grad()
+ def render_traj(self, step: int):
+ """Entry for trajectory rendering."""
+ if self.cfg.disable_video:
+ return
+ print("Running trajectory rendering...")
+ cfg = self.cfg
+ device = self.device
+
+ camtoworlds_all = self.parser.camtoworlds[5:-5]
+ if cfg.render_traj_path == "interp":
+ camtoworlds_all = generate_interpolated_path(
+ camtoworlds_all, 1
+ ) # [N, 3, 4]
+ elif cfg.render_traj_path == "ellipse":
+ height = camtoworlds_all[:, 2, 3].mean()
+ camtoworlds_all = generate_ellipse_path_z(
+ camtoworlds_all, height=height
+ ) # [N, 3, 4]
+ elif cfg.render_traj_path == "spiral":
+ camtoworlds_all = generate_spiral_path(
+ camtoworlds_all,
+ bounds=self.parser.bounds * self.scene_scale,
+ spiral_scale_r=self.parser.extconf["spiral_radius_scale"],
+ )
+ else:
+ raise ValueError(
+ f"Render trajectory type not supported: {cfg.render_traj_path}"
+ )
+
+ camtoworlds_all = np.concatenate(
+ [
+ camtoworlds_all,
+ np.repeat(
+ np.array([[[0.0, 0.0, 0.0, 1.0]]]), len(camtoworlds_all), axis=0
+ ),
+ ],
+ axis=1,
+ ) # [N, 4, 4]
+
+ camtoworlds_all = torch.from_numpy(camtoworlds_all).float().to(device)
+ K = torch.from_numpy(list(self.parser.Ks_dict.values())[0]).float().to(device)
+ width, height = list(self.parser.imsize_dict.values())[0]
+
+ # save to video
+ video_dir = f"{cfg.result_dir}/videos"
+ os.makedirs(video_dir, exist_ok=True)
+ writer = imageio.get_writer(f"{video_dir}/traj_{step}.mp4", fps=30)
+ for i in tqdm.trange(len(camtoworlds_all), desc="Rendering trajectory"):
+ camtoworlds = camtoworlds_all[i : i + 1]
+ Ks = K[None]
+
+ renders, _, _ = self.rasterize_splats(
+ camtoworlds=camtoworlds,
+ Ks=Ks,
+ width=width,
+ height=height,
+ sh_degree=cfg.sh_degree,
+ near_plane=cfg.near_plane,
+ far_plane=cfg.far_plane,
+ render_mode="RGB+ED",
+ ) # [1, H, W, 4]
+ colors = torch.clamp(renders[..., 0:3], 0.0, 1.0) # [1, H, W, 3]
+ depths = renders[..., 3:4] # [1, H, W, 1]
+ depths = (depths - depths.min()) / (depths.max() - depths.min())
+ canvas_list = [colors, depths.repeat(1, 1, 1, 3)]
+
+ # write images
+ canvas = torch.cat(canvas_list, dim=2).squeeze(0).cpu().numpy()
+ canvas = (canvas * 255).astype(np.uint8)
+ writer.append_data(canvas)
+ writer.close()
+ print(f"Video saved to {video_dir}/traj_{step}.mp4")
+
+ @torch.no_grad()
+ def run_compression(self, step: int):
+ """Entry for running compression."""
+ print("Running compression...")
+ world_rank = self.world_rank
+
+ compress_dir = f"{cfg.result_dir}/compression/rank{world_rank}"
+ os.makedirs(compress_dir, exist_ok=True)
+
+ self.compression_method.compress(compress_dir, self.splats)
+
+ # evaluate compression
+ splats_c = self.compression_method.decompress(compress_dir)
+ for k in splats_c.keys():
+ self.splats[k].data = splats_c[k].to(self.device)
+ self.eval(step=step, stage="compress")
+
+ @torch.no_grad()
+ def _viewer_render_fn(
+ self, camera_state: CameraState, render_tab_state: RenderTabState
+ ):
+ assert isinstance(render_tab_state, GsplatRenderTabState)
+ if render_tab_state.preview_render:
+ width = render_tab_state.render_width
+ height = render_tab_state.render_height
+ else:
+ width = render_tab_state.viewer_width
+ height = render_tab_state.viewer_height
+ c2w = camera_state.c2w
+ K = camera_state.get_K((width, height))
+ c2w = torch.from_numpy(c2w).float().to(self.device)
+ K = torch.from_numpy(K).float().to(self.device)
+
+ RENDER_MODE_MAP = {
+ "rgb": "RGB",
+ "depth(accumulated)": "D",
+ "depth(expected)": "ED",
+ "alpha": "RGB",
+ }
+
+ render_colors, render_alphas, info = self.rasterize_splats(
+ camtoworlds=c2w[None],
+ Ks=K[None],
+ width=width,
+ height=height,
+ sh_degree=min(render_tab_state.max_sh_degree, self.cfg.sh_degree),
+ near_plane=render_tab_state.near_plane,
+ far_plane=render_tab_state.far_plane,
+ radius_clip=render_tab_state.radius_clip,
+ # radius_clip=0.1,
+ eps2d=render_tab_state.eps2d,
+ backgrounds=torch.tensor([render_tab_state.backgrounds], device=self.device)
+ / 255.0,
+ render_mode=RENDER_MODE_MAP[render_tab_state.render_mode],
+ rasterize_mode=render_tab_state.rasterize_mode,
+ camera_model=render_tab_state.camera_model,
+ ) # [1, H, W, 3]
+ render_tab_state.total_gs_count = len(self.splats["means"])
+ render_tab_state.rendered_gs_count = (info["radii"] > 0).all(-1).sum().item()
+
+ if render_tab_state.render_mode == "rgb":
+ # colors represented with sh are not guranteed to be in [0, 1]
+ render_colors = render_colors[0, ..., 0:3].clamp(0, 1)
+ renders = render_colors.cpu().numpy()
+ elif render_tab_state.render_mode in ["depth(accumulated)", "depth(expected)"]:
+ # normalize depth to [0, 1]
+ depth = render_colors[0, ..., 0:1]
+ if render_tab_state.normalize_nearfar:
+ near_plane = render_tab_state.near_plane
+ far_plane = render_tab_state.far_plane
+ else:
+ near_plane = depth.min()
+ far_plane = depth.max()
+ depth_norm = (depth - near_plane) / (far_plane - near_plane + 1e-10)
+ depth_norm = torch.clip(depth_norm, 0, 1)
+ if render_tab_state.inverse:
+ depth_norm = 1 - depth_norm
+ renders = (
+ apply_float_colormap(depth_norm, render_tab_state.colormap)
+ .cpu()
+ .numpy()
+ )
+ elif render_tab_state.render_mode == "alpha":
+ alpha = render_alphas[0, ..., 0:1]
+ if render_tab_state.inverse:
+ alpha = 1 - alpha
+ renders = (
+ apply_float_colormap(alpha, render_tab_state.colormap).cpu().numpy()
+ )
+ return renders
+
+
+def main(local_rank: int, world_rank, world_size: int, cfg: Config):
+ if world_size > 1 and not cfg.disable_viewer:
+ cfg.disable_viewer = True
+ if world_rank == 0:
+ print("Viewer is disabled in distributed training.")
+
+ runner = Runner(local_rank, world_rank, world_size, cfg)
+
+ if cfg.ckpt is not None:
+ # run eval only
+ ckpts = [
+ torch.load(file, map_location=runner.device, weights_only=True)
+ for file in cfg.ckpt
+ ]
+ for k in runner.splats.keys():
+ runner.splats[k].data = torch.cat([ckpt["splats"][k] for ckpt in ckpts])
+ step = ckpts[0]["step"]
+ runner.eval(step=step)
+ # runner.render_traj(step=step)
+ if cfg.compression is not None:
+ runner.run_compression(step=step)
+ else:
+ runner.train()
+ runner.eval(step=runner.cfg.max_steps)
+ # runner.render_traj(step=runner.cfg.max_steps)
+ print("Training complete.")
+ # runner.viewer.complete()
+ # if not cfg.disable_viewer:
+ # print("Viewer running... Ctrl+C to exit.")
+ # time.sleep(1000000)
+
+
+if __name__ == "__main__":
+ """
+ Usage:
+
+ ```bash
+ # Single GPU training
+ CUDA_VISIBLE_DEVICES=9 python -m examples.simple_trainer default
+
+ # Distributed training on 4 GPUs: Effectively 4x batch size so run 4x less steps.
+ CUDA_VISIBLE_DEVICES=0,1,2,3 python simple_trainer.py default --steps_scaler 0.25
+
+ """
+
+ # Config objects we can choose between.
+ # Each is a tuple of (CLI description, config object).
+ configs = {
+ "default": (
+ "Gaussian splatting training using densification heuristics from the original paper.",
+ Config(
+ strategy=DefaultStrategy(verbose=True),
+ ),
+ ),
+ "mcmc": (
+ "Gaussian splatting training using densification from the paper '3D Gaussian Splatting as Markov Chain Monte Carlo'.",
+ Config(
+ init_opa=0.5,
+ init_scale=0.1,
+ opacity_reg=0.01,
+ scale_reg=0.01,
+ strategy=MCMCStrategy(verbose=True),
+ ),
+ ),
+ }
+ cfg = tyro.extras.overridable_config_cli(configs)
+ cfg.adjust_steps(cfg.steps_scaler)
+
+ # try import extra dependencies
+ if cfg.compression == "png":
+ try:
+ import plas
+ import torchpq
+ except:
+ raise ImportError(
+ "To use PNG compression, you need to install "
+ "torchpq (instruction at https://github.com/DeMoriarty/TorchPQ?tab=readme-ov-file#install) "
+ "and plas (via 'pip install git+https://github.com/fraunhoferhhi/PLAS.git') "
+ )
+
+ cli(main, cfg, verbose=True)
diff --git a/src/post_opt/simple_viewer.py b/src/post_opt/simple_viewer.py
new file mode 100644
index 0000000000000000000000000000000000000000..fd1191e82d4660194ba6d107b75a9e0c18dcfc06
--- /dev/null
+++ b/src/post_opt/simple_viewer.py
@@ -0,0 +1,275 @@
+import argparse
+import math
+import os
+import time
+
+import imageio
+import numpy as np
+import torch
+import torch.nn.functional as F
+import tqdm
+import viser
+from pathlib import Path
+from gsplat._helper import load_test_data
+from gsplat.distributed import cli
+from gsplat.rendering import rasterization
+
+from nerfview import CameraState, RenderTabState, apply_float_colormap
+from examples.gsplat_viewer import GsplatViewer, GsplatRenderTabState
+
+
+def main(local_rank: int, world_rank, world_size: int, args):
+ torch.manual_seed(42)
+ device = torch.device("cuda", local_rank)
+
+ if args.ckpt is None:
+ (
+ means,
+ quats,
+ scales,
+ opacities,
+ colors,
+ viewmats,
+ Ks,
+ width,
+ height,
+ ) = load_test_data(device=device, scene_grid=args.scene_grid)
+
+ assert world_size <= 2
+ means = means[world_rank::world_size].contiguous()
+ means.requires_grad = True
+ quats = quats[world_rank::world_size].contiguous()
+ quats.requires_grad = True
+ scales = scales[world_rank::world_size].contiguous()
+ scales.requires_grad = True
+ opacities = opacities[world_rank::world_size].contiguous()
+ opacities.requires_grad = True
+ colors = colors[world_rank::world_size].contiguous()
+ colors.requires_grad = True
+
+ viewmats = viewmats[world_rank::world_size][:1].contiguous()
+ Ks = Ks[world_rank::world_size][:1].contiguous()
+
+ sh_degree = None
+ C = len(viewmats)
+ N = len(means)
+ print("rank", world_rank, "Number of Gaussians:", N, "Number of Cameras:", C)
+
+ # batched render
+ for _ in tqdm.trange(1):
+ render_colors, render_alphas, meta = rasterization(
+ means, # [N, 3]
+ quats, # [N, 4]
+ scales, # [N, 3]
+ opacities, # [N]
+ colors, # [N, S, 3]
+ viewmats, # [C, 4, 4]
+ Ks, # [C, 3, 3]
+ width,
+ height,
+ render_mode="RGB+D",
+ packed=False,
+ distributed=world_size > 1,
+ )
+ C = render_colors.shape[0]
+ assert render_colors.shape == (C, height, width, 4)
+ assert render_alphas.shape == (C, height, width, 1)
+ render_colors.sum().backward()
+
+ render_rgbs = render_colors[..., 0:3]
+ render_depths = render_colors[..., 3:4]
+ render_depths = render_depths / render_depths.max()
+
+ # dump batch images
+ os.makedirs(args.output_dir, exist_ok=True)
+ canvas = (
+ torch.cat(
+ [
+ render_rgbs.reshape(C * height, width, 3),
+ render_depths.reshape(C * height, width, 1).expand(-1, -1, 3),
+ render_alphas.reshape(C * height, width, 1).expand(-1, -1, 3),
+ ],
+ dim=1,
+ )
+ .detach()
+ .cpu()
+ .numpy()
+ )
+ imageio.imsave(
+ f"{args.output_dir}/render_rank{world_rank}.png",
+ (canvas * 255).astype(np.uint8),
+ )
+ else:
+ means, quats, scales, opacities, sh0, shN = [], [], [], [], [], []
+ for ckpt_path in args.ckpt:
+ ckpt = torch.load(ckpt_path, map_location=device)["splats"]
+ means.append(ckpt["means"])
+ quats.append(F.normalize(ckpt["quats"], p=2, dim=-1))
+ scales.append(torch.exp(ckpt["scales"]))
+ opacities.append(torch.sigmoid(ckpt["opacities"]))
+ sh0.append(ckpt["sh0"])
+ shN.append(ckpt["shN"])
+ means = torch.cat(means, dim=0)
+ quats = torch.cat(quats, dim=0)
+ scales = torch.cat(scales, dim=0)
+ opacities = torch.cat(opacities, dim=0)
+ sh0 = torch.cat(sh0, dim=0)
+ shN = torch.cat(shN, dim=0)
+ colors = torch.cat([sh0, shN], dim=-2)
+ sh_degree = int(math.sqrt(colors.shape[-2]) - 1)
+
+ # # crop
+ # aabb = torch.tensor((-1.0, -1.0, -1.0, 1.0, 1.0, 0.7), device=device)
+ # edges = aabb[3:] - aabb[:3]
+ # sel = ((means >= aabb[:3]) & (means <= aabb[3:])).all(dim=-1)
+ # sel = torch.where(sel)[0]
+ # means, quats, scales, colors, opacities = (
+ # means[sel],
+ # quats[sel],
+ # scales[sel],
+ # colors[sel],
+ # opacities[sel],
+ # )
+
+ # # repeat the scene into a grid (to mimic a large-scale setting)
+ # repeats = args.scene_grid
+ # gridx, gridy = torch.meshgrid(
+ # [
+ # torch.arange(-(repeats // 2), repeats // 2 + 1, device=device),
+ # torch.arange(-(repeats // 2), repeats // 2 + 1, device=device),
+ # ],
+ # indexing="ij",
+ # )
+ # grid = torch.stack([gridx, gridy, torch.zeros_like(gridx)], dim=-1).reshape(
+ # -1, 3
+ # )
+ # means = means[None, :, :] + grid[:, None, :] * edges[None, None, :]
+ # means = means.reshape(-1, 3)
+ # quats = quats.repeat(repeats**2, 1)
+ # scales = scales.repeat(repeats**2, 1)
+ # colors = colors.repeat(repeats**2, 1, 1)
+ # opacities = opacities.repeat(repeats**2)
+ print("Number of Gaussians:", len(means))
+
+ # register and open viewer
+ @torch.no_grad()
+ def viewer_render_fn(camera_state: CameraState, render_tab_state: RenderTabState):
+ assert isinstance(render_tab_state, GsplatRenderTabState)
+ if render_tab_state.preview_render:
+ width = render_tab_state.render_width
+ height = render_tab_state.render_height
+ else:
+ width = render_tab_state.viewer_width
+ height = render_tab_state.viewer_height
+ c2w = camera_state.c2w
+ K = camera_state.get_K((width, height))
+ c2w = torch.from_numpy(c2w).float().to(device)
+ K = torch.from_numpy(K).float().to(device)
+ viewmat = c2w.inverse()
+
+ RENDER_MODE_MAP = {
+ "rgb": "RGB",
+ "depth(accumulated)": "D",
+ "depth(expected)": "ED",
+ "alpha": "RGB",
+ }
+
+ render_colors, render_alphas, info = rasterization(
+ means, # [N, 3]
+ quats, # [N, 4]
+ scales, # [N, 3]
+ opacities, # [N]
+ colors, # [N, S, 3]
+ viewmat[None], # [1, 4, 4]
+ K[None], # [1, 3, 3]
+ width,
+ height,
+ sh_degree=(
+ min(render_tab_state.max_sh_degree, sh_degree)
+ if sh_degree is not None
+ else None
+ ),
+ near_plane=render_tab_state.near_plane,
+ far_plane=render_tab_state.far_plane,
+ radius_clip=render_tab_state.radius_clip,
+ eps2d=render_tab_state.eps2d,
+ backgrounds=torch.tensor([render_tab_state.backgrounds], device=device)
+ / 255.0,
+ render_mode=RENDER_MODE_MAP[render_tab_state.render_mode],
+ rasterize_mode=render_tab_state.rasterize_mode,
+ camera_model=render_tab_state.camera_model,
+ )
+ render_tab_state.total_gs_count = len(means)
+ render_tab_state.rendered_gs_count = (info["radii"] > 0).all(-1).sum().item()
+
+ if render_tab_state.render_mode == "rgb":
+ # colors represented with sh are not guranteed to be in [0, 1]
+ render_colors = render_colors[0, ..., 0:3].clamp(0, 1)
+ renders = render_colors.cpu().numpy()
+ elif render_tab_state.render_mode in ["depth(accumulated)", "depth(expected)"]:
+ # normalize depth to [0, 1]
+ depth = render_colors[0, ..., 0:1]
+ if render_tab_state.normalize_nearfar:
+ near_plane = render_tab_state.near_plane
+ far_plane = render_tab_state.far_plane
+ else:
+ near_plane = depth.min()
+ far_plane = depth.max()
+ depth_norm = (depth - near_plane) / (far_plane - near_plane + 1e-10)
+ depth_norm = torch.clip(depth_norm, 0, 1)
+ if render_tab_state.inverse:
+ depth_norm = 1 - depth_norm
+ renders = (
+ apply_float_colormap(depth_norm, render_tab_state.colormap)
+ .cpu()
+ .numpy()
+ )
+ elif render_tab_state.render_mode == "alpha":
+ alpha = render_alphas[0, ..., 0:1]
+ if render_tab_state.inverse:
+ alpha = 1 - alpha
+ renders = (
+ apply_float_colormap(alpha, render_tab_state.colormap).cpu().numpy()
+ )
+ return renders
+
+ server = viser.ViserServer(port=args.port, verbose=False)
+ _ = GsplatViewer(
+ server=server,
+ render_fn=viewer_render_fn,
+ output_dir=Path(args.output_dir),
+ mode="rendering",
+ )
+ print("Viewer running... Ctrl+C to exit.")
+ time.sleep(100000)
+
+
+if __name__ == "__main__":
+ """
+ # Use single GPU to view the scene
+ CUDA_VISIBLE_DEVICES=9 python -m simple_viewer \
+ --ckpt results/garden/ckpts/ckpt_6999_rank0.pt \
+ --output_dir results/garden/ \
+ --port 8082
+
+ CUDA_VISIBLE_DEVICES=9 python -m simple_viewer \
+ --output_dir results/garden/ \
+ --port 8082
+ """
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--output_dir", type=str, default="results/", help="where to dump outputs"
+ )
+ parser.add_argument(
+ "--scene_grid", type=int, default=1, help="repeat the scene into a grid of NxN"
+ )
+ parser.add_argument(
+ "--ckpt", type=str, nargs="+", default=None, help="path to the .pt file"
+ )
+ parser.add_argument(
+ "--port", type=int, default=8080, help="port for the viewer server"
+ )
+ args = parser.parse_args()
+ assert args.scene_grid % 2 == 1, "scene_grid must be odd"
+
+ cli(main, args, verbose=True)
diff --git a/src/post_opt/utils.py b/src/post_opt/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..b79f4244bd6bfcc1c8aff85fd7bfe59608410e41
--- /dev/null
+++ b/src/post_opt/utils.py
@@ -0,0 +1,224 @@
+import random
+
+import numpy as np
+import torch
+from sklearn.neighbors import NearestNeighbors
+from torch import Tensor
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+from matplotlib import colormaps
+
+
+class CameraOptModule(torch.nn.Module):
+ """Camera pose optimization module."""
+
+ def __init__(self, n: int):
+ super().__init__()
+ # Delta positions (3D) + Delta rotations (6D)
+ self.embeds = torch.nn.Embedding(n, 9)
+ # Identity rotation in 6D representation
+ self.register_buffer("identity", torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0]))
+
+ def zero_init(self):
+ torch.nn.init.zeros_(self.embeds.weight)
+
+ def random_init(self, std: float):
+ torch.nn.init.normal_(self.embeds.weight, std=std)
+
+ def forward(self, camtoworlds: Tensor, embed_ids: Tensor) -> Tensor:
+ """Adjust camera pose based on deltas.
+
+ Args:
+ camtoworlds: (..., 4, 4)
+ embed_ids: (...,)
+
+ Returns:
+ updated camtoworlds: (..., 4, 4)
+ """
+ assert camtoworlds.shape[:-2] == embed_ids.shape
+ batch_shape = camtoworlds.shape[:-2]
+ pose_deltas = self.embeds(embed_ids) # (..., 9)
+ dx, drot = pose_deltas[..., :3], pose_deltas[..., 3:]
+ rot = rotation_6d_to_matrix(
+ drot + self.identity.expand(*batch_shape, -1)
+ ) # (..., 3, 3)
+ transform = torch.eye(4, device=pose_deltas.device).repeat((*batch_shape, 1, 1))
+ transform[..., :3, :3] = rot
+ transform[..., :3, 3] = dx
+ return torch.matmul(camtoworlds, transform)
+
+
+class AppearanceOptModule(torch.nn.Module):
+ """Appearance optimization module."""
+
+ def __init__(
+ self,
+ n: int,
+ feature_dim: int,
+ embed_dim: int = 16,
+ sh_degree: int = 3,
+ mlp_width: int = 64,
+ mlp_depth: int = 2,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.sh_degree = sh_degree
+ self.embeds = torch.nn.Embedding(n, embed_dim)
+ layers = []
+ layers.append(
+ torch.nn.Linear(embed_dim + feature_dim + (sh_degree + 1) ** 2, mlp_width)
+ )
+ layers.append(torch.nn.ReLU(inplace=True))
+ for _ in range(mlp_depth - 1):
+ layers.append(torch.nn.Linear(mlp_width, mlp_width))
+ layers.append(torch.nn.ReLU(inplace=True))
+ layers.append(torch.nn.Linear(mlp_width, 3))
+ self.color_head = torch.nn.Sequential(*layers)
+
+ def forward(
+ self, features: Tensor, embed_ids: Tensor, dirs: Tensor, sh_degree: int
+ ) -> Tensor:
+ """Adjust appearance based on embeddings.
+
+ Args:
+ features: (N, feature_dim)
+ embed_ids: (C,)
+ dirs: (C, N, 3)
+
+ Returns:
+ colors: (C, N, 3)
+ """
+ from gsplat.cuda._torch_impl import _eval_sh_bases_fast
+
+ C, N = dirs.shape[:2]
+ # Camera embeddings
+ if embed_ids is None:
+ embeds = torch.zeros(C, self.embed_dim, device=features.device)
+ else:
+ embeds = self.embeds(embed_ids) # [C, D2]
+ embeds = embeds[:, None, :].expand(-1, N, -1) # [C, N, D2]
+ # GS features
+ features = features[None, :, :].expand(C, -1, -1) # [C, N, D1]
+ # View directions
+ dirs = F.normalize(dirs, dim=-1) # [C, N, 3]
+ num_bases_to_use = (sh_degree + 1) ** 2
+ num_bases = (self.sh_degree + 1) ** 2
+ sh_bases = torch.zeros(C, N, num_bases, device=features.device) # [C, N, K]
+ sh_bases[:, :, :num_bases_to_use] = _eval_sh_bases_fast(num_bases_to_use, dirs)
+ # Get colors
+ if self.embed_dim > 0:
+ h = torch.cat([embeds, features, sh_bases], dim=-1) # [C, N, D1 + D2 + K]
+ else:
+ h = torch.cat([features, sh_bases], dim=-1)
+ colors = self.color_head(h)
+ return colors
+
+
+def rotation_6d_to_matrix(d6: Tensor) -> Tensor:
+ """
+ Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
+ using Gram--Schmidt orthogonalization per Section B of [1]. Adapted from pytorch3d.
+ Args:
+ d6: 6D rotation representation, of size (*, 6)
+
+ Returns:
+ batch of rotation matrices of size (*, 3, 3)
+
+ [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+ On the Continuity of Rotation Representations in Neural Networks.
+ IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+ Retrieved from http://arxiv.org/abs/1812.07035
+ """
+
+ a1, a2 = d6[..., :3], d6[..., 3:]
+ b1 = F.normalize(a1, dim=-1)
+ b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
+ b2 = F.normalize(b2, dim=-1)
+ b3 = torch.cross(b1, b2, dim=-1)
+ return torch.stack((b1, b2, b3), dim=-2)
+
+
+def knn(x: Tensor, K: int = 4) -> Tensor:
+ x_np = x.cpu().numpy()
+ model = NearestNeighbors(n_neighbors=K, metric="euclidean").fit(x_np)
+ distances, _ = model.kneighbors(x_np)
+ return torch.from_numpy(distances).to(x)
+
+
+def rgb_to_sh(rgb: Tensor) -> Tensor:
+ C0 = 0.28209479177387814
+ return (rgb - 0.5) / C0
+
+
+def set_random_seed(seed: int):
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+
+
+# ref: https://github.com/hbb1/2d-gaussian-splatting/blob/main/utils/general_utils.py#L163
+def colormap(img, cmap="jet"):
+ W, H = img.shape[:2]
+ dpi = 300
+ fig, ax = plt.subplots(1, figsize=(H / dpi, W / dpi), dpi=dpi)
+ im = ax.imshow(img, cmap=cmap)
+ ax.set_axis_off()
+ fig.colorbar(im, ax=ax)
+ fig.tight_layout()
+ fig.canvas.draw()
+ data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+ data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+ img = torch.from_numpy(data).float().permute(2, 0, 1)
+ plt.close()
+ return img
+
+
+def apply_float_colormap(img: torch.Tensor, colormap: str = "turbo") -> torch.Tensor:
+ """Convert single channel to a color img.
+
+ Args:
+ img (torch.Tensor): (..., 1) float32 single channel image.
+ colormap (str): Colormap for img.
+
+ Returns:
+ (..., 3) colored img with colors in [0, 1].
+ """
+ img = torch.nan_to_num(img, 0)
+ if colormap == "gray":
+ return img.repeat(1, 1, 3)
+ img_long = (img * 255).long()
+ img_long_min = torch.min(img_long)
+ img_long_max = torch.max(img_long)
+ assert img_long_min >= 0, f"the min value is {img_long_min}"
+ assert img_long_max <= 255, f"the max value is {img_long_max}"
+ return torch.tensor(
+ colormaps[colormap].colors, # type: ignore
+ device=img.device,
+ )[img_long[..., 0]]
+
+
+def apply_depth_colormap(
+ depth: torch.Tensor,
+ acc: torch.Tensor = None,
+ near_plane: float = None,
+ far_plane: float = None,
+) -> torch.Tensor:
+ """Converts a depth image to color for easier analysis.
+
+ Args:
+ depth (torch.Tensor): (..., 1) float32 depth.
+ acc (torch.Tensor | None): (..., 1) optional accumulation mask.
+ near_plane: Closest depth to consider. If None, use min image value.
+ far_plane: Furthest depth to consider. If None, use max image value.
+
+ Returns:
+ (..., 3) colored depth image with colors in [0, 1].
+ """
+ near_plane = near_plane or float(torch.min(depth))
+ far_plane = far_plane or float(torch.max(depth))
+ depth = (depth - near_plane) / (far_plane - near_plane + 1e-10)
+ depth = torch.clip(depth, 0.0, 1.0)
+ img = apply_float_colormap(depth, colormap="turbo")
+ if acc is not None:
+ img = img * acc + (1.0 - acc)
+ return img
diff --git a/src/utils/__init.py b/src/utils/__init.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/src/utils/ba.py b/src/utils/ba.py
new file mode 100644
index 0000000000000000000000000000000000000000..6844f8fcc332c3ad484ba0816428d5471379413d
--- /dev/null
+++ b/src/utils/ba.py
@@ -0,0 +1,407 @@
+import pycolmap
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from src.model.encoder.vggt.utils.pose_enc import pose_encoding_to_extri_intri
+from src.model.encoder.vggt.utils.geometry import closed_form_inverse_se3, unproject_depth_map_to_point_map
+from lightglue import ALIKED, SuperPoint, SIFT
+from src.utils.tensor_to_pycolmap import batch_matrix_to_pycolmap, pycolmap_to_batch_matrix
+
+
+_RESNET_MEAN = [0.485, 0.456, 0.406]
+_RESNET_STD = [0.229, 0.224, 0.225]
+
+
+def generate_rank_by_dino(
+ images, query_frame_num, image_size=518, model_name="dinov2_vitb14_reg", device="cuda", spatial_similarity=True
+):
+ """
+ Generate a ranking of frames using DINO ViT features.
+
+ Args:
+ images: Tensor of shape (S, 3, H, W) with values in range [0, 1]
+ query_frame_num: Number of frames to select
+ image_size: Size to resize images to before processing
+ model_name: Name of the DINO model to use
+ device: Device to run the model on
+ spatial_similarity: Whether to use spatial token similarity or CLS token similarity
+
+ Returns:
+ List of frame indices ranked by their representativeness
+ """
+ dino_v2_model = torch.hub.load('facebookresearch/dinov2', model_name)
+ dino_v2_model.eval()
+ dino_v2_model = dino_v2_model.to(device)
+
+ resnet_mean = torch.tensor(_RESNET_MEAN, device=device).view(1, 3, 1, 1)
+ resnet_std = torch.tensor(_RESNET_STD, device=device).view(1, 3, 1, 1)
+ images_resnet_norm = (images - resnet_mean) / resnet_std
+
+ with torch.no_grad():
+ frame_feat = dino_v2_model(images_resnet_norm, is_training=True)
+
+ if spatial_similarity:
+ frame_feat = frame_feat["x_norm_patchtokens"]
+ frame_feat_norm = F.normalize(frame_feat, p=2, dim=1)
+
+ # Compute the similarity matrix
+ frame_feat_norm = frame_feat_norm.permute(1, 0, 2)
+ similarity_matrix = torch.bmm(
+ frame_feat_norm, frame_feat_norm.transpose(-1, -2)
+ )
+ similarity_matrix = similarity_matrix.mean(dim=0)
+ else:
+ frame_feat = frame_feat["x_norm_clstoken"]
+ frame_feat_norm = F.normalize(frame_feat, p=2, dim=1)
+ similarity_matrix = torch.mm(
+ frame_feat_norm, frame_feat_norm.transpose(-1, -2)
+ )
+
+ distance_matrix = 100 - similarity_matrix.clone()
+
+ # Ignore self-pairing
+ similarity_matrix.fill_diagonal_(-100)
+ similarity_sum = similarity_matrix.sum(dim=1)
+
+ # Find the most common frame
+ most_common_frame_index = torch.argmax(similarity_sum).item()
+
+ # Conduct FPS sampling starting from the most common frame
+ fps_idx = farthest_point_sampling(
+ distance_matrix, query_frame_num, most_common_frame_index
+ )
+
+ return fps_idx
+
+
+def farthest_point_sampling(
+ distance_matrix, num_samples, most_common_frame_index=0
+):
+ """
+ Farthest point sampling algorithm to select diverse frames.
+
+ Args:
+ distance_matrix: Matrix of distances between frames
+ num_samples: Number of frames to select
+ most_common_frame_index: Index of the first frame to select
+
+ Returns:
+ List of selected frame indices
+ """
+ distance_matrix = distance_matrix.clamp(min=0)
+ N = distance_matrix.size(0)
+
+ # Initialize with the most common frame
+ selected_indices = [most_common_frame_index]
+ check_distances = distance_matrix[selected_indices]
+
+ while len(selected_indices) < num_samples:
+ # Find the farthest point from the current set of selected points
+ farthest_point = torch.argmax(check_distances)
+ selected_indices.append(farthest_point.item())
+
+ check_distances = distance_matrix[farthest_point]
+ # Mark already selected points to avoid selecting them again
+ check_distances[selected_indices] = 0
+
+ # Break if all points have been selected
+ if len(selected_indices) == N:
+ break
+
+ return selected_indices
+
+
+def calculate_index_mappings(query_index, S, device=None):
+ """
+ Construct an order that switches [query_index] and [0]
+ so that the content of query_index would be placed at [0].
+
+ Args:
+ query_index: Index to swap with 0
+ S: Total number of elements
+ device: Device to place the tensor on
+
+ Returns:
+ Tensor of indices with the swapped order
+ """
+ new_order = torch.arange(S)
+ new_order[0] = query_index
+ new_order[query_index] = 0
+ if device is not None:
+ new_order = new_order.to(device)
+ return new_order
+
+
+def switch_tensor_order(tensors, order, dim=1):
+ """
+ Reorder tensors along a specific dimension according to the given order.
+
+ Args:
+ tensors: List of tensors to reorder
+ order: Tensor of indices specifying the new order
+ dim: Dimension along which to reorder
+
+ Returns:
+ List of reordered tensors
+ """
+ return [
+ torch.index_select(tensor, dim, order) if tensor is not None else None
+ for tensor in tensors
+ ]
+
+
+def predict_track(model, images, query_points, dtype=torch.bfloat16, use_tf32_for_track=True, iters=4):
+ """
+ Predict tracks for query points across frames.
+
+ Args:
+ model: VGGT model
+ images: Tensor of images of shape (S, 3, H, W)
+ query_points: Query points to track
+ dtype: Data type for computation
+ use_tf32_for_track: Whether to use TF32 precision for tracking
+ iters: Number of iterations for tracking
+
+ Returns:
+ Predicted tracks, visibility scores, and confidence scores
+ """
+ with torch.no_grad():
+ with torch.cuda.amp.autocast(dtype=dtype):
+ images = images[None] # add batch dimension
+ aggregated_tokens_list, ps_idx = model.aggregator(images)
+
+ if not use_tf32_for_track:
+ track_list, vis_score, conf_score = model.track_head(
+ aggregated_tokens_list, images, ps_idx, query_points=query_points, iters=iters
+ )
+
+ if use_tf32_for_track:
+ with torch.cuda.amp.autocast(enabled=False):
+ track_list, vis_score, conf_score = model.track_head(
+ aggregated_tokens_list, images, ps_idx, query_points=query_points, iters=iters
+ )
+
+ pred_track = track_list[-1]
+
+ return pred_track.squeeze(0), vis_score.squeeze(0), conf_score.squeeze(0)
+
+
+def initialize_feature_extractors(max_query_num, det_thres, extractor_method="aliked", device="cuda"):
+ """
+ Initialize feature extractors that can be reused based on a method string.
+
+ Args:
+ max_query_num: Maximum number of keypoints to extract
+ det_thres: Detection threshold for keypoint extraction
+ extractor_method: String specifying which extractors to use (e.g., "aliked", "sp+sift", "aliked+sp+sift")
+ device: Device to run extraction on
+
+ Returns:
+ Dictionary of initialized extractors
+ """
+ extractors = {}
+ methods = extractor_method.lower().split('+')
+ active_extractors = len(methods)
+
+ for method in methods:
+ method = method.strip()
+ if method == "aliked":
+ aliked_max_points = max_query_num // active_extractors
+ aliked_extractor = ALIKED(max_num_keypoints=aliked_max_points, detection_threshold=det_thres)
+ extractors['aliked'] = aliked_extractor.to(device).eval()
+ elif method == "sp":
+ sp_max_points = max_query_num // active_extractors
+ sp_extractor = SuperPoint(max_num_keypoints=sp_max_points, detection_threshold=det_thres)
+ extractors['sp'] = sp_extractor.to(device).eval()
+ elif method == "sift":
+ sift_max_points = max_query_num // active_extractors
+ sift_extractor = SIFT(max_num_keypoints=sift_max_points)
+ extractors['sift'] = sift_extractor.to(device).eval()
+ else:
+ print(f"Warning: Unknown feature extractor '{method}', ignoring.")
+
+ if not extractors:
+ print(f"Warning: No valid extractors found in '{extractor_method}'. Using ALIKED by default.")
+ aliked_extractor = ALIKED(max_num_keypoints=max_query_num, detection_threshold=det_thres)
+ extractors['aliked'] = aliked_extractor.to(device).eval()
+
+ return extractors
+
+
+def extract_keypoints(query_image, extractors):
+ """
+ Extract keypoints using pre-initialized feature extractors.
+
+ Args:
+ query_image: Input image tensor (3xHxW, range [0, 1])
+ extractors: Dictionary of initialized extractors
+
+ Returns:
+ Tensor of keypoint coordinates (1xNx2)
+ """
+ query_points_round = None
+
+ with torch.no_grad():
+ for extractor_name, extractor in extractors.items():
+ query_points_data = extractor.extract(query_image)
+ extractor_points = query_points_data["keypoints"].round()
+
+ if query_points_round is not None:
+ query_points_round = torch.cat([query_points_round, extractor_points], dim=1)
+ else:
+ query_points_round = extractor_points
+
+ return query_points_round
+
+
+def run_vggt_with_ba(model, images, image_names=None, dtype=torch.bfloat16,
+ max_query_num=2048, det_thres=0.005, query_frame_num=3,
+ extractor_method="aliked+sp+sift",
+ max_reproj_error=4,
+ shared_camera=True,
+ ):
+ """
+ Run VGGT with bundle adjustment for pose estimation.
+
+ Args:
+ model: VGGT model
+ images: Tensor of images of shape (S, 3, H, W)
+ image_names: Optional list of image names
+ dtype: Data type for computation
+
+ Returns:
+ Predicted extrinsic camera parameters
+
+ TODO:
+ - [ ] Use VGGT's vit instead of dinov2 for rank generation
+ """
+ device = images.device
+ frame_num = images.shape[0]
+
+ # TODO: use vggt's vit instead of dinov2
+ # Select representative frames for feature extraction
+ query_frame_indexes = generate_rank_by_dino(
+ images, query_frame_num, image_size=518,
+ model_name="dinov2_vitb14_reg", device=device,
+ spatial_similarity=False
+ )
+
+ # Add the first image to the front if not already present
+ if 0 in query_frame_indexes:
+ query_frame_indexes.remove(0)
+ query_frame_indexes = [0, *query_frame_indexes]
+
+ # Get initial pose and depth predictions
+
+ with torch.no_grad(), torch.cuda.amp.autocast(enabled=False, dtype=torch.bfloat16):
+ aggregated_tokens_list, patch_start_idx = model.aggregator(images, intermediate_layer_idx=model.cfg.intermediate_layer_idx)
+ with torch.cuda.amp.autocast(enabled=False):
+ fp32_tokens = [token.float() for token in aggregated_tokens_list]
+ pred_all_pose_enc = model.camera_head(fp32_tokens)[-1]
+ pred_all_extrinsic, pred_all_intrinsic = pose_encoding_to_extri_intri(pred_all_pose_enc, images.shape[-2:])
+ pred_extrinsic = pred_all_extrinsic[0]
+ pred_intrinsic = pred_all_intrinsic[0]
+ depth_map, depth_conf = model.depth_head(aggregated_tokens_list, images=images, patch_start_idx=patch_start_idx)
+
+ world_points = unproject_depth_map_to_point_map(depth_map, pred_extrinsic, pred_intrinsic)
+ world_points = torch.from_numpy(world_points).to(device)
+ world_points_conf = depth_conf.to(device)
+
+ torch.cuda.empty_cache()
+
+ # Lists to store predictions
+ pred_tracks = []
+ pred_vis_scores = []
+ pred_conf_scores = []
+ pred_world_points = []
+ pred_world_points_conf = []
+
+ # Initialize feature extractors
+ extractors = initialize_feature_extractors(max_query_num, det_thres, extractor_method, device)
+
+ # Process each query frame
+ for query_index in query_frame_indexes:
+ query_image = images[query_index]
+ query_points_round = extract_keypoints(query_image, extractors)
+
+ # Reorder images to put query image first
+ reorder_index = calculate_index_mappings(query_index, frame_num, device=device)
+ reorder_images = switch_tensor_order([images], reorder_index, dim=0)[0]
+
+ # Track points across frames
+ reorder_tracks, reorder_vis_score, reorder_conf_score = predict_track(
+ model, reorder_images, query_points_round, dtype=dtype, use_tf32_for_track=True, iters=4
+ )
+
+ # Restore original order
+ pred_track, pred_vis, pred_score = switch_tensor_order(
+ [reorder_tracks, reorder_vis_score, reorder_conf_score], reorder_index, dim=0
+ )
+
+ pred_tracks.append(pred_track)
+ pred_vis_scores.append(pred_vis)
+ pred_conf_scores.append(pred_score)
+
+ # Get corresponding 3D points
+ query_points_round_long = query_points_round.squeeze(0).long()
+ query_world_points = world_points[query_index][
+ query_points_round_long[:, 1], query_points_round_long[:, 0]
+ ]
+ query_world_points_conf = world_points_conf[query_index][
+ query_points_round_long[:, 1], query_points_round_long[:, 0]
+ ]
+
+ pred_world_points.append(query_world_points)
+ pred_world_points_conf.append(query_world_points_conf)
+
+ # Concatenate prediction lists
+ pred_tracks = torch.cat(pred_tracks, dim=1)
+ pred_vis_scores = torch.cat(pred_vis_scores, dim=1)
+ pred_conf_scores = torch.cat(pred_conf_scores, dim=1)
+ pred_world_points = torch.cat(pred_world_points, dim=0)
+ pred_world_points_conf = torch.cat(pred_world_points_conf, dim=0)
+
+ # Filter points by confidence
+ filtered_flag = pred_world_points_conf > 1.5
+
+ if filtered_flag.sum() > 1024:
+ # well if the number of points is too small, we will not filter
+ pred_world_points = pred_world_points[filtered_flag]
+ pred_world_points_conf = pred_world_points_conf[filtered_flag]
+
+ pred_tracks = pred_tracks[:, filtered_flag]
+ pred_vis_scores = pred_vis_scores[:, filtered_flag]
+ pred_conf_scores = pred_conf_scores[:, filtered_flag]
+
+ torch.cuda.empty_cache()
+
+ # Bundle adjustment parameters
+ S, _, H, W = images.shape
+ image_size = torch.tensor([W, H], dtype=pred_tracks.dtype, device=device)
+
+ # Run bundle adjustment
+ reconstruction = batch_matrix_to_pycolmap(
+ pred_world_points,
+ pred_extrinsic,
+ pred_intrinsic,
+ pred_tracks,
+ image_size,
+ max_reproj_error=max_reproj_error,
+ shared_camera=shared_camera
+ )
+
+ ba_options = pycolmap.BundleAdjustmentOptions()
+ pycolmap.bundle_adjustment(reconstruction, ba_options)
+ _, updated_extrinsic, _, _ = pycolmap_to_batch_matrix(
+ reconstruction, device=device, camera_type="SIMPLE_PINHOLE"
+ )
+
+ return updated_extrinsic
+
+
+
+
+
+
+
diff --git a/src/utils/cropping.py b/src/utils/cropping.py
new file mode 100644
index 0000000000000000000000000000000000000000..1e635174fe869d9da83c7428fd5c7e81020cff84
--- /dev/null
+++ b/src/utils/cropping.py
@@ -0,0 +1,163 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# croppping utilities
+# --------------------------------------------------------
+import PIL.Image
+import os
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+import cv2 # noqa
+import numpy as np # noqa
+from src.utils.geometry import colmap_to_opencv_intrinsics, opencv_to_colmap_intrinsics # noqa
+try:
+ lanczos = PIL.Image.Resampling.LANCZOS
+ bicubic = PIL.Image.Resampling.BICUBIC
+except AttributeError:
+ lanczos = PIL.Image.LANCZOS
+ bicubic = PIL.Image.BICUBIC
+
+
+class ImageList:
+ """ Convenience class to aply the same operation to a whole set of images.
+ """
+
+ def __init__(self, images):
+ if not isinstance(images, (tuple, list, set)):
+ images = [images]
+ self.images = []
+ for image in images:
+ if not isinstance(image, PIL.Image.Image):
+ image = PIL.Image.fromarray(image)
+ self.images.append(image)
+
+ def __len__(self):
+ return len(self.images)
+
+ def to_pil(self):
+ return tuple(self.images) if len(self.images) > 1 else self.images[0]
+
+ @property
+ def size(self):
+ sizes = [im.size for im in self.images]
+ assert all(sizes[0] == s for s in sizes)
+ return sizes[0]
+
+ def resize(self, *args, **kwargs):
+ return ImageList(self._dispatch('resize', *args, **kwargs))
+
+ def crop(self, *args, **kwargs):
+ return ImageList(self._dispatch('crop', *args, **kwargs))
+
+ def _dispatch(self, func, *args, **kwargs):
+ return [getattr(im, func)(*args, **kwargs) for im in self.images]
+
+
+def rescale_image(image, camera_intrinsics, output_resolution, force=True):
+ """ Jointly rescale a (image, depthmap)
+ so that (out_width, out_height) >= output_res
+ """
+ image = ImageList(image)
+ input_resolution = np.array(image.size) # (W,H)
+ output_resolution = np.array(output_resolution)
+ # define output resolution
+ assert output_resolution.shape == (2,)
+ scale_final = max(output_resolution / image.size) + 1e-8
+ if scale_final >= 1 and not force: # image is already smaller than what is asked
+ return (image.to_pil(), camera_intrinsics)
+ output_resolution = np.floor(input_resolution * scale_final).astype(int)
+
+ # first rescale the image so that it contains the crop
+ image = image.resize(tuple(output_resolution), resample=lanczos if scale_final < 1 else bicubic)
+
+ # no offset here; simple rescaling
+ camera_intrinsics = camera_matrix_of_crop(
+ camera_intrinsics, input_resolution, output_resolution, scaling=scale_final)
+
+ return image.to_pil(), camera_intrinsics
+
+
+def camera_matrix_of_crop(input_camera_matrix, input_resolution, output_resolution, scaling=1, offset_factor=0.5, offset=None):
+ # Margins to offset the origin
+ margins = np.asarray(input_resolution) * scaling - output_resolution
+ assert np.all(margins >= 0.0)
+ if offset is None:
+ offset = offset_factor * margins
+
+ # Generate new camera parameters
+ output_camera_matrix_colmap = opencv_to_colmap_intrinsics(input_camera_matrix)
+ output_camera_matrix_colmap[:2, :] *= scaling
+ output_camera_matrix_colmap[:2, 2] -= offset
+ output_camera_matrix = colmap_to_opencv_intrinsics(output_camera_matrix_colmap)
+
+ return output_camera_matrix
+
+
+def crop_image(image, camera_intrinsics, crop_bbox):
+ """
+ Return a crop of the input view.
+ """
+ image = ImageList(image)
+ l, t, r, b = crop_bbox
+
+ image = image.crop((l, t, r, b))
+
+ camera_intrinsics = camera_intrinsics.copy()
+ camera_intrinsics[0, 2] -= l
+ camera_intrinsics[1, 2] -= t
+
+ return image.to_pil(), camera_intrinsics
+
+
+def bbox_from_intrinsics_in_out(input_camera_matrix, output_camera_matrix, output_resolution):
+ out_width, out_height = output_resolution
+ l, t = np.int32(np.round(input_camera_matrix[:2, 2] - output_camera_matrix[:2, 2]))
+ crop_bbox = (l, t, l + out_width, t + out_height)
+ return crop_bbox
+
+def rescale_image_depthmap(image, depthmap, camera_intrinsics, output_resolution, force=True):
+ """ Jointly rescale a (image, depthmap)
+ so that (out_width, out_height) >= output_res
+ """
+ image = ImageList(image)
+ input_resolution = np.array(image.size) # (W,H)
+ output_resolution = np.array(output_resolution)
+ if depthmap is not None:
+ # can also use this with masks instead of depthmaps
+ assert tuple(depthmap.shape[:2]) == image.size[::-1]
+
+ # define output resolution
+ assert output_resolution.shape == (2,)
+ scale_final = max(output_resolution / image.size) + 1e-8
+ if scale_final >= 1 and not force: # image is already smaller than what is asked
+ return (image.to_pil(), depthmap, camera_intrinsics)
+ output_resolution = np.floor(input_resolution * scale_final).astype(int)
+
+ # first rescale the image so that it contains the crop
+ # breakpoint()
+ image = image.resize(tuple(output_resolution), resample=lanczos if scale_final < 1 else bicubic)
+ if depthmap is not None:
+ depthmap = cv2.resize(depthmap, tuple(output_resolution), fx=scale_final,
+ fy=scale_final, interpolation=cv2.INTER_NEAREST)
+
+ # no offset here; simple rescaling
+ camera_intrinsics = camera_matrix_of_crop(
+ camera_intrinsics, input_resolution, tuple(output_resolution), scaling=scale_final)
+
+ return image.to_pil(), depthmap, camera_intrinsics
+
+def crop_image_depthmap(image, depthmap, camera_intrinsics, crop_bbox):
+ """
+ Return a crop of the input view.
+ """
+ image = ImageList(image)
+ l, t, r, b = crop_bbox
+
+ image = image.crop((l, t, r, b))
+ depthmap = depthmap[t:b, l:r]
+
+ camera_intrinsics = camera_intrinsics.copy()
+ camera_intrinsics[0, 2] -= l
+ camera_intrinsics[1, 2] -= t
+
+ return image.to_pil(), depthmap, camera_intrinsics
\ No newline at end of file
diff --git a/src/utils/device.py b/src/utils/device.py
new file mode 100644
index 0000000000000000000000000000000000000000..bae8a503be3d75b60258135a5134dfdd7154ad59
--- /dev/null
+++ b/src/utils/device.py
@@ -0,0 +1,95 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# utilitary functions for DUSt3R
+# --------------------------------------------------------
+import numpy as np
+import torch
+
+
+def todevice(batch, device, callback=None, non_blocking=False):
+ """Transfer some variables to another device (i.e. GPU, CPU:torch, CPU:numpy).
+
+ batch: list, tuple, dict of tensors or other things
+ device: pytorch device or 'numpy'
+ callback: function that would be called on every sub-elements.
+ """
+ if callback:
+ batch = callback(batch)
+
+ if isinstance(batch, dict):
+ return {k: todevice(v, device) for k, v in batch.items()}
+
+ if isinstance(batch, (tuple, list)):
+ return type(batch)(todevice(x, device) for x in batch)
+
+ x = batch
+ if device == "numpy":
+ if isinstance(x, torch.Tensor):
+ x = x.detach().cpu().numpy()
+ elif x is not None:
+ if isinstance(x, np.ndarray):
+ x = torch.from_numpy(x)
+ if torch.is_tensor(x):
+ x = x.to(device, non_blocking=non_blocking)
+ return x
+
+
+to_device = todevice # alias
+
+
+def to_numpy(x):
+ return todevice(x, "numpy")
+
+
+def to_cpu(x):
+ return todevice(x, "cpu")
+
+
+def to_cuda(x):
+ return todevice(x, "cuda")
+
+
+def collate_with_cat(whatever, lists=False):
+ if isinstance(whatever, dict):
+ return {k: collate_with_cat(vals, lists=lists) for k, vals in whatever.items()}
+
+ elif isinstance(whatever, (tuple, list)):
+ if len(whatever) == 0:
+ return whatever
+ elem = whatever[0]
+ T = type(whatever)
+
+ if elem is None:
+ return None
+ if isinstance(elem, (bool, float, int, str)):
+ return whatever
+ if isinstance(elem, tuple):
+ return T(collate_with_cat(x, lists=lists) for x in zip(*whatever))
+ if isinstance(elem, dict):
+ return {
+ k: collate_with_cat([e[k] for e in whatever], lists=lists) for k in elem
+ }
+
+ if isinstance(elem, torch.Tensor):
+ return listify(whatever) if lists else torch.cat(whatever)
+ if isinstance(elem, np.ndarray):
+ return (
+ listify(whatever)
+ if lists
+ else torch.cat([torch.from_numpy(x) for x in whatever])
+ )
+
+ # otherwise, we just chain lists
+ return sum(whatever, T())
+
+
+def listify(elems):
+ return [x for e in elems for x in e]
diff --git a/src/utils/geometry.py b/src/utils/geometry.py
new file mode 100644
index 0000000000000000000000000000000000000000..2124b269b641716489ee20767efc77534005a2fd
--- /dev/null
+++ b/src/utils/geometry.py
@@ -0,0 +1,428 @@
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# geometry utilitary functions
+# --------------------------------------------------------
+import torch
+import numpy as np
+from scipy.spatial import cKDTree as KDTree
+
+def invalid_to_nans(arr, valid_mask, ndim=999):
+ if valid_mask is not None:
+ arr = arr.clone()
+ arr[~valid_mask] = float('nan')
+ if arr.ndim > ndim:
+ arr = arr.flatten(-2 - (arr.ndim - ndim), -2)
+ return arr
+
+
+def invalid_to_zeros(arr, valid_mask, ndim=999):
+ if valid_mask is not None:
+ arr = arr.clone()
+ arr[~valid_mask] = 0
+ nnz = valid_mask.view(len(valid_mask), -1).sum(1)
+ else:
+ nnz = arr.numel() // len(arr) if len(arr) else 0 # number of point per image
+ if arr.ndim > ndim:
+ arr = arr.flatten(-2 - (arr.ndim - ndim), -2)
+ return arr, nnz
+
+
+def xy_grid(W, H, device=None, origin=(0, 0), unsqueeze=None, cat_dim=-1, homogeneous=False, **arange_kw):
+ """ Output a (H,W,2) array of int32
+ with output[j,i,0] = i + origin[0]
+ output[j,i,1] = j + origin[1]
+ """
+ if device is None:
+ # numpy
+ arange, meshgrid, stack, ones = np.arange, np.meshgrid, np.stack, np.ones
+ else:
+ # torch
+ arange = lambda *a, **kw: torch.arange(*a, device=device, **kw)
+ meshgrid, stack = torch.meshgrid, torch.stack
+ ones = lambda *a: torch.ones(*a, device=device)
+
+ tw, th = [arange(o, o + s, **arange_kw) for s, o in zip((W, H), origin)]
+ grid = meshgrid(tw, th, indexing='xy')
+ if homogeneous:
+ grid = grid + (ones((H, W)),)
+ if unsqueeze is not None:
+ grid = (grid[0].unsqueeze(unsqueeze), grid[1].unsqueeze(unsqueeze))
+ if cat_dim is not None:
+ grid = stack(grid, cat_dim)
+ return grid
+
+
+def geotrf(Trf, pts, ncol=None, norm=False):
+ """ Apply a geometric transformation to a list of 3-D points.
+
+ H: 3x3 or 4x4 projection matrix (typically a Homography)
+ p: numpy/torch/tuple of coordinates. Shape must be (...,2) or (...,3)
+
+ ncol: int. number of columns of the result (2 or 3)
+ norm: float. if != 0, the resut is projected on the z=norm plane.
+
+ Returns an array of projected 2d points.
+ """
+ assert Trf.ndim >= 2
+ if isinstance(Trf, np.ndarray):
+ pts = np.asarray(pts)
+ elif isinstance(Trf, torch.Tensor):
+ pts = torch.as_tensor(pts, dtype=Trf.dtype)
+
+ # adapt shape if necessary
+ output_reshape = pts.shape[:-1]
+ ncol = ncol or pts.shape[-1]
+
+ # optimized code
+ if (isinstance(Trf, torch.Tensor) and isinstance(pts, torch.Tensor) and
+ Trf.ndim == 3 and pts.ndim == 4):
+ d = pts.shape[3]
+ if Trf.shape[-1] == d:
+ pts = torch.einsum("bij, bhwj -> bhwi", Trf, pts)
+ elif Trf.shape[-1] == d + 1:
+ pts = torch.einsum("bij, bhwj -> bhwi", Trf[:, :d, :d], pts) + Trf[:, None, None, :d, d]
+ else:
+ raise ValueError(f'bad shape, not ending with 3 or 4, for {pts.shape=}')
+ else:
+ if Trf.ndim >= 3:
+ n = Trf.ndim - 2
+ assert Trf.shape[:n] == pts.shape[:n], 'batch size does not match'
+ Trf = Trf.reshape(-1, Trf.shape[-2], Trf.shape[-1])
+
+ if pts.ndim > Trf.ndim:
+ # Trf == (B,d,d) & pts == (B,H,W,d) --> (B, H*W, d)
+ pts = pts.reshape(Trf.shape[0], -1, pts.shape[-1])
+ elif pts.ndim == 2:
+ # Trf == (B,d,d) & pts == (B,d) --> (B, 1, d)
+ pts = pts[:, None, :]
+
+ if pts.shape[-1] + 1 == Trf.shape[-1]:
+ Trf = Trf.swapaxes(-1, -2) # transpose Trf
+ pts = pts @ Trf[..., :-1, :] + Trf[..., -1:, :]
+ elif pts.shape[-1] == Trf.shape[-1]:
+ Trf = Trf.swapaxes(-1, -2) # transpose Trf
+ pts = pts @ Trf
+ else:
+ pts = Trf @ pts.T
+ if pts.ndim >= 2:
+ pts = pts.swapaxes(-1, -2)
+
+ if norm:
+ pts = pts / pts[..., -1:] # DONT DO /= BECAUSE OF WEIRD PYTORCH BUG
+ if norm != 1:
+ pts *= norm
+
+ res = pts[..., :ncol].reshape(*output_reshape, ncol)
+ return res
+
+
+def inv(mat):
+ """ Invert a torch or numpy matrix
+ """
+ if isinstance(mat, torch.Tensor):
+ return torch.linalg.inv(mat)
+ if isinstance(mat, np.ndarray):
+ return np.linalg.inv(mat)
+ raise ValueError(f'bad matrix type = {type(mat)}')
+
+
+def depthmap_to_pts3d(depth, pseudo_focal, pp=None, **_):
+ """
+ Args:
+ - depthmap (BxHxW array):
+ - pseudo_focal: [B,H,W] ; [B,2,H,W] or [B,1,H,W]
+ Returns:
+ pointmap of absolute coordinates (BxHxWx3 array)
+ """
+
+ if len(depth.shape) == 4:
+ B, H, W, n = depth.shape
+ else:
+ B, H, W = depth.shape
+ n = None
+
+ if len(pseudo_focal.shape) == 3: # [B,H,W]
+ pseudo_focalx = pseudo_focaly = pseudo_focal
+ elif len(pseudo_focal.shape) == 4: # [B,2,H,W] or [B,1,H,W]
+ pseudo_focalx = pseudo_focal[:, 0]
+ if pseudo_focal.shape[1] == 2:
+ pseudo_focaly = pseudo_focal[:, 1]
+ else:
+ pseudo_focaly = pseudo_focalx
+ else:
+ raise NotImplementedError("Error, unknown input focal shape format.")
+
+ assert pseudo_focalx.shape == depth.shape[:3]
+ assert pseudo_focaly.shape == depth.shape[:3]
+ grid_x, grid_y = xy_grid(W, H, cat_dim=0, device=depth.device)[:, None]
+
+ # set principal point
+ if pp is None:
+ grid_x = grid_x - (W - 1) / 2
+ grid_y = grid_y - (H - 1) / 2
+ else:
+ grid_x = grid_x.expand(B, -1, -1) - pp[:, 0, None, None]
+ grid_y = grid_y.expand(B, -1, -1) - pp[:, 1, None, None]
+
+ if n is None:
+ pts3d = torch.empty((B, H, W, 3), device=depth.device)
+ pts3d[..., 0] = depth * grid_x / pseudo_focalx
+ pts3d[..., 1] = depth * grid_y / pseudo_focaly
+ pts3d[..., 2] = depth
+ else:
+ pts3d = torch.empty((B, H, W, 3, n), device=depth.device)
+ pts3d[..., 0, :] = depth * (grid_x / pseudo_focalx)[..., None]
+ pts3d[..., 1, :] = depth * (grid_y / pseudo_focaly)[..., None]
+ pts3d[..., 2, :] = depth
+ return pts3d
+
+
+def depthmap_to_camera_coordinates(depthmap, camera_intrinsics, pseudo_focal=None):
+ """
+ Args:
+ - depthmap (HxW array):
+ - camera_intrinsics: a 3x3 matrix
+ Returns:
+ pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels.
+ """
+ camera_intrinsics = np.float32(camera_intrinsics)
+ H, W = depthmap.shape
+
+ # Compute 3D ray associated with each pixel
+ # Strong assumption: there are no skew terms
+ assert camera_intrinsics[0, 1] == 0.0
+ assert camera_intrinsics[1, 0] == 0.0
+ if pseudo_focal is None:
+ fu = camera_intrinsics[0, 0]
+ fv = camera_intrinsics[1, 1]
+ else:
+ assert pseudo_focal.shape == (H, W)
+ fu = fv = pseudo_focal
+ cu = camera_intrinsics[0, 2]
+ cv = camera_intrinsics[1, 2]
+
+ u, v = np.meshgrid(np.arange(W), np.arange(H))
+ z_cam = depthmap
+ x_cam = (u - cu) * z_cam / fu
+ y_cam = (v - cv) * z_cam / fv
+ X_cam = np.stack((x_cam, y_cam, z_cam), axis=-1).astype(np.float32)
+
+ # Mask for valid coordinates
+ valid_mask = (depthmap > 0.0)
+ return X_cam, valid_mask
+
+
+def depthmap_to_absolute_camera_coordinates(depthmap, camera_intrinsics, camera_pose=None, **kw):
+ """
+ Args:
+ - depthmap (HxW array):
+ - camera_intrinsics: a 3x3 matrix
+ - camera_pose: a 4x3 or 4x4 cam2world matrix
+ Returns:
+ pointmap of absolute coordinates (HxWx3 array), and a mask specifying valid pixels."""
+ X_cam, valid_mask = depthmap_to_camera_coordinates(depthmap, camera_intrinsics)
+
+ X_world = X_cam # default
+ if camera_pose is not None:
+ # R_cam2world = np.float32(camera_params["R_cam2world"])
+ # t_cam2world = np.float32(camera_params["t_cam2world"]).squeeze()
+ R_cam2world = camera_pose[:3, :3]
+ t_cam2world = camera_pose[:3, 3]
+
+ # Express in absolute coordinates (invalid depth values)
+ X_world = np.einsum("ik, vuk -> vui", R_cam2world, X_cam) + t_cam2world[None, None, :]
+
+ return X_world, valid_mask
+
+
+def colmap_to_opencv_intrinsics(K):
+ """
+ Modify camera intrinsics to follow a different convention.
+ Coordinates of the center of the top-left pixels are by default:
+ - (0.5, 0.5) in Colmap
+ - (0,0) in OpenCV
+ """
+ K = K.copy()
+ K[0, 2] -= 0.5
+ K[1, 2] -= 0.5
+ return K
+
+
+def opencv_to_colmap_intrinsics(K):
+ """
+ Modify camera intrinsics to follow a different convention.
+ Coordinates of the center of the top-left pixels are by default:
+ - (0.5, 0.5) in Colmap
+ - (0,0) in OpenCV
+ """
+ K = K.copy()
+ K[0, 2] += 0.5
+ K[1, 2] += 0.5
+ return K
+
+
+def normalize_pointcloud(pts1, pts2, norm_mode='avg_dis', valid1=None, valid2=None, ret_factor=False):
+ """ renorm pointmaps pts1, pts2 with norm_mode
+ """
+ assert pts1.ndim >= 3 and pts1.shape[-1] == 3
+ assert pts2 is None or (pts2.ndim >= 3 and pts2.shape[-1] == 3)
+ norm_mode, dis_mode = norm_mode.split('_')
+
+ if norm_mode == 'avg':
+ # gather all points together (joint normalization)
+ nan_pts1, nnz1 = invalid_to_zeros(pts1, valid1, ndim=3)
+ nan_pts2, nnz2 = invalid_to_zeros(pts2, valid2, ndim=3) if pts2 is not None else (None, 0)
+ all_pts = torch.cat((nan_pts1, nan_pts2), dim=1) if pts2 is not None else nan_pts1
+
+ # compute distance to origin
+ all_dis = all_pts.norm(dim=-1)
+ if dis_mode == 'dis':
+ pass # do nothing
+ elif dis_mode == 'log1p':
+ all_dis = torch.log1p(all_dis)
+ elif dis_mode == 'warp-log1p':
+ # actually warp input points before normalizing them
+ log_dis = torch.log1p(all_dis)
+ warp_factor = log_dis / all_dis.clip(min=1e-8)
+ H1, W1 = pts1.shape[1:-1]
+ pts1 = pts1 * warp_factor[:, :W1 * H1].view(-1, H1, W1, 1)
+ if pts2 is not None:
+ H2, W2 = pts2.shape[1:-1]
+ pts2 = pts2 * warp_factor[:, W1 * H1:].view(-1, H2, W2, 1)
+ all_dis = log_dis # this is their true distance afterwards
+ else:
+ raise ValueError(f'bad {dis_mode=}')
+
+ norm_factor = all_dis.sum(dim=1) / (nnz1 + nnz2 + 1e-8)
+ else:
+ # gather all points together (joint normalization)
+ nan_pts1 = invalid_to_nans(pts1, valid1, ndim=3)
+ nan_pts2 = invalid_to_nans(pts2, valid2, ndim=3) if pts2 is not None else None
+ all_pts = torch.cat((nan_pts1, nan_pts2), dim=1) if pts2 is not None else nan_pts1
+
+ # compute distance to origin
+ all_dis = all_pts.norm(dim=-1)
+
+ if norm_mode == 'avg':
+ norm_factor = all_dis.nanmean(dim=1)
+ elif norm_mode == 'median':
+ norm_factor = all_dis.nanmedian(dim=1).values.detach()
+ elif norm_mode == 'sqrt':
+ norm_factor = all_dis.sqrt().nanmean(dim=1)**2
+ else:
+ raise ValueError(f'bad {norm_mode=}')
+
+ norm_factor = norm_factor.clip(min=1e-8)
+ while norm_factor.ndim < pts1.ndim:
+ norm_factor.unsqueeze_(-1)
+
+ res = pts1 / norm_factor
+ if pts2 is not None:
+ res = (res, pts2 / norm_factor)
+ if ret_factor:
+ res = res + (norm_factor,)
+ return res
+
+
+@torch.no_grad()
+def get_joint_pointcloud_depth(z1, z2, valid_mask1, valid_mask2=None, quantile=0.5):
+ # set invalid points to NaN
+ _z1 = invalid_to_nans(z1, valid_mask1).reshape(len(z1), -1)
+ _z2 = invalid_to_nans(z2, valid_mask2).reshape(len(z2), -1) if z2 is not None else None
+ _z = torch.cat((_z1, _z2), dim=-1) if z2 is not None else _z1
+
+ # compute median depth overall (ignoring nans)
+ if quantile == 0.5:
+ shift_z = torch.nanmedian(_z, dim=-1).values
+ else:
+ shift_z = torch.nanquantile(_z, quantile, dim=-1)
+ return shift_z # (B,)
+
+
+@torch.no_grad()
+def get_joint_pointcloud_center_scale(pts1, pts2, valid_mask1=None, valid_mask2=None, z_only=False, center=True):
+ # set invalid points to NaN
+ _pts1 = invalid_to_nans(pts1, valid_mask1).reshape(len(pts1), -1, 3)
+ _pts2 = invalid_to_nans(pts2, valid_mask2).reshape(len(pts2), -1, 3) if pts2 is not None else None
+ _pts = torch.cat((_pts1, _pts2), dim=1) if pts2 is not None else _pts1
+
+ # compute median center
+ _center = torch.nanmedian(_pts, dim=1, keepdim=True).values # (B,1,3)
+ if z_only:
+ _center[..., :2] = 0 # do not center X and Y
+
+ # compute median norm
+ _norm = ((_pts - _center) if center else _pts).norm(dim=-1)
+ scale = torch.nanmedian(_norm, dim=1).values
+ return _center[:, None, :, :], scale[:, None, None, None]
+
+
+def find_reciprocal_matches(P1, P2):
+ """
+ returns 3 values:
+ 1 - reciprocal_in_P2: a boolean array of size P2.shape[0], a "True" value indicates a match
+ 2 - nn2_in_P1: a int array of size P2.shape[0], it contains the indexes of the closest points in P1
+ 3 - reciprocal_in_P2.sum(): the number of matches
+ """
+ tree1 = KDTree(P1)
+ tree2 = KDTree(P2)
+
+ _, nn1_in_P2 = tree2.query(P1, workers=8)
+ _, nn2_in_P1 = tree1.query(P2, workers=8)
+
+ reciprocal_in_P1 = (nn2_in_P1[nn1_in_P2] == np.arange(len(nn1_in_P2)))
+ reciprocal_in_P2 = (nn1_in_P2[nn2_in_P1] == np.arange(len(nn2_in_P1)))
+ assert reciprocal_in_P1.sum() == reciprocal_in_P2.sum()
+ return reciprocal_in_P2, nn2_in_P1, reciprocal_in_P2.sum()
+
+
+def get_med_dist_between_poses(poses):
+ from scipy.spatial.distance import pdist
+ return np.median(pdist([p[:3, 3].detach().cpu().numpy() for p in poses]))
+
+def get_rel_pos(_rel_xyz, free_space, voxel_size):
+ """
+ Args:
+ _rel_xyz: torch.tensor
+ shape [..., 3]
+ free_space: str
+ something like 'soft'
+ grid: GridBatch
+ fvdb grid
+
+ Returns:
+ rel_pos: torch.tensor
+ shape [..., 3], the position within a voxel (compared with the voxel center)
+ """
+ if free_space == "hard":
+ rel_pos = torch.sigmoid(_rel_xyz) * voxel_size
+ elif free_space == "soft":
+ # free space [-1, 2]
+ rel_pos = (torch.sigmoid(_rel_xyz) * 3 - 1) * voxel_size
+ elif free_space == "soft-2":
+ # free space [-2, 3]
+ rel_pos = (torch.sigmoid(_rel_xyz) * 5 - 2) * voxel_size
+ elif free_space == "soft-3":
+ # free space [-3, 4]
+ rel_pos = (torch.sigmoid(_rel_xyz) * 7 - 3) * voxel_size
+ elif free_space == "soft-4":
+ # free space [-4, 5]
+ rel_pos = (torch.sigmoid(_rel_xyz) * 9 - 4) * voxel_size
+ elif free_space == "soft-5":
+ # free space [-5, 6]
+ rel_pos = (torch.sigmoid(_rel_xyz) * 11 - 5) * voxel_size
+ elif free_space == "tanh-3":
+ # free space [-2.5, 3.5]
+ rel_pos = (torch.tanh(_rel_xyz) * 3 + 0.5) * voxel_size
+ elif free_space == "free-1":
+ rel_pos = _rel_xyz * voxel_size
+ elif free_space == "free-2":
+ rel_pos = _rel_xyz
+ elif free_space == "center":
+ rel_pos = (torch.zeros_like(_rel_xyz) + 0.5) * voxel_size
+ else:
+ raise NotImplementedError
+
+ return rel_pos
diff --git a/src/utils/image.py b/src/utils/image.py
new file mode 100644
index 0000000000000000000000000000000000000000..494502fd2b9ad5dcd0f7baf7929156d002c70c93
--- /dev/null
+++ b/src/utils/image.py
@@ -0,0 +1,185 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# utilitary functions about images (loading/converting...)
+# --------------------------------------------------------
+import os
+
+import numpy as np
+import PIL.Image
+import torch
+import torchvision.transforms as tvf
+from PIL.ImageOps import exif_transpose
+from PIL import Image
+import torchvision
+
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+import cv2
+
+try:
+ from pillow_heif import register_heif_opener
+
+ register_heif_opener()
+ heif_support_enabled = True
+except ImportError:
+ heif_support_enabled = False
+
+ImgNorm = tvf.Compose([tvf.ToTensor(), tvf.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+
+
+def imread_cv2(path, options=cv2.IMREAD_COLOR):
+ """Open an image or a depthmap with opencv-python."""
+ if path.endswith((".exr", "EXR")):
+ options = cv2.IMREAD_ANYDEPTH
+ img = cv2.imread(path, options)
+ if img is None:
+ raise IOError(f"Could not load image={path} with {options=}")
+ if img.ndim == 3:
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+ return img
+
+
+def rgb(ftensor, true_shape=None):
+ if isinstance(ftensor, list):
+ return [rgb(x, true_shape=true_shape) for x in ftensor]
+ if isinstance(ftensor, torch.Tensor):
+ ftensor = ftensor.detach().cpu().numpy() # H,W,3
+ if ftensor.ndim == 3 and ftensor.shape[0] == 3:
+ ftensor = ftensor.transpose(1, 2, 0)
+ elif ftensor.ndim == 4 and ftensor.shape[1] == 3:
+ ftensor = ftensor.transpose(0, 2, 3, 1)
+ if true_shape is not None:
+ H, W = true_shape
+ ftensor = ftensor[:H, :W]
+ if ftensor.dtype == np.uint8:
+ img = np.float32(ftensor) / 255
+ else:
+ img = (ftensor * 0.5) + 0.5
+ return img.clip(min=0, max=1)
+
+
+def _resize_pil_image(img, long_edge_size):
+ S = max(img.size)
+ if S > long_edge_size:
+ interp = PIL.Image.LANCZOS
+ elif S <= long_edge_size:
+ interp = PIL.Image.BICUBIC
+ new_size = tuple(int(round(x * long_edge_size / S)) for x in img.size)
+ return img.resize(new_size, interp)
+
+
+def load_images(folder_or_list, size, square_ok=False, verbose=True, rotate_clockwise_90=False, crop_to_landscape=False):
+ """open and convert all images in a list or folder to proper input format for DUSt3R"""
+ if isinstance(folder_or_list, str):
+ if verbose:
+ print(f">> Loading images from {folder_or_list}")
+ root, folder_content = folder_or_list, sorted(os.listdir(folder_or_list))
+
+ elif isinstance(folder_or_list, list):
+ if verbose:
+ print(f">> Loading a list of {len(folder_or_list)} images")
+ root, folder_content = "", folder_or_list
+
+ else:
+ raise ValueError(f"bad {folder_or_list=} ({type(folder_or_list)})")
+
+ supported_images_extensions = [".jpg", ".jpeg", ".png"]
+ if heif_support_enabled:
+ supported_images_extensions += [".heic", ".heif"]
+ supported_images_extensions = tuple(supported_images_extensions)
+
+ imgs = []
+ for path in folder_content:
+ if not path.lower().endswith(supported_images_extensions):
+ continue
+ img = exif_transpose(PIL.Image.open(os.path.join(root, path))).convert("RGB")
+ if rotate_clockwise_90:
+ img = img.rotate(-90, expand=True)
+ if crop_to_landscape:
+ # Crop to a landscape aspect ratio (e.g., 16:9)
+ desired_aspect_ratio = 4 / 3
+ width, height = img.size
+ current_aspect_ratio = width / height
+
+ if current_aspect_ratio > desired_aspect_ratio:
+ # Wider than landscape: crop width
+ new_width = int(height * desired_aspect_ratio)
+ left = (width - new_width) // 2
+ right = left + new_width
+ top = 0
+ bottom = height
+ else:
+ # Taller than landscape: crop height
+ new_height = int(width / desired_aspect_ratio)
+ top = (height - new_height) // 2
+ bottom = top + new_height
+ left = 0
+ right = width
+
+ img = img.crop((left, top, right, bottom))
+
+ W1, H1 = img.size
+ if size == 224:
+ # resize short side to 224 (then crop)
+ img = _resize_pil_image(img, round(size * max(W1 / H1, H1 / W1)))
+ else:
+ # resize long side to 512
+ img = _resize_pil_image(img, size)
+ W, H = img.size
+ cx, cy = W // 2, H // 2
+ if size == 224:
+ half = min(cx, cy)
+ img = img.crop((cx - half, cy - half, cx + half, cy + half))
+ else:
+ halfw, halfh = ((2 * cx) // 16) * 8, ((2 * cy) // 16) * 8
+ if not (square_ok) and W == H:
+ halfh = 3 * halfw / 4
+ img = img.crop((cx - halfw, cy - halfh, cx + halfw, cy + halfh))
+
+ W2, H2 = img.size
+ if verbose:
+ print(f" - adding {path} with resolution {W1}x{H1} --> {W2}x{H2}")
+ imgs.append(
+ dict(
+ img=ImgNorm(img)[None],
+ true_shape=np.int32([img.size[::-1]]),
+ idx=len(imgs),
+ instance=str(len(imgs)),
+ )
+ )
+
+ assert imgs, "no images foud at " + root
+ if verbose:
+ print(f" (Found {len(imgs)} images)")
+ return imgs
+
+def process_image(img_path):
+ img = Image.open(img_path)
+ if img.mode == 'RGBA':
+ # Convert RGBA to RGB by removing alpha channel
+ img = img.convert('RGB')
+ # Resize to maintain aspect ratio and then center crop to 448x448
+ width, height = img.size
+ if width > height:
+ new_height = 448
+ new_width = int(width * (new_height / height))
+ else:
+ new_width = 448
+ new_height = int(height * (new_width / width))
+ img = img.resize((new_width, new_height))
+
+ # Center crop
+ left = (new_width - 448) // 2
+ top = (new_height - 448) // 2
+ right = left + 448
+ bottom = top + 448
+ img = img.crop((left, top, right, bottom))
+ img_tensor = torchvision.transforms.ToTensor()(img) * 2.0 - 1.0 # [-1, 1]
+ return img_tensor
\ No newline at end of file
diff --git a/src/utils/misc.py b/src/utils/misc.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c34192786d91aeb2d67e8b5a8cbf7412a1c682d
--- /dev/null
+++ b/src/utils/misc.py
@@ -0,0 +1,531 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# utilitary functions for CroCo
+# --------------------------------------------------------
+# References:
+# MAE: https://github.com/facebookresearch/mae
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+
+import builtins
+import datetime
+import json
+import math
+import os
+import time
+from collections import defaultdict, deque
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch import inf
+
+
+class SmoothedValue(object):
+ """Track a series of values and provide access to smoothed values over a
+ window or the global series average.
+ """
+
+ def __init__(self, window_size=20, fmt=None):
+ if fmt is None:
+ fmt = "{median:.4f} ({global_avg:.4f})"
+ self.deque = deque(maxlen=window_size)
+ self.total = 0.0
+ self.count = 0
+ self.fmt = fmt
+
+ def update(self, value, n=1):
+ self.deque.append(value)
+ self.count += n
+ self.total += value * n
+
+ def synchronize_between_processes(self):
+ """
+ Warning: does not synchronize the deque!
+ """
+ if not is_dist_avail_and_initialized():
+ return
+ t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+ dist.barrier()
+ dist.all_reduce(t)
+ t = t.tolist()
+ self.count = int(t[0])
+ self.total = t[1]
+
+ @property
+ def median(self):
+ d = torch.tensor(list(self.deque))
+ return d.median().item()
+
+ @property
+ def avg(self):
+ d = torch.tensor(list(self.deque), dtype=torch.float32)
+ return d.mean().item()
+
+ @property
+ def global_avg(self):
+ return self.total / self.count
+
+ @property
+ def max(self):
+ return max(self.deque)
+
+ @property
+ def value(self):
+ return self.deque[-1]
+
+ def __str__(self):
+ return self.fmt.format(
+ median=self.median,
+ avg=self.avg,
+ global_avg=self.global_avg,
+ max=self.max,
+ value=self.value,
+ )
+
+
+class MetricLogger(object):
+ def __init__(self, delimiter="\t"):
+ self.meters = defaultdict(SmoothedValue)
+ self.delimiter = delimiter
+
+ def update(self, **kwargs):
+ for k, v in kwargs.items():
+ if v is None:
+ continue
+ if isinstance(v, torch.Tensor):
+ v = v.item()
+ assert isinstance(v, (float, int))
+ self.meters[k].update(v)
+
+ def __getattr__(self, attr):
+ if attr in self.meters:
+ return self.meters[attr]
+ if attr in self.__dict__:
+ return self.__dict__[attr]
+ raise AttributeError(
+ "'{}' object has no attribute '{}'".format(type(self).__name__, attr)
+ )
+
+ def __str__(self):
+ loss_str = []
+ for name, meter in self.meters.items():
+ loss_str.append("{}: {}".format(name, str(meter)))
+ return self.delimiter.join(loss_str)
+
+ def synchronize_between_processes(self):
+ for meter in self.meters.values():
+ meter.synchronize_between_processes()
+
+ def add_meter(self, name, meter):
+ self.meters[name] = meter
+
+ def log_every(self, iterable, print_freq, header=None, max_iter=None):
+ i = 0
+ if not header:
+ header = ""
+ start_time = time.time()
+ end = time.time()
+ iter_time = SmoothedValue(fmt="{avg:.4f}")
+ data_time = SmoothedValue(fmt="{avg:.4f}")
+ len_iterable = min(len(iterable), max_iter) if max_iter else len(iterable)
+ space_fmt = ":" + str(len(str(len_iterable))) + "d"
+ log_msg = [
+ header,
+ "[{0" + space_fmt + "}/{1}]",
+ "eta: {eta}",
+ "{meters}",
+ "time: {time}",
+ "data: {data}",
+ ]
+ if torch.cuda.is_available():
+ log_msg.append("max mem: {memory:.0f}")
+ log_msg = self.delimiter.join(log_msg)
+ MB = 1024.0 * 1024.0
+ for it, obj in enumerate(iterable):
+ data_time.update(time.time() - end)
+ yield obj
+ iter_time.update(time.time() - end)
+ if i % print_freq == 0 or i == len_iterable - 1:
+ eta_seconds = iter_time.global_avg * (len_iterable - i)
+ eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+ if torch.cuda.is_available():
+ print(
+ log_msg.format(
+ i,
+ len_iterable,
+ eta=eta_string,
+ meters=str(self),
+ time=str(iter_time),
+ data=str(data_time),
+ memory=torch.cuda.max_memory_allocated() / MB,
+ )
+ )
+ else:
+ print(
+ log_msg.format(
+ i,
+ len_iterable,
+ eta=eta_string,
+ meters=str(self),
+ time=str(iter_time),
+ data=str(data_time),
+ )
+ )
+ i += 1
+ end = time.time()
+ if max_iter and it >= max_iter:
+ break
+ total_time = time.time() - start_time
+ total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+ print(
+ "{} Total time: {} ({:.4f} s / it)".format(
+ header, total_time_str, total_time / len_iterable
+ )
+ )
+
+
+def setup_for_distributed(is_master):
+ """
+ This function disables printing when not in master process
+ """
+ builtin_print = builtins.print
+
+ def print(*args, **kwargs):
+ force = kwargs.pop("force", False)
+ force = force or (get_world_size() > 8)
+ if is_master or force:
+ now = datetime.datetime.now().time()
+ builtin_print("[{}] ".format(now), end="") # print with time stamp
+ builtin_print(*args, **kwargs)
+
+ builtins.print = print
+
+
+def is_dist_avail_and_initialized():
+ if not dist.is_available():
+ return False
+ if not dist.is_initialized():
+ return False
+ return True
+
+
+def get_world_size():
+ if not is_dist_avail_and_initialized():
+ return 1
+ return dist.get_world_size()
+
+
+def get_rank():
+ if not is_dist_avail_and_initialized():
+ return 0
+ return dist.get_rank()
+
+
+def is_main_process():
+ return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+ if is_main_process():
+ torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+ nodist = args.nodist if hasattr(args, "nodist") else False
+ if "RANK" in os.environ and "WORLD_SIZE" in os.environ and not nodist:
+ args.rank = int(os.environ["RANK"])
+ args.world_size = int(os.environ["WORLD_SIZE"])
+ args.gpu = int(os.environ["LOCAL_RANK"])
+ else:
+ print("Not using distributed mode")
+ setup_for_distributed(is_master=True) # hack
+ args.distributed = False
+ return
+
+ args.distributed = True
+
+ torch.cuda.set_device(args.gpu)
+ args.dist_backend = "nccl"
+ print(
+ "| distributed init (rank {}): {}, gpu {}".format(
+ args.rank, args.dist_url, args.gpu
+ ),
+ flush=True,
+ )
+ torch.distributed.init_process_group(
+ backend=args.dist_backend,
+ init_method=args.dist_url,
+ world_size=args.world_size,
+ rank=args.rank,
+ )
+ torch.distributed.barrier()
+ setup_for_distributed(args.rank == 0)
+
+
+class NativeScalerWithGradNormCount:
+ state_dict_key = "amp_scaler"
+
+ def __init__(self, enabled=True):
+ self._scaler = torch.cuda.amp.GradScaler(enabled=enabled)
+
+ def __call__(
+ self,
+ loss,
+ optimizer,
+ clip_grad=None,
+ parameters=None,
+ create_graph=False,
+ update_grad=True,
+ ):
+ self._scaler.scale(loss).backward(create_graph=create_graph)
+ if update_grad:
+ if clip_grad is not None:
+ assert parameters is not None
+ self._scaler.unscale_(
+ optimizer
+ ) # unscale the gradients of optimizer's assigned params in-place
+ norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+ else:
+ self._scaler.unscale_(optimizer)
+ # TODO: FIXME: get_grad_norm_ is a poorly-implemented func that is very slow, and its return not even used
+ # norm = get_grad_norm_(parameters)
+ norm = None
+ self._scaler.step(optimizer)
+ self._scaler.update()
+ else:
+ norm = None
+ return norm
+
+ def state_dict(self):
+ return self._scaler.state_dict()
+
+ def load_state_dict(self, state_dict):
+ self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+ if isinstance(parameters, torch.Tensor):
+ parameters = [parameters]
+ parameters = [p for p in parameters if p.grad is not None]
+ norm_type = float(norm_type)
+ if len(parameters) == 0:
+ return torch.tensor(0.0)
+ device = parameters[0].grad.device
+ if norm_type == inf:
+ total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+ else:
+ total_norm = torch.norm(
+ torch.stack(
+ [torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]
+ ),
+ norm_type,
+ )
+ return total_norm
+
+
+def save_model(
+ args, epoch, model_without_ddp, optimizer, loss_scaler, fname=None, best_so_far=None
+):
+ output_dir = Path(args.output_dir)
+ if fname is None:
+ fname = str(epoch)
+ checkpoint_path = output_dir / ("checkpoint-%s.pth" % fname)
+ to_save = {
+ "model": model_without_ddp.state_dict(),
+ "optimizer": optimizer.state_dict(),
+ "scaler": loss_scaler.state_dict(),
+ "args": args,
+ "epoch": epoch,
+ }
+ if best_so_far is not None:
+ to_save["best_so_far"] = best_so_far
+ print(f">> Saving model to {checkpoint_path} ...")
+ save_on_master(to_save, checkpoint_path)
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler):
+ args.start_epoch = 0
+ best_so_far = None
+ if args.resume is not None:
+ if args.resume.startswith("https"):
+ checkpoint = torch.hub.load_state_dict_from_url(
+ args.resume, map_location="cpu", check_hash=True
+ )
+ else:
+ checkpoint = torch.load(args.resume, map_location="cpu")
+ print("Resume checkpoint %s" % args.resume)
+ model_without_ddp.load_state_dict(checkpoint["model"], strict=False)
+ args.start_epoch = checkpoint["epoch"] + 1
+ optimizer.load_state_dict(checkpoint["optimizer"])
+ if "scaler" in checkpoint:
+ loss_scaler.load_state_dict(checkpoint["scaler"])
+ if "best_so_far" in checkpoint:
+ best_so_far = checkpoint["best_so_far"]
+ print(" & best_so_far={:g}".format(best_so_far))
+ else:
+ print("")
+ print("With optim & sched! start_epoch={:d}".format(args.start_epoch), end="")
+ return best_so_far
+
+
+def all_reduce_mean(x):
+ world_size = get_world_size()
+ if world_size > 1:
+ x_reduce = torch.tensor(x).cuda()
+ dist.all_reduce(x_reduce)
+ x_reduce /= world_size
+ return x_reduce.item()
+ else:
+ return x
+
+
+def _replace(text, src, tgt, rm=""):
+ """Advanced string replacement.
+ Given a text:
+ - replace all elements in src by the corresponding element in tgt
+ - remove all elements in rm
+ """
+ if len(tgt) == 1:
+ tgt = tgt * len(src)
+ assert len(src) == len(tgt), f"'{src}' and '{tgt}' should have the same len"
+ for s, t in zip(src, tgt):
+ text = text.replace(s, t)
+ for c in rm:
+ text = text.replace(c, "")
+ return text
+
+
+def filename(obj):
+ """transform a python obj or cmd into a proper filename.
+ - \1 gets replaced by slash '/'
+ - \2 gets replaced by comma ','
+ """
+ if not isinstance(obj, str):
+ obj = repr(obj)
+ obj = str(obj).replace("()", "")
+ obj = _replace(obj, "_,(*/\1\2", "-__x%/,", rm=" )'\"")
+ assert all(len(s) < 256 for s in obj.split(os.sep)), (
+ "filename too long (>256 characters):\n" + obj
+ )
+ return obj
+
+
+def _get_num_layer_for_vit(var_name, enc_depth, dec_depth):
+ if var_name in ("cls_token", "mask_token", "pos_embed", "global_tokens"):
+ return 0
+ elif var_name.startswith("patch_embed"):
+ return 0
+ elif var_name.startswith("enc_blocks"):
+ layer_id = int(var_name.split(".")[1])
+ return layer_id + 1
+ elif var_name.startswith("decoder_embed") or var_name.startswith(
+ "enc_norm"
+ ): # part of the last black
+ return enc_depth
+ elif var_name.startswith("dec_blocks"):
+ layer_id = int(var_name.split(".")[1])
+ return enc_depth + layer_id + 1
+ elif var_name.startswith("dec_norm"): # part of the last block
+ return enc_depth + dec_depth
+ elif any(var_name.startswith(k) for k in ["head", "prediction_head"]):
+ return enc_depth + dec_depth + 1
+ else:
+ raise NotImplementedError(var_name)
+
+
+def get_parameter_groups(
+ model, weight_decay, layer_decay=1.0, skip_list=(), no_lr_scale_list=[]
+):
+ parameter_group_names = {}
+ parameter_group_vars = {}
+ enc_depth, dec_depth = None, None
+ # prepare layer decay values
+ assert layer_decay == 1.0 or 0.0 < layer_decay < 1.0
+ if layer_decay < 1.0:
+ enc_depth = model.enc_depth
+ dec_depth = model.dec_depth if hasattr(model, "dec_blocks") else 0
+ num_layers = enc_depth + dec_depth
+ layer_decay_values = list(
+ layer_decay ** (num_layers + 1 - i) for i in range(num_layers + 2)
+ )
+
+ for name, param in model.named_parameters():
+ if not param.requires_grad:
+ continue # frozen weights
+
+ # Assign weight decay values
+ if len(param.shape) == 1 or name.endswith(".bias") or name in skip_list:
+ group_name = "no_decay"
+ this_weight_decay = 0.0
+ else:
+ group_name = "decay"
+ this_weight_decay = weight_decay
+
+ # Assign layer ID for LR scaling
+ if layer_decay < 1.0:
+ skip_scale = False
+ layer_id = _get_num_layer_for_vit(name, enc_depth, dec_depth)
+ group_name = "layer_%d_%s" % (layer_id, group_name)
+ if name in no_lr_scale_list:
+ skip_scale = True
+ group_name = f"{group_name}_no_lr_scale"
+ else:
+ layer_id = 0
+ skip_scale = True
+
+ if group_name not in parameter_group_names:
+ if not skip_scale:
+ scale = layer_decay_values[layer_id]
+ else:
+ scale = 1.0
+
+ parameter_group_names[group_name] = {
+ "weight_decay": this_weight_decay,
+ "params": [],
+ "lr_scale": scale,
+ }
+ parameter_group_vars[group_name] = {
+ "weight_decay": this_weight_decay,
+ "params": [],
+ "lr_scale": scale,
+ }
+
+ parameter_group_vars[group_name]["params"].append(param)
+ parameter_group_names[group_name]["params"].append(name)
+ print("Param groups = %s" % json.dumps(parameter_group_names, indent=2))
+ return list(parameter_group_vars.values())
+
+
+def adjust_learning_rate(optimizer, epoch, args):
+ """Decay the learning rate with half-cycle cosine after warmup"""
+
+ if epoch < args.warmup_epochs:
+ lr = args.lr * epoch / args.warmup_epochs
+ else:
+ lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * (
+ 1.0
+ + math.cos(
+ math.pi
+ * (epoch - args.warmup_epochs)
+ / (args.epochs - args.warmup_epochs)
+ )
+ )
+
+ for param_group in optimizer.param_groups:
+ if "lr_scale" in param_group:
+ param_group["lr"] = lr * param_group["lr_scale"]
+ else:
+ param_group["lr"] = lr
+
+ return lr
diff --git a/src/utils/point.py b/src/utils/point.py
new file mode 100644
index 0000000000000000000000000000000000000000..5eb126006427139c38ca607fcc650f2427567e6d
--- /dev/null
+++ b/src/utils/point.py
@@ -0,0 +1,47 @@
+import torch
+from torch import Tensor
+
+def get_normal_map(depth_map: torch.Tensor, intrinsic: torch.Tensor) -> torch.Tensor:
+ """
+ Convert a depth map to camera coordinates.
+
+ Args:
+ depth_map (torch.Tensor): Depth map of shape (H, W).
+ intrinsic (torch.Tensor): Camera intrinsic matrix of shape (3, 3).
+
+ Returns:
+ tuple[torch.Tensor, torch.Tensor]: Camera coordinates (H, W, 3)
+ """
+ B, H, W = depth_map.shape
+ assert intrinsic.shape == (B, 3, 3), "Intrinsic matrix must be Bx3x3"
+ assert (intrinsic[:, 0, 1] == 0).all() and (intrinsic[:, 1, 0] == 0).all(), "Intrinsic matrix must have zero skew"
+
+ # Intrinsic parameters
+ fu = intrinsic[:, 0, 0] * W # (B,)
+ fv = intrinsic[:, 1, 1] * H # (B,)
+ cu = intrinsic[:, 0, 2] * W # (B,)
+ cv = intrinsic[:, 1, 2] * H # (B,)
+
+ # Generate grid of pixel coordinates
+ u = torch.arange(W, device=depth_map.device)[None, None, :].expand(B, H, W)
+ v = torch.arange(H, device=depth_map.device)[None, :, None].expand(B, H, W)
+
+ # Unproject to camera coordinates (B, H, W)
+ x_cam = (u - cu[:, None, None]) * depth_map / fu[:, None, None]
+ y_cam = (v - cv[:, None, None]) * depth_map / fv[:, None, None]
+ z_cam = depth_map
+
+ # Stack to form camera coordinates (B, H, W, 3)
+ cam_coords = torch.stack((x_cam, y_cam, z_cam), dim=-1).to(dtype=torch.float32)
+
+ output = torch.zeros_like(cam_coords)
+ # Calculate dx using batch dimension (B, H-2, W-2, 3)
+ dx = cam_coords[:, 2:, 1:-1] - cam_coords[:, :-2, 1:-1]
+ # Calculate dy using batch dimension (B, H-2, W-2, 3)
+ dy = cam_coords[:, 1:-1, 2:] - cam_coords[:, 1:-1, :-2]
+ # Cross product and normalization (B, H-2, W-2, 3)
+ normal_map = torch.nn.functional.normalize(torch.cross(dx, dy, dim=-1), dim=-1)
+ # Assign the computed normal map to the output tensor
+ output[:, 1:-1, 1:-1, :] = normal_map
+
+ return output
\ No newline at end of file
diff --git a/src/utils/pose.py b/src/utils/pose.py
new file mode 100644
index 0000000000000000000000000000000000000000..2bdba20257c1922a4cd3241c152c23a2382bdb9d
--- /dev/null
+++ b/src/utils/pose.py
@@ -0,0 +1,218 @@
+import torch
+import numpy as np
+from src.model.encoder.vggt.utils.rotation import mat_to_quat
+from src.model.encoder.vggt.utils.geometry import closed_form_inverse_se3, unproject_depth_map_to_point_map
+
+
+def convert_pt3d_RT_to_opencv(Rot, Trans):
+ """
+ Convert Point3D extrinsic matrices to OpenCV convention.
+
+ Args:
+ Rot: 3D rotation matrix in Point3D format
+ Trans: 3D translation vector in Point3D format
+
+ Returns:
+ extri_opencv: 3x4 extrinsic matrix in OpenCV format
+ """
+ rot_pt3d = np.array(Rot)
+ trans_pt3d = np.array(Trans)
+
+ trans_pt3d[:2] *= -1
+ rot_pt3d[:, :2] *= -1
+ rot_pt3d = rot_pt3d.transpose(1, 0)
+ extri_opencv = np.hstack((rot_pt3d, trans_pt3d[:, None]))
+ return extri_opencv
+
+
+def build_pair_index(N, B=1):
+ """
+ Build indices for all possible pairs of frames.
+
+ Args:
+ N: Number of frames
+ B: Batch size
+
+ Returns:
+ i1, i2: Indices for all possible pairs
+ """
+ i1_, i2_ = torch.combinations(torch.arange(N), 2, with_replacement=False).unbind(-1)
+ i1, i2 = [(i[None] + torch.arange(B)[:, None] * N).reshape(-1) for i in [i1_, i2_]]
+ return i1, i2
+
+
+def rotation_angle(rot_gt, rot_pred, batch_size=None, eps=1e-15):
+ """
+ Calculate rotation angle error between ground truth and predicted rotations.
+
+ Args:
+ rot_gt: Ground truth rotation matrices
+ rot_pred: Predicted rotation matrices
+ batch_size: Batch size for reshaping the result
+ eps: Small value to avoid numerical issues
+
+ Returns:
+ Rotation angle error in degrees
+ """
+ q_pred = mat_to_quat(rot_pred)
+ q_gt = mat_to_quat(rot_gt)
+
+ loss_q = (1 - (q_pred * q_gt).sum(dim=1) ** 2).clamp(min=eps)
+ err_q = torch.arccos(1 - 2 * loss_q)
+
+ rel_rangle_deg = err_q * 180 / np.pi
+
+ if batch_size is not None:
+ rel_rangle_deg = rel_rangle_deg.reshape(batch_size, -1)
+
+ return rel_rangle_deg
+
+
+def translation_angle(tvec_gt, tvec_pred, batch_size=None, ambiguity=True):
+ """
+ Calculate translation angle error between ground truth and predicted translations.
+
+ Args:
+ tvec_gt: Ground truth translation vectors
+ tvec_pred: Predicted translation vectors
+ batch_size: Batch size for reshaping the result
+ ambiguity: Whether to handle direction ambiguity
+
+ Returns:
+ Translation angle error in degrees
+ """
+ rel_tangle_deg = compare_translation_by_angle(tvec_gt, tvec_pred)
+ rel_tangle_deg = rel_tangle_deg * 180.0 / np.pi
+
+ if ambiguity:
+ rel_tangle_deg = torch.min(rel_tangle_deg, (180 - rel_tangle_deg).abs())
+
+ if batch_size is not None:
+ rel_tangle_deg = rel_tangle_deg.reshape(batch_size, -1)
+
+ return rel_tangle_deg
+
+
+def compare_translation_by_angle(t_gt, t, eps=1e-15, default_err=1e6):
+ """
+ Normalize the translation vectors and compute the angle between them.
+
+ Args:
+ t_gt: Ground truth translation vectors
+ t: Predicted translation vectors
+ eps: Small value to avoid division by zero
+ default_err: Default error value for invalid cases
+
+ Returns:
+ Angular error between translation vectors in radians
+ """
+ t_norm = torch.norm(t, dim=1, keepdim=True)
+ t = t / (t_norm + eps)
+
+ t_gt_norm = torch.norm(t_gt, dim=1, keepdim=True)
+ t_gt = t_gt / (t_gt_norm + eps)
+
+ loss_t = torch.clamp_min(1.0 - torch.sum(t * t_gt, dim=1) ** 2, eps)
+ err_t = torch.acos(torch.sqrt(1 - loss_t))
+
+ err_t[torch.isnan(err_t) | torch.isinf(err_t)] = default_err
+ return err_t
+
+
+def calculate_auc(r_error, t_error, max_threshold=30, return_list=False):
+ """
+ Calculate the Area Under the Curve (AUC) for the given error arrays using PyTorch.
+
+ Args:
+ r_error: torch.Tensor representing R error values (Degree)
+ t_error: torch.Tensor representing T error values (Degree)
+ max_threshold: Maximum threshold value for binning the histogram
+ return_list: Whether to return the normalized histogram as well
+
+ Returns:
+ AUC value, and optionally the normalized histogram
+ """
+ error_matrix = torch.stack((r_error, t_error), dim=1)
+ max_errors, _ = torch.max(error_matrix, dim=1)
+ histogram = torch.histc(
+ max_errors, bins=max_threshold + 1, min=0, max=max_threshold
+ )
+ num_pairs = float(max_errors.size(0))
+ normalized_histogram = histogram / num_pairs
+
+ if return_list:
+ return (
+ torch.cumsum(normalized_histogram, dim=0).mean(),
+ normalized_histogram,
+ )
+ return torch.cumsum(normalized_histogram, dim=0).mean()
+
+
+def calculate_auc_np(r_error, t_error, max_threshold=30):
+ """
+ Calculate the Area Under the Curve (AUC) for the given error arrays using NumPy.
+
+ Args:
+ r_error: numpy array representing R error values (Degree)
+ t_error: numpy array representing T error values (Degree)
+ max_threshold: Maximum threshold value for binning the histogram
+
+ Returns:
+ AUC value and the normalized histogram
+ """
+ error_matrix = np.concatenate((r_error[:, None], t_error[:, None]), axis=1)
+ max_errors = np.max(error_matrix, axis=1)
+ bins = np.arange(max_threshold + 1)
+ histogram, _ = np.histogram(max_errors, bins=bins)
+ num_pairs = float(len(max_errors))
+ normalized_histogram = histogram.astype(float) / num_pairs
+ return np.mean(np.cumsum(normalized_histogram)), normalized_histogram
+
+
+def se3_to_relative_pose_error(pred_se3, gt_se3, num_frames):
+ """
+ Compute rotation and translation errors between predicted and ground truth poses.
+
+ Args:
+ pred_se3: Predicted SE(3) transformations
+ gt_se3: Ground truth SE(3) transformations
+ num_frames: Number of frames
+
+ Returns:
+ Rotation and translation angle errors in degrees
+ """
+ pair_idx_i1, pair_idx_i2 = build_pair_index(num_frames)
+
+ # Compute relative camera poses between pairs
+ # We use closed_form_inverse to avoid potential numerical loss by torch.inverse()
+ relative_pose_gt = closed_form_inverse_se3(gt_se3[pair_idx_i1]).bmm(
+ gt_se3[pair_idx_i2]
+ )
+ relative_pose_pred = closed_form_inverse_se3(pred_se3[pair_idx_i1]).bmm(
+ pred_se3[pair_idx_i2]
+ )
+
+ # Compute the difference in rotation and translation
+ rel_rangle_deg = rotation_angle(
+ relative_pose_gt[:, :3, :3], relative_pose_pred[:, :3, :3]
+ )
+ rel_tangle_deg = translation_angle(
+ relative_pose_gt[:, :3, 3], relative_pose_pred[:, :3, 3]
+ )
+
+ return rel_rangle_deg, rel_tangle_deg
+
+
+def align_to_first_camera(camera_poses):
+ """
+ Align all camera poses to the first camera's coordinate frame.
+
+ Args:
+ camera_poses: Tensor of shape (N, 4, 4) containing camera poses as SE3 transformations
+
+ Returns:
+ Tensor of shape (N, 4, 4) containing aligned camera poses
+ """
+ first_cam_extrinsic_inv = closed_form_inverse_se3(camera_poses[0][None])
+ aligned_poses = torch.matmul(camera_poses, first_cam_extrinsic_inv)
+ return aligned_poses
\ No newline at end of file
diff --git a/src/utils/render.py b/src/utils/render.py
new file mode 100644
index 0000000000000000000000000000000000000000..63e576230d33ab76b5dba10916c768b098ca83fb
--- /dev/null
+++ b/src/utils/render.py
@@ -0,0 +1,284 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import os
+import enum
+import types
+from typing import List, Mapping, Optional, Text, Tuple, Union
+import copy
+from PIL import Image
+# import mediapy as media
+from matplotlib import cm
+from tqdm import tqdm
+
+import torch
+
+def normalize(x: np.ndarray) -> np.ndarray:
+ """Normalization helper function."""
+ return x / np.linalg.norm(x)
+
+def pad_poses(p: np.ndarray) -> np.ndarray:
+ """Pad [..., 3, 4] pose matrices with a homogeneous bottom row [0,0,0,1]."""
+ bottom = np.broadcast_to([0, 0, 0, 1.], p[..., :1, :4].shape)
+ return np.concatenate([p[..., :3, :4], bottom], axis=-2)
+
+
+def unpad_poses(p: np.ndarray) -> np.ndarray:
+ """Remove the homogeneous bottom row from [..., 4, 4] pose matrices."""
+ return p[..., :3, :4]
+
+
+def recenter_poses(poses: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+ """Recenter poses around the origin."""
+ cam2world = average_pose(poses)
+ transform = np.linalg.inv(pad_poses(cam2world))
+ poses = transform @ pad_poses(poses)
+ return unpad_poses(poses), transform
+
+
+def average_pose(poses: np.ndarray) -> np.ndarray:
+ """New pose using average position, z-axis, and up vector of input poses."""
+ position = poses[:, :3, 3].mean(0)
+ z_axis = poses[:, :3, 2].mean(0)
+ up = poses[:, :3, 1].mean(0)
+ cam2world = viewmatrix(z_axis, up, position)
+ return cam2world
+
+def viewmatrix(lookdir: np.ndarray, up: np.ndarray,
+ position: np.ndarray) -> np.ndarray:
+ """Construct lookat view matrix."""
+ vec2 = normalize(lookdir)
+ vec0 = normalize(np.cross(up, vec2))
+ vec1 = normalize(np.cross(vec2, vec0))
+ m = np.stack([vec0, vec1, vec2, position], axis=1)
+ return m
+
+def focus_point_fn(poses: np.ndarray) -> np.ndarray:
+ """Calculate nearest point to all focal axes in poses."""
+ directions, origins = poses[:, :3, 2:3], poses[:, :3, 3:4]
+ m = np.eye(3) - directions * np.transpose(directions, [0, 2, 1])
+ mt_m = np.transpose(m, [0, 2, 1]) @ m
+ focus_pt = np.linalg.inv(mt_m.mean(0)) @ (mt_m @ origins).mean(0)[:, 0]
+ return focus_pt
+
+def transform_poses_pca(poses: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+ """Transforms poses so principal components lie on XYZ axes.
+
+ Args:
+ poses: a (N, 3, 4) array containing the cameras' camera to world transforms.
+
+ Returns:
+ A tuple (poses, transform), with the transformed poses and the applied
+ camera_to_world transforms.
+ """
+ t = poses[:, :3, 3]
+ t_mean = t.mean(axis=0)
+ t = t - t_mean
+
+ eigval, eigvec = np.linalg.eig(t.T @ t)
+ # Sort eigenvectors in order of largest to smallest eigenvalue.
+ inds = np.argsort(eigval)[::-1]
+ eigvec = eigvec[:, inds]
+ rot = eigvec.T
+ if np.linalg.det(rot) < 0:
+ rot = np.diag(np.array([1, 1, -1])) @ rot
+
+ transform = np.concatenate([rot, rot @ -t_mean[:, None]], -1)
+ poses_recentered = unpad_poses(transform @ pad_poses(poses))
+ transform = np.concatenate([transform, np.eye(4)[3:]], axis=0)
+
+ # Flip coordinate system if z component of y-axis is negative
+ if poses_recentered.mean(axis=0)[2, 1] < 0:
+ poses_recentered = np.diag(np.array([1, -1, -1])) @ poses_recentered
+ transform = np.diag(np.array([1, -1, -1, 1])) @ transform
+
+ return poses_recentered, transform
+ # points = np.random.rand(3,100)
+ # points_h = np.concatenate((points,np.ones_like(points[:1])), axis=0)
+ # (poses_recentered @ points_h)[0]
+ # (transform @ pad_poses(poses) @ points_h)[0,:3]
+ # import pdb; pdb.set_trace()
+
+ # # Just make sure it's it in the [-1, 1]^3 cube
+ # scale_factor = 1. / np.max(np.abs(poses_recentered[:, :3, 3]))
+ # poses_recentered[:, :3, 3] *= scale_factor
+ # transform = np.diag(np.array([scale_factor] * 3 + [1])) @ transform
+
+ # return poses_recentered, transform
+
+def generate_ellipse_path(poses: np.ndarray,
+ n_frames: int = 120,
+ const_speed: bool = True,
+ z_variation: float = 0.,
+ z_phase: float = 0.) -> np.ndarray:
+ """Generate an elliptical render path based on the given poses."""
+ # Calculate the focal point for the path (cameras point toward this).
+ center = focus_point_fn(poses)
+ # Path height sits at z=0 (in middle of zero-mean capture pattern).
+ offset = np.array([center[0], center[1], 0])
+
+ # Calculate scaling for ellipse axes based on input camera positions.
+ sc = np.percentile(np.abs(poses[:, :3, 3] - offset), 90, axis=0)
+ # Use ellipse that is symmetric about the focal point in xy.
+ low = -sc + offset
+ high = sc + offset
+ # Optional height variation need not be symmetric
+ z_low = np.percentile((poses[:, :3, 3]), 10, axis=0)
+ z_high = np.percentile((poses[:, :3, 3]), 90, axis=0)
+
+ def get_positions(theta):
+ # Interpolate between bounds with trig functions to get ellipse in x-y.
+ # Optionally also interpolate in z to change camera height along path.
+ return np.stack([
+ low[0] + (high - low)[0] * (np.cos(theta) * .5 + .5),
+ low[1] + (high - low)[1] * (np.sin(theta) * .5 + .5),
+ z_variation * (z_low[2] + (z_high - z_low)[2] *
+ (np.cos(theta + 2 * np.pi * z_phase) * .5 + .5)),
+ ], -1)
+
+ theta = np.linspace(0, 2. * np.pi, n_frames + 1, endpoint=True)
+ positions = get_positions(theta)
+
+ #if const_speed:
+
+ # # Resample theta angles so that the velocity is closer to constant.
+ # lengths = np.linalg.norm(positions[1:] - positions[:-1], axis=-1)
+ # theta = stepfun.sample(None, theta, np.log(lengths), n_frames + 1)
+ # positions = get_positions(theta)
+
+ # Throw away duplicated last position.
+ positions = positions[:-1]
+
+ # Set path's up vector to axis closest to average of input pose up vectors.
+ avg_up = poses[:, :3, 1].mean(0)
+ avg_up = avg_up / np.linalg.norm(avg_up)
+ ind_up = np.argmax(np.abs(avg_up))
+ up = np.eye(3)[ind_up] * np.sign(avg_up[ind_up])
+
+ return np.stack([viewmatrix(p - center, up, p) for p in positions])
+
+
+def generate_path(viewpoint_cameras, n_frames=480):
+ # c2ws = np.array([np.linalg.inv(np.asarray((cam.world_view_transform.T).cpu().numpy())) for cam in viewpoint_cameras])
+ c2ws = viewpoint_cameras.cpu().numpy()
+ pose = c2ws[:,:3,:] @ np.diag([1, -1, -1, 1])
+ pose_recenter, colmap_to_world_transform = transform_poses_pca(pose)
+
+ # generate new poses
+ new_poses = generate_ellipse_path(poses=pose_recenter, n_frames=n_frames)
+ # warp back to orignal scale
+ new_poses = np.linalg.inv(colmap_to_world_transform) @ pad_poses(new_poses)
+
+ return new_poses
+
+ # traj = []
+ # for c2w in new_poses:
+ # c2w = c2w @ np.diag([1, -1, -1, 1])
+ # cam = copy.deepcopy(viewpoint_cameras[0])
+ # cam.image_height = int(cam.image_height / 2) * 2
+ # cam.image_width = int(cam.image_width / 2) * 2
+ # cam.world_view_transform = torch.from_numpy(np.linalg.inv(c2w).T).float().cuda()
+ # cam.full_proj_transform = (cam.world_view_transform.unsqueeze(0).bmm(cam.projection_matrix.unsqueeze(0))).squeeze(0)
+ # cam.camera_center = cam.world_view_transform.inverse()[3, :3]
+ # traj.append(cam)
+
+ # return traj
+
+def load_img(pth: str) -> np.ndarray:
+ """Load an image and cast to float32."""
+ with open(pth, 'rb') as f:
+ image = np.array(Image.open(f), dtype=np.float32)
+ return image
+
+
+def create_videos(base_dir, input_dir, out_name, num_frames=480):
+ """Creates videos out of the images saved to disk."""
+ # Last two parts of checkpoint path are experiment name and scene name.
+ video_prefix = f'{out_name}'
+
+ zpad = max(5, len(str(num_frames - 1)))
+ idx_to_str = lambda idx: str(idx).zfill(zpad)
+
+ os.makedirs(base_dir, exist_ok=True)
+ render_dist_curve_fn = np.log
+
+ # Load one example frame to get image shape and depth range.
+ depth_file = os.path.join(input_dir, 'vis', f'depth_{idx_to_str(0)}.tiff')
+ depth_frame = load_img(depth_file)
+ shape = depth_frame.shape
+ p = 3
+ distance_limits = np.percentile(depth_frame.flatten(), [p, 100 - p])
+ lo, hi = [render_dist_curve_fn(x) for x in distance_limits]
+ print(f'Video shape is {shape[:2]}')
+
+ video_kwargs = {
+ 'shape': shape[:2],
+ 'codec': 'h264',
+ 'fps': 60,
+ 'crf': 18,
+ }
+
+ for k in ['depth', 'normal', 'color']:
+ video_file = os.path.join(base_dir, f'{video_prefix}_{k}.mp4')
+ input_format = 'gray' if k == 'alpha' else 'rgb'
+
+
+ file_ext = 'png' if k in ['color', 'normal'] else 'tiff'
+ idx = 0
+
+ if k == 'color':
+ file0 = os.path.join(input_dir, 'renders', f'{idx_to_str(0)}.{file_ext}')
+ else:
+ file0 = os.path.join(input_dir, 'vis', f'{k}_{idx_to_str(0)}.{file_ext}')
+
+ if not os.path.exists(file0):
+ print(f'Images missing for tag {k}')
+ continue
+ print(f'Making video {video_file}...')
+ with media.VideoWriter(
+ video_file, **video_kwargs, input_format=input_format) as writer:
+ for idx in tqdm(range(num_frames)):
+ # img_file = os.path.join(input_dir, f'{k}_{idx_to_str(idx)}.{file_ext}')
+ if k == 'color':
+ img_file = os.path.join(input_dir, 'renders', f'{idx_to_str(idx)}.{file_ext}')
+ else:
+ img_file = os.path.join(input_dir, 'vis', f'{k}_{idx_to_str(idx)}.{file_ext}')
+
+ if not os.path.exists(img_file):
+ ValueError(f'Image file {img_file} does not exist.')
+ img = load_img(img_file)
+ if k in ['color', 'normal']:
+ img = img / 255.
+ elif k.startswith('depth'):
+ img = render_dist_curve_fn(img)
+ img = np.clip((img - np.minimum(lo, hi)) / np.abs(hi - lo), 0, 1)
+ img = cm.get_cmap('turbo')(img)[..., :3]
+
+ frame = (np.clip(np.nan_to_num(img), 0., 1.) * 255.).astype(np.uint8)
+ writer.add_image(frame)
+ idx += 1
+
+def save_img_u8(img, pth):
+ """Save an image (probably RGB) in [0, 1] to disk as a uint8 PNG."""
+ with open(pth, 'wb') as f:
+ Image.fromarray(
+ (np.clip(np.nan_to_num(img), 0., 1.) * 255.).astype(np.uint8)).save(
+ f, 'PNG')
+
+
+def save_img_f32(depthmap, pth):
+ """Save an image (probably a depthmap) to disk as a float32 TIFF."""
+ with open(pth, 'wb') as f:
+ Image.fromarray(np.nan_to_num(depthmap).astype(np.float32)).save(f, 'TIFF')
\ No newline at end of file
diff --git a/src/utils/tensor_to_pycolmap.py b/src/utils/tensor_to_pycolmap.py
new file mode 100644
index 0000000000000000000000000000000000000000..1332f6617b933134557555c99c1b040eb1569514
--- /dev/null
+++ b/src/utils/tensor_to_pycolmap.py
@@ -0,0 +1,320 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+import pycolmap
+
+# TODO: frame_idx should start from 1 instead of 0 in colmap
+def batch_matrix_to_pycolmap(
+ points3d,
+ extrinsics,
+ intrinsics,
+ tracks,
+ image_size,
+ masks=None,
+ max_reproj_error=None,
+ max_points3D_val=3000,
+ shared_camera=False,
+ camera_type="SIMPLE_PINHOLE",
+ extra_params=None,
+):
+ """
+ Convert Batched Pytorch Tensors to PyCOLMAP
+
+ Check https://github.com/colmap/pycolmap for more details about its format
+ """
+
+ # points3d: Px3
+ # extrinsics: Nx3x4
+ # intrinsics: Nx3x3
+ # tracks: NxPx2
+ # masks: NxP
+ # image_size: 2, assume all the frames have been padded to the same size
+ # where N is the number of frames and P is the number of tracks
+
+ N, P, _ = tracks.shape
+ assert len(extrinsics) == N
+ assert len(intrinsics) == N
+ assert len(points3d) == P
+ assert image_size.shape[0] == 2
+
+ projected_points_2d, projected_points_cam = project_3D_points(points3d, extrinsics, intrinsics, return_points_cam=True)
+ projected_diff = (projected_points_2d - tracks).norm(dim=-1)
+ projected_points_2d[projected_points_cam[:, -1] <= 0] = 1e6
+ reproj_mask = projected_diff < max_reproj_error
+
+ if masks is not None:
+ masks = torch.logical_and(masks, reproj_mask)
+ else:
+ masks = reproj_mask
+
+ extrinsics = extrinsics.cpu().numpy()
+ intrinsics = intrinsics.cpu().numpy()
+
+ if extra_params is not None:
+ extra_params = extra_params.cpu().numpy()
+
+
+ tracks = tracks.cpu().numpy()
+ points3d = points3d.cpu().numpy()
+ image_size = image_size.cpu().numpy()
+
+ # Reconstruction object, following the format of PyCOLMAP/COLMAP
+ reconstruction = pycolmap.Reconstruction()
+
+ masks = masks.cpu().numpy()
+
+ inlier_num = masks.sum(0)
+ valid_mask = inlier_num >= 2 # a track is invalid if without two inliers
+ valid_idx = np.nonzero(valid_mask)[0]
+
+ # Only add 3D points that have sufficient 2D points
+ for vidx in valid_idx:
+ reconstruction.add_point3D(
+ points3d[vidx], pycolmap.Track(), np.zeros(3)
+ )
+
+ num_points3D = len(valid_idx)
+
+ camera = None
+ # frame idx
+ for fidx in range(N):
+ # set camera
+ if camera is None or (not shared_camera):
+ if camera_type == "SIMPLE_RADIAL":
+ focal = (intrinsics[fidx][0, 0] + intrinsics[fidx][1, 1]) / 2
+ pycolmap_intri = np.array(
+ [
+ focal,
+ intrinsics[fidx][0, 2],
+ intrinsics[fidx][1, 2],
+ extra_params[fidx][0],
+ ]
+ )
+ elif camera_type == "SIMPLE_PINHOLE":
+ focal = (intrinsics[fidx][0, 0] + intrinsics[fidx][1, 1]) / 2
+ pycolmap_intri = np.array(
+ [
+ focal,
+ intrinsics[fidx][0, 2],
+ intrinsics[fidx][1, 2],
+ ]
+ )
+ else:
+ raise ValueError(
+ f"Camera type {camera_type} is not supported yet"
+ )
+
+ camera = pycolmap.Camera(
+ model=camera_type,
+ width=image_size[0],
+ height=image_size[1],
+ params=pycolmap_intri,
+ camera_id=fidx,
+ )
+
+ # add camera
+ reconstruction.add_camera(camera)
+
+ # set image
+ cam_from_world = pycolmap.Rigid3d(
+ pycolmap.Rotation3d(extrinsics[fidx][:3, :3]),
+ extrinsics[fidx][:3, 3],
+ ) # Rot and Trans
+ image = pycolmap.Image(
+ id=fidx,
+ name=f"image_{fidx}",
+ camera_id=camera.camera_id,
+ cam_from_world=cam_from_world,
+ )
+
+ points2D_list = []
+
+ point2D_idx = 0
+ # NOTE point3D_id start by 1
+ for point3D_id in range(1, num_points3D + 1):
+ original_track_idx = valid_idx[point3D_id - 1]
+
+ if (
+ reconstruction.points3D[point3D_id].xyz < max_points3D_val
+ ).all():
+ if masks[fidx][original_track_idx]:
+ # It seems we don't need +0.5 for BA
+ point2D_xy = tracks[fidx][original_track_idx]
+ # Please note when adding the Point2D object
+ # It not only requires the 2D xy location, but also the id to 3D point
+ points2D_list.append(
+ pycolmap.Point2D(point2D_xy, point3D_id)
+ )
+
+ # add element
+ track = reconstruction.points3D[point3D_id].track
+ track.add_element(fidx, point2D_idx)
+ point2D_idx += 1
+
+ assert point2D_idx == len(points2D_list)
+
+ try:
+ image.points2D = pycolmap.ListPoint2D(points2D_list)
+ image.registered = True
+ except:
+ print(f"frame {fidx} is out of BA")
+ image.registered = False
+
+ # add image
+ reconstruction.add_image(image)
+
+ return reconstruction
+
+
+def pycolmap_to_batch_matrix(
+ reconstruction, device="cuda", camera_type="SIMPLE_PINHOLE"
+):
+ """
+ Convert a PyCOLMAP Reconstruction Object to batched PyTorch tensors.
+
+ Args:
+ reconstruction (pycolmap.Reconstruction): The reconstruction object from PyCOLMAP.
+ device (str): The device to place the tensors on (default: "cuda").
+ camera_type (str): The type of camera model used (default: "SIMPLE_PINHOLE").
+
+ Returns:
+ tuple: A tuple containing points3D, extrinsics, intrinsics, and optionally extra_params.
+ """
+
+ num_images = len(reconstruction.images)
+ max_points3D_id = max(reconstruction.point3D_ids())
+ points3D = np.zeros((max_points3D_id, 3))
+
+ for point3D_id in reconstruction.points3D:
+ points3D[point3D_id - 1] = reconstruction.points3D[point3D_id].xyz
+ points3D = torch.from_numpy(points3D).to(device)
+
+ extrinsics = []
+ intrinsics = []
+
+ extra_params = [] if camera_type == "SIMPLE_RADIAL" else None
+
+ for i in range(num_images):
+ # Extract and append extrinsics
+ pyimg = reconstruction.images[i]
+ pycam = reconstruction.cameras[pyimg.camera_id]
+ matrix = pyimg.cam_from_world.matrix()
+ extrinsics.append(matrix)
+
+ # Extract and append intrinsics
+ calibration_matrix = pycam.calibration_matrix()
+ intrinsics.append(calibration_matrix)
+
+ if camera_type == "SIMPLE_RADIAL":
+ extra_params.append(pycam.params[-1])
+
+ # Convert lists to torch tensors
+ extrinsics = torch.from_numpy(np.stack(extrinsics)).to(device)
+
+ intrinsics = torch.from_numpy(np.stack(intrinsics)).to(device)
+
+ if camera_type == "SIMPLE_RADIAL":
+ extra_params = torch.from_numpy(np.stack(extra_params)).to(device)
+ extra_params = extra_params[:, None]
+
+ return points3D, extrinsics, intrinsics, extra_params
+
+
+
+
+
+def project_3D_points(
+ points3D,
+ extrinsics,
+ intrinsics=None,
+ extra_params=None,
+ return_points_cam=False,
+ default=0,
+ only_points_cam=False,
+):
+ """
+ Transforms 3D points to 2D using extrinsic and intrinsic parameters.
+ Args:
+ points3D (torch.Tensor): 3D points of shape Px3.
+ extrinsics (torch.Tensor): Extrinsic parameters of shape Bx3x4.
+ intrinsics (torch.Tensor): Intrinsic parameters of shape Bx3x3.
+ extra_params (torch.Tensor): Extra parameters of shape BxN, which is used for radial distortion.
+ Returns:
+ torch.Tensor: Transformed 2D points of shape BxNx2.
+ """
+ with torch.cuda.amp.autocast(dtype=torch.double):
+ N = points3D.shape[0] # Number of points
+ B = extrinsics.shape[0] # Batch size, i.e., number of cameras
+ points3D_homogeneous = torch.cat(
+ [points3D, torch.ones_like(points3D[..., 0:1])], dim=1
+ ) # Nx4
+ # Reshape for batch processing
+ points3D_homogeneous = points3D_homogeneous.unsqueeze(0).expand(
+ B, -1, -1
+ ) # BxNx4
+
+ # Step 1: Apply extrinsic parameters
+ # Transform 3D points to camera coordinate system for all cameras
+ points_cam = torch.bmm(
+ extrinsics, points3D_homogeneous.transpose(-1, -2)
+ )
+
+ if only_points_cam:
+ return points_cam
+
+ # Step 2: Apply intrinsic parameters and (optional) distortion
+ points2D = img_from_cam(intrinsics, points_cam, extra_params)
+
+ if return_points_cam:
+ return points2D, points_cam
+ return points2D
+
+
+def img_from_cam(intrinsics, points_cam, extra_params=None, default=0.0):
+ """
+ Applies intrinsic parameters and optional distortion to the given 3D points.
+
+ Args:
+ intrinsics (torch.Tensor): Intrinsic camera parameters of shape Bx3x3.
+ points_cam (torch.Tensor): 3D points in camera coordinates of shape Bx3xN.
+ extra_params (torch.Tensor, optional): Distortion parameters of shape BxN, where N can be 1, 2, or 4.
+ default (float, optional): Default value to replace NaNs in the output.
+
+ Returns:
+ points2D (torch.Tensor): 2D points in pixel coordinates of shape BxNx2.
+ """
+
+ # Normalize by the third coordinate (homogeneous division)
+ points_cam = points_cam / points_cam[:, 2:3, :]
+ # Extract uv
+ uv = points_cam[:, :2, :]
+
+ # Apply distortion if extra_params are provided
+ if extra_params is not None:
+ uu, vv = apply_distortion(extra_params, uv[:, 0], uv[:, 1])
+ uv = torch.stack([uu, vv], dim=1)
+
+ # Prepare points_cam for batch matrix multiplication
+ points_cam_homo = torch.cat(
+ (uv, torch.ones_like(uv[:, :1, :])), dim=1
+ ) # Bx3xN
+ # Apply intrinsic parameters using batch matrix multiplication
+ points2D_homo = torch.bmm(intrinsics, points_cam_homo) # Bx3xN
+
+ # Extract x and y coordinates
+ points2D = points2D_homo[:, :2, :] # Bx2xN
+
+ # Replace NaNs with default value
+ points2D = torch.nan_to_num(points2D, nan=default)
+
+ return points2D.transpose(1, 2) # BxNx2
+
diff --git a/src/utils/transforms.py b/src/utils/transforms.py
new file mode 100644
index 0000000000000000000000000000000000000000..809802e3c3c2f290b00fb1cd80d724535b554ec8
--- /dev/null
+++ b/src/utils/transforms.py
@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# DUST3R default transforms
+# --------------------------------------------------------
+import torchvision.transforms as tvf
+
+from model.dataset.utils.image import ImgNorm
+
+# define the standard image transforms
+ColorJitter = tvf.Compose([tvf.ColorJitter(0.5, 0.5, 0.5, 0.1), ImgNorm])
diff --git a/src/utils/viz.py b/src/utils/viz.py
new file mode 100644
index 0000000000000000000000000000000000000000..c757e8c835626dabed6ad497338e96465e7dd2e1
--- /dev/null
+++ b/src/utils/viz.py
@@ -0,0 +1,377 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Visualization utilities using trimesh
+# --------------------------------------------------------
+import numpy as np
+import PIL.Image
+import torch
+from scipy.spatial.transform import Rotation
+
+from model.dataset.utils.device import to_numpy
+from model.dataset.utils.geometry import geotrf, get_med_dist_between_poses
+from model.dataset.utils.image import rgb
+
+try:
+ import trimesh
+except ImportError:
+ print("/!\\ module trimesh is not installed, cannot visualize results /!\\")
+
+
+def cat_3d(vecs):
+ if isinstance(vecs, (np.ndarray, torch.Tensor)):
+ vecs = [vecs]
+ return np.concatenate([p.reshape(-1, 3) for p in to_numpy(vecs)])
+
+
+def show_raw_pointcloud(pts3d, colors, point_size=2):
+ scene = trimesh.Scene()
+
+ pct = trimesh.PointCloud(cat_3d(pts3d), colors=cat_3d(colors))
+ scene.add_geometry(pct)
+
+ scene.show(line_settings={"point_size": point_size})
+
+
+def pts3d_to_trimesh(img, pts3d, valid=None):
+ H, W, THREE = img.shape
+ assert THREE == 3
+ assert img.shape == pts3d.shape
+
+ vertices = pts3d.reshape(-1, 3)
+
+ # make squares: each pixel == 2 triangles
+ idx = np.arange(len(vertices)).reshape(H, W)
+ idx1 = idx[:-1, :-1].ravel() # top-left corner
+ idx2 = idx[:-1, +1:].ravel() # right-left corner
+ idx3 = idx[+1:, :-1].ravel() # bottom-left corner
+ idx4 = idx[+1:, +1:].ravel() # bottom-right corner
+ faces = np.concatenate(
+ (
+ np.c_[idx1, idx2, idx3],
+ np.c_[
+ idx3, idx2, idx1
+ ], # same triangle, but backward (cheap solution to cancel face culling)
+ np.c_[idx2, idx3, idx4],
+ np.c_[
+ idx4, idx3, idx2
+ ], # same triangle, but backward (cheap solution to cancel face culling)
+ ),
+ axis=0,
+ )
+
+ # prepare triangle colors
+ face_colors = np.concatenate(
+ (
+ img[:-1, :-1].reshape(-1, 3),
+ img[:-1, :-1].reshape(-1, 3),
+ img[+1:, +1:].reshape(-1, 3),
+ img[+1:, +1:].reshape(-1, 3),
+ ),
+ axis=0,
+ )
+
+ # remove invalid faces
+ if valid is not None:
+ assert valid.shape == (H, W)
+ valid_idxs = valid.ravel()
+ valid_faces = valid_idxs[faces].all(axis=-1)
+ faces = faces[valid_faces]
+ face_colors = face_colors[valid_faces]
+
+ assert len(faces) == len(face_colors)
+ return dict(vertices=vertices, face_colors=face_colors, faces=faces)
+
+
+def cat_meshes(meshes):
+ vertices, faces, colors = zip(
+ *[(m["vertices"], m["faces"], m["face_colors"]) for m in meshes]
+ )
+ n_vertices = np.cumsum([0] + [len(v) for v in vertices])
+ for i in range(len(faces)):
+ faces[i][:] += n_vertices[i]
+
+ vertices = np.concatenate(vertices)
+ colors = np.concatenate(colors)
+ faces = np.concatenate(faces)
+ return dict(vertices=vertices, face_colors=colors, faces=faces)
+
+
+def show_duster_pairs(view1, view2, pred1, pred2):
+ import matplotlib.pyplot as pl
+
+ pl.ion()
+
+ for e in range(len(view1["instance"])):
+ i = view1["idx"][e]
+ j = view2["idx"][e]
+ img1 = rgb(view1["img"][e])
+ img2 = rgb(view2["img"][e])
+ conf1 = pred1["conf"][e].squeeze()
+ conf2 = pred2["conf"][e].squeeze()
+ score = conf1.mean() * conf2.mean()
+ print(f">> Showing pair #{e} {i}-{j} {score=:g}")
+ pl.clf()
+ pl.subplot(221).imshow(img1)
+ pl.subplot(223).imshow(img2)
+ pl.subplot(222).imshow(conf1, vmin=1, vmax=30)
+ pl.subplot(224).imshow(conf2, vmin=1, vmax=30)
+ pts1 = pred1["pts3d"][e]
+ pts2 = pred2["pts3d_in_other_view"][e]
+ pl.subplots_adjust(0, 0, 1, 1, 0, 0)
+ if input("show pointcloud? (y/n) ") == "y":
+ show_raw_pointcloud(cat(pts1, pts2), cat(img1, img2), point_size=5)
+
+
+def auto_cam_size(im_poses):
+ return 0.1 * get_med_dist_between_poses(im_poses)
+
+
+class SceneViz:
+ def __init__(self):
+ self.scene = trimesh.Scene()
+
+ def add_pointcloud(self, pts3d, color, mask=None):
+ pts3d = to_numpy(pts3d)
+ mask = to_numpy(mask)
+ if mask is None:
+ mask = [slice(None)] * len(pts3d)
+ pts = np.concatenate([p[m] for p, m in zip(pts3d, mask)])
+ pct = trimesh.PointCloud(pts.reshape(-1, 3))
+
+ if isinstance(color, (list, np.ndarray, torch.Tensor)):
+ color = to_numpy(color)
+ col = np.concatenate([p[m] for p, m in zip(color, mask)])
+ assert col.shape == pts.shape
+ pct.visual.vertex_colors = uint8(col.reshape(-1, 3))
+ else:
+ assert len(color) == 3
+ pct.visual.vertex_colors = np.broadcast_to(uint8(color), pts.shape)
+
+ self.scene.add_geometry(pct)
+ return self
+
+ def add_camera(
+ self,
+ pose_c2w,
+ focal=None,
+ color=(0, 0, 0),
+ image=None,
+ imsize=None,
+ cam_size=0.03,
+ ):
+ pose_c2w, focal, color, image = to_numpy((pose_c2w, focal, color, image))
+ add_scene_cam(self.scene, pose_c2w, color, image, focal, screen_width=cam_size)
+ return self
+
+ def add_cameras(
+ self, poses, focals=None, images=None, imsizes=None, colors=None, **kw
+ ):
+ def get(arr, idx):
+ return None if arr is None else arr[idx]
+
+ for i, pose_c2w in enumerate(poses):
+ self.add_camera(
+ pose_c2w,
+ get(focals, i),
+ image=get(images, i),
+ color=get(colors, i),
+ imsize=get(imsizes, i),
+ **kw,
+ )
+ return self
+
+ def show(self, point_size=2, viewer=None):
+ return self.scene.show(viewer=viewer, line_settings={"point_size": point_size})
+
+
+def show_raw_pointcloud_with_cams(
+ imgs, pts3d, mask, focals, cams2world, point_size=2, cam_size=0.05, cam_color=None
+):
+ """Visualization of a pointcloud with cameras
+ imgs = (N, H, W, 3) or N-size list of [(H,W,3), ...]
+ pts3d = (N, H, W, 3) or N-size list of [(H,W,3), ...]
+ focals = (N,) or N-size list of [focal, ...]
+ cams2world = (N,4,4) or N-size list of [(4,4), ...]
+ """
+ assert len(pts3d) == len(mask) <= len(imgs) <= len(cams2world) == len(focals)
+ pts3d = to_numpy(pts3d)
+ imgs = to_numpy(imgs)
+ focals = to_numpy(focals)
+ cams2world = to_numpy(cams2world)
+
+ scene = trimesh.Scene()
+
+ # full pointcloud
+ pts = np.concatenate([p[m] for p, m in zip(pts3d, mask)])
+ col = np.concatenate([p[m] for p, m in zip(imgs, mask)])
+ pct = trimesh.PointCloud(pts.reshape(-1, 3), colors=col.reshape(-1, 3))
+ scene.add_geometry(pct)
+
+ # add each camera
+ for i, pose_c2w in enumerate(cams2world):
+ if isinstance(cam_color, list):
+ camera_edge_color = cam_color[i]
+ else:
+ camera_edge_color = cam_color or CAM_COLORS[i % len(CAM_COLORS)]
+ add_scene_cam(
+ scene,
+ pose_c2w,
+ camera_edge_color,
+ imgs[i] if i < len(imgs) else None,
+ focals[i],
+ screen_width=cam_size,
+ )
+
+ scene.show(line_settings={"point_size": point_size})
+
+
+def add_scene_cam(
+ scene, pose_c2w, edge_color, image=None, focal=None, imsize=None, screen_width=0.03
+):
+ if image is not None:
+ H, W, THREE = image.shape
+ assert THREE == 3
+ if image.dtype != np.uint8:
+ image = np.uint8(255 * image)
+ elif imsize is not None:
+ W, H = imsize
+ elif focal is not None:
+ H = W = focal / 1.1
+ else:
+ H = W = 1
+
+ if focal is None:
+ focal = min(H, W) * 1.1 # default value
+ elif isinstance(focal, np.ndarray):
+ focal = focal[0]
+
+ # create fake camera
+ height = focal * screen_width / H
+ width = screen_width * 0.5**0.5
+ rot45 = np.eye(4)
+ rot45[:3, :3] = Rotation.from_euler("z", np.deg2rad(45)).as_matrix()
+ rot45[2, 3] = -height # set the tip of the cone = optical center
+ aspect_ratio = np.eye(4)
+ aspect_ratio[0, 0] = W / H
+ transform = pose_c2w @ OPENGL @ aspect_ratio @ rot45
+ cam = trimesh.creation.cone(width, height, sections=4) # , transform=transform)
+
+ # this is the image
+ if image is not None:
+ vertices = geotrf(transform, cam.vertices[[4, 5, 1, 3]])
+ faces = np.array([[0, 1, 2], [0, 2, 3], [2, 1, 0], [3, 2, 0]])
+ img = trimesh.Trimesh(vertices=vertices, faces=faces)
+ uv_coords = np.float32([[0, 0], [1, 0], [1, 1], [0, 1]])
+ img.visual = trimesh.visual.TextureVisuals(
+ uv_coords, image=PIL.Image.fromarray(image)
+ )
+ scene.add_geometry(img)
+
+ # this is the camera mesh
+ rot2 = np.eye(4)
+ rot2[:3, :3] = Rotation.from_euler("z", np.deg2rad(2)).as_matrix()
+ vertices = np.r_[cam.vertices, 0.95 * cam.vertices, geotrf(rot2, cam.vertices)]
+ vertices = geotrf(transform, vertices)
+ faces = []
+ for face in cam.faces:
+ if 0 in face:
+ continue
+ a, b, c = face
+ a2, b2, c2 = face + len(cam.vertices)
+ a3, b3, c3 = face + 2 * len(cam.vertices)
+
+ # add 3 pseudo-edges
+ faces.append((a, b, b2))
+ faces.append((a, a2, c))
+ faces.append((c2, b, c))
+
+ faces.append((a, b, b3))
+ faces.append((a, a3, c))
+ faces.append((c3, b, c))
+
+ # no culling
+ faces += [(c, b, a) for a, b, c in faces]
+
+ cam = trimesh.Trimesh(vertices=vertices, faces=faces)
+ cam.visual.face_colors[:, :3] = edge_color
+ scene.add_geometry(cam)
+
+
+def cat(a, b):
+ return np.concatenate((a.reshape(-1, 3), b.reshape(-1, 3)))
+
+
+OPENGL = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
+
+
+CAM_COLORS = [
+ (255, 0, 0),
+ (0, 0, 255),
+ (0, 255, 0),
+ (255, 0, 255),
+ (255, 204, 0),
+ (0, 204, 204),
+ (128, 255, 255),
+ (255, 128, 255),
+ (255, 255, 128),
+ (0, 0, 0),
+ (128, 128, 128),
+]
+
+
+def uint8(colors):
+ if not isinstance(colors, np.ndarray):
+ colors = np.array(colors)
+ if np.issubdtype(colors.dtype, np.floating):
+ colors *= 255
+ assert 0 <= colors.min() and colors.max() < 256
+ return np.uint8(colors)
+
+
+def segment_sky(image):
+ import cv2
+ from scipy import ndimage
+
+ # Convert to HSV
+ image = to_numpy(image)
+ if np.issubdtype(image.dtype, np.floating):
+ image = np.uint8(255 * image.clip(min=0, max=1))
+ hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+
+ # Define range for blue color and create mask
+ lower_blue = np.array([0, 0, 100])
+ upper_blue = np.array([30, 255, 255])
+ mask = cv2.inRange(hsv, lower_blue, upper_blue).view(bool)
+
+ # add luminous gray
+ mask |= (hsv[:, :, 1] < 10) & (hsv[:, :, 2] > 150)
+ mask |= (hsv[:, :, 1] < 30) & (hsv[:, :, 2] > 180)
+ mask |= (hsv[:, :, 1] < 50) & (hsv[:, :, 2] > 220)
+
+ # Morphological operations
+ kernel = np.ones((5, 5), np.uint8)
+ mask2 = ndimage.binary_opening(mask, structure=kernel)
+
+ # keep only largest CC
+ _, labels, stats, _ = cv2.connectedComponentsWithStats(
+ mask2.view(np.uint8), connectivity=8
+ )
+ cc_sizes = stats[1:, cv2.CC_STAT_AREA]
+ order = cc_sizes.argsort()[::-1] # bigger first
+ i = 0
+ selection = []
+ while i < len(order) and cc_sizes[order[i]] > cc_sizes[order[0]] / 2:
+ selection.append(1 + order[i])
+ i += 1
+ mask3 = np.in1d(labels, selection).reshape(labels.shape)
+
+ # Apply mask
+ return torch.from_numpy(mask3)
diff --git a/src/visualization/annotation.py b/src/visualization/annotation.py
new file mode 100644
index 0000000000000000000000000000000000000000..ebcbdbb668e0a32cfc05ee9c986b029a7ff50a34
--- /dev/null
+++ b/src/visualization/annotation.py
@@ -0,0 +1,49 @@
+from pathlib import Path
+from string import ascii_letters, digits, punctuation
+
+import numpy as np
+import torch
+from einops import rearrange
+from jaxtyping import Float
+from PIL import Image, ImageDraw, ImageFont
+from torch import Tensor
+
+from .layout import vcat
+
+EXPECTED_CHARACTERS = digits + punctuation + ascii_letters
+
+
+def draw_label(
+ text: str,
+ font: Path,
+ font_size: int,
+ device: torch.device = torch.device("cpu"),
+) -> Float[Tensor, "3 height width"]:
+ """Draw a black label on a white background with no border."""
+ try:
+ font = ImageFont.truetype(str(font), font_size)
+ except OSError:
+ font = ImageFont.load_default()
+ left, _, right, _ = font.getbbox(text)
+ width = right - left
+ _, top, _, bottom = font.getbbox(EXPECTED_CHARACTERS)
+ height = bottom - top
+ image = Image.new("RGB", (width, height), color="white")
+ draw = ImageDraw.Draw(image)
+ draw.text((0, 0), text, font=font, fill="black")
+ image = torch.tensor(np.array(image) / 255, dtype=torch.float32, device=device)
+ return rearrange(image, "h w c -> c h w")
+
+
+def add_label(
+ image: Float[Tensor, "3 width height"],
+ label: str,
+ font: Path = Path("assets/Inter-Regular.otf"),
+ font_size: int = 24,
+) -> Float[Tensor, "3 width_with_label height_with_label"]:
+ return vcat(
+ draw_label(label, font, font_size, image.device),
+ image,
+ align="left",
+ gap=4,
+ )
diff --git a/src/visualization/camera_trajectory/interpolation.py b/src/visualization/camera_trajectory/interpolation.py
new file mode 100644
index 0000000000000000000000000000000000000000..f1c151441480e889a6e16bb8e705adbfb94ee60d
--- /dev/null
+++ b/src/visualization/camera_trajectory/interpolation.py
@@ -0,0 +1,255 @@
+import torch
+from einops import einsum, rearrange, reduce
+from jaxtyping import Float
+from scipy.spatial.transform import Rotation as R
+from torch import Tensor
+
+
+def interpolate_intrinsics(
+ initial: Float[Tensor, "*#batch 3 3"],
+ final: Float[Tensor, "*#batch 3 3"],
+ t: Float[Tensor, " time_step"],
+) -> Float[Tensor, "*batch time_step 3 3"]:
+ initial = rearrange(initial, "... i j -> ... () i j")
+ final = rearrange(final, "... i j -> ... () i j")
+ t = rearrange(t, "t -> t () ()")
+ return initial + (final - initial) * t
+
+
+def intersect_rays(
+ a_origins: Float[Tensor, "*#batch dim"],
+ a_directions: Float[Tensor, "*#batch dim"],
+ b_origins: Float[Tensor, "*#batch dim"],
+ b_directions: Float[Tensor, "*#batch dim"],
+) -> Float[Tensor, "*batch dim"]:
+ """Compute the least-squares intersection of rays. Uses the math from here:
+ https://math.stackexchange.com/a/1762491/286022
+ """
+
+ # Broadcast and stack the tensors.
+ a_origins, a_directions, b_origins, b_directions = torch.broadcast_tensors(
+ a_origins, a_directions, b_origins, b_directions
+ )
+ origins = torch.stack((a_origins, b_origins), dim=-2)
+ directions = torch.stack((a_directions, b_directions), dim=-2)
+
+ # Compute n_i * n_i^T - eye(3) from the equation.
+ n = einsum(directions, directions, "... n i, ... n j -> ... n i j")
+ n = n - torch.eye(3, dtype=origins.dtype, device=origins.device)
+
+ # Compute the left-hand side of the equation.
+ lhs = reduce(n, "... n i j -> ... i j", "sum")
+
+ # Compute the right-hand side of the equation.
+ rhs = einsum(n, origins, "... n i j, ... n j -> ... n i")
+ rhs = reduce(rhs, "... n i -> ... i", "sum")
+
+ # Left-matrix-multiply both sides by the inverse of lhs to find p.
+ return torch.linalg.lstsq(lhs, rhs).solution
+
+
+def normalize(a: Float[Tensor, "*#batch dim"]) -> Float[Tensor, "*#batch dim"]:
+ return a / a.norm(dim=-1, keepdim=True)
+
+
+def generate_coordinate_frame(
+ y: Float[Tensor, "*#batch 3"],
+ z: Float[Tensor, "*#batch 3"],
+) -> Float[Tensor, "*batch 3 3"]:
+ """Generate a coordinate frame given perpendicular, unit-length Y and Z vectors."""
+ y, z = torch.broadcast_tensors(y, z)
+ return torch.stack([y.cross(z), y, z], dim=-1)
+
+
+def generate_rotation_coordinate_frame(
+ a: Float[Tensor, "*#batch 3"],
+ b: Float[Tensor, "*#batch 3"],
+ eps: float = 1e-4,
+) -> Float[Tensor, "*batch 3 3"]:
+ """Generate a coordinate frame where the Y direction is normal to the plane defined
+ by unit vectors a and b. The other axes are arbitrary."""
+ device = a.device
+
+ # Replace every entry in b that's parallel to the corresponding entry in a with an
+ # arbitrary vector.
+ b = b.detach().clone()
+ parallel = (einsum(a, b, "... i, ... i -> ...").abs() - 1).abs() < eps
+ b[parallel] = torch.tensor([0, 0, 1], dtype=b.dtype, device=device)
+ parallel = (einsum(a, b, "... i, ... i -> ...").abs() - 1).abs() < eps
+ b[parallel] = torch.tensor([0, 1, 0], dtype=b.dtype, device=device)
+
+ # Generate the coordinate frame. The initial cross product defines the plane.
+ return generate_coordinate_frame(normalize(a.cross(b)), a)
+
+
+def matrix_to_euler(
+ rotations: Float[Tensor, "*batch 3 3"],
+ pattern: str,
+) -> Float[Tensor, "*batch 3"]:
+ *batch, _, _ = rotations.shape
+ rotations = rotations.reshape(-1, 3, 3)
+ angles_np = R.from_matrix(rotations.detach().cpu().numpy()).as_euler(pattern)
+ rotations = torch.tensor(angles_np, dtype=rotations.dtype, device=rotations.device)
+ return rotations.reshape(*batch, 3)
+
+
+def euler_to_matrix(
+ rotations: Float[Tensor, "*batch 3"],
+ pattern: str,
+) -> Float[Tensor, "*batch 3 3"]:
+ *batch, _ = rotations.shape
+ rotations = rotations.reshape(-1, 3)
+ matrix_np = R.from_euler(pattern, rotations.detach().cpu().numpy()).as_matrix()
+ rotations = torch.tensor(matrix_np, dtype=rotations.dtype, device=rotations.device)
+ return rotations.reshape(*batch, 3, 3)
+
+
+def extrinsics_to_pivot_parameters(
+ extrinsics: Float[Tensor, "*#batch 4 4"],
+ pivot_coordinate_frame: Float[Tensor, "*#batch 3 3"],
+ pivot_point: Float[Tensor, "*#batch 3"],
+) -> Float[Tensor, "*batch 5"]:
+ """Convert the extrinsics to a representation with 5 degrees of freedom:
+ 1. Distance from pivot point in the "X" (look cross pivot axis) direction.
+ 2. Distance from pivot point in the "Y" (pivot axis) direction.
+ 3. Distance from pivot point in the Z (look) direction
+ 4. Angle in plane
+ 5. Twist (rotation not in plane)
+ """
+
+ # The pivot coordinate frame's Z axis is normal to the plane.
+ pivot_axis = pivot_coordinate_frame[..., :, 1]
+
+ # Compute the translation elements of the pivot parametrization.
+ translation_frame = generate_coordinate_frame(pivot_axis, extrinsics[..., :3, 2])
+ origin = extrinsics[..., :3, 3]
+ delta = pivot_point - origin
+ translation = einsum(translation_frame, delta, "... i j, ... i -> ... j")
+
+ # Add the rotation elements of the pivot parametrization.
+ inverted = pivot_coordinate_frame.inverse() @ extrinsics[..., :3, :3]
+ y, _, z = matrix_to_euler(inverted, "YXZ").unbind(dim=-1)
+
+ return torch.cat([translation, y[..., None], z[..., None]], dim=-1)
+
+
+def pivot_parameters_to_extrinsics(
+ parameters: Float[Tensor, "*#batch 5"],
+ pivot_coordinate_frame: Float[Tensor, "*#batch 3 3"],
+ pivot_point: Float[Tensor, "*#batch 3"],
+) -> Float[Tensor, "*batch 4 4"]:
+ translation, y, z = parameters.split((3, 1, 1), dim=-1)
+
+ euler = torch.cat((y, torch.zeros_like(y), z), dim=-1)
+ rotation = pivot_coordinate_frame @ euler_to_matrix(euler, "YXZ")
+
+ # The pivot coordinate frame's Z axis is normal to the plane.
+ pivot_axis = pivot_coordinate_frame[..., :, 1]
+
+ translation_frame = generate_coordinate_frame(pivot_axis, rotation[..., :3, 2])
+ delta = einsum(translation_frame, translation, "... i j, ... j -> ... i")
+ origin = pivot_point - delta
+
+ *batch, _ = origin.shape
+ extrinsics = torch.eye(4, dtype=parameters.dtype, device=parameters.device)
+ extrinsics = extrinsics.broadcast_to((*batch, 4, 4)).clone()
+ extrinsics[..., 3, 3] = 1
+ extrinsics[..., :3, :3] = rotation
+ extrinsics[..., :3, 3] = origin
+ return extrinsics
+
+
+def interpolate_circular(
+ a: Float[Tensor, "*#batch"],
+ b: Float[Tensor, "*#batch"],
+ t: Float[Tensor, "*#batch"],
+) -> Float[Tensor, " *batch"]:
+ a, b, t = torch.broadcast_tensors(a, b, t)
+
+ tau = 2 * torch.pi
+ a = a % tau
+ b = b % tau
+
+ # Consider piecewise edge cases.
+ d = (b - a).abs()
+ a_left = a - tau
+ d_left = (b - a_left).abs()
+ a_right = a + tau
+ d_right = (b - a_right).abs()
+ use_d = (d < d_left) & (d < d_right)
+ use_d_left = (d_left < d_right) & (~use_d)
+ use_d_right = (~use_d) & (~use_d_left)
+
+ result = a + (b - a) * t
+ result[use_d_left] = (a_left + (b - a_left) * t)[use_d_left]
+ result[use_d_right] = (a_right + (b - a_right) * t)[use_d_right]
+
+ return result
+
+
+def interpolate_pivot_parameters(
+ initial: Float[Tensor, "*#batch 5"],
+ final: Float[Tensor, "*#batch 5"],
+ t: Float[Tensor, " time_step"],
+) -> Float[Tensor, "*batch time_step 5"]:
+ initial = rearrange(initial, "... d -> ... () d")
+ final = rearrange(final, "... d -> ... () d")
+ t = rearrange(t, "t -> t ()")
+ ti, ri = initial.split((3, 2), dim=-1)
+ tf, rf = final.split((3, 2), dim=-1)
+
+ t_lerp = ti + (tf - ti) * t
+ r_lerp = interpolate_circular(ri, rf, t)
+
+ return torch.cat((t_lerp, r_lerp), dim=-1)
+
+
+@torch.no_grad()
+def interpolate_extrinsics(
+ initial: Float[Tensor, "*#batch 4 4"],
+ final: Float[Tensor, "*#batch 4 4"],
+ t: Float[Tensor, " time_step"],
+ eps: float = 1e-4,
+) -> Float[Tensor, "*batch time_step 4 4"]:
+ """Interpolate extrinsics by rotating around their "focus point," which is the
+ least-squares intersection between the look vectors of the initial and final
+ extrinsics.
+ """
+
+ initial = initial.type(torch.float64)
+ final = final.type(torch.float64)
+ t = t.type(torch.float64)
+
+ # Based on the dot product between the look vectors, pick from one of two cases:
+ # 1. Look vectors are parallel: interpolate about their origins' midpoint.
+ # 3. Look vectors aren't parallel: interpolate about their focus point.
+ initial_look = initial[..., :3, 2]
+ final_look = final[..., :3, 2]
+ dot_products = einsum(initial_look, final_look, "... i, ... i -> ...")
+ parallel_mask = (dot_products.abs() - 1).abs() < eps
+
+ # Pick focus points.
+ initial_origin = initial[..., :3, 3]
+ final_origin = final[..., :3, 3]
+ pivot_point = 0.5 * (initial_origin + final_origin)
+ pivot_point[~parallel_mask] = intersect_rays(
+ initial_origin[~parallel_mask],
+ initial_look[~parallel_mask],
+ final_origin[~parallel_mask],
+ final_look[~parallel_mask],
+ )
+
+ # Convert to pivot parameters.
+ pivot_frame = generate_rotation_coordinate_frame(initial_look, final_look, eps=eps)
+ initial_params = extrinsics_to_pivot_parameters(initial, pivot_frame, pivot_point)
+ final_params = extrinsics_to_pivot_parameters(final, pivot_frame, pivot_point)
+
+ # Interpolate the pivot parameters.
+ interpolated_params = interpolate_pivot_parameters(initial_params, final_params, t)
+
+ # Convert back.
+ return pivot_parameters_to_extrinsics(
+ interpolated_params.type(torch.float32),
+ rearrange(pivot_frame, "... i j -> ... () i j").type(torch.float32),
+ rearrange(pivot_point, "... xyz -> ... () xyz").type(torch.float32),
+ )
diff --git a/src/visualization/camera_trajectory/spin.py b/src/visualization/camera_trajectory/spin.py
new file mode 100644
index 0000000000000000000000000000000000000000..eadddcbc9b33075e902d610a3fc40ef0ba7bde2e
--- /dev/null
+++ b/src/visualization/camera_trajectory/spin.py
@@ -0,0 +1,37 @@
+import numpy as np
+import torch
+from einops import repeat
+from jaxtyping import Float
+from scipy.spatial.transform import Rotation as R
+from torch import Tensor
+
+
+def generate_spin(
+ num_frames: int,
+ device: torch.device,
+ elevation: float,
+ radius: float,
+) -> Float[Tensor, "frame 4 4"]:
+ # Translate back along the camera's look vector.
+ tf_translation = torch.eye(4, dtype=torch.float32, device=device)
+ tf_translation[:2] *= -1
+ tf_translation[2, 3] = -radius
+
+ # Generate the transformation for the azimuth.
+ phi = 2 * np.pi * (np.arange(num_frames) / num_frames)
+ rotation_vectors = np.stack([np.zeros_like(phi), phi, np.zeros_like(phi)], axis=-1)
+
+ azimuth = R.from_rotvec(rotation_vectors).as_matrix()
+ azimuth = torch.tensor(azimuth, dtype=torch.float32, device=device)
+ tf_azimuth = torch.eye(4, dtype=torch.float32, device=device)
+ tf_azimuth = repeat(tf_azimuth, "i j -> b i j", b=num_frames).clone()
+ tf_azimuth[:, :3, :3] = azimuth
+
+ # Generate the transformation for the elevation.
+ deg_elevation = np.deg2rad(elevation)
+ elevation = R.from_rotvec(np.array([deg_elevation, 0, 0], dtype=np.float32))
+ elevation = torch.tensor(elevation.as_matrix())
+ tf_elevation = torch.eye(4, dtype=torch.float32, device=device)
+ tf_elevation[:3, :3] = elevation
+
+ return tf_azimuth @ tf_elevation @ tf_translation
diff --git a/src/visualization/camera_trajectory/wobble.py b/src/visualization/camera_trajectory/wobble.py
new file mode 100644
index 0000000000000000000000000000000000000000..5fd71c8f8e3d808561894f993f6cd4988e469b39
--- /dev/null
+++ b/src/visualization/camera_trajectory/wobble.py
@@ -0,0 +1,32 @@
+import torch
+from einops import rearrange
+from jaxtyping import Float
+from torch import Tensor
+
+
+@torch.no_grad()
+def generate_wobble_transformation(
+ radius: Float[Tensor, "*#batch"],
+ t: Float[Tensor, " time_step"],
+ num_rotations: int = 1,
+ scale_radius_with_t: bool = True,
+) -> Float[Tensor, "*batch time_step 4 4"]:
+ # Generate a translation in the image plane.
+ tf = torch.eye(4, dtype=torch.float32, device=t.device)
+ tf = tf.broadcast_to((*radius.shape, t.shape[0], 4, 4)).clone()
+ radius = radius[..., None]
+ if scale_radius_with_t:
+ radius = radius * t
+ tf[..., 0, 3] = torch.sin(2 * torch.pi * num_rotations * t) * radius
+ tf[..., 1, 3] = -torch.cos(2 * torch.pi * num_rotations * t) * radius
+ return tf
+
+
+@torch.no_grad()
+def generate_wobble(
+ extrinsics: Float[Tensor, "*#batch 4 4"],
+ radius: Float[Tensor, "*#batch"],
+ t: Float[Tensor, " time_step"],
+) -> Float[Tensor, "*batch time_step 4 4"]:
+ tf = generate_wobble_transformation(radius, t)
+ return rearrange(extrinsics, "... i j -> ... () i j") @ tf
diff --git a/src/visualization/color_map.py b/src/visualization/color_map.py
new file mode 100644
index 0000000000000000000000000000000000000000..2926fffd896610ebbc1b5786170dcba4cbdb1a73
--- /dev/null
+++ b/src/visualization/color_map.py
@@ -0,0 +1,48 @@
+import torch
+from colorspacious import cspace_convert
+from einops import rearrange
+from jaxtyping import Float
+from matplotlib import cm
+from torch import Tensor
+
+
+def apply_color_map(
+ x: Float[Tensor, " *batch"],
+ color_map: str = "inferno",
+) -> Float[Tensor, "*batch 3"]:
+ cmap = cm.get_cmap(color_map)
+
+ # Convert to NumPy so that Matplotlib color maps can be used.
+ mapped = cmap(x.detach().clip(min=0, max=1).cpu().numpy())[..., :3]
+
+ # Convert back to the original format.
+ return torch.tensor(mapped, device=x.device, dtype=torch.float32)
+
+
+def apply_color_map_to_image(
+ image: Float[Tensor, "*batch height width"],
+ color_map: str = "inferno",
+) -> Float[Tensor, "*batch 3 height with"]:
+ image = apply_color_map(image, color_map)
+ return rearrange(image, "... h w c -> ... c h w")
+
+
+def apply_color_map_2d(
+ x: Float[Tensor, "*#batch"],
+ y: Float[Tensor, "*#batch"],
+) -> Float[Tensor, "*batch 3"]:
+ red = cspace_convert((189, 0, 0), "sRGB255", "CIELab")
+ blue = cspace_convert((0, 45, 255), "sRGB255", "CIELab")
+ white = cspace_convert((255, 255, 255), "sRGB255", "CIELab")
+ x_np = x.detach().clip(min=0, max=1).cpu().numpy()[..., None]
+ y_np = y.detach().clip(min=0, max=1).cpu().numpy()[..., None]
+
+ # Interpolate between red and blue on the x axis.
+ interpolated = x_np * red + (1 - x_np) * blue
+
+ # Interpolate between color and white on the y axis.
+ interpolated = y_np * interpolated + (1 - y_np) * white
+
+ # Convert to RGB.
+ rgb = cspace_convert(interpolated, "CIELab", "sRGB1")
+ return torch.tensor(rgb, device=x.device, dtype=torch.float32).clip(min=0, max=1)
diff --git a/src/visualization/colors.py b/src/visualization/colors.py
new file mode 100644
index 0000000000000000000000000000000000000000..ef82964e0bf71f147d41d95a6e389b1f22ba57b5
--- /dev/null
+++ b/src/visualization/colors.py
@@ -0,0 +1,32 @@
+from PIL import ImageColor
+
+# https://sashamaps.net/docs/resources/20-colors/
+DISTINCT_COLORS = [
+ "#e6194b",
+ "#3cb44b",
+ "#ffe119",
+ "#4363d8",
+ "#f58231",
+ "#911eb4",
+ "#46f0f0",
+ "#f032e6",
+ "#bcf60c",
+ "#fabebe",
+ "#008080",
+ "#e6beff",
+ "#9a6324",
+ "#fffac8",
+ "#800000",
+ "#aaffc3",
+ "#808000",
+ "#ffd8b1",
+ "#000075",
+ "#808080",
+ "#ffffff",
+ "#000000",
+]
+
+
+def get_distinct_color(index: int) -> tuple[float, float, float]:
+ hex = DISTINCT_COLORS[index % len(DISTINCT_COLORS)]
+ return tuple(x / 255 for x in ImageColor.getcolor(hex, "RGB"))
diff --git a/src/visualization/drawing/cameras.py b/src/visualization/drawing/cameras.py
new file mode 100644
index 0000000000000000000000000000000000000000..6a60beb53a1a24931cd4b822e5ebe65b6425a6d6
--- /dev/null
+++ b/src/visualization/drawing/cameras.py
@@ -0,0 +1,195 @@
+from typing import Optional
+
+import torch
+from einops import einsum, rearrange, repeat
+from jaxtyping import Float
+from torch import Tensor
+
+from ...geometry.projection import unproject
+from ..annotation import add_label
+from .lines import draw_lines
+from .types import Scalar, sanitize_scalar
+
+
+def draw_cameras(
+ resolution: int,
+ extrinsics: Float[Tensor, "batch 4 4"],
+ intrinsics: Float[Tensor, "batch 3 3"],
+ color: Float[Tensor, "batch 3"],
+ near: Optional[Scalar] = None,
+ far: Optional[Scalar] = None,
+ margin: float = 0.1, # relative to AABB
+ frustum_scale: float = 0.05, # relative to image resolution
+) -> Float[Tensor, "3 3 height width"]:
+ device = extrinsics.device
+
+ # Compute scene bounds.
+ minima, maxima = compute_aabb(extrinsics, intrinsics, near, far)
+ scene_minima, scene_maxima = compute_equal_aabb_with_margin(
+ minima, maxima, margin=margin
+ )
+ span = (scene_maxima - scene_minima).max()
+
+ # Compute frustum locations.
+ corner_depth = (span * frustum_scale)[None]
+ frustum_corners = unproject_frustum_corners(extrinsics, intrinsics, corner_depth)
+ if near is not None:
+ near_corners = unproject_frustum_corners(extrinsics, intrinsics, near)
+ if far is not None:
+ far_corners = unproject_frustum_corners(extrinsics, intrinsics, far)
+
+ # Project the cameras onto each axis-aligned plane.
+ projections = []
+ for projected_axis in range(3):
+ image = torch.zeros(
+ (3, resolution, resolution),
+ dtype=torch.float32,
+ device=device,
+ )
+ image_x_axis = (projected_axis + 1) % 3
+ image_y_axis = (projected_axis + 2) % 3
+
+ def project(points: Float[Tensor, "*batch 3"]) -> Float[Tensor, "*batch 2"]:
+ x = points[..., image_x_axis]
+ y = points[..., image_y_axis]
+ return torch.stack([x, y], dim=-1)
+
+ x_range, y_range = torch.stack(
+ (project(scene_minima), project(scene_maxima)), dim=-1
+ )
+
+ # Draw near and far planes.
+ if near is not None:
+ projected_near_corners = project(near_corners)
+ image = draw_lines(
+ image,
+ rearrange(projected_near_corners, "b p xy -> (b p) xy"),
+ rearrange(projected_near_corners.roll(1, 1), "b p xy -> (b p) xy"),
+ color=0.25,
+ width=2,
+ x_range=x_range,
+ y_range=y_range,
+ )
+ if far is not None:
+ projected_far_corners = project(far_corners)
+ image = draw_lines(
+ image,
+ rearrange(projected_far_corners, "b p xy -> (b p) xy"),
+ rearrange(projected_far_corners.roll(1, 1), "b p xy -> (b p) xy"),
+ color=0.25,
+ width=2,
+ x_range=x_range,
+ y_range=y_range,
+ )
+ if near is not None and far is not None:
+ image = draw_lines(
+ image,
+ rearrange(projected_near_corners, "b p xy -> (b p) xy"),
+ rearrange(projected_far_corners, "b p xy -> (b p) xy"),
+ color=0.25,
+ width=2,
+ x_range=x_range,
+ y_range=y_range,
+ )
+
+ # Draw the camera frustums themselves.
+ projected_origins = project(extrinsics[:, :3, 3])
+ projected_frustum_corners = project(frustum_corners)
+ start = [
+ repeat(projected_origins, "b xy -> (b p) xy", p=4),
+ rearrange(projected_frustum_corners.roll(1, 1), "b p xy -> (b p) xy"),
+ ]
+ start = rearrange(torch.cat(start, dim=0), "(r b p) xy -> (b r p) xy", r=2, p=4)
+ image = draw_lines(
+ image,
+ start,
+ repeat(projected_frustum_corners, "b p xy -> (b r p) xy", r=2),
+ color=repeat(color, "b c -> (b r p) c", r=2, p=4),
+ width=2,
+ x_range=x_range,
+ y_range=y_range,
+ )
+
+ x_name = "XYZ"[image_x_axis]
+ y_name = "XYZ"[image_y_axis]
+ image = add_label(image, f"{x_name}{y_name} Projection")
+
+ # TODO: Draw axis indicators.
+ projections.append(image)
+
+ return torch.stack(projections)
+
+
+def compute_aabb(
+ extrinsics: Float[Tensor, "batch 4 4"],
+ intrinsics: Float[Tensor, "batch 3 3"],
+ near: Optional[Scalar] = None,
+ far: Optional[Scalar] = None,
+) -> tuple[
+ Float[Tensor, "3"], # minima of the scene
+ Float[Tensor, "3"], # maxima of the scene
+]:
+ """Compute an axis-aligned bounding box for the camera frustums."""
+
+ device = extrinsics.device
+
+ # These points are included in the AABB.
+ points = [extrinsics[:, :3, 3]]
+
+ if near is not None:
+ near = sanitize_scalar(near, device)
+ corners = unproject_frustum_corners(extrinsics, intrinsics, near)
+ points.append(rearrange(corners, "b p xyz -> (b p) xyz"))
+
+ if far is not None:
+ far = sanitize_scalar(far, device)
+ corners = unproject_frustum_corners(extrinsics, intrinsics, far)
+ points.append(rearrange(corners, "b p xyz -> (b p) xyz"))
+
+ points = torch.cat(points, dim=0)
+ return points.min(dim=0).values, points.max(dim=0).values
+
+
+def compute_equal_aabb_with_margin(
+ minima: Float[Tensor, "*#batch 3"],
+ maxima: Float[Tensor, "*#batch 3"],
+ margin: float = 0.1,
+) -> tuple[
+ Float[Tensor, "*batch 3"], # minima of the scene
+ Float[Tensor, "*batch 3"], # maxima of the scene
+]:
+ midpoint = (maxima + minima) * 0.5
+ span = (maxima - minima).max() * (1 + margin)
+ scene_minima = midpoint - 0.5 * span
+ scene_maxima = midpoint + 0.5 * span
+ return scene_minima, scene_maxima
+
+
+def unproject_frustum_corners(
+ extrinsics: Float[Tensor, "batch 4 4"],
+ intrinsics: Float[Tensor, "batch 3 3"],
+ depth: Float[Tensor, "#batch"],
+) -> Float[Tensor, "batch 4 3"]:
+ device = extrinsics.device
+
+ # Get coordinates for the corners. Following them in a circle makes a rectangle.
+ xy = torch.linspace(0, 1, 2, device=device)
+ xy = torch.stack(torch.meshgrid(xy, xy, indexing="xy"), dim=-1)
+ xy = rearrange(xy, "i j xy -> (i j) xy")
+ xy = xy[torch.tensor([0, 1, 3, 2], device=device)]
+
+ # Get ray directions in camera space.
+ directions = unproject(
+ xy,
+ torch.ones(1, dtype=torch.float32, device=device),
+ rearrange(intrinsics, "b i j -> b () i j"),
+ )
+
+ # Divide by the z coordinate so that multiplying by depth will produce orthographic
+ # depth (z depth) as opposed to Euclidean depth (distance from the camera).
+ directions = directions / directions[..., -1:]
+ directions = einsum(extrinsics[..., :3, :3], directions, "b i j, b r j -> b r i")
+
+ origins = rearrange(extrinsics[:, :3, 3], "b xyz -> b () xyz")
+ depth = rearrange(depth, "b -> b () ()")
+ return origins + depth * directions
diff --git a/src/visualization/drawing/coordinate_conversion.py b/src/visualization/drawing/coordinate_conversion.py
new file mode 100644
index 0000000000000000000000000000000000000000..42ac5aca7cf903b61e6e9f9018c8aa1811300fa3
--- /dev/null
+++ b/src/visualization/drawing/coordinate_conversion.py
@@ -0,0 +1,44 @@
+from typing import Optional, Protocol, runtime_checkable
+
+import torch
+from jaxtyping import Float
+from torch import Tensor
+
+from .types import Pair, sanitize_pair
+
+
+@runtime_checkable
+class ConversionFunction(Protocol):
+ def __call__(
+ self,
+ xy: Float[Tensor, "*batch 2"],
+ ) -> Float[Tensor, "*batch 2"]:
+ pass
+
+
+def generate_conversions(
+ shape: tuple[int, int],
+ device: torch.device,
+ x_range: Optional[Pair] = None,
+ y_range: Optional[Pair] = None,
+) -> tuple[
+ ConversionFunction, # conversion from world coordinates to pixel coordinates
+ ConversionFunction, # conversion from pixel coordinates to world coordinates
+]:
+ h, w = shape
+ x_range = sanitize_pair((0, w) if x_range is None else x_range, device)
+ y_range = sanitize_pair((0, h) if y_range is None else y_range, device)
+ minima, maxima = torch.stack((x_range, y_range), dim=-1)
+ wh = torch.tensor((w, h), dtype=torch.float32, device=device)
+
+ def convert_world_to_pixel(
+ xy: Float[Tensor, "*batch 2"],
+ ) -> Float[Tensor, "*batch 2"]:
+ return (xy - minima) / (maxima - minima) * wh
+
+ def convert_pixel_to_world(
+ xy: Float[Tensor, "*batch 2"],
+ ) -> Float[Tensor, "*batch 2"]:
+ return xy / wh * (maxima - minima) + minima
+
+ return convert_world_to_pixel, convert_pixel_to_world
diff --git a/src/visualization/drawing/lines.py b/src/visualization/drawing/lines.py
new file mode 100644
index 0000000000000000000000000000000000000000..85ce39825f456f5511049b2b1237836ded600416
--- /dev/null
+++ b/src/visualization/drawing/lines.py
@@ -0,0 +1,83 @@
+from typing import Literal, Optional
+
+import torch
+from einops import einsum, repeat
+from jaxtyping import Float
+from torch import Tensor
+
+from .coordinate_conversion import generate_conversions
+from .rendering import render_over_image
+from .types import Pair, Scalar, Vector, sanitize_scalar, sanitize_vector
+
+
+def draw_lines(
+ image: Float[Tensor, "3 height width"],
+ start: Vector,
+ end: Vector,
+ color: Vector,
+ width: Scalar,
+ cap: Literal["butt", "round", "square"] = "round",
+ num_msaa_passes: int = 1,
+ x_range: Optional[Pair] = None,
+ y_range: Optional[Pair] = None,
+) -> Float[Tensor, "3 height width"]:
+ device = image.device
+ start = sanitize_vector(start, 2, device)
+ end = sanitize_vector(end, 2, device)
+ color = sanitize_vector(color, 3, device)
+ width = sanitize_scalar(width, device)
+ (num_lines,) = torch.broadcast_shapes(
+ start.shape[0],
+ end.shape[0],
+ color.shape[0],
+ width.shape,
+ )
+
+ # Convert world-space points to pixel space.
+ _, h, w = image.shape
+ world_to_pixel, _ = generate_conversions((h, w), device, x_range, y_range)
+ start = world_to_pixel(start)
+ end = world_to_pixel(end)
+
+ def color_function(
+ xy: Float[Tensor, "point 2"],
+ ) -> Float[Tensor, "point 4"]:
+ # Define a vector between the start and end points.
+ delta = end - start
+ delta_norm = delta.norm(dim=-1, keepdim=True)
+ u_delta = delta / delta_norm
+
+ # Define a vector between each sample and the start point.
+ indicator = xy - start[:, None]
+
+ # Determine whether each sample is inside the line in the parallel direction.
+ extra = 0.5 * width[:, None] if cap == "square" else 0
+ parallel = einsum(u_delta, indicator, "l xy, l s xy -> l s")
+ parallel_inside_line = (parallel <= delta_norm + extra) & (parallel > -extra)
+
+ # Determine whether each sample is inside the line perpendicularly.
+ perpendicular = indicator - parallel[..., None] * u_delta[:, None]
+ perpendicular_inside_line = perpendicular.norm(dim=-1) < 0.5 * width[:, None]
+
+ inside_line = parallel_inside_line & perpendicular_inside_line
+
+ # Compute round caps.
+ if cap == "round":
+ near_start = indicator.norm(dim=-1) < 0.5 * width[:, None]
+ inside_line |= near_start
+ end_indicator = indicator = xy - end[:, None]
+ near_end = end_indicator.norm(dim=-1) < 0.5 * width[:, None]
+ inside_line |= near_end
+
+ # Determine the sample's color.
+ selectable_color = color.broadcast_to((num_lines, 3))
+ arrangement = inside_line * torch.arange(num_lines, device=device)[:, None]
+ top_color = selectable_color.gather(
+ dim=0,
+ index=repeat(arrangement.argmax(dim=0), "s -> s c", c=3),
+ )
+ rgba = torch.cat((top_color, inside_line.any(dim=0).float()[:, None]), dim=-1)
+
+ return rgba
+
+ return render_over_image(image, color_function, device, num_passes=num_msaa_passes)
diff --git a/src/visualization/drawing/points.py b/src/visualization/drawing/points.py
new file mode 100644
index 0000000000000000000000000000000000000000..671db100d34cd9121cb2778dcdb7252ec915bb2d
--- /dev/null
+++ b/src/visualization/drawing/points.py
@@ -0,0 +1,59 @@
+from typing import Optional
+
+import torch
+from einops import repeat
+from jaxtyping import Float
+from torch import Tensor
+
+from .coordinate_conversion import generate_conversions
+from .rendering import render_over_image
+from .types import Pair, Scalar, Vector, sanitize_scalar, sanitize_vector
+
+
+def draw_points(
+ image: Float[Tensor, "3 height width"],
+ points: Vector,
+ color: Vector = [1, 1, 1],
+ radius: Scalar = 1,
+ inner_radius: Scalar = 0,
+ num_msaa_passes: int = 1,
+ x_range: Optional[Pair] = None,
+ y_range: Optional[Pair] = None,
+) -> Float[Tensor, "3 height width"]:
+ device = image.device
+ points = sanitize_vector(points, 2, device)
+ color = sanitize_vector(color, 3, device)
+ radius = sanitize_scalar(radius, device)
+ inner_radius = sanitize_scalar(inner_radius, device)
+ (num_points,) = torch.broadcast_shapes(
+ points.shape[0],
+ color.shape[0],
+ radius.shape,
+ inner_radius.shape,
+ )
+
+ # Convert world-space points to pixel space.
+ _, h, w = image.shape
+ world_to_pixel, _ = generate_conversions((h, w), device, x_range, y_range)
+ points = world_to_pixel(points)
+
+ def color_function(
+ xy: Float[Tensor, "point 2"],
+ ) -> Float[Tensor, "point 4"]:
+ # Define a vector between the start and end points.
+ delta = xy[:, None] - points[None]
+ delta_norm = delta.norm(dim=-1)
+ mask = (delta_norm >= inner_radius[None]) & (delta_norm <= radius[None])
+
+ # Determine the sample's color.
+ selectable_color = color.broadcast_to((num_points, 3))
+ arrangement = mask * torch.arange(num_points, device=device)
+ top_color = selectable_color.gather(
+ dim=0,
+ index=repeat(arrangement.argmax(dim=1), "s -> s c", c=3),
+ )
+ rgba = torch.cat((top_color, mask.any(dim=1).float()[:, None]), dim=-1)
+
+ return rgba
+
+ return render_over_image(image, color_function, device, num_passes=num_msaa_passes)
diff --git a/src/visualization/drawing/rendering.py b/src/visualization/drawing/rendering.py
new file mode 100644
index 0000000000000000000000000000000000000000..65842a8ec5b9c6b109d19ef9509bc04ca0437ea7
--- /dev/null
+++ b/src/visualization/drawing/rendering.py
@@ -0,0 +1,152 @@
+from typing import Protocol, runtime_checkable
+
+import torch
+from einops import rearrange, reduce
+from jaxtyping import Bool, Float
+from torch import Tensor
+
+
+@runtime_checkable
+class ColorFunction(Protocol):
+ def __call__(
+ self,
+ xy: Float[Tensor, "point 2"],
+ ) -> Float[Tensor, "point 4"]: # RGBA color
+ pass
+
+
+def generate_sample_grid(
+ shape: tuple[int, int],
+ device: torch.device,
+) -> Float[Tensor, "height width 2"]:
+ h, w = shape
+ x = torch.arange(w, device=device) + 0.5
+ y = torch.arange(h, device=device) + 0.5
+ x, y = torch.meshgrid(x, y, indexing="xy")
+ return torch.stack([x, y], dim=-1)
+
+
+def detect_msaa_pixels(
+ image: Float[Tensor, "batch 4 height width"],
+) -> Bool[Tensor, "batch height width"]:
+ b, _, h, w = image.shape
+
+ mask = torch.zeros((b, h, w), dtype=torch.bool, device=image.device)
+
+ # Detect horizontal differences.
+ horizontal = (image[:, :, :, 1:] != image[:, :, :, :-1]).any(dim=1)
+ mask[:, :, 1:] |= horizontal
+ mask[:, :, :-1] |= horizontal
+
+ # Detect vertical differences.
+ vertical = (image[:, :, 1:, :] != image[:, :, :-1, :]).any(dim=1)
+ mask[:, 1:, :] |= vertical
+ mask[:, :-1, :] |= vertical
+
+ # Detect diagonal (top left to bottom right) differences.
+ tlbr = (image[:, :, 1:, 1:] != image[:, :, :-1, :-1]).any(dim=1)
+ mask[:, 1:, 1:] |= tlbr
+ mask[:, :-1, :-1] |= tlbr
+
+ # Detect diagonal (top right to bottom left) differences.
+ trbl = (image[:, :, :-1, 1:] != image[:, :, 1:, :-1]).any(dim=1)
+ mask[:, :-1, 1:] |= trbl
+ mask[:, 1:, :-1] |= trbl
+
+ return mask
+
+
+def reduce_straight_alpha(
+ rgba: Float[Tensor, "batch 4 height width"],
+) -> Float[Tensor, "batch 4"]:
+ color, alpha = rgba.split((3, 1), dim=1)
+
+ # Color becomes a weighted average of color (weighted by alpha).
+ weighted_color = reduce(color * alpha, "b c h w -> b c", "sum")
+ alpha_sum = reduce(alpha, "b c h w -> b c", "sum")
+ color = weighted_color / (alpha_sum + 1e-10)
+
+ # Alpha becomes mean alpha.
+ alpha = reduce(alpha, "b c h w -> b c", "mean")
+
+ return torch.cat((color, alpha), dim=-1)
+
+
+@torch.no_grad()
+def run_msaa_pass(
+ xy: Float[Tensor, "batch height width 2"],
+ color_function: ColorFunction,
+ scale: float,
+ subdivision: int,
+ remaining_passes: int,
+ device: torch.device,
+ batch_size: int = int(2**16),
+) -> Float[Tensor, "batch 4 height width"]: # color (RGBA with straight alpha)
+ # Sample the color function.
+ b, h, w, _ = xy.shape
+ color = [
+ color_function(batch)
+ for batch in rearrange(xy, "b h w xy -> (b h w) xy").split(batch_size)
+ ]
+ color = torch.cat(color, dim=0)
+ color = rearrange(color, "(b h w) c -> b c h w", b=b, h=h, w=w)
+
+ # If any MSAA passes remain, subdivide.
+ if remaining_passes > 0:
+ mask = detect_msaa_pixels(color)
+ batch_index, row_index, col_index = torch.where(mask)
+ xy = xy[batch_index, row_index, col_index]
+
+ offsets = generate_sample_grid((subdivision, subdivision), device)
+ offsets = (offsets / subdivision - 0.5) * scale
+
+ color_fine = run_msaa_pass(
+ xy[:, None, None] + offsets,
+ color_function,
+ scale / subdivision,
+ subdivision,
+ remaining_passes - 1,
+ device,
+ batch_size=batch_size,
+ )
+ color[batch_index, :, row_index, col_index] = reduce_straight_alpha(color_fine)
+
+ return color
+
+
+@torch.no_grad()
+def render(
+ shape: tuple[int, int],
+ color_function: ColorFunction,
+ device: torch.device,
+ subdivision: int = 8,
+ num_passes: int = 2,
+) -> Float[Tensor, "4 height width"]: # color (RGBA with straight alpha)
+ xy = generate_sample_grid(shape, device)
+ return run_msaa_pass(
+ xy[None],
+ color_function,
+ 1.0,
+ subdivision,
+ num_passes,
+ device,
+ )[0]
+
+
+def render_over_image(
+ image: Float[Tensor, "3 height width"],
+ color_function: ColorFunction,
+ device: torch.device,
+ subdivision: int = 8,
+ num_passes: int = 1,
+) -> Float[Tensor, "3 height width"]:
+ _, h, w = image.shape
+ overlay = render(
+ (h, w),
+ color_function,
+ device,
+ subdivision=subdivision,
+ num_passes=num_passes,
+ )
+ color, alpha = overlay.split((3, 1), dim=0)
+ return image * (1 - alpha) + color * alpha
diff --git a/src/visualization/drawing/types.py b/src/visualization/drawing/types.py
new file mode 100644
index 0000000000000000000000000000000000000000..4bbbe3a068052ac389a8332304ddbcea78762b86
--- /dev/null
+++ b/src/visualization/drawing/types.py
@@ -0,0 +1,67 @@
+from typing import Iterable, Union
+
+import torch
+from einops import repeat
+from jaxtyping import Float, Shaped
+from torch import Tensor
+
+Real = Union[float, int]
+
+Vector = Union[
+ Real,
+ Iterable[Real],
+ Shaped[Tensor, "dim"],
+ Shaped[Tensor, "batch dim"],
+]
+
+
+def sanitize_vector(
+ vector: Vector,
+ dim: int,
+ device: torch.device,
+) -> Float[Tensor, "*#batch dim"]:
+ if isinstance(vector, Tensor):
+ vector = vector.type(torch.float32).to(device)
+ else:
+ vector = torch.tensor(vector, dtype=torch.float32, device=device)
+ while vector.ndim < 2:
+ vector = vector[None]
+ if vector.shape[-1] == 1:
+ vector = repeat(vector, "... () -> ... c", c=dim)
+ assert vector.shape[-1] == dim
+ assert vector.ndim == 2
+ return vector
+
+
+Scalar = Union[
+ Real,
+ Iterable[Real],
+ Shaped[Tensor, ""],
+ Shaped[Tensor, " batch"],
+]
+
+
+def sanitize_scalar(scalar: Scalar, device: torch.device) -> Float[Tensor, "*#batch"]:
+ if isinstance(scalar, Tensor):
+ scalar = scalar.type(torch.float32).to(device)
+ else:
+ scalar = torch.tensor(scalar, dtype=torch.float32, device=device)
+ while scalar.ndim < 1:
+ scalar = scalar[None]
+ assert scalar.ndim == 1
+ return scalar
+
+
+Pair = Union[
+ Iterable[Real],
+ Shaped[Tensor, "2"],
+]
+
+
+def sanitize_pair(pair: Pair, device: torch.device) -> Float[Tensor, "2"]:
+ if isinstance(pair, Tensor):
+ pair = pair.type(torch.float32).to(device)
+ else:
+ pair = torch.tensor(pair, dtype=torch.float32, device=device)
+ assert pair.shape == (2,)
+ return pair
diff --git a/src/visualization/layout.py b/src/visualization/layout.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca44b8a0814f8e601d82aaea9dd737d8881c972c
--- /dev/null
+++ b/src/visualization/layout.py
@@ -0,0 +1,228 @@
+"""This file contains useful layout utilities for images. They are:
+
+- add_border: Add a border to an image.
+- cat/hcat/vcat: Join images by arranging them in a line. If the images have different
+ sizes, they are aligned as specified (start, end, center). Allows you to specify a gap
+ between images.
+
+Images are assumed to be float32 tensors with shape (channel, height, width).
+"""
+
+from typing import Any, Generator, Iterable, Literal, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from jaxtyping import Float
+from torch import Tensor
+
+Alignment = Literal["start", "center", "end"]
+Axis = Literal["horizontal", "vertical"]
+Color = Union[
+ int,
+ float,
+ Iterable[int],
+ Iterable[float],
+ Float[Tensor, "#channel"],
+ Float[Tensor, ""],
+]
+
+
+def _sanitize_color(color: Color) -> Float[Tensor, "#channel"]:
+ # Convert tensor to list (or individual item).
+ if isinstance(color, torch.Tensor):
+ color = color.tolist()
+
+ # Turn iterators and individual items into lists.
+ if isinstance(color, Iterable):
+ color = list(color)
+ else:
+ color = [color]
+
+ return torch.tensor(color, dtype=torch.float32)
+
+
+def _intersperse(iterable: Iterable, delimiter: Any) -> Generator[Any, None, None]:
+ it = iter(iterable)
+ yield next(it)
+ for item in it:
+ yield delimiter
+ yield item
+
+
+def _get_main_dim(main_axis: Axis) -> int:
+ return {
+ "horizontal": 2,
+ "vertical": 1,
+ }[main_axis]
+
+
+def _get_cross_dim(main_axis: Axis) -> int:
+ return {
+ "horizontal": 1,
+ "vertical": 2,
+ }[main_axis]
+
+
+def _compute_offset(base: int, overlay: int, align: Alignment) -> slice:
+ assert base >= overlay
+ offset = {
+ "start": 0,
+ "center": (base - overlay) // 2,
+ "end": base - overlay,
+ }[align]
+ return slice(offset, offset + overlay)
+
+
+def overlay(
+ base: Float[Tensor, "channel base_height base_width"],
+ overlay: Float[Tensor, "channel overlay_height overlay_width"],
+ main_axis: Axis,
+ main_axis_alignment: Alignment,
+ cross_axis_alignment: Alignment,
+) -> Float[Tensor, "channel base_height base_width"]:
+ # The overlay must be smaller than the base.
+ _, base_height, base_width = base.shape
+ _, overlay_height, overlay_width = overlay.shape
+ assert base_height >= overlay_height and base_width >= overlay_width
+
+ # Compute spacing on the main dimension.
+ main_dim = _get_main_dim(main_axis)
+ main_slice = _compute_offset(
+ base.shape[main_dim], overlay.shape[main_dim], main_axis_alignment
+ )
+
+ # Compute spacing on the cross dimension.
+ cross_dim = _get_cross_dim(main_axis)
+ cross_slice = _compute_offset(
+ base.shape[cross_dim], overlay.shape[cross_dim], cross_axis_alignment
+ )
+
+ # Combine the slices and paste the overlay onto the base accordingly.
+ selector = [..., None, None]
+ selector[main_dim] = main_slice
+ selector[cross_dim] = cross_slice
+ result = base.clone()
+ result[selector] = overlay
+ return result
+
+
+def cat(
+ main_axis: Axis,
+ *images: Iterable[Float[Tensor, "channel _ _"]],
+ align: Alignment = "center",
+ gap: int = 8,
+ gap_color: Color = 1,
+) -> Float[Tensor, "channel height width"]:
+ """Arrange images in a line. The interface resembles a CSS div with flexbox."""
+ device = images[0].device
+ gap_color = _sanitize_color(gap_color).to(device)
+
+ # Find the maximum image side length in the cross axis dimension.
+ cross_dim = _get_cross_dim(main_axis)
+ cross_axis_length = max(image.shape[cross_dim] for image in images)
+
+ # Pad the images.
+ padded_images = []
+ for image in images:
+ # Create an empty image with the correct size.
+ padded_shape = list(image.shape)
+ padded_shape[cross_dim] = cross_axis_length
+ base = torch.ones(padded_shape, dtype=torch.float32, device=device)
+ base = base * gap_color[:, None, None]
+ padded_images.append(overlay(base, image, main_axis, "start", align))
+
+ # Intersperse separators if necessary.
+ if gap > 0:
+ # Generate a separator.
+ c, _, _ = images[0].shape
+ separator_size = [gap, gap]
+ separator_size[cross_dim - 1] = cross_axis_length
+ separator = torch.ones((c, *separator_size), dtype=torch.float32, device=device)
+ separator = separator * gap_color[:, None, None]
+
+ # Intersperse the separator between the images.
+ padded_images = list(_intersperse(padded_images, separator))
+
+ return torch.cat(padded_images, dim=_get_main_dim(main_axis))
+
+
+def hcat(
+ *images: Iterable[Float[Tensor, "channel _ _"]],
+ align: Literal["start", "center", "end", "top", "bottom"] = "start",
+ gap: int = 8,
+ gap_color: Color = 1,
+):
+ """Shorthand for a horizontal linear concatenation."""
+ return cat(
+ "horizontal",
+ *images,
+ align={
+ "start": "start",
+ "center": "center",
+ "end": "end",
+ "top": "start",
+ "bottom": "end",
+ }[align],
+ gap=gap,
+ gap_color=gap_color,
+ )
+
+
+def vcat(
+ *images: Iterable[Float[Tensor, "channel _ _"]],
+ align: Literal["start", "center", "end", "left", "right"] = "start",
+ gap: int = 8,
+ gap_color: Color = 1,
+):
+ """Shorthand for a horizontal linear concatenation."""
+ return cat(
+ "vertical",
+ *images,
+ align={
+ "start": "start",
+ "center": "center",
+ "end": "end",
+ "left": "start",
+ "right": "end",
+ }[align],
+ gap=gap,
+ gap_color=gap_color,
+ )
+
+
+def add_border(
+ image: Float[Tensor, "channel height width"],
+ border: int = 8,
+ color: Color = 1,
+) -> Float[Tensor, "channel new_height new_width"]:
+ color = _sanitize_color(color).to(image)
+ c, h, w = image.shape
+ result = torch.empty(
+ (c, h + 2 * border, w + 2 * border), dtype=torch.float32, device=image.device
+ )
+ result[:] = color[:, None, None]
+ result[:, border : h + border, border : w + border] = image
+ return result
+
+
+def resize(
+ image: Float[Tensor, "channel height width"],
+ shape: Optional[tuple[int, int]] = None,
+ width: Optional[int] = None,
+ height: Optional[int] = None,
+) -> Float[Tensor, "channel new_height new_width"]:
+ assert (shape is not None) + (width is not None) + (height is not None) == 1
+ _, h, w = image.shape
+
+ if width is not None:
+ shape = (int(h * width / w), width)
+ elif height is not None:
+ shape = (height, int(w * height / h))
+
+ return F.interpolate(
+ image[None],
+ shape,
+ mode="bilinear",
+ align_corners=False,
+ antialias="bilinear",
+ )[0]
diff --git a/src/visualization/validation_in_3d.py b/src/visualization/validation_in_3d.py
new file mode 100644
index 0000000000000000000000000000000000000000..0fceafe9493cb33260c0446d7fd07fd2a6a27bb1
--- /dev/null
+++ b/src/visualization/validation_in_3d.py
@@ -0,0 +1,115 @@
+import torch
+from jaxtyping import Float, Shaped
+from torch import Tensor
+
+from ..model.decoder.cuda_splatting import render_cuda_orthographic
+from ..model.types import Gaussians
+from ..visualization.annotation import add_label
+from ..visualization.drawing.cameras import draw_cameras
+from .drawing.cameras import compute_equal_aabb_with_margin
+
+
+def pad(images: list[Shaped[Tensor, "..."]]) -> list[Shaped[Tensor, "..."]]:
+ shapes = torch.stack([torch.tensor(x.shape) for x in images])
+ padded_shape = shapes.max(dim=0)[0]
+ results = [
+ torch.ones(padded_shape.tolist(), dtype=x.dtype, device=x.device)
+ for x in images
+ ]
+ for image, result in zip(images, results):
+ slices = [slice(0, x) for x in image.shape]
+ result[slices] = image[slices]
+ return results
+
+
+def render_projections(
+ gaussians: Gaussians,
+ resolution: int,
+ margin: float = 0.1,
+ draw_label: bool = True,
+ extra_label: str = "",
+) -> Float[Tensor, "batch 3 3 height width"]:
+ device = gaussians.means.device
+ b, _, _ = gaussians.means.shape
+
+ # Compute the minima and maxima of the scene.
+ minima = gaussians.means.min(dim=1).values
+ maxima = gaussians.means.max(dim=1).values
+ scene_minima, scene_maxima = compute_equal_aabb_with_margin(
+ minima, maxima, margin=margin
+ )
+
+ projections = []
+ for look_axis in range(3):
+ right_axis = (look_axis + 1) % 3
+ down_axis = (look_axis + 2) % 3
+
+ # Define the extrinsics for rendering.
+ extrinsics = torch.zeros((b, 4, 4), dtype=torch.float32, device=device)
+ extrinsics[:, right_axis, 0] = 1
+ extrinsics[:, down_axis, 1] = 1
+ extrinsics[:, look_axis, 2] = 1
+ extrinsics[:, right_axis, 3] = 0.5 * (
+ scene_minima[:, right_axis] + scene_maxima[:, right_axis]
+ )
+ extrinsics[:, down_axis, 3] = 0.5 * (
+ scene_minima[:, down_axis] + scene_maxima[:, down_axis]
+ )
+ extrinsics[:, look_axis, 3] = scene_minima[:, look_axis]
+ extrinsics[:, 3, 3] = 1
+
+ # Define the intrinsics for rendering.
+ extents = scene_maxima - scene_minima
+ far = extents[:, look_axis]
+ near = torch.zeros_like(far)
+ width = extents[:, right_axis]
+ height = extents[:, down_axis]
+
+ projection = render_cuda_orthographic(
+ extrinsics,
+ width,
+ height,
+ near,
+ far,
+ (resolution, resolution),
+ torch.zeros((b, 3), dtype=torch.float32, device=device),
+ gaussians.means,
+ gaussians.covariances,
+ gaussians.harmonics,
+ gaussians.opacities,
+ fov_degrees=10.0,
+ )
+ if draw_label:
+ right_axis_name = "XYZ"[right_axis]
+ down_axis_name = "XYZ"[down_axis]
+ label = f"{right_axis_name}{down_axis_name} Projection {extra_label}"
+ projection = torch.stack([add_label(x, label) for x in projection])
+
+ projections.append(projection)
+
+ return torch.stack(pad(projections), dim=1)
+
+
+def render_cameras(batch: dict, resolution: int) -> Float[Tensor, "3 3 height width"]:
+ # Define colors for context and target views.
+ num_context_views = batch["context"]["extrinsics"].shape[1]
+ num_target_views = batch["target"]["extrinsics"].shape[1]
+ color = torch.ones(
+ (num_target_views + num_context_views, 3),
+ dtype=torch.float32,
+ device=batch["target"]["extrinsics"].device,
+ )
+ color[num_context_views:, 1:] = 0
+
+ return draw_cameras(
+ resolution,
+ torch.cat(
+ (batch["context"]["extrinsics"][0], batch["target"]["extrinsics"][0])
+ ),
+ torch.cat(
+ (batch["context"]["intrinsics"][0], batch["target"]["intrinsics"][0])
+ ),
+ color,
+ torch.cat((batch["context"]["near"][0], batch["target"]["near"][0])),
+ torch.cat((batch["context"]["far"][0], batch["target"]["far"][0])),
+ )
diff --git a/src/visualization/video_render.py b/src/visualization/video_render.py
new file mode 100644
index 0000000000000000000000000000000000000000..b7a6153e8d5d98b947f533fae19087122e56bf6f
--- /dev/null
+++ b/src/visualization/video_render.py
@@ -0,0 +1,162 @@
+from typing import Protocol, runtime_checkable
+
+import cv2
+import torch
+from einops import repeat, pack
+from jaxtyping import Float
+from torch import Tensor
+
+from .camera_trajectory.interpolation import interpolate_extrinsics, interpolate_intrinsics
+from .camera_trajectory.wobble import generate_wobble, generate_wobble_transformation
+from .layout import vcat
+from ..dataset.types import BatchedExample
+from ..misc.image_io import save_video
+from ..misc.utils import vis_depth_map
+from ..model.decoder import Decoder
+from ..model.types import Gaussians
+
+
+@runtime_checkable
+class TrajectoryFn(Protocol):
+ def __call__(
+ self,
+ t: Float[Tensor, " t"],
+ ) -> tuple[
+ Float[Tensor, "batch view 4 4"], # extrinsics
+ Float[Tensor, "batch view 3 3"], # intrinsics
+ ]:
+ pass
+
+
+def render_video_wobble(
+ gaussians: Gaussians,
+ decoder: Decoder,
+ batch: BatchedExample,
+ num_frames: int = 60,
+ smooth: bool = True,
+ loop_reverse: bool = True,
+ add_depth: bool = False,
+) -> Tensor:
+ # Two views are needed to get the wobble radius，use the first and the last view
+ _, v, _, _ = batch["context"]["extrinsics"].shape
+
+ def trajectory_fn(t):
+ origin_a = batch["context"]["extrinsics"][:, 0, :3, 3]
+ origin_b = batch["context"]["extrinsics"][:, -1, :3, 3]
+ delta = (origin_a - origin_b).norm(dim=-1)
+ extrinsics = generate_wobble(
+ batch["context"]["extrinsics"][:, 0],
+ delta * 0.25,
+ t,
+ )
+ intrinsics = repeat(
+ batch["context"]["intrinsics"][:, 0],
+ "b i j -> b v i j",
+ v=t.shape[0],
+ )
+ return extrinsics, intrinsics
+
+ return render_video_generic(gaussians, decoder, batch, trajectory_fn, num_frames, smooth, loop_reverse, add_depth)
+
+
+def render_video_interpolation(
+ gaussians: Gaussians,
+ decoder: Decoder,
+ batch: BatchedExample,
+ num_frames: int = 60,
+ smooth: bool = True,
+ loop_reverse: bool = True,
+ add_depth: bool = False,
+) -> Tensor:
+ _, v, _, _ = batch["context"]["extrinsics"].shape
+
+ def trajectory_fn(t):
+ extrinsics = interpolate_extrinsics(
+ batch["context"]["extrinsics"][0, 0],
+ batch["context"]["extrinsics"][0, -1],
+ t,
+ )
+ intrinsics = interpolate_intrinsics(
+ batch["context"]["intrinsics"][0, 0],
+ batch["context"]["intrinsics"][0, -1],
+ t,
+ )
+ return extrinsics[None], intrinsics[None]
+
+ return render_video_generic(gaussians, decoder, batch, trajectory_fn, num_frames, smooth, loop_reverse, add_depth)
+
+
+def render_video_interpolation_exaggerated(
+ gaussians: Gaussians,
+ decoder: Decoder,
+ batch: BatchedExample,
+ num_frames: int = 300,
+ smooth: bool = False,
+ loop_reverse: bool = False,
+ add_depth: bool = False,
+) -> Tensor:
+ # Two views are needed to get the wobble radius.
+ _, v, _, _ = batch["context"]["extrinsics"].shape
+
+ def trajectory_fn(t):
+ origin_a = batch["context"]["extrinsics"][:, 0, :3, 3]
+ origin_b = batch["context"]["extrinsics"][:, -1, :3, 3]
+ delta = (origin_a - origin_b).norm(dim=-1)
+ tf = generate_wobble_transformation(
+ delta * 0.5,
+ t,
+ 5,
+ scale_radius_with_t=False,
+ )
+ extrinsics = interpolate_extrinsics(
+ batch["context"]["extrinsics"][0, 0],
+ batch["context"]["extrinsics"][0, -1],
+ t * 5 - 2,
+ )
+ intrinsics = interpolate_intrinsics(
+ batch["context"]["intrinsics"][0, 0],
+ batch["context"]["extrinsics"][0, -1],
+ t * 5 - 2,
+ )
+ return extrinsics @ tf, intrinsics[None]
+
+ return render_video_generic(gaussians, decoder, batch, trajectory_fn, num_frames, smooth, loop_reverse, add_depth)
+
+
+def render_video_generic(
+ gaussians: Gaussians,
+ decoder: Decoder,
+ batch: BatchedExample,
+ trajectory_fn: TrajectoryFn,
+ num_frames: int = 30,
+ smooth: bool = True,
+ loop_reverse: bool = True,
+ add_depth: bool = False,
+) -> Tensor:
+ device = gaussians.means.device
+
+ t = torch.linspace(0, 1, num_frames, dtype=torch.float32, device=device)
+ if smooth:
+ t = (torch.cos(torch.pi * (t + 1)) + 1) / 2
+
+ extrinsics, intrinsics = trajectory_fn(t)
+
+ _, _, _, h, w = batch["context"]["image"].shape
+
+ near = repeat(batch["context"]["near"][:, 0], "b -> b v", v=num_frames)
+ far = repeat(batch["context"]["far"][:, 0], "b -> b v", v=num_frames)
+ output = decoder.forward(
+ gaussians, extrinsics, intrinsics, near, far, (h, w), "depth"
+ )
+ images = [
+ vcat(rgb, depth) if add_depth else rgb
+ for rgb, depth in zip(output.color[0], vis_depth_map(output.depth[0]))
+ ]
+
+ video = torch.stack(images)
+ # video = (video.clip(min=0, max=1) * 255).type(torch.uint8).cpu().numpy()
+ if loop_reverse:
+ # video = pack([video, video[::-1][1:-1]], "* c h w")[0]
+ video = pack([video, video.flip(dims=(0,))[1:-1]], "* c h w")[0]
+
+ return video