fix wandb and image indexing

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.

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

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)
 

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
         )

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(

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

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 (

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:

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
+# 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

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")],