chores: rebase commits

threestudio/data/__init__.py

@@ -0,0 +1 @@
+from . import image, uncond

threestudio/data/image.py

@@ -0,0 +1,351 @@
+import bisect
+import math
+import os
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset, IterableDataset
+
+import threestudio
+from threestudio import register
+from threestudio.data.uncond import (
+    RandomCameraDataModuleConfig,
+    RandomCameraDataset,
+    RandomCameraIterableDataset,
+)
+from threestudio.utils.base import Updateable
+from threestudio.utils.config import parse_structured
+from threestudio.utils.misc import get_rank
+from threestudio.utils.ops import (
+    get_mvp_matrix,
+    get_projection_matrix,
+    get_ray_directions,
+    get_rays,
+)
+from threestudio.utils.typing import *
+
+
+@dataclass
+class SingleImageDataModuleConfig:
+    # height and width should be Union[int, List[int]]
+    # but OmegaConf does not support Union of containers
+    height: Any = 96
+    width: Any = 96
+    resolution_milestones: List[int] = field(default_factory=lambda: [])
+    default_elevation_deg: float = 0.0
+    default_azimuth_deg: float = -180.0
+    default_camera_distance: float = 1.2
+    default_fovy_deg: float = 60.0
+    image_path: str = ""
+    use_random_camera: bool = True
+    random_camera: dict = field(default_factory=dict)
+    rays_noise_scale: float = 2e-3
+    batch_size: int = 1
+    requires_depth: bool = False
+    requires_normal: bool = False
+    rays_d_normalize: bool = True
+    use_mixed_camera_config: bool = False
+
+
+class SingleImageDataBase:
+    def setup(self, cfg, split):
+        self.split = split
+        self.rank = get_rank()
+        self.cfg: SingleImageDataModuleConfig = cfg
+
+        if self.cfg.use_random_camera:
+            random_camera_cfg = parse_structured(
+                RandomCameraDataModuleConfig, self.cfg.get("random_camera", {})
+            )
+            # FIXME: 
+            if self.cfg.use_mixed_camera_config:
+                if self.rank % 2 == 0:
+                    random_camera_cfg.camera_distance_range=[self.cfg.default_camera_distance, self.cfg.default_camera_distance]
+                    random_camera_cfg.fovy_range=[self.cfg.default_fovy_deg, self.cfg.default_fovy_deg]
+                    self.fixed_camera_intrinsic = True
+                else:
+                    self.fixed_camera_intrinsic = False
+            if split == "train":
+                self.random_pose_generator = RandomCameraIterableDataset(
+                    random_camera_cfg
+                )
+            else:
+                self.random_pose_generator = RandomCameraDataset(
+                    random_camera_cfg, split
+                )
+
+        elevation_deg = torch.FloatTensor([self.cfg.default_elevation_deg])
+        azimuth_deg = torch.FloatTensor([self.cfg.default_azimuth_deg])
+        camera_distance = torch.FloatTensor([self.cfg.default_camera_distance])
+
+        elevation = elevation_deg * math.pi / 180
+        azimuth = azimuth_deg * math.pi / 180
+        camera_position: Float[Tensor, "1 3"] = torch.stack(
+            [
+                camera_distance * torch.cos(elevation) * torch.cos(azimuth),
+                camera_distance * torch.cos(elevation) * torch.sin(azimuth),
+                camera_distance * torch.sin(elevation),
+            ],
+            dim=-1,
+        )
+
+        center: Float[Tensor, "1 3"] = torch.zeros_like(camera_position)
+        up: Float[Tensor, "1 3"] = torch.as_tensor([0, 0, 1], dtype=torch.float32)[None]
+
+        light_position: Float[Tensor, "1 3"] = camera_position
+        lookat: Float[Tensor, "1 3"] = F.normalize(center - camera_position, dim=-1)
+        right: Float[Tensor, "1 3"] = F.normalize(torch.cross(lookat, up), dim=-1)
+        up = F.normalize(torch.cross(right, lookat), dim=-1)
+        self.c2w: Float[Tensor, "1 3 4"] = torch.cat(
+            [torch.stack([right, up, -lookat], dim=-1), camera_position[:, :, None]],
+            dim=-1,
+        )
+        self.c2w4x4: Float[Tensor, "B 4 4"] = torch.cat(
+            [self.c2w, torch.zeros_like(self.c2w[:, :1])], dim=1
+        )
+        self.c2w4x4[:, 3, 3] = 1.0
+
+        self.camera_position = camera_position
+        self.light_position = light_position
+        self.elevation_deg, self.azimuth_deg = elevation_deg, azimuth_deg
+        self.camera_distance = camera_distance
+        self.fovy = torch.deg2rad(torch.FloatTensor([self.cfg.default_fovy_deg]))
+
+        self.heights: List[int] = (
+            [self.cfg.height] if isinstance(self.cfg.height, int) else self.cfg.height
+        )
+        self.widths: List[int] = (
+            [self.cfg.width] if isinstance(self.cfg.width, int) else self.cfg.width
+        )
+        assert len(self.heights) == len(self.widths)
+        self.resolution_milestones: List[int]
+        if len(self.heights) == 1 and len(self.widths) == 1:
+            if len(self.cfg.resolution_milestones) > 0:
+                threestudio.warn(
+                    "Ignoring resolution_milestones since height and width are not changing"
+                )
+            self.resolution_milestones = [-1]
+        else:
+            assert len(self.heights) == len(self.cfg.resolution_milestones) + 1
+            self.resolution_milestones = [-1] + self.cfg.resolution_milestones
+
+        self.directions_unit_focals = [
+            get_ray_directions(H=height, W=width, focal=1.0)
+            for (height, width) in zip(self.heights, self.widths)
+        ]
+        self.focal_lengths = [
+            0.5 * height / torch.tan(0.5 * self.fovy) for height in self.heights
+        ]
+
+        self.height: int = self.heights[0]
+        self.width: int = self.widths[0]
+        self.directions_unit_focal = self.directions_unit_focals[0]
+        self.focal_length = self.focal_lengths[0]
+        self.set_rays()
+        self.load_images()
+        self.prev_height = self.height
+
+    def set_rays(self):
+        # get directions by dividing directions_unit_focal by focal length
+        directions: Float[Tensor, "1 H W 3"] = self.directions_unit_focal[None]
+        directions[:, :, :, :2] = directions[:, :, :, :2] / self.focal_length
+
+        rays_o, rays_d = get_rays(
+            directions,
+            self.c2w,
+            keepdim=True,
+            noise_scale=self.cfg.rays_noise_scale,
+            normalize=self.cfg.rays_d_normalize,
+        )
+
+        proj_mtx: Float[Tensor, "4 4"] = get_projection_matrix(
+            self.fovy, self.width / self.height, 0.01, 100.0
+        )  # FIXME: hard-coded near and far
+        mvp_mtx: Float[Tensor, "4 4"] = get_mvp_matrix(self.c2w, proj_mtx)
+
+        self.rays_o, self.rays_d = rays_o, rays_d
+        self.mvp_mtx = mvp_mtx
+
+    def load_images(self):
+        # load image
+        assert os.path.exists(
+            self.cfg.image_path
+        ), f"Could not find image {self.cfg.image_path}!"
+        rgba = cv2.cvtColor(
+            cv2.imread(self.cfg.image_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA
+        )
+        rgba = (
+            cv2.resize(
+                rgba, (self.width, self.height), interpolation=cv2.INTER_AREA
+            ).astype(np.float32)
+            / 255.0
+        )
+        rgb = rgba[..., :3]
+        self.rgb: Float[Tensor, "1 H W 3"] = (
+            torch.from_numpy(rgb).unsqueeze(0).contiguous().to(self.rank)
+        )
+        self.mask: Float[Tensor, "1 H W 1"] = (
+            torch.from_numpy(rgba[..., 3:] > 0.5).unsqueeze(0).to(self.rank)
+        )
+        print(
+            f"[INFO] single image dataset: load image {self.cfg.image_path} {self.rgb.shape}"
+        )
+
+        # load depth
+        if self.cfg.requires_depth:
+            depth_path = self.cfg.image_path.replace("_rgba.png", "_depth.png")
+            assert os.path.exists(depth_path)
+            depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
+            depth = cv2.resize(
+                depth, (self.width, self.height), interpolation=cv2.INTER_AREA
+            )
+            self.depth: Float[Tensor, "1 H W 1"] = (
+                torch.from_numpy(depth.astype(np.float32) / 255.0)
+                .unsqueeze(0)
+                .to(self.rank)
+            )
+            print(
+                f"[INFO] single image dataset: load depth {depth_path} {self.depth.shape}"
+            )
+        else:
+            self.depth = None
+
+        # load normal
+        if self.cfg.requires_normal:
+            normal_path = self.cfg.image_path.replace("_rgba.png", "_normal.png")
+            assert os.path.exists(normal_path)
+            normal = cv2.imread(normal_path, cv2.IMREAD_UNCHANGED)
+            normal = cv2.resize(
+                normal, (self.width, self.height), interpolation=cv2.INTER_AREA
+            )
+            self.normal: Float[Tensor, "1 H W 3"] = (
+                torch.from_numpy(normal.astype(np.float32) / 255.0)
+                .unsqueeze(0)
+                .to(self.rank)
+            )
+            print(
+                f"[INFO] single image dataset: load normal {normal_path} {self.normal.shape}"
+            )
+        else:
+            self.normal = None
+
+    def get_all_images(self):
+        return self.rgb
+
+    def update_step_(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        size_ind = bisect.bisect_right(self.resolution_milestones, global_step) - 1
+        self.height = self.heights[size_ind]
+        if self.height == self.prev_height:
+            return
+
+        self.prev_height = self.height
+        self.width = self.widths[size_ind]
+        self.directions_unit_focal = self.directions_unit_focals[size_ind]
+        self.focal_length = self.focal_lengths[size_ind]
+        threestudio.debug(f"Training height: {self.height}, width: {self.width}")
+        self.set_rays()
+        self.load_images()
+
+
+class SingleImageIterableDataset(IterableDataset, SingleImageDataBase, Updateable):
+    def __init__(self, cfg: Any, split: str) -> None:
+        super().__init__()
+        self.setup(cfg, split)
+
+    def collate(self, batch) -> Dict[str, Any]:
+        batch = {
+            "rays_o": self.rays_o,
+            "rays_d": self.rays_d,
+            "mvp_mtx": self.mvp_mtx,
+            "camera_positions": self.camera_position,
+            "light_positions": self.light_position,
+            "elevation": self.elevation_deg,
+            "azimuth": self.azimuth_deg,
+            "camera_distances": self.camera_distance,
+            "rgb": self.rgb,
+            "ref_depth": self.depth,
+            "ref_normal": self.normal,
+            "mask": self.mask,
+            "height": self.cfg.height,
+            "width": self.cfg.width,
+            "c2w": self.c2w,
+            "c2w4x4": self.c2w4x4,
+        }
+        if self.cfg.use_random_camera:
+            batch["random_camera"] = self.random_pose_generator.collate(None)
+
+        return batch
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        self.update_step_(epoch, global_step, on_load_weights)
+        self.random_pose_generator.update_step(epoch, global_step, on_load_weights)
+
+    def __iter__(self):
+        while True:
+            yield {}
+
+
+class SingleImageDataset(Dataset, SingleImageDataBase):
+    def __init__(self, cfg: Any, split: str) -> None:
+        super().__init__()
+        self.setup(cfg, split)
+
+    def __len__(self):
+        return len(self.random_pose_generator)
+
+    def __getitem__(self, index):
+        batch = self.random_pose_generator[index]
+        batch.update(
+            {
+            "height": self.random_pose_generator.cfg.eval_height,
+            "width": self.random_pose_generator.cfg.eval_width,
+            "mvp_mtx_ref": self.mvp_mtx[0],
+            "c2w_ref": self.c2w4x4,
+            }
+        )
+        return batch
+
+
+@register("single-image-datamodule")
+class SingleImageDataModule(pl.LightningDataModule):
+    cfg: SingleImageDataModuleConfig
+
+    def __init__(self, cfg: Optional[Union[dict, DictConfig]] = None) -> None:
+        super().__init__()
+        self.cfg = parse_structured(SingleImageDataModuleConfig, cfg)
+
+    def setup(self, stage=None) -> None:
+        if stage in [None, "fit"]:
+            self.train_dataset = SingleImageIterableDataset(self.cfg, "train")
+        if stage in [None, "fit", "validate"]:
+            self.val_dataset = SingleImageDataset(self.cfg, "val")
+        if stage in [None, "test", "predict"]:
+            self.test_dataset = SingleImageDataset(self.cfg, "test")
+
+    def prepare_data(self):
+        pass
+
+    def general_loader(self, dataset, batch_size, collate_fn=None) -> DataLoader:
+        return DataLoader(
+            dataset, num_workers=0, batch_size=batch_size, collate_fn=collate_fn
+        )
+
+    def train_dataloader(self) -> DataLoader:
+        return self.general_loader(
+            self.train_dataset,
+            batch_size=self.cfg.batch_size,
+            collate_fn=self.train_dataset.collate,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        return self.general_loader(self.val_dataset, batch_size=1)
+
+    def test_dataloader(self) -> DataLoader:
+        return self.general_loader(self.test_dataset, batch_size=1)
+
+    def predict_dataloader(self) -> DataLoader:
+        return self.general_loader(self.test_dataset, batch_size=1)

threestudio/data/images.py

@@ -0,0 +1,351 @@
+import bisect
+import math
+import os
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset, IterableDataset
+
+import threestudio
+from threestudio import register
+from threestudio.data.uncond import (
+    RandomCameraDataModuleConfig,
+    RandomCameraDataset,
+    RandomCameraIterableDataset,
+)
+from threestudio.utils.base import Updateable
+from threestudio.utils.config import parse_structured
+from threestudio.utils.misc import get_rank
+from threestudio.utils.ops import (
+    get_mvp_matrix,
+    get_projection_matrix,
+    get_ray_directions,
+    get_rays,
+)
+from threestudio.utils.typing import *
+
+
+@dataclass
+class SingleImageDataModuleConfig:
+    # height and width should be Union[int, List[int]]
+    # but OmegaConf does not support Union of containers
+    height: Any = 96
+    width: Any = 96
+    resolution_milestones: List[int] = field(default_factory=lambda: [])
+    default_elevation_deg: float = 0.0
+    default_azimuth_deg: float = -180.0
+    default_camera_distance: float = 1.2
+    default_fovy_deg: float = 60.0
+    image_path: str = ""
+    use_random_camera: bool = True
+    random_camera: dict = field(default_factory=dict)
+    rays_noise_scale: float = 2e-3
+    batch_size: int = 1
+    requires_depth: bool = False
+    requires_normal: bool = False
+    rays_d_normalize: bool = True
+    use_mixed_camera_config: bool = False
+
+
+class SingleImageDataBase:
+    def setup(self, cfg, split):
+        self.split = split
+        self.rank = get_rank()
+        self.cfg: SingleImageDataModuleConfig = cfg
+
+        if self.cfg.use_random_camera:
+            random_camera_cfg = parse_structured(
+                RandomCameraDataModuleConfig, self.cfg.get("random_camera", {})
+            )
+            # FIXME: 
+            if self.cfg.use_mixed_camera_config:
+                if self.rank % 2 == 0:
+                    random_camera_cfg.camera_distance_range=[self.cfg.default_camera_distance, self.cfg.default_camera_distance]
+                    random_camera_cfg.fovy_range=[self.cfg.default_fovy_deg, self.cfg.default_fovy_deg]
+                    self.fixed_camera_intrinsic = True
+                else:
+                    self.fixed_camera_intrinsic = False
+            if split == "train":
+                self.random_pose_generator = RandomCameraIterableDataset(
+                    random_camera_cfg
+                )
+            else:
+                self.random_pose_generator = RandomCameraDataset(
+                    random_camera_cfg, split
+                )
+
+        elevation_deg = torch.FloatTensor([self.cfg.default_elevation_deg])
+        azimuth_deg = torch.FloatTensor([self.cfg.default_azimuth_deg])
+        camera_distance = torch.FloatTensor([self.cfg.default_camera_distance])
+
+        elevation = elevation_deg * math.pi / 180
+        azimuth = azimuth_deg * math.pi / 180
+        camera_position: Float[Tensor, "1 3"] = torch.stack(
+            [
+                camera_distance * torch.cos(elevation) * torch.cos(azimuth),
+                camera_distance * torch.cos(elevation) * torch.sin(azimuth),
+                camera_distance * torch.sin(elevation),
+            ],
+            dim=-1,
+        )
+
+        center: Float[Tensor, "1 3"] = torch.zeros_like(camera_position)
+        up: Float[Tensor, "1 3"] = torch.as_tensor([0, 0, 1], dtype=torch.float32)[None]
+
+        light_position: Float[Tensor, "1 3"] = camera_position
+        lookat: Float[Tensor, "1 3"] = F.normalize(center - camera_position, dim=-1)
+        right: Float[Tensor, "1 3"] = F.normalize(torch.cross(lookat, up), dim=-1)
+        up = F.normalize(torch.cross(right, lookat), dim=-1)
+        self.c2w: Float[Tensor, "1 3 4"] = torch.cat(
+            [torch.stack([right, up, -lookat], dim=-1), camera_position[:, :, None]],
+            dim=-1,
+        )
+        self.c2w4x4: Float[Tensor, "B 4 4"] = torch.cat(
+            [self.c2w, torch.zeros_like(self.c2w[:, :1])], dim=1
+        )
+        self.c2w4x4[:, 3, 3] = 1.0
+
+        self.camera_position = camera_position
+        self.light_position = light_position
+        self.elevation_deg, self.azimuth_deg = elevation_deg, azimuth_deg
+        self.camera_distance = camera_distance
+        self.fovy = torch.deg2rad(torch.FloatTensor([self.cfg.default_fovy_deg]))
+
+        self.heights: List[int] = (
+            [self.cfg.height] if isinstance(self.cfg.height, int) else self.cfg.height
+        )
+        self.widths: List[int] = (
+            [self.cfg.width] if isinstance(self.cfg.width, int) else self.cfg.width
+        )
+        assert len(self.heights) == len(self.widths)
+        self.resolution_milestones: List[int]
+        if len(self.heights) == 1 and len(self.widths) == 1:
+            if len(self.cfg.resolution_milestones) > 0:
+                threestudio.warn(
+                    "Ignoring resolution_milestones since height and width are not changing"
+                )
+            self.resolution_milestones = [-1]
+        else:
+            assert len(self.heights) == len(self.cfg.resolution_milestones) + 1
+            self.resolution_milestones = [-1] + self.cfg.resolution_milestones
+
+        self.directions_unit_focals = [
+            get_ray_directions(H=height, W=width, focal=1.0)
+            for (height, width) in zip(self.heights, self.widths)
+        ]
+        self.focal_lengths = [
+            0.5 * height / torch.tan(0.5 * self.fovy) for height in self.heights
+        ]
+
+        self.height: int = self.heights[0]
+        self.width: int = self.widths[0]
+        self.directions_unit_focal = self.directions_unit_focals[0]
+        self.focal_length = self.focal_lengths[0]
+        self.set_rays()
+        self.load_images()
+        self.prev_height = self.height
+
+    def set_rays(self):
+        # get directions by dividing directions_unit_focal by focal length
+        directions: Float[Tensor, "1 H W 3"] = self.directions_unit_focal[None]
+        directions[:, :, :, :2] = directions[:, :, :, :2] / self.focal_length
+
+        rays_o, rays_d = get_rays(
+            directions,
+            self.c2w,
+            keepdim=True,
+            noise_scale=self.cfg.rays_noise_scale,
+            normalize=self.cfg.rays_d_normalize,
+        )
+
+        proj_mtx: Float[Tensor, "4 4"] = get_projection_matrix(
+            self.fovy, self.width / self.height, 0.01, 100.0
+        )  # FIXME: hard-coded near and far
+        mvp_mtx: Float[Tensor, "4 4"] = get_mvp_matrix(self.c2w, proj_mtx)
+
+        self.rays_o, self.rays_d = rays_o, rays_d
+        self.mvp_mtx = mvp_mtx
+
+    def load_images(self):
+        # load image
+        assert os.path.exists(
+            self.cfg.image_path
+        ), f"Could not find image {self.cfg.image_path}!"
+        rgba = cv2.cvtColor(
+            cv2.imread(self.cfg.image_path, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGRA2RGBA
+        )
+        rgba = (
+            cv2.resize(
+                rgba, (self.width, self.height), interpolation=cv2.INTER_AREA
+            ).astype(np.float32)
+            / 255.0
+        )
+        rgb = rgba[..., :3]
+        self.rgb: Float[Tensor, "1 H W 3"] = (
+            torch.from_numpy(rgb).unsqueeze(0).contiguous().to(self.rank)
+        )
+        self.mask: Float[Tensor, "1 H W 1"] = (
+            torch.from_numpy(rgba[..., 3:] > 0.5).unsqueeze(0).to(self.rank)
+        )
+        print(
+            f"[INFO] single image dataset: load image {self.cfg.image_path} {self.rgb.shape}"
+        )
+
+        # load depth
+        if self.cfg.requires_depth:
+            depth_path = self.cfg.image_path.replace("_rgba.png", "_depth.png")
+            assert os.path.exists(depth_path)
+            depth = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
+            depth = cv2.resize(
+                depth, (self.width, self.height), interpolation=cv2.INTER_AREA
+            )
+            self.depth: Float[Tensor, "1 H W 1"] = (
+                torch.from_numpy(depth.astype(np.float32) / 255.0)
+                .unsqueeze(0)
+                .to(self.rank)
+            )
+            print(
+                f"[INFO] single image dataset: load depth {depth_path} {self.depth.shape}"
+            )
+        else:
+            self.depth = None
+
+        # load normal
+        if self.cfg.requires_normal:
+            normal_path = self.cfg.image_path.replace("_rgba.png", "_normal.png")
+            assert os.path.exists(normal_path)
+            normal = cv2.imread(normal_path, cv2.IMREAD_UNCHANGED)
+            normal = cv2.resize(
+                normal, (self.width, self.height), interpolation=cv2.INTER_AREA
+            )
+            self.normal: Float[Tensor, "1 H W 3"] = (
+                torch.from_numpy(normal.astype(np.float32) / 255.0)
+                .unsqueeze(0)
+                .to(self.rank)
+            )
+            print(
+                f"[INFO] single image dataset: load normal {normal_path} {self.normal.shape}"
+            )
+        else:
+            self.normal = None
+
+    def get_all_images(self):
+        return self.rgb
+
+    def update_step_(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        size_ind = bisect.bisect_right(self.resolution_milestones, global_step) - 1
+        self.height = self.heights[size_ind]
+        if self.height == self.prev_height:
+            return
+
+        self.prev_height = self.height
+        self.width = self.widths[size_ind]
+        self.directions_unit_focal = self.directions_unit_focals[size_ind]
+        self.focal_length = self.focal_lengths[size_ind]
+        threestudio.debug(f"Training height: {self.height}, width: {self.width}")
+        self.set_rays()
+        self.load_images()
+
+
+class SingleImageIterableDataset(IterableDataset, SingleImageDataBase, Updateable):
+    def __init__(self, cfg: Any, split: str) -> None:
+        super().__init__()
+        self.setup(cfg, split)
+
+    def collate(self, batch) -> Dict[str, Any]:
+        batch = {
+            "rays_o": self.rays_o,
+            "rays_d": self.rays_d,
+            "mvp_mtx": self.mvp_mtx,
+            "camera_positions": self.camera_position,
+            "light_positions": self.light_position,
+            "elevation": self.elevation_deg,
+            "azimuth": self.azimuth_deg,
+            "camera_distances": self.camera_distance,
+            "rgb": self.rgb,
+            "ref_depth": self.depth,
+            "ref_normal": self.normal,
+            "mask": self.mask,
+            "height": self.cfg.height,
+            "width": self.cfg.width,
+            "c2w": self.c2w,
+            "c2w4x4": self.c2w4x4,
+        }
+        if self.cfg.use_random_camera:
+            batch["random_camera"] = self.random_pose_generator.collate(None)
+
+        return batch
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        self.update_step_(epoch, global_step, on_load_weights)
+        self.random_pose_generator.update_step(epoch, global_step, on_load_weights)
+
+    def __iter__(self):
+        while True:
+            yield {}
+
+
+class SingleImageDataset(Dataset, SingleImageDataBase):
+    def __init__(self, cfg: Any, split: str) -> None:
+        super().__init__()
+        self.setup(cfg, split)
+
+    def __len__(self):
+        return len(self.random_pose_generator)
+
+    def __getitem__(self, index):
+        batch = self.random_pose_generator[index]
+        batch.update(
+            {
+            "height": self.random_pose_generator.cfg.eval_height,
+            "width": self.random_pose_generator.cfg.eval_width,
+            "mvp_mtx_ref": self.mvp_mtx[0],
+            "c2w_ref": self.c2w4x4,
+            }
+        )
+        return batch
+
+
+@register("single-image-datamodule")
+class SingleImageDataModule(pl.LightningDataModule):
+    cfg: SingleImageDataModuleConfig
+
+    def __init__(self, cfg: Optional[Union[dict, DictConfig]] = None) -> None:
+        super().__init__()
+        self.cfg = parse_structured(SingleImageDataModuleConfig, cfg)
+
+    def setup(self, stage=None) -> None:
+        if stage in [None, "fit"]:
+            self.train_dataset = SingleImageIterableDataset(self.cfg, "train")
+        if stage in [None, "fit", "validate"]:
+            self.val_dataset = SingleImageDataset(self.cfg, "val")
+        if stage in [None, "test", "predict"]:
+            self.test_dataset = SingleImageDataset(self.cfg, "test")
+
+    def prepare_data(self):
+        pass
+
+    def general_loader(self, dataset, batch_size, collate_fn=None) -> DataLoader:
+        return DataLoader(
+            dataset, num_workers=0, batch_size=batch_size, collate_fn=collate_fn
+        )
+
+    def train_dataloader(self) -> DataLoader:
+        return self.general_loader(
+            self.train_dataset,
+            batch_size=self.cfg.batch_size,
+            collate_fn=self.train_dataset.collate,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        return self.general_loader(self.val_dataset, batch_size=1)
+
+    def test_dataloader(self) -> DataLoader:
+        return self.general_loader(self.test_dataset, batch_size=1)
+
+    def predict_dataloader(self) -> DataLoader:
+        return self.general_loader(self.test_dataset, batch_size=1)

threestudio/data/uncond.py

@@ -0,0 +1,518 @@
+import bisect
+import math
+import random
+from dataclasses import dataclass, field
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset, IterableDataset
+
+import threestudio
+from threestudio import register
+from threestudio.utils.base import Updateable
+from threestudio.utils.config import parse_structured
+from threestudio.utils.misc import get_device
+from threestudio.utils.ops import (
+    get_full_projection_matrix,
+    get_mvp_matrix,
+    get_projection_matrix,
+    get_ray_directions,
+    get_rays,
+)
+from threestudio.utils.typing import *
+
+
+@dataclass
+class RandomCameraDataModuleConfig:
+    # height, width, and batch_size should be Union[int, List[int]]
+    # but OmegaConf does not support Union of containers
+    height: Any = 64
+    width: Any = 64
+    batch_size: Any = 1
+    resolution_milestones: List[int] = field(default_factory=lambda: [])
+    eval_height: int = 512
+    eval_width: int = 512
+    eval_batch_size: int = 1
+    n_val_views: int = 1
+    n_test_views: int = 120
+    elevation_range: Tuple[float, float] = (-10, 90)
+    azimuth_range: Tuple[float, float] = (-180, 180)
+    camera_distance_range: Tuple[float, float] = (1, 1.5)
+    fovy_range: Tuple[float, float] = (
+        40,
+        70,
+    )  # in degrees, in vertical direction (along height)
+    camera_perturb: float = 0.1
+    center_perturb: float = 0.2
+    up_perturb: float = 0.02
+    light_position_perturb: float = 1.0
+    light_distance_range: Tuple[float, float] = (0.8, 1.5)
+    eval_elevation_deg: float = 15.0
+    eval_camera_distance: float = 1.5
+    eval_fovy_deg: float = 70.0
+    light_sample_strategy: str = "dreamfusion"
+    batch_uniform_azimuth: bool = True
+    progressive_until: int = 0  # progressive ranges for elevation, azimuth, r, fovy
+
+    rays_d_normalize: bool = True
+
+
+class RandomCameraIterableDataset(IterableDataset, Updateable):
+    def __init__(self, cfg: Any) -> None:
+        super().__init__()
+        self.cfg: RandomCameraDataModuleConfig = cfg
+        self.heights: List[int] = (
+            [self.cfg.height] if isinstance(self.cfg.height, int) else self.cfg.height
+        )
+        self.widths: List[int] = (
+            [self.cfg.width] if isinstance(self.cfg.width, int) else self.cfg.width
+        )
+        self.batch_sizes: List[int] = (
+            [self.cfg.batch_size]
+            if isinstance(self.cfg.batch_size, int)
+            else self.cfg.batch_size
+        )
+        assert len(self.heights) == len(self.widths) == len(self.batch_sizes)
+        self.resolution_milestones: List[int]
+        if (
+            len(self.heights) == 1
+            and len(self.widths) == 1
+            and len(self.batch_sizes) == 1
+        ):
+            if len(self.cfg.resolution_milestones) > 0:
+                threestudio.warn(
+                    "Ignoring resolution_milestones since height and width are not changing"
+                )
+            self.resolution_milestones = [-1]
+        else:
+            assert len(self.heights) == len(self.cfg.resolution_milestones) + 1
+            self.resolution_milestones = [-1] + self.cfg.resolution_milestones
+
+        self.directions_unit_focals = [
+            get_ray_directions(H=height, W=width, focal=1.0)
+            for (height, width) in zip(self.heights, self.widths)
+        ]
+        self.height: int = self.heights[0]
+        self.width: int = self.widths[0]
+        self.batch_size: int = self.batch_sizes[0]
+        self.directions_unit_focal = self.directions_unit_focals[0]
+        self.elevation_range = self.cfg.elevation_range
+        self.azimuth_range = self.cfg.azimuth_range
+        self.camera_distance_range = self.cfg.camera_distance_range
+        self.fovy_range = self.cfg.fovy_range
+
+    def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
+        size_ind = bisect.bisect_right(self.resolution_milestones, global_step) - 1
+        self.height = self.heights[size_ind]
+        self.width = self.widths[size_ind]
+        self.batch_size = self.batch_sizes[size_ind]
+        self.directions_unit_focal = self.directions_unit_focals[size_ind]
+        threestudio.debug(
+            f"Training height: {self.height}, width: {self.width}, batch_size: {self.batch_size}"
+        )
+        # progressive view
+        self.progressive_view(global_step)
+
+    def __iter__(self):
+        while True:
+            yield {}
+
+    def progressive_view(self, global_step):
+        r = min(1.0, global_step / (self.cfg.progressive_until + 1))
+        self.elevation_range = [
+            (1 - r) * self.cfg.eval_elevation_deg + r * self.cfg.elevation_range[0],
+            (1 - r) * self.cfg.eval_elevation_deg + r * self.cfg.elevation_range[1],
+        ]
+        self.azimuth_range = [
+            (1 - r) * 0.0 + r * self.cfg.azimuth_range[0],
+            (1 - r) * 0.0 + r * self.cfg.azimuth_range[1],
+        ]
+        # self.camera_distance_range = [
+        #     (1 - r) * self.cfg.eval_camera_distance
+        #     + r * self.cfg.camera_distance_range[0],
+        #     (1 - r) * self.cfg.eval_camera_distance
+        #     + r * self.cfg.camera_distance_range[1],
+        # ]
+        # self.fovy_range = [
+        #     (1 - r) * self.cfg.eval_fovy_deg + r * self.cfg.fovy_range[0],
+        #     (1 - r) * self.cfg.eval_fovy_deg + r * self.cfg.fovy_range[1],
+        # ]
+
+    def collate(self, batch) -> Dict[str, Any]:
+        # sample elevation angles
+        elevation_deg: Float[Tensor, "B"]
+        elevation: Float[Tensor, "B"]
+        if random.random() < 0.5:
+            # sample elevation angles uniformly with a probability 0.5 (biased towards poles)
+            elevation_deg = (
+                torch.rand(self.batch_size)
+                * (self.elevation_range[1] - self.elevation_range[0])
+                + self.elevation_range[0]
+            )
+            elevation = elevation_deg * math.pi / 180
+        else:
+            # otherwise sample uniformly on sphere
+            elevation_range_percent = [
+                self.elevation_range[0] / 180.0 * math.pi,
+                self.elevation_range[1] / 180.0 * math.pi,
+            ]
+            # inverse transform sampling
+            elevation = torch.asin(
+                (
+                    torch.rand(self.batch_size)
+                    * (
+                        math.sin(elevation_range_percent[1])
+                        - math.sin(elevation_range_percent[0])
+                    )
+                    + math.sin(elevation_range_percent[0])
+                )
+            )
+            elevation_deg = elevation / math.pi * 180.0
+
+        # sample azimuth angles from a uniform distribution bounded by azimuth_range
+        azimuth_deg: Float[Tensor, "B"]
+        if self.cfg.batch_uniform_azimuth:
+            # ensures sampled azimuth angles in a batch cover the whole range
+            azimuth_deg = (
+                torch.rand(self.batch_size) + torch.arange(self.batch_size)
+            ) / self.batch_size * (
+                self.azimuth_range[1] - self.azimuth_range[0]
+            ) + self.azimuth_range[
+                0
+            ]
+        else:
+            # simple random sampling
+            azimuth_deg = (
+                torch.rand(self.batch_size)
+                * (self.azimuth_range[1] - self.azimuth_range[0])
+                + self.azimuth_range[0]
+            )
+        azimuth = azimuth_deg * math.pi / 180
+
+        # sample distances from a uniform distribution bounded by distance_range
+        camera_distances: Float[Tensor, "B"] = (
+            torch.rand(self.batch_size)
+            * (self.camera_distance_range[1] - self.camera_distance_range[0])
+            + self.camera_distance_range[0]
+        )
+
+        # convert spherical coordinates to cartesian coordinates
+        # right hand coordinate system, x back, y right, z up
+        # elevation in (-90, 90), azimuth from +x to +y in (-180, 180)
+        camera_positions: Float[Tensor, "B 3"] = torch.stack(
+            [
+                camera_distances * torch.cos(elevation) * torch.cos(azimuth),
+                camera_distances * torch.cos(elevation) * torch.sin(azimuth),
+                camera_distances * torch.sin(elevation),
+            ],
+            dim=-1,
+        )
+
+        # default scene center at origin
+        center: Float[Tensor, "B 3"] = torch.zeros_like(camera_positions)
+        # default camera up direction as +z
+        up: Float[Tensor, "B 3"] = torch.as_tensor([0, 0, 1], dtype=torch.float32)[
+            None, :
+        ].repeat(self.batch_size, 1)
+
+        # sample camera perturbations from a uniform distribution [-camera_perturb, camera_perturb]
+        camera_perturb: Float[Tensor, "B 3"] = (
+            torch.rand(self.batch_size, 3) * 2 * self.cfg.camera_perturb
+            - self.cfg.camera_perturb
+        )
+        camera_positions = camera_positions + camera_perturb
+        # sample center perturbations from a normal distribution with mean 0 and std center_perturb
+        center_perturb: Float[Tensor, "B 3"] = (
+            torch.randn(self.batch_size, 3) * self.cfg.center_perturb
+        )
+        center = center + center_perturb
+        # sample up perturbations from a normal distribution with mean 0 and std up_perturb
+        up_perturb: Float[Tensor, "B 3"] = (
+            torch.randn(self.batch_size, 3) * self.cfg.up_perturb
+        )
+        up = up + up_perturb
+
+        # sample fovs from a uniform distribution bounded by fov_range
+        fovy_deg: Float[Tensor, "B"] = (
+            torch.rand(self.batch_size) * (self.fovy_range[1] - self.fovy_range[0])
+            + self.fovy_range[0]
+        )
+        fovy = fovy_deg * math.pi / 180
+
+        # sample light distance from a uniform distribution bounded by light_distance_range
+        light_distances: Float[Tensor, "B"] = (
+            torch.rand(self.batch_size)
+            * (self.cfg.light_distance_range[1] - self.cfg.light_distance_range[0])
+            + self.cfg.light_distance_range[0]
+        )
+
+        if self.cfg.light_sample_strategy == "dreamfusion":
+            # sample light direction from a normal distribution with mean camera_position and std light_position_perturb
+            light_direction: Float[Tensor, "B 3"] = F.normalize(
+                camera_positions
+                + torch.randn(self.batch_size, 3) * self.cfg.light_position_perturb,
+                dim=-1,
+            )
+            # get light position by scaling light direction by light distance
+            light_positions: Float[Tensor, "B 3"] = (
+                light_direction * light_distances[:, None]
+            )
+        elif self.cfg.light_sample_strategy == "magic3d":
+            # sample light direction within restricted angle range (pi/3)
+            local_z = F.normalize(camera_positions, dim=-1)
+            local_x = F.normalize(
+                torch.stack(
+                    [local_z[:, 1], -local_z[:, 0], torch.zeros_like(local_z[:, 0])],
+                    dim=-1,
+                ),
+                dim=-1,
+            )
+            local_y = F.normalize(torch.cross(local_z, local_x, dim=-1), dim=-1)
+            rot = torch.stack([local_x, local_y, local_z], dim=-1)
+            light_azimuth = (
+                torch.rand(self.batch_size) * math.pi * 2 - math.pi
+            )  # [-pi, pi]
+            light_elevation = (
+                torch.rand(self.batch_size) * math.pi / 3 + math.pi / 6
+            )  # [pi/6, pi/2]
+            light_positions_local = torch.stack(
+                [
+                    light_distances
+                    * torch.cos(light_elevation)
+                    * torch.cos(light_azimuth),
+                    light_distances
+                    * torch.cos(light_elevation)
+                    * torch.sin(light_azimuth),
+                    light_distances * torch.sin(light_elevation),
+                ],
+                dim=-1,
+            )
+            light_positions = (rot @ light_positions_local[:, :, None])[:, :, 0]
+        else:
+            raise ValueError(
+                f"Unknown light sample strategy: {self.cfg.light_sample_strategy}"
+            )
+
+        lookat: Float[Tensor, "B 3"] = F.normalize(center - camera_positions, dim=-1)
+        right: Float[Tensor, "B 3"] = F.normalize(torch.cross(lookat, up), dim=-1)
+        up = F.normalize(torch.cross(right, lookat), dim=-1)
+        c2w3x4: Float[Tensor, "B 3 4"] = torch.cat(
+            [torch.stack([right, up, -lookat], dim=-1), camera_positions[:, :, None]],
+            dim=-1,
+        )
+        c2w: Float[Tensor, "B 4 4"] = torch.cat(
+            [c2w3x4, torch.zeros_like(c2w3x4[:, :1])], dim=1
+        )
+        c2w[:, 3, 3] = 1.0
+
+        # get directions by dividing directions_unit_focal by focal length
+        focal_length: Float[Tensor, "B"] = 0.5 * self.height / torch.tan(0.5 * fovy)
+        directions: Float[Tensor, "B H W 3"] = self.directions_unit_focal[
+            None, :, :, :
+        ].repeat(self.batch_size, 1, 1, 1)
+        directions[:, :, :, :2] = (
+            directions[:, :, :, :2] / focal_length[:, None, None, None]
+        )
+
+        # Importance note: the returned rays_d MUST be normalized!
+        rays_o, rays_d = get_rays(
+            directions, c2w, keepdim=True, normalize=self.cfg.rays_d_normalize
+        )
+
+        self.proj_mtx: Float[Tensor, "B 4 4"] = get_projection_matrix(
+            fovy, self.width / self.height, 0.1, 1000.0
+        )  # FIXME: hard-coded near and far
+        mvp_mtx: Float[Tensor, "B 4 4"] = get_mvp_matrix(c2w, self.proj_mtx)
+        self.fovy = fovy
+
+        return {
+            "rays_o": rays_o,
+            "rays_d": rays_d,
+            "mvp_mtx": mvp_mtx,
+            "camera_positions": camera_positions,
+            "c2w": c2w,
+            "light_positions": light_positions,
+            "elevation": elevation_deg,
+            "azimuth": azimuth_deg,
+            "camera_distances": camera_distances,
+            "height": self.height,
+            "width": self.width,
+            "fovy": self.fovy,
+            "proj_mtx": self.proj_mtx,
+        }
+
+
+class RandomCameraDataset(Dataset):
+    def __init__(self, cfg: Any, split: str) -> None:
+        super().__init__()
+        self.cfg: RandomCameraDataModuleConfig = cfg
+        self.split = split
+
+        if split == "val":
+            self.n_views = self.cfg.n_val_views
+        else:
+            self.n_views = self.cfg.n_test_views
+
+        azimuth_deg: Float[Tensor, "B"]
+        if self.split == "val":
+            # make sure the first and last view are not the same
+            azimuth_deg = torch.linspace(0, 360.0, self.n_views + 1)[: self.n_views]
+        else:
+            azimuth_deg = torch.linspace(0, 360.0, self.n_views)
+        elevation_deg: Float[Tensor, "B"] = torch.full_like(
+            azimuth_deg, self.cfg.eval_elevation_deg
+        )
+        camera_distances: Float[Tensor, "B"] = torch.full_like(
+            elevation_deg, self.cfg.eval_camera_distance
+        )
+
+        elevation = elevation_deg * math.pi / 180
+        azimuth = azimuth_deg * math.pi / 180
+
+        # convert spherical coordinates to cartesian coordinates
+        # right hand coordinate system, x back, y right, z up
+        # elevation in (-90, 90), azimuth from +x to +y in (-180, 180)
+        camera_positions: Float[Tensor, "B 3"] = torch.stack(
+            [
+                camera_distances * torch.cos(elevation) * torch.cos(azimuth),
+                camera_distances * torch.cos(elevation) * torch.sin(azimuth),
+                camera_distances * torch.sin(elevation),
+            ],
+            dim=-1,
+        )
+
+        # default scene center at origin
+        center: Float[Tensor, "B 3"] = torch.zeros_like(camera_positions)
+        # default camera up direction as +z
+        up: Float[Tensor, "B 3"] = torch.as_tensor([0, 0, 1], dtype=torch.float32)[
+            None, :
+        ].repeat(self.cfg.eval_batch_size, 1)
+
+        fovy_deg: Float[Tensor, "B"] = torch.full_like(
+            elevation_deg, self.cfg.eval_fovy_deg
+        )
+        fovy = fovy_deg * math.pi / 180
+        light_positions: Float[Tensor, "B 3"] = camera_positions
+
+        lookat: Float[Tensor, "B 3"] = F.normalize(center - camera_positions, dim=-1)
+        right: Float[Tensor, "B 3"] = F.normalize(torch.cross(lookat, up), dim=-1)
+        up = F.normalize(torch.cross(right, lookat), dim=-1)
+        c2w3x4: Float[Tensor, "B 3 4"] = torch.cat(
+            [torch.stack([right, up, -lookat], dim=-1), camera_positions[:, :, None]],
+            dim=-1,
+        )
+        c2w: Float[Tensor, "B 4 4"] = torch.cat(
+            [c2w3x4, torch.zeros_like(c2w3x4[:, :1])], dim=1
+        )
+        c2w[:, 3, 3] = 1.0
+
+        # get directions by dividing directions_unit_focal by focal length
+        focal_length: Float[Tensor, "B"] = (
+            0.5 * self.cfg.eval_height / torch.tan(0.5 * fovy)
+        )
+        directions_unit_focal = get_ray_directions(
+            H=self.cfg.eval_height, W=self.cfg.eval_width, focal=1.0
+        )
+        directions: Float[Tensor, "B H W 3"] = directions_unit_focal[
+            None, :, :, :
+        ].repeat(self.n_views, 1, 1, 1)
+        directions[:, :, :, :2] = (
+            directions[:, :, :, :2] / focal_length[:, None, None, None]
+        )
+
+        rays_o, rays_d = get_rays(
+            directions, c2w, keepdim=True, normalize=self.cfg.rays_d_normalize
+        )
+        self.proj_mtx: Float[Tensor, "B 4 4"] = get_projection_matrix(
+            fovy, self.cfg.eval_width / self.cfg.eval_height, 0.1, 1000.0
+        )  # FIXME: hard-coded near and far
+        mvp_mtx: Float[Tensor, "B 4 4"] = get_mvp_matrix(c2w, self.proj_mtx)
+
+        self.rays_o, self.rays_d = rays_o, rays_d
+        self.mvp_mtx = mvp_mtx
+        self.c2w = c2w
+        self.camera_positions = camera_positions
+        self.light_positions = light_positions
+        self.elevation, self.azimuth = elevation, azimuth
+        self.elevation_deg, self.azimuth_deg = elevation_deg, azimuth_deg
+        self.camera_distances = camera_distances
+        self.fovy = fovy
+
+    def __len__(self):
+        return self.n_views
+
+    def __getitem__(self, index):
+        return {
+            "index": index,
+            "rays_o": self.rays_o[index],
+            "rays_d": self.rays_d[index],
+            "mvp_mtx": self.mvp_mtx[index],
+            "c2w": self.c2w[index],
+            "camera_positions": self.camera_positions[index],
+            "light_positions": self.light_positions[index],
+            "elevation": self.elevation_deg[index],
+            "azimuth": self.azimuth_deg[index],
+            "camera_distances": self.camera_distances[index],
+            "height": self.cfg.eval_height,
+            "width": self.cfg.eval_width,
+            "fovy": self.fovy[index],
+            "proj_mtx": self.proj_mtx[index],
+        }
+
+    def collate(self, batch):
+        batch = torch.utils.data.default_collate(batch)
+        batch.update({"height": self.cfg.eval_height, "width": self.cfg.eval_width})
+        return batch
+
+
+@register("random-camera-datamodule")
+class RandomCameraDataModule(pl.LightningDataModule):
+    cfg: RandomCameraDataModuleConfig
+
+    def __init__(self, cfg: Optional[Union[dict, DictConfig]] = None) -> None:
+        super().__init__()
+        self.cfg = parse_structured(RandomCameraDataModuleConfig, cfg)
+
+    def setup(self, stage=None) -> None:
+        if stage in [None, "fit"]:
+            self.train_dataset = RandomCameraIterableDataset(self.cfg)
+        if stage in [None, "fit", "validate"]:
+            self.val_dataset = RandomCameraDataset(self.cfg, "val")
+        if stage in [None, "test", "predict"]:
+            self.test_dataset = RandomCameraDataset(self.cfg, "test")
+
+    def prepare_data(self):
+        pass
+
+    def general_loader(self, dataset, batch_size, collate_fn=None) -> DataLoader:
+        return DataLoader(
+            dataset,
+            # very important to disable multi-processing if you want to change self attributes at runtime!
+            # (for example setting self.width and self.height in update_step)
+            num_workers=0,  # type: ignore
+            batch_size=batch_size,
+            collate_fn=collate_fn,
+        )
+
+    def train_dataloader(self) -> DataLoader:
+        return self.general_loader(
+            self.train_dataset, batch_size=None, collate_fn=self.train_dataset.collate
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        return self.general_loader(
+            self.val_dataset, batch_size=1, collate_fn=self.val_dataset.collate
+        )
+        # return self.general_loader(self.train_dataset, batch_size=None, collate_fn=self.train_dataset.collate)
+
+    def test_dataloader(self) -> DataLoader:
+        return self.general_loader(
+            self.test_dataset, batch_size=1, collate_fn=self.test_dataset.collate
+        )
+
+    def predict_dataloader(self) -> DataLoader:
+        return self.general_loader(
+            self.test_dataset, batch_size=1, collate_fn=self.test_dataset.collate
+        )