chores: rebase commits

threestudio/models/background/__init__.py

@@ -0,0 +1,6 @@
+from . import (
+    base,
+    neural_environment_map_background,
+    solid_color_background,
+    textured_background,
+)

threestudio/models/background/base.py

@@ -0,0 +1,24 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.utils.base import BaseModule
+from threestudio.utils.typing import *
+
+
+class BaseBackground(BaseModule):
+    @dataclass
+    class Config(BaseModule.Config):
+        pass
+
+    cfg: Config
+
+    def configure(self):
+        pass
+
+    def forward(self, dirs: Float[Tensor, "B H W 3"]) -> Float[Tensor, "B H W Nc"]:
+        raise NotImplementedError

threestudio/models/background/neural_environment_map_background.py

@@ -0,0 +1,71 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.models.networks import get_encoding, get_mlp
+from threestudio.utils.ops import get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("neural-environment-map-background")
+class NeuralEnvironmentMapBackground(BaseBackground):
+    @dataclass
+    class Config(BaseBackground.Config):
+        n_output_dims: int = 3
+        color_activation: str = "sigmoid"
+        dir_encoding_config: dict = field(
+            default_factory=lambda: {"otype": "SphericalHarmonics", "degree": 3}
+        )
+        mlp_network_config: dict = field(
+            default_factory=lambda: {
+                "otype": "VanillaMLP",
+                "activation": "ReLU",
+                "n_neurons": 16,
+                "n_hidden_layers": 2,
+            }
+        )
+        random_aug: bool = False
+        random_aug_prob: float = 0.5
+        eval_color: Optional[Tuple[float, float, float]] = None
+
+        # multi-view diffusion
+        share_aug_bg: bool = False
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.encoding = get_encoding(3, self.cfg.dir_encoding_config)
+        self.network = get_mlp(
+            self.encoding.n_output_dims,
+            self.cfg.n_output_dims,
+            self.cfg.mlp_network_config,
+        )
+
+    def forward(self, dirs: Float[Tensor, "B H W 3"]) -> Float[Tensor, "B H W Nc"]:
+        if not self.training and self.cfg.eval_color is not None:
+            return torch.ones(*dirs.shape[:-1], self.cfg.n_output_dims).to(
+                dirs
+            ) * torch.as_tensor(self.cfg.eval_color).to(dirs)
+        # viewdirs must be normalized before passing to this function
+        dirs = (dirs + 1.0) / 2.0  # (-1, 1) => (0, 1)
+        dirs_embd = self.encoding(dirs.view(-1, 3))
+        color = self.network(dirs_embd).view(*dirs.shape[:-1], self.cfg.n_output_dims)
+        color = get_activation(self.cfg.color_activation)(color)
+        if (
+            self.training
+            and self.cfg.random_aug
+            and random.random() < self.cfg.random_aug_prob
+        ):
+            # use random background color with probability random_aug_prob
+            n_color = 1 if self.cfg.share_aug_bg else dirs.shape[0]
+            color = color * 0 + (  # prevent checking for unused parameters in DDP
+                torch.rand(n_color, 1, 1, self.cfg.n_output_dims)
+                .to(dirs)
+                .expand(*dirs.shape[:-1], -1)
+            )
+        return color

threestudio/models/background/solid_color_background.py

@@ -0,0 +1,51 @@
+import random
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.utils.typing import *
+
+
+@threestudio.register("solid-color-background")
+class SolidColorBackground(BaseBackground):
+    @dataclass
+    class Config(BaseBackground.Config):
+        n_output_dims: int = 3
+        color: Tuple = (1.0, 1.0, 1.0)
+        learned: bool = False
+        random_aug: bool = False
+        random_aug_prob: float = 0.5
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.env_color: Float[Tensor, "Nc"]
+        if self.cfg.learned:
+            self.env_color = nn.Parameter(
+                torch.as_tensor(self.cfg.color, dtype=torch.float32)
+            )
+        else:
+            self.register_buffer(
+                "env_color", torch.as_tensor(self.cfg.color, dtype=torch.float32)
+            )
+
+    def forward(self, dirs: Float[Tensor, "B H W 3"]) -> Float[Tensor, "B H W Nc"]:
+        color = torch.ones(*dirs.shape[:-1], self.cfg.n_output_dims).to(
+            dirs
+        ) * self.env_color.to(dirs)
+        if (
+            self.training
+            and self.cfg.random_aug
+            and random.random() < self.cfg.random_aug_prob
+        ):
+            # use random background color with probability random_aug_prob
+            color = color * 0 + (  # prevent checking for unused parameters in DDP
+                torch.rand(dirs.shape[0], 1, 1, self.cfg.n_output_dims)
+                .to(dirs)
+                .expand(*dirs.shape[:-1], -1)
+            )
+        return color

threestudio/models/background/textured_background.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import threestudio
+from threestudio.models.background.base import BaseBackground
+from threestudio.utils.ops import get_activation
+from threestudio.utils.typing import *
+
+
+@threestudio.register("textured-background")
+class TexturedBackground(BaseBackground):
+    @dataclass
+    class Config(BaseBackground.Config):
+        n_output_dims: int = 3
+        height: int = 64
+        width: int = 64
+        color_activation: str = "sigmoid"
+
+    cfg: Config
+
+    def configure(self) -> None:
+        self.texture = nn.Parameter(
+            torch.randn((1, self.cfg.n_output_dims, self.cfg.height, self.cfg.width))
+        )
+
+    def spherical_xyz_to_uv(self, dirs: Float[Tensor, "*B 3"]) -> Float[Tensor, "*B 2"]:
+        x, y, z = dirs[..., 0], dirs[..., 1], dirs[..., 2]
+        xy = (x**2 + y**2) ** 0.5
+        u = torch.atan2(xy, z) / torch.pi
+        v = torch.atan2(y, x) / (torch.pi * 2) + 0.5
+        uv = torch.stack([u, v], -1)
+        return uv
+
+    def forward(self, dirs: Float[Tensor, "*B 3"]) -> Float[Tensor, "*B Nc"]:
+        dirs_shape = dirs.shape[:-1]
+        uv = self.spherical_xyz_to_uv(dirs.reshape(-1, dirs.shape[-1]))
+        uv = 2 * uv - 1  # rescale to [-1, 1] for grid_sample
+        uv = uv.reshape(1, -1, 1, 2)
+        color = (
+            F.grid_sample(
+                self.texture,
+                uv,
+                mode="bilinear",
+                padding_mode="reflection",
+                align_corners=False,
+            )
+            .reshape(self.cfg.n_output_dims, -1)
+            .T.reshape(*dirs_shape, self.cfg.n_output_dims)
+        )
+        color = get_activation(self.cfg.color_activation)(color)
+        return color