From 72da747fc2fd00bd71f2b6ad8297aa94797f8006 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Tue, 1 Aug 2023 15:50:22 -0700 Subject: [PATCH 1/4] Remove system-level CUDA install from install instructions (#2289) * Remove system-level CUDA install from Linux install instructions * Remove system-level CUDA from Windows install instructions * Un-flip CUDA 11.7 / 11.8 --- docs/quickstart/installation.md | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md index 30d75ebc3..271c7caec 100644 --- a/docs/quickstart/installation.md +++ b/docs/quickstart/installation.md @@ -5,23 +5,22 @@ ::::::{tab-set} :::::{tab-item} Linux -Install CUDA. This library has been tested with version 11.8. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html). +Nerfstudio requires `python >= 3.8`. We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before proceeding. ::::: :::::{tab-item} Windows + Install [Git](https://git-scm.com/downloads). Install Visual Studio 2022. This must be done before installing CUDA. The necessary components are included in the `Desktop Development with C++` workflow (also called `C++ Build Tools` in the BuildTools edition). -Install CUDA 11.8. You can find CUDA download links [here](https://developer.nvidia.com/cuda-toolkit-archive) and more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html). +Nerfstudio requires `python >= 3.8`. We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before proceeding. ::::: :::::: ## Create environment -Nerfstudio requires `python >= 3.8`. We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before proceeding. - ```bash conda create --name nerfstudio -y python=3.8 conda activate nerfstudio @@ -55,7 +54,7 @@ To build the necessary CUDA extensions, `cuda-toolkit` is also required. We recommend installing with conda: ```bash -conda install -c "nvidia/label/cuda-11.7.1" cuda-toolkit +conda install -c "nvidia/label/cuda-11.8.0" cuda-toolkit ``` ::: @@ -71,7 +70,7 @@ To build the necessary CUDA extensions, `cuda-toolkit` is also required. We recommend installing with conda: ```bash -conda install -c "nvidia/label/cuda-11.8.0" cuda-toolkit +conda install -c "nvidia/label/cuda-11.7.1" cuda-toolkit ``` ::: From 71c6b9cf3db4541bedaa17940e4392ef557d14c1 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Tue, 1 Aug 2023 21:52:47 -0400 Subject: [PATCH 2/4] Monodepth Regularization (#2284) * monodepth implemented * device * pair pixel sampler * docstring * removing unused imports * fix bugs * cache in cache dir, no longer edit json file * moving to depth nerfacto * fixed bug using median depth by accident * removing depthdataset from nerfacto * added pixelsamplerconfig * added equirectangular warning back * bug fix * delete sample_config.py + small fixes * rename to pixel_sampler * cleaning * test * passing num_rays_per_batch properly * fixed transforms typo * renaming in nerfacto * remove depth files from test data * removed defaulting to base pixel sampler with mask * adding patch_size back --------- Co-authored-by: AdamRashid96 Co-authored-by: Matthew Tancik --- nerfstudio/configs/method_configs.py | 2 + .../data/datamanagers/base_datamanager.py | 23 +- nerfstudio/data/datasets/depth_dataset.py | 86 ++++++- nerfstudio/data/pixel_samplers.py | 231 +++++++++++++----- nerfstudio/model_components/losses.py | 14 ++ nerfstudio/models/depth_nerfacto.py | 13 +- nerfstudio/models/nerfacto.py | 10 +- pyproject.toml | 3 +- 8 files changed, 293 insertions(+), 89 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 72e4f20c8..a92c042bf 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -22,6 +22,7 @@ from collections import OrderedDict from typing import Dict import tyro +from nerfstudio.data.pixel_samplers import PairPixelSamplerConfig from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig @@ -209,6 +210,7 @@ method_configs["depth-nerfacto"] = TrainerConfig( pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( _target=VanillaDataManager[DepthDataset], + pixel_sampler=PairPixelSamplerConfig(), dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index b3f16af1c..e115949d5 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -55,9 +55,9 @@ from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.pixel_samplers import ( - EquirectangularPixelSampler, - PatchPixelSampler, PixelSampler, + PixelSamplerConfig, + PatchPixelSamplerConfig, ) from nerfstudio.data.utils.dataloaders import ( CacheDataloader, @@ -348,6 +348,8 @@ class VanillaDataManagerConfig(DataManagerConfig): """ patch_size: int = 1 """Size of patch to sample from. If >1, patch-based sampling will be used.""" + pixel_sampler: PixelSamplerConfig = PixelSamplerConfig() + """Specifies the pixel sampler used to sample pixels from images.""" TDataset = TypeVar("TDataset", bound=InputDataset, default=InputDataset) @@ -457,19 +459,18 @@ class VanillaDataManager(DataManager, Generic[TDataset]): scale_factor=self.config.camera_res_scale_factor, ) - def _get_pixel_sampler(self, dataset: TDataset, *args: Any, **kwargs: Any) -> PixelSampler: + def _get_pixel_sampler(self, dataset: TDataset, num_rays_per_batch: int) -> PixelSampler: """Infer pixel sampler to use.""" if self.config.patch_size > 1: - return PatchPixelSampler(*args, **kwargs, patch_size=self.config.patch_size) - - # If all images are equirectangular, use equirectangular pixel sampler - is_equirectangular = dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value - if is_equirectangular.all(): - return EquirectangularPixelSampler(*args, **kwargs) - # Otherwise, use the default pixel sampler + return PatchPixelSamplerConfig().setup( + patch_size=self.config.patch_size, num_rays_per_batch=num_rays_per_batch + ) + is_equirectangular = (dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value).all() if is_equirectangular.any(): CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.") - return PixelSampler(*args, **kwargs) + return self.config.pixel_sampler.setup( + is_equirectangular=is_equirectangular, num_rays_per_batch=num_rays_per_batch + ) def setup_train(self): """Sets up the data loaders for training""" diff --git a/nerfstudio/data/datasets/depth_dataset.py b/nerfstudio/data/datasets/depth_dataset.py index 1d64cdf1c..85d20ce7c 100644 --- a/nerfstudio/data/datasets/depth_dataset.py +++ b/nerfstudio/data/datasets/depth_dataset.py @@ -18,31 +18,95 @@ Depth dataset. from typing import Dict +import numpy as np + from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs +from nerfstudio.model_components import losses from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.utils.data_utils import get_depth_image_from_path +from nerfstudio.utils.rich_utils import CONSOLE + +from typing import Union +from PIL import Image +import torch +from rich.progress import track +from pathlib import Path +import json class DepthDataset(InputDataset): - """Dataset that returns images and depths. + """Dataset that returns images and depths. If no depths are found, then we generate them with Zoe Depth. Args: dataparser_outputs: description of where and how to read input images. scale_factor: The scaling factor for the dataparser outputs. """ - exclude_batch_keys_from_device = InputDataset.exclude_batch_keys_from_device + ["depth_image"] - def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0): super().__init__(dataparser_outputs, scale_factor) - assert ( - "depth_filenames" in dataparser_outputs.metadata.keys() - and dataparser_outputs.metadata["depth_filenames"] is not None - ) + # if there are no depth images than we want to generate them all with zoe depth + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + if len(dataparser_outputs.image_filenames) > 0 and ( + "depth_filenames" not in dataparser_outputs.metadata.keys() + or dataparser_outputs.metadata["depth_filenames"] is None + ): + losses.DEPTH_METRIC = 0 + CONSOLE.print("[bold yellow] No depth data found!") + cache = dataparser_outputs.image_filenames[0].parent / "depths.npy" + if cache.exists(): + CONSOLE.print("[bold yellow] Loading cache!") + # load all the depths + self.depths = np.load(cache) + self.depths = torch.from_numpy(self.depths).to(device) + else: + depth_tensors = [] + transforms = self._find_transform(dataparser_outputs.image_filenames[0]) + data = dataparser_outputs.image_filenames[0].parent + if transforms is not None: + meta = json.load(open(transforms, "r")) + frames = meta["frames"] + filenames = [data / frames[j]["file_path"].split("/")[-1] for j in range(len(frames))] + else: + meta = None + frames = None + filenames = dataparser_outputs.image_filenames + + repo = "isl-org/ZoeDepth" + self.zoe = torch.compile(torch.hub.load(repo, "ZoeD_NK", pretrained=True).to(device)) + + for i in track(range(len(filenames)), description="Generating depth images"): + image_filename = filenames[i] + pil_image = Image.open(image_filename) + image = np.array(pil_image, dtype="uint8") # shape is (h, w) or (h, w, 3 or 4) + if len(image.shape) == 2: + image = image[:, :, None].repeat(3, axis=2) + image = torch.from_numpy(image.astype("float32") / 255.0) + + with torch.no_grad(): + image = torch.permute(image, (2, 0, 1)).unsqueeze(0).to(device) + if image.shape[1] == 4: + image = image[:, :3, :, :] + depth_tensor = self.zoe.infer(image).squeeze().unsqueeze(-1) + + depth_tensors.append(depth_tensor) + + self.depths = torch.stack(depth_tensors) + np.save(cache, self.depths.cpu().numpy()) + dataparser_outputs.metadata["depth_filenames"] = None + dataparser_outputs.metadata["depth_unit_scale_factor"] = 1.0 + self.metadata["depth_filenames"] = None + self.metadata["depth_unit_scale_factor"] = 1.0 + + else: + losses.DEPTH_METRIC = 0 + self.depth_filenames = self.metadata["depth_filenames"] self.depth_unit_scale_factor = self.metadata["depth_unit_scale_factor"] def get_metadata(self, data: Dict) -> Dict: + if self.depth_filenames is None: + return {"depth_image": self.depths[data["image_idx"]]} + filepath = self.depth_filenames[data["image_idx"]] height = int(self._dataparser_outputs.cameras.height[data["image_idx"]]) width = int(self._dataparser_outputs.cameras.width[data["image_idx"]]) @@ -54,3 +118,11 @@ class DepthDataset(InputDataset): ) return {"depth_image": depth_image} + + def _find_transform(self, image_path: Path) -> Union[Path, None]: + while image_path.parent != image_path: + transform_path = image_path.parent / "transforms.json" + if transform_path.exists(): + return transform_path + image_path = image_path.parent + return None diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index a145dce46..bb5e9d5c5 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -17,25 +17,55 @@ Code for sampling pixels. """ import random -from typing import Dict, Optional, Union import torch from jaxtyping import Int from torch import Tensor +from dataclasses import dataclass, field +from typing import ( + Dict, + Optional, + Type, + Union, +) + +from nerfstudio.configs.base_config import ( + InstantiateConfig, +) + + +@dataclass +class PixelSamplerConfig(InstantiateConfig): + """Configuration for pixel sampler instantiation.""" + + _target: Type = field(default_factory=lambda: PixelSampler) + """Target class to instantiate.""" + num_rays_per_batch: int = 4096 + """Number of rays to sample per batch.""" + keep_full_image: bool = False + """Whether or not to include a reference to the full image in returned batch.""" + is_equirectangular: bool = False + """List of whether or not camera i is equirectangular.""" + class PixelSampler: """Samples 'pixel_batch's from 'image_batch's. Args: - num_rays_per_batch: number of rays to sample per batch - keep_full_image: whether or not to include a reference to the full image in returned batch + config: the DataManagerConfig used to instantiate class """ - def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False, **kwargs) -> None: + config: PixelSamplerConfig + + def __init__(self, config: PixelSamplerConfig, **kwargs) -> None: self.kwargs = kwargs - self.num_rays_per_batch = num_rays_per_batch - self.keep_full_image = keep_full_image + self.config = config + # Possibly override some values if they are present in the kwargs dictionary + self.config.num_rays_per_batch = self.kwargs.get("num_rays_per_batch", self.config.num_rays_per_batch) + self.config.keep_full_image = self.kwargs.get("keep_full_image", self.config.keep_full_image) + self.config.is_equirectangular = self.kwargs.get("is_equirectangular", self.config.is_equirectangular) + self.set_num_rays_per_batch(self.config.num_rays_per_batch) def set_num_rays_per_batch(self, num_rays_per_batch: int): """Set the number of rays to sample per batch. @@ -74,6 +104,36 @@ class PixelSampler: return indices + def sample_method_equirectangular( + self, + batch_size: int, + num_images: int, + image_height: int, + image_width: int, + mask: Optional[Tensor] = None, + device: Union[torch.device, str] = "cpu", + ) -> Int[Tensor, "batch_size 3"]: + if isinstance(mask, torch.Tensor): + # Note: if there is a mask, sampling reduces back to uniform sampling, which gives more + # sampling weight to the poles of the image than the equators. + # TODO(kevinddchen): implement the correct mask-sampling method. + + indices = self.sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) + else: + # We sample theta uniformly in [0, 2*pi] + # We sample phi in [0, pi] according to the PDF f(phi) = sin(phi) / 2. + # This is done by inverse transform sampling. + # http://corysimon.github.io/articles/uniformdistn-on-sphere/ + num_images_rand = torch.rand(batch_size, device=device) + phi_rand = torch.acos(1 - 2 * torch.rand(batch_size, device=device)) / torch.pi + theta_rand = torch.rand(batch_size, device=device) + indices = torch.floor( + torch.stack((num_images_rand, phi_rand, theta_rand), dim=-1) + * torch.tensor([num_images, image_height, image_width], device=device) + ).long() + + return indices + def collate_image_dataset_batch(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): """ Operates on a batch of images and samples pixels to use for generating rays. @@ -90,11 +150,21 @@ class PixelSampler: num_images, image_height, image_width, _ = batch["image"].shape if "mask" in batch: - indices = self.sample_method( - num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device - ) + if self.config.is_equirectangular: + indices = self.sample_method_equirectangular( + num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device + ) + else: + indices = self.sample_method( + num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device + ) else: - indices = self.sample_method(num_rays_per_batch, num_images, image_height, image_width, device=device) + if self.config.is_equirectangular: + indices = self.sample_method_equirectangular( + num_rays_per_batch, num_images, image_height, image_width, device=device + ) + else: + indices = self.sample_method(num_rays_per_batch, num_images, image_height, image_width, device=device) c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) c, y, x = c.cpu(), y.cpu(), x.cpu() @@ -156,7 +226,12 @@ class PixelSampler: image_height, image_width, _ = batch["image"][i].shape if i == num_images - 1: num_rays_in_batch = num_rays_per_batch - (num_images - 1) * num_rays_in_batch - indices = self.sample_method(num_rays_in_batch, 1, image_height, image_width, device=device) + if self.config.is_equirectangular: + indices = self.sample_method_equirectangular( + num_rays_in_batch, 1, image_height, image_width, device=device + ) + else: + indices = self.sample_method(num_rays_in_batch, 1, image_height, image_width, device=device) indices[:, 0] = i all_indices.append(indices) all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]]) @@ -192,56 +267,26 @@ class PixelSampler: if isinstance(image_batch["image"], list): image_batch = dict(image_batch.items()) # copy the dictionary so we don't modify the original pixel_batch = self.collate_image_dataset_batch_list( - image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image + image_batch, self.num_rays_per_batch, keep_full_image=self.config.keep_full_image ) elif isinstance(image_batch["image"], torch.Tensor): pixel_batch = self.collate_image_dataset_batch( - image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image + image_batch, self.num_rays_per_batch, keep_full_image=self.config.keep_full_image ) else: raise ValueError("image_batch['image'] must be a list or torch.Tensor") return pixel_batch -class EquirectangularPixelSampler(PixelSampler): - """Samples 'pixel_batch's from 'image_batch's. Assumes images are - equirectangular and the sampling is done uniformly on the sphere. +@dataclass +class PatchPixelSamplerConfig(PixelSamplerConfig): + """Config dataclass for PatchPixelSampler.""" - Args: - num_rays_per_batch: number of rays to sample per batch - keep_full_image: whether or not to include a reference to the full image in returned batch - """ - - # overrides base method - def sample_method( - self, - batch_size: int, - num_images: int, - image_height: int, - image_width: int, - mask: Optional[Tensor] = None, - device: Union[torch.device, str] = "cpu", - ) -> Int[Tensor, "batch_size 3"]: - if isinstance(mask, torch.Tensor): - # Note: if there is a mask, sampling reduces back to uniform sampling, which gives more - # sampling weight to the poles of the image than the equators. - # TODO(kevinddchen): implement the correct mask-sampling method. - - indices = super().sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) - else: - # We sample theta uniformly in [0, 2*pi] - # We sample phi in [0, pi] according to the PDF f(phi) = sin(phi) / 2. - # This is done by inverse transform sampling. - # http://corysimon.github.io/articles/uniformdistn-on-sphere/ - num_images_rand = torch.rand(batch_size, device=device) - phi_rand = torch.acos(1 - 2 * torch.rand(batch_size, device=device)) / torch.pi - theta_rand = torch.rand(batch_size, device=device) - indices = torch.floor( - torch.stack((num_images_rand, phi_rand, theta_rand), dim=-1) - * torch.tensor([num_images, image_height, image_width], device=device) - ).long() - - return indices + _target: Type = field(default_factory=lambda: PatchPixelSampler) + """Target class to instantiate.""" + patch_size: int = 2 + """Side length of patch. This must be consistent in the method + config in order for samples to be reshaped into patches correctly.""" class PatchPixelSampler(PixelSampler): @@ -249,16 +294,10 @@ class PatchPixelSampler(PixelSampler): from the images randomly. Useful for patch-based losses. Args: - num_rays_per_batch: number of rays to sample per batch - keep_full_image: whether or not to include a reference to the full image in returned batch - patch_size: side length of patch. This must be consistent in the method - config in order for samples to be reshaped into patches correctly. + config: the PatchPixelSamplerConfig used to instantiate class """ - def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False, **kwargs) -> None: - self.patch_size = kwargs["patch_size"] - num_rays = (num_rays_per_batch // (self.patch_size**2)) * (self.patch_size**2) - super().__init__(num_rays, keep_full_image, **kwargs) + config: PatchPixelSamplerConfig def set_num_rays_per_batch(self, num_rays_per_batch: int): """Set the number of rays to sample per batch. Overridden to deal with patch-based sampling. @@ -266,7 +305,7 @@ class PatchPixelSampler(PixelSampler): Args: num_rays_per_batch: number of rays to sample per batch """ - self.num_rays_per_batch = (num_rays_per_batch // (self.patch_size**2)) * (self.patch_size**2) + self.num_rays_per_batch = (num_rays_per_batch // (self.config.patch_size**2)) * (self.config.patch_size**2) # overrides base method def sample_method( @@ -279,19 +318,23 @@ class PatchPixelSampler(PixelSampler): device: Union[torch.device, str] = "cpu", ) -> Int[Tensor, "batch_size 3"]: if isinstance(mask, Tensor): - # Note: if there is a mask, sampling reduces back to uniform sampling - indices = super().sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) + # Note: if there is a mask, should switch to the base PixelSampler class + raise NotImplementedError() else: - sub_bs = batch_size // (self.patch_size**2) + sub_bs = batch_size // (self.config.patch_size**2) indices = torch.rand((sub_bs, 3), device=device) * torch.tensor( - [num_images, image_height - self.patch_size, image_width - self.patch_size], + [num_images, image_height - self.config.patch_size, image_width - self.config.patch_size], device=device, ) - indices = indices.view(sub_bs, 1, 1, 3).broadcast_to(sub_bs, self.patch_size, self.patch_size, 3).clone() + indices = ( + indices.view(sub_bs, 1, 1, 3) + .broadcast_to(sub_bs, self.config.patch_size, self.config.patch_size, 3) + .clone() + ) yys, xxs = torch.meshgrid( - torch.arange(self.patch_size, device=device), torch.arange(self.patch_size, device=device) + torch.arange(self.config.patch_size, device=device), torch.arange(self.config.patch_size, device=device) ) indices[:, ..., 1] += yys indices[:, ..., 2] += xxs @@ -300,3 +343,59 @@ class PatchPixelSampler(PixelSampler): indices = indices.flatten(0, 2) return indices + + +@dataclass +class PairPixelSamplerConfig(PixelSamplerConfig): + """Config dataclass for PatchPixelSampler.""" + + _target: Type = field(default_factory=lambda: PairPixelSampler) + """Target class to instantiate.""" + radius: int = 2 + """max distance between pairs of pixels.""" + + +class PairPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods + """Samples pair of pixels from 'image_batch's. Samples pairs of pixels from + from the images randomly within a 'radius' distance apart. Useful for pair-based losses. + + Args: + config: the PairPixelSamplerConfig used to instantiate class + """ + + def __init__(self, config: PairPixelSamplerConfig, **kwargs) -> None: + self.config = config + self.radius = self.config.radius + self.rays_to_sample = self.config.num_rays_per_batch // 2 + super().__init__(self.config, **kwargs) + + # overrides base method + def sample_method( # pylint: disable=no-self-use + self, + batch_size: int, + num_images: int, + image_height: int, + image_width: int, + mask: Optional[Tensor] = None, + device: Union[torch.device, str] = "cpu", + ) -> Int[Tensor, "batch_size 3"]: + if mask: + # Note: if there is a mask, should switch to the base PixelSampler class + + raise NotImplementedError() + else: + s = (self.rays_to_sample, 1) + ns = torch.randint(0, num_images, s, dtype=torch.long, device=device) + hs = torch.randint(self.radius, image_height - self.radius, s, dtype=torch.long, device=device) + ws = torch.randint(self.radius, image_width - self.radius, s, dtype=torch.long, device=device) + indices = torch.concat((ns, hs, ws), dim=1) + + pair_indices = torch.hstack( + ( + torch.zeros(self.rays_to_sample, 1, device=device, dtype=torch.long), + torch.randint(-self.radius, self.radius, (self.rays_to_sample, 2), device=device, dtype=torch.long), + ) + ) + pair_indices += indices + indices = torch.hstack((indices, pair_indices)).view(self.rays_to_sample * 2, 3) + return indices diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index 44fefafd1..8c53c24f1 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -35,6 +35,7 @@ EPS = 1.0e-7 # Sigma scale factor from Urban Radiance Fields (Rematas et al., 2022) URF_SIGMA_SCALE_FACTOR = 3.0 +DEPTH_METRIC = 1 class DepthLossType(Enum): @@ -561,3 +562,16 @@ def scale_gradients_by_distance_squared( for key, value in field_outputs.items(): out[key], _ = cast(Tuple[Tensor, Tensor], _GradientScaler.apply(value, scaling)) return out + + +def depth_ranking_loss(rendered_depth, gt_depth): + """ + Depth ranking loss as described in the SparseNeRF paper + Assumes that the layout of the batch comes from a PairPixelSampler, so that adjacent samples in the gt_depth + and rendered_depth are from pixels with a radius of each other + """ + m = 1e-4 + dpt_diff = gt_depth[::2, :] - gt_depth[1::2, :] + out_diff = rendered_depth[::2, :] - rendered_depth[1::2, :] + m + differing_signs = torch.sign(dpt_diff) != torch.sign(out_diff) + return torch.nanmean((out_diff[differing_signs] * torch.sign(out_diff[differing_signs]))) diff --git a/nerfstudio/models/depth_nerfacto.py b/nerfstudio/models/depth_nerfacto.py index e3328fd6d..f3f791e80 100644 --- a/nerfstudio/models/depth_nerfacto.py +++ b/nerfstudio/models/depth_nerfacto.py @@ -22,9 +22,11 @@ from dataclasses import dataclass, field from typing import Dict, Tuple, Type import torch +import numpy as np from nerfstudio.cameras.rays import RayBundle -from nerfstudio.model_components.losses import DepthLossType, depth_loss +from nerfstudio.model_components import losses +from nerfstudio.model_components.losses import DepthLossType, depth_loss, depth_ranking_loss from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig from nerfstudio.utils import colormaps @@ -98,7 +100,14 @@ class DepthNerfactoModel(NerfactoModel): loss_dict = super().get_loss_dict(outputs, batch, metrics_dict) if self.training: assert metrics_dict is not None and "depth_loss" in metrics_dict - loss_dict["depth_loss"] = self.config.depth_loss_mult * metrics_dict["depth_loss"] + # If real depth used + if losses.DEPTH_METRIC: + loss_dict["depth_loss"] = self.config.depth_loss_mult * metrics_dict["depth_loss"] + else: + loss_dict["depth_ranking"] = np.interp(self.step, [0, 2000], [0, 0.2]) * depth_ranking_loss( + outputs["expected_depth"], batch["depth_image"] + ) + # Else use depth ranking loss return loss_dict diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 065fb571e..1122d5c20 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -220,7 +220,8 @@ class NerfactoModel(Model): # renderers self.renderer_rgb = RGBRenderer(background_color=self.config.background_color) self.renderer_accumulation = AccumulationRenderer() - self.renderer_depth = DepthRenderer() + self.renderer_depth = DepthRenderer(method="median") + self.renderer_expected_depth = DepthRenderer(method="expected") self.renderer_normals = NormalsRenderer() # shaders @@ -233,6 +234,7 @@ class NerfactoModel(Model): self.psnr = PeakSignalNoiseRatio(data_range=1.0) self.ssim = structural_similarity_index_measure self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) + self.step = 0 def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} @@ -251,6 +253,7 @@ class NerfactoModel(Model): def set_anneal(step): # https://arxiv.org/pdf/2111.12077.pdf eq. 18 train_frac = np.clip(step / N, 0, 1) + self.step = step def bias(x, b): return b * x / ((b - 1) * x + 1) @@ -286,13 +289,16 @@ class NerfactoModel(Model): ray_samples_list.append(ray_samples) rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights) - depth = self.renderer_depth(weights=weights, ray_samples=ray_samples) + with torch.no_grad(): + depth = self.renderer_depth(weights=weights, ray_samples=ray_samples) + expected_depth = self.renderer_expected_depth(weights=weights, ray_samples=ray_samples) accumulation = self.renderer_accumulation(weights=weights) outputs = { "rgb": rgb, "accumulation": accumulation, "depth": depth, + "expected_depth": expected_depth, } if self.config.predict_normals: diff --git a/pyproject.toml b/pyproject.toml index 804dab015..61e04bc1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,7 +56,8 @@ dependencies = [ "nuscenes-devkit>=1.1.1", "wandb>=0.13.3", "xatlas", - "trimesh>=3.20.2" + "trimesh>=3.20.2", + "timm==0.6.7" ] [project.urls] From 5b4abc41e61805afcf849efa4289a27c72f29d01 Mon Sep 17 00:00:00 2001 From: machenmusik Date: Wed, 2 Aug 2023 02:35:27 -0400 Subject: [PATCH 3/4] Allow same dimensions optimizations to be disabled (#2292) Use ns-process-data images --no-same-dimensions ... to handle images sets with different dimensions properly. (Fast downscaling assumes image dimensions are the same.) --- .../colmap_converter_to_nerfstudio_dataset.py | 2 + .../images_to_nerfstudio_dataset.py | 1 + nerfstudio/process_data/process_data_utils.py | 63 ++++++++++++------- 3 files changed, 42 insertions(+), 24 deletions(-) diff --git a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py index fb2fb88b6..13ed7da66 100644 --- a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py +++ b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py @@ -96,6 +96,8 @@ class ColmapConverterToNerfstudioDataset(BaseConverterToNerfstudioDataset): """If True, export and use depth maps induced from SfM points.""" include_depth_debug: bool = False """If --use-sfm-depth and this flag is True, also export debug images showing Sf overlaid upon input images.""" + same_dimensions: bool = True + """Whether to assume all images are same dimensions and so to use fast downscaling with no autorotation.""" @staticmethod def default_colmap_path() -> Path: diff --git a/nerfstudio/process_data/images_to_nerfstudio_dataset.py b/nerfstudio/process_data/images_to_nerfstudio_dataset.py index dc613412a..3f47b70ac 100644 --- a/nerfstudio/process_data/images_to_nerfstudio_dataset.py +++ b/nerfstudio/process_data/images_to_nerfstudio_dataset.py @@ -65,6 +65,7 @@ class ImagesToNerfstudioDataset(ColmapConverterToNerfstudioDataset): crop_factor=self.crop_factor, verbose=self.verbose, num_downscales=self.num_downscales, + same_dimensions=self.same_dimensions, ) image_rename_map = dict((a.name, b.name) for a, b in image_rename_map_paths.items()) num_frames = len(image_rename_map) diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index 355e5bc1c..d25ed0dfa 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -215,6 +215,7 @@ def copy_images_list( verbose: bool = False, upscale_factor: Optional[int] = None, nearest_neighbor: bool = False, + same_dimensions: bool = True, ) -> List[Path]: """Copy all images in a list of Paths. Useful for filtering from a directory. Args: @@ -250,8 +251,15 @@ def copy_images_list( rgb = raw.postprocess() imageio.imsave(copied_image_path, rgb) image_paths[idx] = copied_image_path - else: + elif same_dimensions: + # Fast path; just copy the file shutil.copy(image_path, copied_image_path) + else: + # Slow path; let ffmpeg perform autorotation (and clear metadata) + ffmpeg_cmd = f"ffmpeg -y -i {image_path} -metadata:s:v:0 rotate=0 {copied_image_path}" + if verbose: + CONSOLE.log(f"... {ffmpeg_cmd}") + run_command(ffmpeg_cmd, verbose=verbose) except shutil.SameFileError: pass copied_image_paths.append(copied_image_path) @@ -259,9 +267,6 @@ def copy_images_list( nn_flag = "" if not nearest_neighbor else ":flags=neighbor" downscale_chains = [f"[t{i}]scale=iw/{2**i}:ih/{2**i}{nn_flag}[out{i}]" for i in range(num_downscales + 1)] downscale_dirs = [Path(str(image_dir) + (f"_{2**i}" if i > 0 else "")) for i in range(num_downscales + 1)] - downscale_paths = [ - downscale_dirs[i] / ("frame_%05d" + copied_image_paths[0].suffix) for i in range(num_downscales + 1) - ] for dir in downscale_dirs: dir.mkdir(parents=True, exist_ok=True) @@ -274,30 +279,38 @@ def copy_images_list( ) num_frames = len(image_paths) - ffmpeg_cmd = f'ffmpeg -y -noautorotate -i "{image_dir / f"frame_%05d{copied_image_paths[0].suffix}"}" -q:v 2 ' + # ffmpeg batch commands assume all images are the same dimensions. + # When this is not the case (e.g. mixed portrait and landscape images), we need to do individually. + # (Unfortunately, that is much slower.) + for framenum in range(1, (1 if same_dimensions else num_frames) + 1): + framename = "frame_%05d" if same_dimensions else f"frame_{framenum:05d}" + ffmpeg_cmd = f'ffmpeg -y -noautorotate -i "{image_dir / f"{framename}{copied_image_paths[0].suffix}"}" -q:v 2 ' - crop_cmd = "" - if crop_border_pixels is not None: - crop_cmd = f"crop=iw-{crop_border_pixels*2}:ih-{crop_border_pixels*2}[cropped];[cropped]" - elif crop_factor != (0.0, 0.0, 0.0, 0.0): - height = 1 - crop_factor[0] - crop_factor[1] - width = 1 - crop_factor[2] - crop_factor[3] - start_x = crop_factor[2] - start_y = crop_factor[0] - crop_cmd = f"crop=w=iw*{width}:h=ih*{height}:x=iw*{start_x}:y=ih*{start_y}[cropped];[cropped]" + crop_cmd = "" + if crop_border_pixels is not None: + crop_cmd = f"crop=iw-{crop_border_pixels*2}:ih-{crop_border_pixels*2}[cropped];[cropped]" + elif crop_factor != (0.0, 0.0, 0.0, 0.0): + height = 1 - crop_factor[0] - crop_factor[1] + width = 1 - crop_factor[2] - crop_factor[3] + start_x = crop_factor[2] + start_y = crop_factor[0] + crop_cmd = f"crop=w=iw*{width}:h=ih*{height}:x=iw*{start_x}:y=ih*{start_y}[cropped];[cropped]" - select_cmd = "[0:v]" - if upscale_factor is not None: - select_cmd = f"[0:v]scale=iw*{upscale_factor}:ih*{upscale_factor}:flags=neighbor[upscaled];[upscaled]" + select_cmd = "[0:v]" + if upscale_factor is not None: + select_cmd = f"[0:v]scale=iw*{upscale_factor}:ih*{upscale_factor}:flags=neighbor[upscaled];[upscaled]" - downscale_cmd = f' -filter_complex "{select_cmd}{crop_cmd}{downscale_chain}"' + "".join( - [f' -map "[out{i}]" "{downscale_paths[i]}"' for i in range(num_downscales + 1)] - ) + downscale_cmd = f' -filter_complex "{select_cmd}{crop_cmd}{downscale_chain}"' + "".join( + [ + f' -map "[out{i}]" "{downscale_dirs[i] / f"{framename}{copied_image_paths[0].suffix}"}"' + for i in range(num_downscales + 1) + ] + ) - ffmpeg_cmd += downscale_cmd - if verbose: - CONSOLE.log(f"... {ffmpeg_cmd}") - run_command(ffmpeg_cmd, verbose=verbose) + ffmpeg_cmd += downscale_cmd + if verbose: + CONSOLE.log(f"... {ffmpeg_cmd}") + run_command(ffmpeg_cmd, verbose=verbose) if num_frames == 0: CONSOLE.log("[bold red]:skull: No usable images in the data folder.") @@ -353,6 +366,7 @@ def copy_images( verbose: bool = False, crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), num_downscales: int = 0, + same_dimensions: bool = True, ) -> OrderedDict[Path, Path]: """Copy images from a directory to a new directory. @@ -377,6 +391,7 @@ def copy_images( crop_factor=crop_factor, verbose=verbose, num_downscales=num_downscales, + same_dimensions=same_dimensions, ) return OrderedDict((original_path, new_path) for original_path, new_path in zip(image_paths, copied_images)) From 388fcb09e11c4bd191b4bd6c319e7e388ccef321 Mon Sep 17 00:00:00 2001 From: machenmusik Date: Wed, 2 Aug 2023 12:44:35 -0400 Subject: [PATCH 4/4] convert raw to JPG not TIFF, to avoid fast-downscale ffmpeg error (#2295) --- nerfstudio/process_data/process_data_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index d25ed0dfa..6a10153e8 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -36,7 +36,7 @@ POLYCAM_UPSCALING_TIMES = 2 """Lowercase suffixes to treat as raw image.""" ALLOWED_RAW_EXTS = [".cr2"] """Suffix to use for converted images from raw.""" -RAW_CONVERTED_SUFFIX = ".tiff" +RAW_CONVERTED_SUFFIX = ".jpg" class CameraModel(Enum): @@ -244,7 +244,7 @@ def copy_images_list( CONSOLE.log(f"Copying image {idx + 1} of {len(image_paths)}...") copied_image_path = image_dir / f"frame_{idx + 1:05d}{image_path.suffix}" try: - # if CR2 raw, we want to read raw and write TIFF, and change the file suffix for downstream processing + # if CR2 raw, we want to read raw and write RAW_CONVERTED_SUFFIX, and change the file suffix for downstream processing if image_path.suffix.lower() in ALLOWED_RAW_EXTS: copied_image_path = image_dir / f"frame_{idx + 1:05d}{RAW_CONVERTED_SUFFIX}" with rawpy.imread(str(image_path)) as raw: