diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index b17e4fdeb..b754ae206 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -18,6 +18,7 @@ Datamanager. from __future__ import annotations +import random from abc import abstractmethod from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple, Type, Union @@ -51,10 +52,9 @@ from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PixelSampler from nerfstudio.data.utils.dataloaders import CacheDataloader from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate -from nerfstudio.utils.misc import get_dict_to_torch from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfstudio.model_components.ray_generators import RayGenerator -from nerfstudio.utils.misc import IterableWrapper +from nerfstudio.utils.misc import IterableWrapper, get_dict_to_torch CONSOLE = Console(width=120) @@ -424,10 +424,10 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method return camera_ray_bundle, batch def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]: - for idx in range(len(self.eval_dataset)): - camera_ray_bundle, batch = self.get_eval_image(idx) - return idx, camera_ray_bundle, batch - raise ValueError("No more eval images.") + # sample a random index from the eval dataset + idx = random.randint(0, len(self.eval_dataset) - 1) + camera_ray_bundle, batch = self.get_eval_image(idx) + return idx, camera_ray_bundle, batch def get_param_groups(self) -> Dict[str, List[Parameter]]: # pylint: disable=no-self-use """Get the param groups for the data manager. diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 2ff1f1c96..af0d78d96 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -37,12 +37,14 @@ def collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, sample_all num_images, image_height, image_width, _ = batch["image"].shape if sample_all_pixels: + image_indices = (torch.ones(num_images, device=device) * batch["image_idx"].to(device)).to(torch.int64) collated_batch = {key: value for key, value in batch.items() if key != "image_idx" and value is not None} collated_batch["indices"] = torch.stack( torch.meshgrid( - torch.arange(num_images, device=device), + image_indices, torch.arange(image_height, device=device), torch.arange(image_width, device=device), + indexing="ij", ), dim=-1, ) @@ -59,7 +61,9 @@ def collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, sample_all ).long() c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) - collated_batch = {key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None} + collated_batch = { + key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None + } assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape @@ -170,9 +174,13 @@ class PixelSampler: # pylint: disable=too-few-public-methods """ if isinstance(image_batch["image"], list): image_batch = dict(image_batch.items()) # copy the dictioary so we don't modify the original - pixel_batch = collate_image_dataset_batch_list(image_batch, self.num_rays_per_batch, sample_all_pixels=sample_all_pixels) + pixel_batch = collate_image_dataset_batch_list( + image_batch, self.num_rays_per_batch, sample_all_pixels=sample_all_pixels + ) elif isinstance(image_batch["image"], torch.Tensor): - pixel_batch = collate_image_dataset_batch(image_batch, self.num_rays_per_batch, sample_all_pixels=sample_all_pixels) + pixel_batch = collate_image_dataset_batch( + image_batch, self.num_rays_per_batch, sample_all_pixels=sample_all_pixels + ) else: raise ValueError("image_batch['image'] must be a list or torch.Tensor") return pixel_batch diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index 765f93793..cbb3695a4 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -32,8 +32,8 @@ from nerfstudio.cameras.cameras import Cameras from nerfstudio.cameras.rays import RayBundle from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate -from nerfstudio.utils.misc import get_dict_to_torch from nerfstudio.model_components.ray_generators import RayGenerator +from nerfstudio.utils.misc import get_dict_to_torch CONSOLE = Console(width=120) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 904aa7577..244b0f411 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -96,7 +96,12 @@ class Trainer: self._check_viewer_warnings() # set up writers/profilers if enabled writer_log_path = self.base_dir / config.logging.relative_log_dir - writer.setup_event_writer(config.is_wandb_enabled(), config.is_tensorboard_enabled(), log_dir=writer_log_path, wandb_name=config.wandb_name) + writer.setup_event_writer( + config.is_wandb_enabled(), + config.is_tensorboard_enabled(), + log_dir=writer_log_path, + wandb_name=config.wandb_name, + ) writer.setup_local_writer( config.logging, max_iter=config.trainer.max_num_iterations, banner_messages=banner_messages ) diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 1a8c425d6..1213503eb 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -310,10 +310,7 @@ class TCNNNerfactoField(Field): x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions) outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x) - h_list = [ - d, - density_embedding.view(-1, self.geo_feat_dim) - ] + h_list = [d, density_embedding.view(-1, self.geo_feat_dim)] if self.use_appearance_embedding: h_list += [embedded_appearance.view(-1, self.appearance_embedding_dim)] h = torch.cat( diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index cff57963e..a67c1ec7d 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -218,7 +218,7 @@ class NGPModel(Model): def get_loss_dict(self, outputs, batch, metrics_dict=None): image = batch["image"].to(self.device) - mask = outputs["alive_ray_mask"] #.squeeze() + mask = outputs["alive_ray_mask"] rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask]) loss_dict = {"rgb_loss": rgb_loss} return loss_dict diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 0bb77a29b..5e02f57e2 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -17,13 +17,13 @@ Abstracts for the Pipeline class. """ from __future__ import annotations +import random import typing from abc import abstractmethod from dataclasses import dataclass, field from time import time from typing import Any, Dict, List, Optional, Type, Union, cast -import random import torch import torch.distributed as dist from rich.progress import ( diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index a057a335c..ad61cb9fd 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -188,7 +188,9 @@ def setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages @check_main_thread -def setup_event_writer(is_wandb_enabled: bool, is_tensorboard_enabled: bool, log_dir: Path, wandb_name: Optional[str] = None) -> None: +def setup_event_writer( + is_wandb_enabled: bool, is_tensorboard_enabled: bool, log_dir: Path, wandb_name: Optional[str] = None +) -> None: """Initialization of all event writers specified in config Args: diff --git a/scripts/datasets/process_mipnerf360.py b/scripts/datasets/process_mipnerf360.py index 9278cd36d..5dea78706 100644 --- a/scripts/datasets/process_mipnerf360.py +++ b/scripts/datasets/process_mipnerf360.py @@ -79,4 +79,4 @@ with Pool() as p: # the default method # DATASET="person" METHOD="nerfacto-default" && ns-train nerfacto --vis wandb --data "data/nerfstudio/${DATASET}" --wandb-name "${METHOD}-${DATASET}" --experiment-name "${METHOD}-${DATASET}" # without pose optimization -# DATASET="person" METHOD="nerfacto-no-pose" && ns-train nerfacto --vis wandb --data "data/nerfstudio/${DATASET}" --wandb-name "${METHOD}-${DATASET}" --experiment-name "${METHOD}-${DATASET}" --pipeline.datamanager.camera-optimizer.mode off \ No newline at end of file +# DATASET="person" METHOD="nerfacto-no-pose" && ns-train nerfacto --vis wandb --data "data/nerfstudio/${DATASET}" --wandb-name "${METHOD}-${DATASET}" --experiment-name "${METHOD}-${DATASET}" --pipeline.datamanager.camera-optimizer.mode off diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index 3926caa00..12a0a0227 100755 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -283,6 +283,7 @@ class PhototourismDownload(DatasetDownload): shutil.rmtree(tmp_path) os.remove(download_path) + Commands = Union[ Annotated[BlenderDownload, tyro.conf.subcommand(name="blender")], Annotated[FriendsDownload, tyro.conf.subcommand(name="friends")],