├── assets
├── demo.gif
├── example_ref.png
└── example_input.png
├── requirements.txt
├── examples
├── gsplat
│ ├── requirements.txt
│ ├── datasets
│ │ ├── normalize.py
│ │ ├── traj.py
│ │ └── colmap.py
│ ├── utils.py
│ └── lib_bilagrid.py
├── nerfstudio
│ ├── pyproject.toml
│ └── difix3d
│ │ ├── difix3d_config.py
│ │ ├── difix3d_datamanager.py
│ │ ├── difix3d_trainer.py
│ │ ├── difix3d.py
│ │ ├── difix3d_pipeline.py
│ │ └── difix3d_field.py
└── utils.py
├── src
├── dataset.py
├── loss.py
├── inference_difix.py
├── model.py
└── train_difix.py
├── README.md
└── LICENSE.txt
/assets/demo.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nv-tlabs/Difix3D/HEAD/assets/demo.gif
--------------------------------------------------------------------------------
/assets/example_ref.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nv-tlabs/Difix3D/HEAD/assets/example_ref.png
--------------------------------------------------------------------------------
/assets/example_input.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nv-tlabs/Difix3D/HEAD/assets/example_input.png
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | torch
2 | torchvision
3 | torchmetrics
4 | wandb
5 | imageio[ffmpeg]
6 | einops
7 | lpips
8 | xformers
9 | peft==0.9.0
10 | diffusers==0.25.1
11 | huggingface-hub==0.25.1
12 | transformers==4.38.0
--------------------------------------------------------------------------------
/examples/gsplat/requirements.txt:
--------------------------------------------------------------------------------
1 | # assume torch is already installed
2 |
3 | # pycolmap for data parsing
4 | git+https://github.com/rmbrualla/pycolmap@cc7ea4b7301720ac29287dbe450952511b32125e
5 | # (optional) nerfacc for torch version rasterization
6 | # git+https://github.com/nerfstudio-project/nerfacc
7 |
8 | viser
9 | nerfview==0.0.2
10 | imageio[ffmpeg]
11 | numpy<2.0.0
12 | scikit-learn
13 | tqdm
14 | torchmetrics[image]
15 | opencv-python
16 | tyro>=0.8.8
17 | Pillow
18 | tensorboard
19 | tensorly
20 | pyyaml
21 | matplotlib
22 | git+https://github.com/rahul-goel/fused-ssim@1272e21a282342e89537159e4bad508b19b34157
23 |
--------------------------------------------------------------------------------
/examples/nerfstudio/pyproject.toml:
--------------------------------------------------------------------------------
1 | [project]
2 | name = "difix3d"
3 | version = "0.1.0"
4 |
5 | dependencies=[
6 | "nerfstudio>=0.3.0",
7 | ]
8 |
9 | # black
10 | [tool.black]
11 | line-length = 120
12 |
13 | # pylint
14 | [tool.pylint.messages_control]
15 | max-line-length = 120
16 | generated-members = ["numpy.*", "torch.*", "cv2.*", "cv.*"]
17 | good-names-rgxs = "^[_a-zA-Z][_a-z0-9]?$"
18 | ignore-paths = ["scripts/colmap2nerf.py"]
19 | jobs = 0
20 | ignored-classes = ["TensorDataclass"]
21 |
22 | disable = [
23 | "duplicate-code",
24 | "fixme",
25 | "logging-fstring-interpolation",
26 | "too-many-arguments",
27 | "too-many-branches",
28 | "too-many-instance-attributes",
29 | "too-many-locals",
30 | "unnecessary-ellipsis",
31 | ]
32 |
33 | [tool.setuptools.packages.find]
34 | include = ["difix3d"]
35 |
36 | [project.entry-points.'nerfstudio.method_configs']
37 | difix3d = 'difix3d.difix3d_config:difix3d_method'
--------------------------------------------------------------------------------
/src/dataset.py:
--------------------------------------------------------------------------------
1 | import json
2 | import torch
3 | from PIL import Image
4 | import torchvision.transforms.functional as F
5 |
6 |
7 | class PairedDataset(torch.utils.data.Dataset):
8 | def __init__(self, dataset_path, split, height=576, width=1024, tokenizer=None):
9 |
10 | super().__init__()
11 | with open(dataset_path, "r") as f:
12 | self.data = json.load(f)[split]
13 | self.img_ids = list(self.data.keys())
14 | self.image_size = (height, width)
15 | self.tokenizer = tokenizer
16 |
17 | def __len__(self):
18 |
19 | return len(self.img_ids)
20 |
21 | def __getitem__(self, idx):
22 |
23 | img_id = self.img_ids[idx]
24 |
25 | input_img = self.data[img_id]["image"]
26 | output_img = self.data[img_id]["target_image"]
27 | ref_img = self.data[img_id]["ref_image"] if "ref_image" in self.data[img_id] else None
28 | caption = self.data[img_id]["prompt"]
29 |
30 | try:
31 | input_img = Image.open(input_img)
32 | output_img = Image.open(output_img)
33 | except:
34 | print("Error loading image:", input_img, output_img)
35 | return self.__getitem__(idx + 1)
36 |
37 | img_t = F.to_tensor(img_t)
38 | img_t = F.resize(img_t, self.image_size)
39 | img_t = F.normalize(img_t, mean=[0.5], std=[0.5])
40 |
41 | output_t = F.to_tensor(output_t)
42 | output_t = F.resize(output_t, self.image_size)
43 | output_t = F.normalize(output_t, mean=[0.5], std=[0.5])
44 |
45 | if ref_img is not None:
46 | ref_img = Image.open(ref_img)
47 | ref_t = F.to_tensor(ref_t)
48 | ref_t = F.resize(ref_t, self.image_size)
49 | ref_t = F.normalize(ref_t, mean=[0.5], std=[0.5])
50 |
51 | img_t = torch.stack([img_t, ref_t], dim=0)
52 | output_t = torch.stack([output_t, ref_t], dim=0)
53 | else:
54 | img_t = img_t.unsqueeze(0)
55 | output_t = output_t.unsqueeze(0)
56 |
57 | out = {
58 | "output_pixel_values": output_t,
59 | "conditioning_pixel_values": img_t,
60 | "caption": caption,
61 | }
62 |
63 | if self.tokenizer is not None:
64 | input_ids = self.tokenizer(
65 | caption, max_length=self.tokenizer.model_max_length,
66 | padding="max_length", truncation=True, return_tensors="pt"
67 | ).input_ids
68 | out["input_ids"] = input_ids
69 |
70 | return out
71 |
--------------------------------------------------------------------------------
/src/loss.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torchvision import transforms, models
3 |
4 | # Define style weights for different layers
5 | STYLE_WEIGHTS = {
6 | 'relu1_2': 1.0 / 2.6,
7 | 'relu2_2': 1.0 / 4.8,
8 | 'relu3_3': 1.0 / 3.7,
9 | 'relu4_3': 1.0 / 5.6,
10 | 'relu5_3': 10.0 / 1.5
11 | }
12 |
13 | def get_features(image, model, layers=None):
14 | """
15 | Extract features from specific layers of a model for a given image.
16 |
17 | Args:
18 | image (torch.Tensor): Input image tensor.
19 | model (torch.nn.Module): Pretrained model (e.g., VGG).
20 | layers (dict): Mapping of layer indices to layer names.
21 |
22 | Returns:
23 | dict: A dictionary of features for the specified layers.
24 | """
25 | if layers is None:
26 | layers = {
27 | '3': 'relu1_2',
28 | '8': 'relu2_2',
29 | '15': 'relu3_3',
30 | '22': 'relu4_3',
31 | '29': 'relu5_3'
32 | }
33 |
34 | features = {}
35 | x = image
36 | for name, layer in model._modules.items():
37 | x = layer(x)
38 | if name in layers:
39 | features[layers[name]] = x
40 | return features
41 |
42 | def gram_matrix(tensor):
43 | """
44 | Compute the Gram matrix for a given tensor.
45 |
46 | Args:
47 | tensor (torch.Tensor): Input tensor of shape (batch_size, depth, height, width).
48 |
49 | Returns:
50 | torch.Tensor: Gram matrix of the input tensor.
51 | """
52 | b, d, h, w = tensor.size()
53 | tensor = tensor.view(b * d, h * w) # Reshape tensor for matrix multiplication
54 | gram = torch.mm(tensor, tensor.t()) # Compute Gram matrix
55 | return gram
56 |
57 | def gram_loss(style, target, model):
58 | """
59 | Compute the Gram loss (style loss) between a style image and a target image.
60 |
61 | Args:
62 | style (torch.Tensor): Style image tensor.
63 | target (torch.Tensor): Target image tensor.
64 | model (torch.nn.Module): Pretrained model (e.g., VGG).
65 |
66 | Returns:
67 | torch.Tensor: The computed Gram loss.
68 | """
69 | # Extract features for the style and target images
70 | style_features = get_features(style, model)
71 | target_features = get_features(target, model)
72 |
73 | # Compute Gram matrices for the style image
74 | style_grams = {layer: gram_matrix(style_features[layer]) for layer in style_features}
75 |
76 | # Initialize total loss
77 | total_loss = 0
78 |
79 | # Compute the weighted Gram loss for each layer
80 | for layer, weight in STYLE_WEIGHTS.items():
81 | target_feature = target_features[layer]
82 | target_gram = gram_matrix(target_feature)
83 | style_gram = style_grams[layer]
84 |
85 | # Compute the layer-specific Gram loss
86 | _, d, h, w = target_feature.shape
87 | layer_loss = weight * torch.mean((target_gram - style_gram) ** 2)
88 | total_loss += layer_loss / (d * h * w)
89 |
90 | return total_loss
--------------------------------------------------------------------------------
/examples/nerfstudio/difix3d/difix3d_config.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 The Nerfstudio Team. All rights reserved.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 |
16 | from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
17 | from nerfstudio.configs.base_config import ViewerConfig
18 | from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
19 | from nerfstudio.engine.optimizers import AdamOptimizerConfig
20 | from nerfstudio.engine.schedulers import ExponentialDecaySchedulerConfig
21 | from nerfstudio.plugins.types import MethodSpecification
22 |
23 | from difix3d.difix3d_datamanager import Difix3DDataManagerConfig
24 | from difix3d.difix3d import Difix3DModelConfig
25 | from difix3d.difix3d_pipeline import Difix3DPipelineConfig
26 | from difix3d.difix3d_trainer import Difix3DTrainerConfig
27 |
28 | difix3d_method = MethodSpecification(
29 | config=Difix3DTrainerConfig(
30 | method_name="difix3d",
31 | # steps_per_eval_batch=1000,
32 | # steps_per_eval_image=100,
33 | # steps_per_save=250,
34 | # max_num_iterations=15000,
35 | # save_only_latest_checkpoint=True,
36 | # mixed_precision=False,
37 | steps_per_eval_batch=500,
38 | steps_per_save=2000,
39 | max_num_iterations=30000,
40 | mixed_precision=True,
41 | pipeline=Difix3DPipelineConfig(
42 | datamanager=Difix3DDataManagerConfig(
43 | dataparser=NerfstudioDataParserConfig(),
44 | train_num_rays_per_batch=16384,
45 | eval_num_rays_per_batch=4096,
46 | ),
47 | model=Difix3DModelConfig(
48 | eval_num_rays_per_chunk=1 << 15,
49 | average_init_density=0.01,
50 | camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"),
51 | ),
52 | ),
53 | optimizers={
54 | "proposal_networks": {
55 | "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
56 | "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
57 | },
58 | "fields": {
59 | "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
60 | "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
61 | },
62 | "camera_opt": {
63 | "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
64 | "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=5000),
65 | },
66 | },
67 | viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
68 | vis="viewer",
69 | ),
70 | description="Difix3D",
71 | )
--------------------------------------------------------------------------------
/examples/nerfstudio/difix3d/difix3d_datamanager.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 The Nerfstudio Team. All rights reserved.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 |
16 | from __future__ import annotations
17 |
18 | from dataclasses import dataclass, field
19 | from typing import Dict, Tuple, Type
20 |
21 | from rich.progress import Console
22 |
23 | from nerfstudio.cameras.rays import RayBundle
24 | from nerfstudio.data.utils.dataloaders import CacheDataloader
25 | from nerfstudio.model_components.ray_generators import RayGenerator
26 | from nerfstudio.data.datamanagers.base_datamanager import (
27 | VanillaDataManager,
28 | VanillaDataManagerConfig,
29 | )
30 |
31 | CONSOLE = Console(width=120)
32 |
33 | @dataclass
34 | class Difix3DDataManagerConfig(VanillaDataManagerConfig):
35 | """Configuration for the Difix3DDataManager."""
36 |
37 | _target: Type = field(default_factory=lambda: Difix3DDataManager)
38 | patch_size: int = 32
39 | """Size of patch to sample from. If >1, patch-based sampling will be used."""
40 |
41 |
42 | class Difix3DDataManager(VanillaDataManager):
43 | """Data manager for Difix3D."""
44 |
45 | config: Difix3DDataManagerConfig
46 |
47 | def setup_train(self):
48 | """Sets up the data loaders for training"""
49 | assert self.train_dataset is not None
50 | CONSOLE.print("Setting up training dataset...")
51 | self.train_image_dataloader = CacheDataloader(
52 | self.train_dataset,
53 | num_images_to_sample_from=self.config.train_num_images_to_sample_from,
54 | num_times_to_repeat_images=self.config.train_num_times_to_repeat_images,
55 | device=self.device,
56 | num_workers=self.world_size * 4,
57 | pin_memory=True,
58 | collate_fn=self.config.collate_fn,
59 | exclude_batch_keys_from_device=self.exclude_batch_keys_from_device,
60 | )
61 | self.iter_train_image_dataloader = iter(self.train_image_dataloader)
62 | self.train_pixel_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch)
63 | self.train_ray_generator = RayGenerator(self.train_dataset.cameras.to(self.device))
64 |
65 | def next_train(self, step: int) -> Tuple[RayBundle, Dict]:
66 | """Returns the next batch of data from the train dataloader."""
67 | self.train_count += 1
68 | image_batch = next(self.iter_train_image_dataloader)
69 | assert self.train_pixel_sampler is not None
70 | assert isinstance(image_batch, dict)
71 | batch = self.train_pixel_sampler.sample(image_batch)
72 | ray_indices = batch["indices"]
73 | ray_bundle = self.train_ray_generator(ray_indices)
74 | return ray_bundle, batch
75 |
--------------------------------------------------------------------------------
/src/inference_difix.py:
--------------------------------------------------------------------------------
1 | import os
2 | import imageio
3 | import argparse
4 | import numpy as np
5 | from PIL import Image
6 | from glob import glob
7 | from tqdm import tqdm
8 | from model import Difix
9 |
10 |
11 | if __name__ == "__main__":
12 | # Argument parser
13 | parser = argparse.ArgumentParser()
14 | parser.add_argument('--input_image', type=str, required=True, help='Path to the input image or directory')
15 | parser.add_argument('--ref_image', type=str, default=None, help='Path to the reference image or directory')
16 | parser.add_argument('--height', type=int, default=576, help='Height of the input image')
17 | parser.add_argument('--width', type=int, default=1024, help='Width of the input image')
18 | parser.add_argument('--prompt', type=str, required=True, help='The prompt to be used')
19 | parser.add_argument('--model_name', type=str, default=None, help='Name of the pretrained model to be used')
20 | parser.add_argument('--model_path', type=str, default=None, help='Path to a model state dict to be used')
21 | parser.add_argument('--output_dir', type=str, default='output', help='Directory to save the output')
22 | parser.add_argument('--seed', type=int, default=42, help='Random seed to be used')
23 | parser.add_argument('--timestep', type=int, default=199, help='Diffusion timestep')
24 | parser.add_argument('--video', action='store_true', help='If the input is a video')
25 | args = parser.parse_args()
26 |
27 | # Create output directory
28 | os.makedirs(args.output_dir, exist_ok=True)
29 |
30 | # Initialize the model
31 | model = Difix(
32 | pretrained_name=args.model_name,
33 | pretrained_path=args.model_path,
34 | timestep=args.timestep,
35 | mv_unet=True if args.ref_image is not None else False,
36 | )
37 | model.set_eval()
38 |
39 | # Load input images
40 | if os.path.isdir(args.input_image):
41 | input_images = sorted(glob(os.path.join(args.input_image, "*.png")))
42 | else:
43 | input_images = [args.input_image]
44 |
45 | # Load reference images if provided
46 | if args.ref_image is not None:
47 | if os.path.isdir(args.ref_image):
48 | ref_images = sorted(glob(os.path.join(args.ref_image, "*")))
49 | else:
50 | ref_images = [args.ref_image]
51 |
52 | assert len(input_images) == len(ref_images), "Number of input images and reference images should be the same"
53 |
54 | # Process images
55 | output_images = []
56 | for i, input_image in enumerate(tqdm(input_images, desc="Processing images")):
57 | image = Image.open(input_image).convert('RGB')
58 | ref_image = Image.open(ref_images[i]).convert('RGB') if args.ref_image is not None else None
59 | output_image = model.sample(
60 | image,
61 | height=args.height,
62 | width=args.width,
63 | ref_image=ref_image,
64 | prompt=args.prompt
65 | )
66 | output_images.append(output_image)
67 |
68 | # Save outputs
69 | if args.video:
70 | # Save as video
71 | video_path = os.path.join(args.output_dir, "output.mp4")
72 | writer = imageio.get_writer(video_path, fps=30)
73 | for output_image in tqdm(output_images, desc="Saving video"):
74 | writer.append_data(np.array(output_image))
75 | writer.close()
76 | else:
77 | # Save as individual images
78 | for i, output_image in enumerate(tqdm(output_images, desc="Saving images")):
79 | output_image.save(os.path.join(args.output_dir, os.path.basename(input_images[i])))
--------------------------------------------------------------------------------
/examples/nerfstudio/difix3d/difix3d_trainer.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 The Nerfstudio Team. All rights reserved.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 |
16 | import dataclasses
17 | from dataclasses import dataclass, field
18 | from typing import Type, Literal
19 | from nerfstudio.engine.trainer import Trainer, TrainerConfig
20 | from nerfstudio.engine.callbacks import TrainingCallbackAttributes
21 | from nerfstudio.viewer.viewer import Viewer as ViewerState
22 | from nerfstudio.utils import profiler, writer
23 |
24 |
25 | @dataclass
26 | class Difix3DTrainerConfig(TrainerConfig):
27 | """Configuration for the Difix3DTrainer."""
28 | _target: Type = field(default_factory=lambda: Difix3DTrainer)
29 |
30 |
31 | class Difix3DTrainer(Trainer):
32 | """Trainer for Difix3D"""
33 |
34 | def __init__(self, config: TrainerConfig, local_rank: int = 0, world_size: int = 1) -> None:
35 |
36 | super().__init__(config, local_rank, world_size)
37 |
38 | def setup(self, test_mode: Literal["test", "val", "inference"] = "val") -> None:
39 | """Setup the Trainer by calling other setup functions.
40 |
41 | Args:
42 | test_mode:
43 | 'val': loads train/val datasets into memory
44 | 'test': loads train/test datasets into memory
45 | 'inference': does not load any dataset into memory
46 | """
47 | self.pipeline = self.config.pipeline.setup(
48 | device=self.device,
49 | test_mode=test_mode,
50 | world_size=self.world_size,
51 | local_rank=self.local_rank,
52 | grad_scaler=self.grad_scaler,
53 | render_dir=self.base_dir / "renders",
54 | )
55 | self.optimizers = self.setup_optimizers()
56 |
57 | # set up viewer if enabled
58 | viewer_log_path = self.base_dir / self.config.viewer.relative_log_filename
59 | self.viewer_state, banner_messages = None, None
60 | if self.config.is_viewer_legacy_enabled() and self.local_rank == 0:
61 | datapath = self.config.data
62 | if datapath is None:
63 | datapath = self.base_dir
64 | self.viewer_state = ViewerLegacyState(
65 | self.config.viewer,
66 | log_filename=viewer_log_path,
67 | datapath=datapath,
68 | pipeline=self.pipeline,
69 | trainer=self,
70 | train_lock=self.train_lock,
71 | )
72 | banner_messages = [f"Legacy viewer at: {self.viewer_state.viewer_url}"]
73 | if self.config.is_viewer_enabled() and self.local_rank == 0:
74 | datapath = self.config.data
75 | if datapath is None:
76 | datapath = self.base_dir
77 | self.viewer_state = ViewerState(
78 | self.config.viewer,
79 | log_filename=viewer_log_path,
80 | datapath=datapath,
81 | pipeline=self.pipeline,
82 | trainer=self,
83 | train_lock=self.train_lock,
84 | share=self.config.viewer.make_share_url,
85 | )
86 | banner_messages = self.viewer_state.viewer_info
87 | self._check_viewer_warnings()
88 |
89 | self._load_checkpoint()
90 |
91 | self.callbacks = self.pipeline.get_training_callbacks(
92 | TrainingCallbackAttributes(
93 | optimizers=self.optimizers, grad_scaler=self.grad_scaler, pipeline=self.pipeline, trainer=self
94 | )
95 | )
96 |
97 | # set up writers/profilers if enabled
98 | writer_log_path = self.base_dir / self.config.logging.relative_log_dir
99 | writer.setup_event_writer(
100 | self.config.is_wandb_enabled(),
101 | self.config.is_tensorboard_enabled(),
102 | self.config.is_comet_enabled(),
103 | log_dir=writer_log_path,
104 | experiment_name=self.config.experiment_name,
105 | project_name=self.config.project_name,
106 | )
107 | writer.setup_local_writer(
108 | self.config.logging, max_iter=self.config.max_num_iterations, banner_messages=banner_messages
109 | )
110 | writer.put_config(name="config", config_dict=dataclasses.asdict(self.config), step=0)
111 | profiler.setup_profiler(self.config.logging, writer_log_path)
--------------------------------------------------------------------------------
/examples/nerfstudio/difix3d/difix3d.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 The Nerfstudio Team. All rights reserved.
2 | #
3 | # Licensed under the Apache License, Version 2.0 (the "License");
4 | # you may not use this file except in compliance with the License.
5 | # You may obtain a copy of the License at
6 | #
7 | # http://www.apache.org/licenses/LICENSE-2.0
8 | #
9 | # Unless required by applicable law or agreed to in writing, software
10 | # distributed under the License is distributed on an "AS IS" BASIS,
11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 | # See the License for the specific language governing permissions and
13 | # limitations under the License.
14 |
15 |
16 | from __future__ import annotations
17 |
18 | from dataclasses import dataclass, field
19 | from typing import Type
20 |
21 | import torch
22 | from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig
23 | from difix3d.difix3d_field import Difix3DField
24 | from nerfstudio.cameras.camera_optimizers import CameraOptimizer
25 | from nerfstudio.field_components.spatial_distortions import SceneContraction
26 | from nerfstudio.fields.density_fields import HashMLPDensityField
27 |
28 |
29 | @dataclass
30 | class Difix3DModelConfig(NerfactoModelConfig):
31 | """Configuration for the Difix3DModel."""
32 | _target: Type = field(default_factory=lambda: Difix3DModel)
33 |
34 | disable_camera_optimizer: bool = False
35 | """Whether to disable the camera optimizer or not."""
36 | freeze_appearance_embedding: bool = False
37 |
38 | class Difix3DModel(NerfactoModel):
39 | """Model for Difix3D."""
40 |
41 | config: Difix3DModelConfig
42 |
43 | def populate_modules(self):
44 | """Set the fields and modules."""
45 | super().populate_modules()
46 |
47 | if self.config.disable_scene_contraction:
48 | scene_contraction = None
49 | else:
50 | scene_contraction = SceneContraction(order=float("inf"))
51 |
52 | appearance_embedding_dim = self.config.appearance_embed_dim if self.config.use_appearance_embedding else 0
53 |
54 | # Fields
55 | self.field = Difix3DField(
56 | self.scene_box.aabb,
57 | hidden_dim=self.config.hidden_dim,
58 | num_levels=self.config.num_levels,
59 | max_res=self.config.max_res,
60 | base_res=self.config.base_res,
61 | features_per_level=self.config.features_per_level,
62 | log2_hashmap_size=self.config.log2_hashmap_size,
63 | hidden_dim_color=self.config.hidden_dim_color,
64 | hidden_dim_transient=self.config.hidden_dim_transient,
65 | spatial_distortion=scene_contraction,
66 | num_images=self.num_train_data,
67 | use_pred_normals=self.config.predict_normals,
68 | use_average_appearance_embedding=self.config.use_average_appearance_embedding,
69 | appearance_embedding_dim=appearance_embedding_dim,
70 | average_init_density=self.config.average_init_density,
71 | implementation=self.config.implementation,
72 | freeze_appearance_embedding=self.config.freeze_appearance_embedding,
73 | )
74 |
75 | self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup(
76 | num_cameras=self.num_train_data, device="cpu"
77 | )
78 | if self.config.disable_camera_optimizer:
79 | self.camera_optimizer.mode = "off"
80 | self.density_fns = []
81 | num_prop_nets = self.config.num_proposal_iterations
82 | # Build the proposal network(s)
83 | self.proposal_networks = torch.nn.ModuleList()
84 | if self.config.use_same_proposal_network:
85 | assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed."
86 | prop_net_args = self.config.proposal_net_args_list[0]
87 | network = HashMLPDensityField(
88 | self.scene_box.aabb,
89 | spatial_distortion=scene_contraction,
90 | **prop_net_args,
91 | average_init_density=self.config.average_init_density,
92 | implementation=self.config.implementation,
93 | )
94 | self.proposal_networks.append(network)
95 | self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)])
96 | else:
97 | for i in range(num_prop_nets):
98 | prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)]
99 | network = HashMLPDensityField(
100 | self.scene_box.aabb,
101 | spatial_distortion=scene_contraction,
102 | **prop_net_args,
103 | average_init_density=self.config.average_init_density,
104 | implementation=self.config.implementation,
105 | )
106 | self.proposal_networks.append(network)
107 | self.density_fns.extend([network.density_fn for network in self.proposal_networks])
108 |
--------------------------------------------------------------------------------
/examples/gsplat/datasets/normalize.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 |
3 |
4 | def similarity_from_cameras(c2w, strict_scaling=False, center_method="focus"):
5 | """
6 | reference: nerf-factory
7 | Get a similarity transform to normalize dataset
8 | from c2w (OpenCV convention) cameras
9 | :param c2w: (N, 4)
10 | :return T (4,4) , scale (float)
11 | """
12 | t = c2w[:, :3, 3]
13 | R = c2w[:, :3, :3]
14 |
15 | # (1) Rotate the world so that z+ is the up axis
16 | # we estimate the up axis by averaging the camera up axes
17 | ups = np.sum(R * np.array([0, -1.0, 0]), axis=-1)
18 | world_up = np.mean(ups, axis=0)
19 | world_up /= np.linalg.norm(world_up)
20 |
21 | up_camspace = np.array([0.0, -1.0, 0.0])
22 | c = (up_camspace * world_up).sum()
23 | cross = np.cross(world_up, up_camspace)
24 | skew = np.array(
25 | [
26 | [0.0, -cross[2], cross[1]],
27 | [cross[2], 0.0, -cross[0]],
28 | [-cross[1], cross[0], 0.0],
29 | ]
30 | )
31 | if c > -1:
32 | R_align = np.eye(3) + skew + (skew @ skew) * 1 / (1 + c)
33 | else:
34 | # In the unlikely case the original data has y+ up axis,
35 | # rotate 180-deg about x axis
36 | R_align = np.array([[-1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
37 |
38 | # R_align = np.eye(3) # DEBUG
39 | R = R_align @ R
40 | fwds = np.sum(R * np.array([0, 0.0, 1.0]), axis=-1)
41 | t = (R_align @ t[..., None])[..., 0]
42 |
43 | # (2) Recenter the scene.
44 | if center_method == "focus":
45 | # find the closest point to the origin for each camera's center ray
46 | nearest = t + (fwds * -t).sum(-1)[:, None] * fwds
47 | translate = -np.median(nearest, axis=0)
48 | elif center_method == "poses":
49 | # use center of the camera positions
50 | translate = -np.median(t, axis=0)
51 | else:
52 | raise ValueError(f"Unknown center_method {center_method}")
53 |
54 | transform = np.eye(4)
55 | transform[:3, 3] = translate
56 | transform[:3, :3] = R_align
57 |
58 | # (3) Rescale the scene using camera distances
59 | scale_fn = np.max if strict_scaling else np.median
60 | scale = 1.0 / scale_fn(np.linalg.norm(t + translate, axis=-1))
61 | transform[:3, :] *= scale
62 |
63 | return transform
64 |
65 |
66 | def align_principle_axes(point_cloud):
67 | # Compute centroid
68 | centroid = np.median(point_cloud, axis=0)
69 |
70 | # Translate point cloud to centroid
71 | translated_point_cloud = point_cloud - centroid
72 |
73 | # Compute covariance matrix
74 | covariance_matrix = np.cov(translated_point_cloud, rowvar=False)
75 |
76 | # Compute eigenvectors and eigenvalues
77 | eigenvalues, eigenvectors = np.linalg.eigh(covariance_matrix)
78 |
79 | # Sort eigenvectors by eigenvalues (descending order) so that the z-axis
80 | # is the principal axis with the smallest eigenvalue.
81 | sort_indices = eigenvalues.argsort()[::-1]
82 | eigenvectors = eigenvectors[:, sort_indices]
83 |
84 | # Check orientation of eigenvectors. If the determinant of the eigenvectors is
85 | # negative, then we need to flip the sign of one of the eigenvectors.
86 | if np.linalg.det(eigenvectors) < 0:
87 | eigenvectors[:, 0] *= -1
88 |
89 | # Create rotation matrix
90 | rotation_matrix = eigenvectors.T
91 |
92 | # Create SE(3) matrix (4x4 transformation matrix)
93 | transform = np.eye(4)
94 | transform[:3, :3] = rotation_matrix
95 | transform[:3, 3] = -rotation_matrix @ centroid
96 |
97 | return transform
98 |
99 |
100 | def transform_points(matrix, points):
101 | """Transform points using an SE(3) matrix.
102 |
103 | Args:
104 | matrix: 4x4 SE(3) matrix
105 | points: Nx3 array of points
106 |
107 | Returns:
108 | Nx3 array of transformed points
109 | """
110 | assert matrix.shape == (4, 4)
111 | assert len(points.shape) == 2 and points.shape[1] == 3
112 | return points @ matrix[:3, :3].T + matrix[:3, 3]
113 |
114 |
115 | def transform_cameras(matrix, camtoworlds):
116 | """Transform cameras using an SE(3) matrix.
117 |
118 | Args:
119 | matrix: 4x4 SE(3) matrix
120 | camtoworlds: Nx4x4 array of camera-to-world matrices
121 |
122 | Returns:
123 | Nx4x4 array of transformed camera-to-world matrices
124 | """
125 | assert matrix.shape == (4, 4)
126 | assert len(camtoworlds.shape) == 3 and camtoworlds.shape[1:] == (4, 4)
127 | camtoworlds = np.einsum("nij, ki -> nkj", camtoworlds, matrix)
128 | scaling = np.linalg.norm(camtoworlds[:, 0, :3], axis=1)
129 | camtoworlds[:, :3, :3] = camtoworlds[:, :3, :3] / scaling[:, None, None]
130 | return camtoworlds
131 |
132 |
133 | def normalize(camtoworlds, points=None):
134 | T1 = similarity_from_cameras(camtoworlds)
135 | camtoworlds = transform_cameras(T1, camtoworlds)
136 | if points is not None:
137 | points = transform_points(T1, points)
138 | T2 = align_principle_axes(points)
139 | camtoworlds = transform_cameras(T2, camtoworlds)
140 | points = transform_points(T2, points)
141 | return camtoworlds, points, T2 @ T1
142 | else:
143 | return camtoworlds, T1
144 |
--------------------------------------------------------------------------------
/examples/utils.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from scipy.spatial.transform import Rotation
3 |
4 |
5 | class CameraPoseInterpolator:
6 | """
7 | A system for interpolating between sets of camera poses with visualization capabilities.
8 | """
9 |
10 | def __init__(self, rotation_weight=1.0, translation_weight=1.0):
11 | """
12 | Initialize the interpolator with weights for pose distance computation.
13 |
14 | Args:
15 | rotation_weight: Weight for rotational distance in pose matching
16 | translation_weight: Weight for translational distance in pose matching
17 | """
18 | self.rotation_weight = rotation_weight
19 | self.translation_weight = translation_weight
20 |
21 | def compute_pose_distance(self, pose1, pose2):
22 | """
23 | Compute weighted distance between two camera poses.
24 |
25 | Args:
26 | pose1, pose2: 4x4 transformation matrices
27 |
28 | Returns:
29 | Combined weighted distance between poses
30 | """
31 | # Translation distance (Euclidean)
32 | t1, t2 = pose1[:3, 3], pose2[:3, 3]
33 | translation_dist = np.linalg.norm(t1 - t2)
34 |
35 | # Rotation distance (angular distance between quaternions)
36 | R1 = Rotation.from_matrix(pose1[:3, :3])
37 | R2 = Rotation.from_matrix(pose2[:3, :3])
38 | q1 = R1.as_quat()
39 | q2 = R2.as_quat()
40 |
41 | # Ensure quaternions are in the same hemisphere
42 | if np.dot(q1, q2) < 0:
43 | q2 = -q2
44 |
45 | rotation_dist = np.arccos(2 * np.dot(q1, q2)**2 - 1)
46 |
47 | return (self.translation_weight * translation_dist +
48 | self.rotation_weight * rotation_dist)
49 |
50 | def find_nearest_assignments(self, training_poses, testing_poses):
51 | """
52 | Find the nearest training camera pose for each testing camera pose.
53 |
54 | Args:
55 | training_poses: [N, 4, 4] array of training camera poses
56 | testing_poses: [M, 4, 4] array of testing camera poses
57 |
58 | Returns:
59 | assignments: list of closest training pose indices for each testing pose
60 | """
61 | M = len(testing_poses)
62 | assignments = []
63 |
64 | for j in range(M):
65 | # Compute distance from each training pose to this testing pose
66 | distances = [self.compute_pose_distance(training_pose, testing_poses[j])
67 | for training_pose in training_poses]
68 | # Find the index of the nearest training pose
69 | nearest_index = np.argmin(distances)
70 | assignments.append(nearest_index)
71 |
72 | return assignments
73 |
74 | def interpolate_rotation(self, R1, R2, t):
75 | """
76 | Interpolate between two rotation matrices using SLERP.
77 | """
78 | q1 = Rotation.from_matrix(R1).as_quat()
79 | q2 = Rotation.from_matrix(R2).as_quat()
80 |
81 | if np.dot(q1, q2) < 0:
82 | q2 = -q2
83 |
84 | # Clamp dot product to avoid invalid values in arccos
85 | dot_product = np.clip(np.dot(q1, q2), -1.0, 1.0)
86 | theta = np.arccos(dot_product)
87 |
88 | if np.abs(theta) < 1e-6:
89 | q_interp = (1 - t) * q1 + t * q2
90 | else:
91 | q_interp = (np.sin((1-t)*theta) * q1 + np.sin(t*theta) * q2) / np.sin(theta)
92 |
93 | q_interp = q_interp / np.linalg.norm(q_interp)
94 | return Rotation.from_quat(q_interp).as_matrix()
95 |
96 | def interpolate_poses(self, training_poses, testing_poses, num_steps=20):
97 | """
98 | Interpolate between camera poses using nearest assignments.
99 |
100 | Args:
101 | training_poses: [N, 4, 4] array of training poses
102 | testing_poses: [M, 4, 4] array of testing poses
103 | num_steps: number of interpolation steps
104 |
105 | Returns:
106 | interpolated_sequences: list of lists of interpolated poses
107 | """
108 | assignments = self.find_nearest_assignments(training_poses, testing_poses)
109 | interpolated_sequences = []
110 |
111 | for test_idx, train_idx in enumerate(assignments):
112 | train_pose = training_poses[train_idx]
113 | test_pose = testing_poses[test_idx]
114 | sequence = []
115 |
116 | for t in np.linspace(0, 1, num_steps):
117 | # Interpolate rotation
118 | R_interp = self.interpolate_rotation(
119 | train_pose[:3, :3],
120 | test_pose[:3, :3],
121 | t
122 | )
123 |
124 | # Interpolate translation
125 | t_interp = (1-t) * train_pose[:3, 3] + t * test_pose[:3, 3]
126 |
127 | # Construct interpolated pose
128 | pose_interp = np.eye(4)
129 | pose_interp[:3, :3] = R_interp
130 | pose_interp[:3, 3] = t_interp
131 |
132 | sequence.append(pose_interp)
133 |
134 | interpolated_sequences.append(sequence)
135 |
136 | return interpolated_sequences
137 |
138 |
139 | def shift_poses(self, training_poses, testing_poses, distance=0.1, threshold=0.1):
140 | """
141 | Shift nearest training poses toward testing poses by a specified distance.
142 |
143 | Args:
144 | training_poses: [N, 4, 4] array of training camera poses
145 | testing_poses: [M, 4, 4] array of testing camera poses
146 | distance: float, the step size to move training pose toward testing pose
147 |
148 | Returns:
149 | novel_poses: [M, 4, 4] array of shifted poses
150 | """
151 | assignments = self.find_nearest_assignments(training_poses, testing_poses)
152 | novel_poses = []
153 |
154 | for test_idx, train_idx in enumerate(assignments):
155 | train_pose = training_poses[train_idx]
156 | test_pose = testing_poses[test_idx]
157 |
158 | if self.compute_pose_distance(train_pose, test_pose) <= distance:
159 | novel_poses.append(test_pose)
160 | continue
161 |
162 | # Calculate translation step if shifting is necessary
163 | t1, t2 = train_pose[:3, 3], test_pose[:3, 3]
164 | translation_direction = t2 - t1
165 | translation_norm = np.linalg.norm(translation_direction)
166 |
167 | if translation_norm > 1e-6:
168 | translation_step = (translation_direction / translation_norm) * distance
169 | new_translation = t1 + translation_step
170 | else:
171 | # If translation direction is too small, use testing pose translation directly
172 | new_translation = t2
173 |
174 | # Check if the new translation would overshoot the testing pose translation
175 | if np.dot(new_translation - t1, t2 - t1) <= 0 or np.linalg.norm(new_translation - t2) <= distance:
176 | new_translation = t2
177 |
178 | # Update rotation
179 | R1 = train_pose[:3, :3]
180 | R2 = test_pose[:3, :3]
181 | if translation_norm > 1e-6:
182 | R_interp = self.interpolate_rotation(R1, R2, min(distance / translation_norm, 1.0))
183 | else:
184 | R_interp = R2 # Use testing rotation if too close
185 |
186 | # Construct shifted pose
187 | shifted_pose = np.eye(4)
188 | shifted_pose[:3, :3] = R_interp
189 | shifted_pose[:3, 3] = new_translation
190 |
191 | novel_poses.append(shifted_pose)
192 |
193 | return np.array(novel_poses)
--------------------------------------------------------------------------------
/examples/gsplat/utils.py:
--------------------------------------------------------------------------------
1 | import random
2 |
3 | import numpy as np
4 | import torch
5 | from sklearn.neighbors import NearestNeighbors
6 | from torch import Tensor
7 | import torch.nn.functional as F
8 | import matplotlib.pyplot as plt
9 | from matplotlib import colormaps
10 |
11 |
12 | class CameraOptModule(torch.nn.Module):
13 | """Camera pose optimization module."""
14 |
15 | def __init__(self, n: int):
16 | super().__init__()
17 | # Delta positions (3D) + Delta rotations (6D)
18 | self.embeds = torch.nn.Embedding(n, 9)
19 | # Identity rotation in 6D representation
20 | self.register_buffer("identity", torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0]))
21 |
22 | def zero_init(self):
23 | torch.nn.init.zeros_(self.embeds.weight)
24 |
25 | def random_init(self, std: float):
26 | torch.nn.init.normal_(self.embeds.weight, std=std)
27 |
28 | def forward(self, camtoworlds: Tensor, embed_ids: Tensor) -> Tensor:
29 | """Adjust camera pose based on deltas.
30 |
31 | Args:
32 | camtoworlds: (..., 4, 4)
33 | embed_ids: (...,)
34 |
35 | Returns:
36 | updated camtoworlds: (..., 4, 4)
37 | """
38 | assert camtoworlds.shape[:-2] == embed_ids.shape
39 | batch_shape = camtoworlds.shape[:-2]
40 | pose_deltas = self.embeds(embed_ids) # (..., 9)
41 | dx, drot = pose_deltas[..., :3], pose_deltas[..., 3:]
42 | rot = rotation_6d_to_matrix(
43 | drot + self.identity.expand(*batch_shape, -1)
44 | ) # (..., 3, 3)
45 | transform = torch.eye(4, device=pose_deltas.device).repeat((*batch_shape, 1, 1))
46 | transform[..., :3, :3] = rot
47 | transform[..., :3, 3] = dx
48 | return torch.matmul(camtoworlds, transform)
49 |
50 |
51 | class AppearanceOptModule(torch.nn.Module):
52 | """Appearance optimization module."""
53 |
54 | def __init__(
55 | self,
56 | n: int,
57 | feature_dim: int,
58 | embed_dim: int = 16,
59 | sh_degree: int = 3,
60 | mlp_width: int = 64,
61 | mlp_depth: int = 2,
62 | ):
63 | super().__init__()
64 | self.embed_dim = embed_dim
65 | self.sh_degree = sh_degree
66 | self.embeds = torch.nn.Embedding(n, embed_dim)
67 | layers = []
68 | layers.append(
69 | torch.nn.Linear(embed_dim + feature_dim + (sh_degree + 1) ** 2, mlp_width)
70 | )
71 | layers.append(torch.nn.ReLU(inplace=True))
72 | for _ in range(mlp_depth - 1):
73 | layers.append(torch.nn.Linear(mlp_width, mlp_width))
74 | layers.append(torch.nn.ReLU(inplace=True))
75 | layers.append(torch.nn.Linear(mlp_width, 3))
76 | self.color_head = torch.nn.Sequential(*layers)
77 |
78 | def forward(
79 | self, features: Tensor, embed_ids: Tensor, dirs: Tensor, sh_degree: int
80 | ) -> Tensor:
81 | """Adjust appearance based on embeddings.
82 |
83 | Args:
84 | features: (N, feature_dim)
85 | embed_ids: (C,)
86 | dirs: (C, N, 3)
87 |
88 | Returns:
89 | colors: (C, N, 3)
90 | """
91 | from gsplat.cuda._torch_impl import _eval_sh_bases_fast
92 |
93 | C, N = dirs.shape[:2]
94 | # Camera embeddings
95 | if embed_ids is None:
96 | embeds = torch.zeros(C, self.embed_dim, device=features.device)
97 | else:
98 | embeds = self.embeds(embed_ids) # [C, D2]
99 | embeds = embeds[:, None, :].expand(-1, N, -1) # [C, N, D2]
100 | # GS features
101 | features = features[None, :, :].expand(C, -1, -1) # [C, N, D1]
102 | # View directions
103 | dirs = F.normalize(dirs, dim=-1) # [C, N, 3]
104 | num_bases_to_use = (sh_degree + 1) ** 2
105 | num_bases = (self.sh_degree + 1) ** 2
106 | sh_bases = torch.zeros(C, N, num_bases, device=features.device) # [C, N, K]
107 | sh_bases[:, :, :num_bases_to_use] = _eval_sh_bases_fast(num_bases_to_use, dirs)
108 | # Get colors
109 | if self.embed_dim > 0:
110 | h = torch.cat([embeds, features, sh_bases], dim=-1) # [C, N, D1 + D2 + K]
111 | else:
112 | h = torch.cat([features, sh_bases], dim=-1)
113 | colors = self.color_head(h)
114 | return colors
115 |
116 |
117 | def rotation_6d_to_matrix(d6: Tensor) -> Tensor:
118 | """
119 | Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
120 | using Gram--Schmidt orthogonalization per Section B of [1]. Adapted from pytorch3d.
121 | Args:
122 | d6: 6D rotation representation, of size (*, 6)
123 |
124 | Returns:
125 | batch of rotation matrices of size (*, 3, 3)
126 |
127 | [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
128 | On the Continuity of Rotation Representations in Neural Networks.
129 | IEEE Conference on Computer Vision and Pattern Recognition, 2019.
130 | Retrieved from http://arxiv.org/abs/1812.07035
131 | """
132 |
133 | a1, a2 = d6[..., :3], d6[..., 3:]
134 | b1 = F.normalize(a1, dim=-1)
135 | b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
136 | b2 = F.normalize(b2, dim=-1)
137 | b3 = torch.cross(b1, b2, dim=-1)
138 | return torch.stack((b1, b2, b3), dim=-2)
139 |
140 |
141 | def knn(x: Tensor, K: int = 4) -> Tensor:
142 | x_np = x.cpu().numpy()
143 | model = NearestNeighbors(n_neighbors=K, metric="euclidean").fit(x_np)
144 | distances, _ = model.kneighbors(x_np)
145 | return torch.from_numpy(distances).to(x)
146 |
147 |
148 | def rgb_to_sh(rgb: Tensor) -> Tensor:
149 | C0 = 0.28209479177387814
150 | return (rgb - 0.5) / C0
151 |
152 |
153 | def set_random_seed(seed: int):
154 | random.seed(seed)
155 | np.random.seed(seed)
156 | torch.manual_seed(seed)
157 |
158 |
159 | # ref: https://github.com/hbb1/2d-gaussian-splatting/blob/main/utils/general_utils.py#L163
160 | def colormap(img, cmap="jet"):
161 | W, H = img.shape[:2]
162 | dpi = 300
163 | fig, ax = plt.subplots(1, figsize=(H / dpi, W / dpi), dpi=dpi)
164 | im = ax.imshow(img, cmap=cmap)
165 | ax.set_axis_off()
166 | fig.colorbar(im, ax=ax)
167 | fig.tight_layout()
168 | fig.canvas.draw()
169 | data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
170 | data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
171 | img = torch.from_numpy(data).float().permute(2, 0, 1)
172 | plt.close()
173 | return img
174 |
175 |
176 | def apply_float_colormap(img: torch.Tensor, colormap: str = "turbo") -> torch.Tensor:
177 | """Convert single channel to a color img.
178 |
179 | Args:
180 | img (torch.Tensor): (..., 1) float32 single channel image.
181 | colormap (str): Colormap for img.
182 |
183 | Returns:
184 | (..., 3) colored img with colors in [0, 1].
185 | """
186 | img = torch.nan_to_num(img, 0)
187 | if colormap == "gray":
188 | return img.repeat(1, 1, 3)
189 | img_long = (img * 255).long()
190 | img_long_min = torch.min(img_long)
191 | img_long_max = torch.max(img_long)
192 | assert img_long_min >= 0, f"the min value is {img_long_min}"
193 | assert img_long_max <= 255, f"the max value is {img_long_max}"
194 | return torch.tensor(
195 | colormaps[colormap].colors, # type: ignore
196 | device=img.device,
197 | )[img_long[..., 0]]
198 |
199 |
200 | def apply_depth_colormap(
201 | depth: torch.Tensor,
202 | acc: torch.Tensor = None,
203 | near_plane: float = None,
204 | far_plane: float = None,
205 | ) -> torch.Tensor:
206 | """Converts a depth image to color for easier analysis.
207 |
208 | Args:
209 | depth (torch.Tensor): (..., 1) float32 depth.
210 | acc (torch.Tensor | None): (..., 1) optional accumulation mask.
211 | near_plane: Closest depth to consider. If None, use min image value.
212 | far_plane: Furthest depth to consider. If None, use max image value.
213 |
214 | Returns:
215 | (..., 3) colored depth image with colors in [0, 1].
216 | """
217 | near_plane = near_plane or float(torch.min(depth))
218 | far_plane = far_plane or float(torch.max(depth))
219 | depth = (depth - near_plane) / (far_plane - near_plane + 1e-10)
220 | depth = torch.clip(depth, 0.0, 1.0)
221 | img = apply_float_colormap(depth, colormap="turbo")
222 | if acc is not None:
223 | img = img * acc + (1.0 - acc)
224 | return img
225 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Difix3D+
2 |
3 | **Difix3D+: Improving 3D Reconstructions with Single-Step Diffusion Models**
4 | [Jay Zhangjie Wu*](https://zhangjiewu.github.io/), [Yuxuan Zhang*](https://scholar.google.com/citations?user=Jt5VvNgAAAAJ&hl=en), [Haithem Turki](https://haithemturki.com/), [Xuanchi Ren](https://xuanchiren.com/), [Jun Gao](https://www.cs.toronto.edu/~jungao/),
5 | [Mike Zheng Shou](https://sites.google.com/view/showlab/home?authuser=0), [Sanja Fidler](https://www.cs.utoronto.ca/~fidler/), [Zan Gojcic†](https://zgojcic.github.io/), [Huan Ling†](https://www.cs.toronto.edu/~linghuan/) _(*/† equal contribution/advising)_
6 | CVPR 2025 (Oral)
7 | [Project Page](https://research.nvidia.com/labs/toronto-ai/difix3d/) | [Paper](https://arxiv.org/abs/2503.01774) | [Model](https://huggingface.co/nvidia/difix) | [Demo](https://huggingface.co/spaces/nvidia/difix)
8 |
9 |
10 |
11 |
38 |
39 | - Low-rank 4D bilateral grid for MR enhancement:
40 |
41 |
42 |
43 |
44 | Below is the API reference.
45 |
46 | """
47 |
48 | import tensorly as tl
49 | import torch
50 | import torch.nn.functional as F
51 | from torch import nn
52 |
53 | tl.set_backend("pytorch")
54 |
55 |
56 | def color_correct(
57 | img: torch.Tensor, ref: torch.Tensor, num_iters: int = 5, eps: float = 0.5 / 255
58 | ) -> torch.Tensor:
59 | """
60 | Warp `img` to match the colors in `ref_img` using iterative color matching.
61 |
62 | This function performs color correction by warping the colors of the input image
63 | to match those of a reference image. It uses a least squares method to find a
64 | transformation that maps the input image's colors to the reference image's colors.
65 |
66 | The algorithm iteratively solves a system of linear equations, updating the set of
67 | unsaturated pixels in each iteration. This approach helps handle non-linear color
68 | transformations and reduces the impact of clipping.
69 |
70 | Args:
71 | img (torch.Tensor): Input image to be color corrected. Shape: [..., num_channels]
72 | ref (torch.Tensor): Reference image to match colors. Shape: [..., num_channels]
73 | num_iters (int, optional): Number of iterations for the color matching process.
74 | Default is 5.
75 | eps (float, optional): Small value to determine the range of unclipped pixels.
76 | Default is 0.5 / 255.
77 |
78 | Returns:
79 | torch.Tensor: Color corrected image with the same shape as the input image.
80 |
81 | Note:
82 | - Both input and reference images should be in the range [0, 1].
83 | - The function works with any number of channels, but typically used with 3 (RGB).
84 | """
85 | if img.shape[-1] != ref.shape[-1]:
86 | raise ValueError(
87 | f"img's {img.shape[-1]} and ref's {ref.shape[-1]} channels must match"
88 | )
89 | num_channels = img.shape[-1]
90 | img_mat = img.reshape([-1, num_channels])
91 | ref_mat = ref.reshape([-1, num_channels])
92 |
93 | def is_unclipped(z):
94 | return (z >= eps) & (z <= 1 - eps) # z \in [eps, 1-eps].
95 |
96 | mask0 = is_unclipped(img_mat)
97 | # Because the set of saturated pixels may change after solving for a
98 | # transformation, we repeatedly solve a system `num_iters` times and update
99 | # our estimate of which pixels are saturated.
100 | for _ in range(num_iters):
101 | # Construct the left hand side of a linear system that contains a quadratic
102 | # expansion of each pixel of `img`.
103 | a_mat = []
104 | for c in range(num_channels):
105 | a_mat.append(img_mat[:, c : (c + 1)] * img_mat[:, c:]) # Quadratic term.
106 | a_mat.append(img_mat) # Linear term.
107 | a_mat.append(torch.ones_like(img_mat[:, :1])) # Bias term.
108 | a_mat = torch.cat(a_mat, dim=-1)
109 | warp = []
110 | for c in range(num_channels):
111 | # Construct the right hand side of a linear system containing each color
112 | # of `ref`.
113 | b = ref_mat[:, c]
114 | # Ignore rows of the linear system that were saturated in the input or are
115 | # saturated in the current corrected color estimate.
116 | mask = mask0[:, c] & is_unclipped(img_mat[:, c]) & is_unclipped(b)
117 | ma_mat = torch.where(mask[:, None], a_mat, torch.zeros_like(a_mat))
118 | mb = torch.where(mask, b, torch.zeros_like(b))
119 | w = torch.linalg.lstsq(ma_mat, mb, rcond=-1)[0]
120 | assert torch.all(torch.isfinite(w))
121 | warp.append(w)
122 | warp = torch.stack(warp, dim=-1)
123 | # Apply the warp to update img_mat.
124 | img_mat = torch.clip(torch.matmul(a_mat, warp), 0, 1)
125 | corrected_img = torch.reshape(img_mat, img.shape)
126 | return corrected_img
127 |
128 |
129 | def bilateral_grid_tv_loss(model, config):
130 | """Computes total variations of bilateral grids."""
131 | total_loss = 0.0
132 |
133 | for bil_grids in model.bil_grids:
134 | total_loss += config.bilgrid_tv_loss_mult * total_variation_loss(
135 | bil_grids.grids
136 | )
137 |
138 | return total_loss
139 |
140 |
141 | def color_affine_transform(affine_mats, rgb):
142 | """Applies color affine transformations.
143 |
144 | Args:
145 | affine_mats (torch.Tensor): Affine transformation matrices. Supported shape: $(..., 3, 4)$.
146 | rgb (torch.Tensor): Input RGB values. Supported shape: $(..., 3)$.
147 |
148 | Returns:
149 | Output transformed colors of shape $(..., 3)$.
150 | """
151 | return (
152 | torch.matmul(affine_mats[..., :3], rgb.unsqueeze(-1)).squeeze(-1)
153 | + affine_mats[..., 3]
154 | )
155 |
156 |
157 | def _num_tensor_elems(t):
158 | return max(torch.prod(torch.tensor(t.size()[1:]).float()).item(), 1.0)
159 |
160 |
161 | def total_variation_loss(x): # noqa: F811
162 | """Returns total variation on multi-dimensional tensors.
163 |
164 | Args:
165 | x (torch.Tensor): The input tensor with shape $(B, C, ...)$, where $B$ is the batch size and $C$ is the channel size.
166 | """
167 | batch_size = x.shape[0]
168 | tv = 0
169 | for i in range(2, len(x.shape)):
170 | n_res = x.shape[i]
171 | idx1 = torch.arange(1, n_res, device=x.device)
172 | idx2 = torch.arange(0, n_res - 1, device=x.device)
173 | x1 = x.index_select(i, idx1)
174 | x2 = x.index_select(i, idx2)
175 | count = _num_tensor_elems(x1)
176 | tv += torch.pow((x1 - x2), 2).sum() / count
177 | return tv / batch_size
178 |
179 |
180 | def slice(bil_grids, xy, rgb, grid_idx):
181 | """Slices a batch of 3D bilateral grids by pixel coordinates `xy` and gray-scale guidances of pixel colors `rgb`.
182 |
183 | Supports 2-D, 3-D, and 4-D input shapes. The first dimension of the input is the batch size
184 | and the last dimension is 2 for `xy`, 3 for `rgb`, and 1 for `grid_idx`.
185 |
186 | The return value is a dictionary containing the affine transformations `affine_mats` sliced from bilateral grids and
187 | the output color `rgb_out` after applying the afffine transformations.
188 |
189 | 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.
190 | Then `affine_mats[i]` can be obtained via slicing the bilateral grid indexed at `grid_idx[i]` by `xy[i, :]` and `rgb2gray(rgb[i, :])`.
191 | For 3-D and 4-D input cases, the behavior of indexing bilateral grids and coordinates is the same with the 2-D case.
192 |
193 | .. note::
194 | This function can be regarded as a wrapper of `color_affine_transform` and `BilateralGrid` with a slight performance improvement.
195 | When `grid_idx` contains a unique index, only a single bilateral grid will used during the slicing. In this case, this function will not
196 | perform tensor indexing to avoid data copy and extra memory
197 | (see [this](https://discuss.pytorch.org/t/does-indexing-a-tensor-return-a-copy-of-it/164905)).
198 |
199 | Args:
200 | bil_grids (`BilateralGrid`): An instance of $N$ bilateral grids.
201 | xy (torch.Tensor): The x-y coordinates of shape $(..., 2)$ in the range of $[0,1]$.
202 | rgb (torch.Tensor): The RGB values of shape $(..., 3)$ for computing the guidance coordinates, ranging in $[0,1]$.
203 | grid_idx (torch.Tensor): The indices of bilateral grids for each slicing. Shape: $(..., 1)$.
204 |
205 | Returns:
206 | A dictionary with keys and values as follows:
207 | ```
208 | {
209 | "rgb": Transformed RGB colors. Shape: (..., 3),
210 | "rgb_affine_mats": The sliced affine transformation matrices from bilateral grids. Shape: (..., 3, 4)
211 | }
212 | ```
213 | """
214 |
215 | sh_ = rgb.shape
216 |
217 | grid_idx_unique = torch.unique(grid_idx)
218 | if len(grid_idx_unique) == 1:
219 | # All pixels are from a single view.
220 | grid_idx = grid_idx_unique # (1,)
221 | xy = xy.unsqueeze(0) # (1, ..., 2)
222 | rgb = rgb.unsqueeze(0) # (1, ..., 3)
223 | else:
224 | # Pixels are randomly sampled from different views.
225 | if len(grid_idx.shape) == 4:
226 | grid_idx = grid_idx[:, 0, 0, 0] # (chunk_size,)
227 | elif len(grid_idx.shape) == 3:
228 | grid_idx = grid_idx[:, 0, 0] # (chunk_size,)
229 | elif len(grid_idx.shape) == 2:
230 | grid_idx = grid_idx[:, 0] # (chunk_size,)
231 | else:
232 | raise ValueError(
233 | "The input to bilateral grid slicing is not supported yet."
234 | )
235 |
236 | affine_mats = bil_grids(xy, rgb, grid_idx)
237 | rgb = color_affine_transform(affine_mats, rgb)
238 |
239 | return {
240 | "rgb": rgb.reshape(*sh_),
241 | "rgb_affine_mats": affine_mats.reshape(
242 | *sh_[:-1], affine_mats.shape[-2], affine_mats.shape[-1]
243 | ),
244 | }
245 |
246 |
247 | class BilateralGrid(nn.Module):
248 | """Class for 3D bilateral grids.
249 |
250 | Holds one or more than one bilateral grids.
251 | """
252 |
253 | def __init__(self, num, grid_X=16, grid_Y=16, grid_W=8):
254 | """
255 | Args:
256 | num (int): The number of bilateral grids (i.e., the number of views).
257 | grid_X (int): Defines grid width $W$.
258 | grid_Y (int): Defines grid height $H$.
259 | grid_W (int): Defines grid guidance dimension $L$.
260 | """
261 | super(BilateralGrid, self).__init__()
262 |
263 | self.grid_width = grid_X
264 | """Grid width. Type: int."""
265 | self.grid_height = grid_Y
266 | """Grid height. Type: int."""
267 | self.grid_guidance = grid_W
268 | """Grid guidance dimension. Type: int."""
269 |
270 | # Initialize grids.
271 | grid = self._init_identity_grid()
272 | self.grids = nn.Parameter(grid.tile(num, 1, 1, 1, 1)) # (N, 12, L, H, W)
273 | """ A 5-D tensor of shape $(N, 12, L, H, W)$."""
274 |
275 | # Weights of BT601 RGB-to-gray.
276 | self.register_buffer("rgb2gray_weight", torch.Tensor([[0.299, 0.587, 0.114]]))
277 | self.rgb2gray = lambda rgb: (rgb @ self.rgb2gray_weight.T) * 2.0 - 1.0
278 | """ A function that converts RGB to gray-scale guidance in $[-1, 1]$."""
279 |
280 | def _init_identity_grid(self):
281 | grid = torch.tensor(
282 | [
283 | 1.0,
284 | 0,
285 | 0,
286 | 0,
287 | 0,
288 | 1.0,
289 | 0,
290 | 0,
291 | 0,
292 | 0,
293 | 1.0,
294 | 0,
295 | ]
296 | ).float()
297 | grid = grid.repeat(
298 | [self.grid_guidance * self.grid_height * self.grid_width, 1]
299 | ) # (L * H * W, 12)
300 | grid = grid.reshape(
301 | 1, self.grid_guidance, self.grid_height, self.grid_width, -1
302 | ) # (1, L, H, W, 12)
303 | grid = grid.permute(0, 4, 1, 2, 3) # (1, 12, L, H, W)
304 | return grid
305 |
306 | def tv_loss(self):
307 | """Computes and returns total variation loss on the bilateral grids."""
308 | return total_variation_loss(self.grids)
309 |
310 | def forward(self, grid_xy, rgb, idx=None):
311 | """Bilateral grid slicing. Supports 2-D, 3-D, 4-D, and 5-D input.
312 | For the 2-D, 3-D, and 4-D cases, please refer to `slice`.
313 | For the 5-D cases, `idx` will be unused and the first dimension of `xy` should be
314 | equal to the number of bilateral grids. Then this function becomes PyTorch's
315 | [`F.grid_sample`](https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html).
316 |
317 | Args:
318 | grid_xy (torch.Tensor): The x-y coordinates in the range of $[0,1]$.
319 | rgb (torch.Tensor): The RGB values in the range of $[0,1]$.
320 | idx (torch.Tensor): The bilateral grid indices.
321 |
322 | Returns:
323 | Sliced affine matrices of shape $(..., 3, 4)$.
324 | """
325 |
326 | grids = self.grids
327 | input_ndims = len(grid_xy.shape)
328 | assert len(rgb.shape) == input_ndims
329 |
330 | if input_ndims > 1 and input_ndims < 5:
331 | # Convert input into 5D
332 | for i in range(5 - input_ndims):
333 | grid_xy = grid_xy.unsqueeze(1)
334 | rgb = rgb.unsqueeze(1)
335 | assert idx is not None
336 | elif input_ndims != 5:
337 | raise ValueError(
338 | "Bilateral grid slicing only takes either 2D, 3D, 4D and 5D inputs"
339 | )
340 |
341 | grids = self.grids
342 | if idx is not None:
343 | grids = grids[idx]
344 | assert grids.shape[0] == grid_xy.shape[0]
345 |
346 | # Generate slicing coordinates.
347 | grid_xy = (grid_xy - 0.5) * 2 # Rescale to [-1, 1].
348 | grid_z = self.rgb2gray(rgb)
349 |
350 | # print(grid_xy.shape, grid_z.shape)
351 | # exit()
352 | grid_xyz = torch.cat([grid_xy, grid_z], dim=-1) # (N, m, h, w, 3)
353 |
354 | affine_mats = F.grid_sample(
355 | grids, grid_xyz, mode="bilinear", align_corners=True, padding_mode="border"
356 | ) # (N, 12, m, h, w)
357 | affine_mats = affine_mats.permute(0, 2, 3, 4, 1) # (N, m, h, w, 12)
358 | affine_mats = affine_mats.reshape(
359 | *affine_mats.shape[:-1], 3, 4
360 | ) # (N, m, h, w, 3, 4)
361 |
362 | for _ in range(5 - input_ndims):
363 | affine_mats = affine_mats.squeeze(1)
364 |
365 | return affine_mats
366 |
367 |
368 | def slice4d(bil_grid4d, xyz, rgb):
369 | """Slices a 4D bilateral grid by point coordinates `xyz` and gray-scale guidances of radiance colors `rgb`.
370 |
371 | Args:
372 | bil_grid4d (`BilateralGridCP4D`): The input 4D bilateral grid.
373 | xyz (torch.Tensor): The xyz coordinates with shape $(..., 3)$.
374 | rgb (torch.Tensor): The RGB values with shape $(..., 3)$.
375 |
376 | Returns:
377 | A dictionary with keys and values as follows:
378 | ```
379 | {
380 | "rgb": Transformed radiance RGB colors. Shape: (..., 3),
381 | "rgb_affine_mats": The sliced affine transformation matrices from the 4D bilateral grid. Shape: (..., 3, 4)
382 | }
383 | ```
384 | """
385 |
386 | affine_mats = bil_grid4d(xyz, rgb)
387 | rgb = color_affine_transform(affine_mats, rgb)
388 |
389 | return {"rgb": rgb, "rgb_affine_mats": affine_mats}
390 |
391 |
392 | class _ScaledTanh(nn.Module):
393 | def __init__(self, s=2.0):
394 | super().__init__()
395 | self.scaler = s
396 |
397 | def forward(self, x):
398 | return torch.tanh(self.scaler * x)
399 |
400 |
401 | class BilateralGridCP4D(nn.Module):
402 | """Class for low-rank 4D bilateral grids."""
403 |
404 | def __init__(
405 | self,
406 | grid_X=16,
407 | grid_Y=16,
408 | grid_Z=16,
409 | grid_W=8,
410 | rank=5,
411 | learn_gray=True,
412 | gray_mlp_width=8,
413 | gray_mlp_depth=2,
414 | init_noise_scale=1e-6,
415 | bound=2.0,
416 | ):
417 | """
418 | Args:
419 | grid_X (int): Defines grid width.
420 | grid_Y (int): Defines grid height.
421 | grid_Z (int): Defines grid depth.
422 | grid_W (int): Defines grid guidance dimension.
423 | rank (int): Rank of the 4D bilateral grid.
424 | learn_gray (bool): If True, an MLP will be learned to convert RGB colors to gray-scale guidances.
425 | gray_mlp_width (int): The MLP width for learnable guidance.
426 | gray_mlp_depth (int): The number of MLP layers for learnable guidance.
427 | init_noise_scale (float): The noise scale of the initialized factors.
428 | bound (float): The bound of the xyz coordinates.
429 | """
430 | super(BilateralGridCP4D, self).__init__()
431 |
432 | self.grid_X = grid_X
433 | """Grid width. Type: int."""
434 | self.grid_Y = grid_Y
435 | """Grid height. Type: int."""
436 | self.grid_Z = grid_Z
437 | """Grid depth. Type: int."""
438 | self.grid_W = grid_W
439 | """Grid guidance dimension. Type: int."""
440 | self.rank = rank
441 | """Rank of the 4D bilateral grid. Type: int."""
442 | self.learn_gray = learn_gray
443 | """Flags of learnable guidance is used. Type: bool."""
444 | self.gray_mlp_width = gray_mlp_width
445 | """The MLP width for learnable guidance. Type: int."""
446 | self.gray_mlp_depth = gray_mlp_depth
447 | """The MLP depth for learnable guidance. Type: int."""
448 | self.init_noise_scale = init_noise_scale
449 | """The noise scale of the initialized factors. Type: float."""
450 | self.bound = bound
451 | """The bound of the xyz coordinates. Type: float."""
452 |
453 | self._init_cp_factors_parafac()
454 |
455 | self.rgb2gray = None
456 | """ A function that converts RGB to gray-scale guidances in $[-1, 1]$.
457 | If `learn_gray` is True, this will be an MLP network."""
458 |
459 | if self.learn_gray:
460 |
461 | def rgb2gray_mlp_linear(layer):
462 | return nn.Linear(
463 | self.gray_mlp_width,
464 | self.gray_mlp_width if layer < self.gray_mlp_depth - 1 else 1,
465 | )
466 |
467 | def rgb2gray_mlp_actfn(_):
468 | return nn.ReLU(inplace=True)
469 |
470 | self.rgb2gray = nn.Sequential(
471 | *(
472 | [nn.Linear(3, self.gray_mlp_width)]
473 | + [
474 | nn_module(layer)
475 | for layer in range(1, self.gray_mlp_depth)
476 | for nn_module in [rgb2gray_mlp_actfn, rgb2gray_mlp_linear]
477 | ]
478 | + [_ScaledTanh(2.0)]
479 | )
480 | )
481 | else:
482 | # Weights of BT601/BT470 RGB-to-gray.
483 | self.register_buffer(
484 | "rgb2gray_weight", torch.Tensor([[0.299, 0.587, 0.114]])
485 | )
486 | self.rgb2gray = lambda rgb: (rgb @ self.rgb2gray_weight.T) * 2.0 - 1.0
487 |
488 | def _init_identity_grid(self):
489 | grid = torch.tensor(
490 | [
491 | 1.0,
492 | 0,
493 | 0,
494 | 0,
495 | 0,
496 | 1.0,
497 | 0,
498 | 0,
499 | 0,
500 | 0,
501 | 1.0,
502 | 0,
503 | ]
504 | ).float()
505 | grid = grid.repeat([self.grid_W * self.grid_Z * self.grid_Y * self.grid_X, 1])
506 | grid = grid.reshape(self.grid_W, self.grid_Z, self.grid_Y, self.grid_X, -1)
507 | grid = grid.permute(4, 0, 1, 2, 3) # (12, grid_W, grid_Z, grid_Y, grid_X)
508 | return grid
509 |
510 | def _init_cp_factors_parafac(self):
511 | # Initialize identity grids.
512 | init_grids = self._init_identity_grid()
513 | # Random noises are added to avoid singularity.
514 | init_grids = torch.randn_like(init_grids) * self.init_noise_scale + init_grids
515 | from tensorly.decomposition import parafac
516 |
517 | # Initialize grid CP factors
518 | _, facs = parafac(init_grids.clone().detach(), rank=self.rank)
519 |
520 | self.num_facs = len(facs)
521 |
522 | self.fac_0 = nn.Linear(facs[0].shape[0], facs[0].shape[1], bias=False)
523 | self.fac_0.weight = nn.Parameter(facs[0]) # (12, rank)
524 |
525 | for i in range(1, self.num_facs):
526 | fac = facs[i].T # (rank, grid_size)
527 | fac = fac.view(1, fac.shape[0], fac.shape[1], 1) # (1, rank, grid_size, 1)
528 | self.register_buffer(f"fac_{i}_init", fac)
529 |
530 | fac_resid = torch.zeros_like(fac)
531 | self.register_parameter(f"fac_{i}", nn.Parameter(fac_resid))
532 |
533 | def tv_loss(self):
534 | """Computes and returns total variation loss on the factors of the low-rank 4D bilateral grids."""
535 |
536 | total_loss = 0
537 | for i in range(1, self.num_facs):
538 | fac = self.get_parameter(f"fac_{i}")
539 | total_loss += total_variation_loss(fac)
540 |
541 | return total_loss
542 |
543 | def forward(self, xyz, rgb):
544 | """Low-rank 4D bilateral grid slicing.
545 |
546 | Args:
547 | xyz (torch.Tensor): The xyz coordinates with shape $(..., 3)$.
548 | rgb (torch.Tensor): The corresponding RGB values with shape $(..., 3)$.
549 |
550 | Returns:
551 | Sliced affine matrices with shape $(..., 3, 4)$.
552 | """
553 | sh_ = xyz.shape
554 | xyz = xyz.reshape(-1, 3) # flatten (N, 3)
555 | rgb = rgb.reshape(-1, 3) # flatten (N, 3)
556 |
557 | xyz = xyz / self.bound
558 | assert self.rgb2gray is not None
559 | gray = self.rgb2gray(rgb)
560 | xyzw = torch.cat([xyz, gray], dim=-1) # (N, 4)
561 | xyzw = xyzw.transpose(0, 1) # (4, N)
562 | coords = torch.stack([torch.zeros_like(xyzw), xyzw], dim=-1) # (4, N, 2)
563 | coords = coords.unsqueeze(1) # (4, 1, N, 2)
564 |
565 | coef = 1.0
566 | for i in range(1, self.num_facs):
567 | fac = self.get_parameter(f"fac_{i}") + self.get_buffer(f"fac_{i}_init")
568 | coef = coef * F.grid_sample(
569 | fac, coords[[i - 1]], align_corners=True, padding_mode="border"
570 | ) # [1, rank, 1, N]
571 | coef = coef.squeeze([0, 2]).transpose(0, 1) # (N, rank) #type: ignore
572 | mat = self.fac_0(coef)
573 | return mat.reshape(*sh_[:-1], 3, 4)
574 |
--------------------------------------------------------------------------------