diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 7adf8cd05..93d17ed70 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -49,12 +49,9 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import ( from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PixelSampler -from nerfstudio.data.utils.dataloaders import ( - CacheDataloader, - FixedIndicesEvalDataloader, - RandIndicesEvalDataloader, -) +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 @@ -384,22 +381,11 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ) self.iter_eval_image_dataloader = iter(self.eval_image_dataloader) self.eval_pixel_sampler = self._get_pixel_sampler(self.eval_dataset, self.config.eval_num_rays_per_batch) - self.eval_ray_generator = RayGenerator( - self.eval_dataset.cameras.to(self.device), - self.train_camera_optimizer, # should be shared between train and eval. - ) - # for loading full images - self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader( - input_dataset=self.eval_dataset, - device=self.device, - num_workers=self.world_size * 4, - ) - self.eval_dataloader = RandIndicesEvalDataloader( - input_dataset=self.eval_dataset, - image_indices=self.config.eval_image_indices, - device=self.device, - num_workers=self.world_size * 4, + self.eval_camera_optimizer = self.config.camera_optimizer.setup( + num_cameras=self.eval_dataset.cameras.size, device=self.device ) + self.eval_ray_generator = RayGenerator(self.eval_dataset.cameras.to(self.device), self.eval_camera_optimizer) + # TODO(ethan): remove use of self.eval_dataloader from the spiral trajectory code def next_train(self, step: int) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the train dataloader.""" @@ -419,12 +405,28 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ray_bundle = self.eval_ray_generator(ray_indices) return ray_bundle, batch + def get_eval_image(self, idx: int) -> Tuple[RayBundle, Dict]: + """Returns the full image at the given index.""" + image_batch = self.eval_dataset[idx] + image_batch = self.config.collate_fn([image_batch]) + image_batch = get_dict_to_torch(image_batch, device=self.device, exclude=["image"]) + # image_batch["image"].shape == (1, H, W, 3) + # this will sample the entire image + batch = self.eval_pixel_sampler.sample(image_batch, sample_all_pixels=True) + ray_indices = batch["indices"] + ray_indices_shape = ray_indices.shape # (1, H, W, 3) + ray_bundle = self.eval_ray_generator(ray_indices.view(-1, 3)) + camera_ray_bundle = ray_bundle.reshape(tuple(ray_indices_shape[1:-1])) # (H, W) + # remove batch dimension from batch + for key in batch: + batch[key] = batch[key][0] + return camera_ray_bundle, batch + def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]: - for camera_ray_bundle, batch in self.eval_dataloader: - assert camera_ray_bundle.camera_indices is not None - image_idx = int(camera_ray_bundle.camera_indices[0, 0, 0]) - return image_idx, camera_ray_bundle, batch - raise ValueError("No more eval images") + 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.") def get_param_groups(self) -> Dict[str, List[Parameter]]: # pylint: disable=no-self-use """Get the param groups for the data manager. @@ -433,11 +435,15 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method """ param_groups = {} - camera_opt_params = list(self.train_camera_optimizer.parameters()) + camera_opt_params_train = list(self.train_camera_optimizer.parameters()) + camera_opt_params_eval = list(self.eval_camera_optimizer.parameters()) if self.config.camera_optimizer.mode != "off": - assert len(camera_opt_params) > 0 - param_groups[self.config.camera_optimizer.param_group] = camera_opt_params + assert len(camera_opt_params_train) > 0 + assert len(camera_opt_params_eval) > 0 + param_groups[self.config.camera_optimizer.param_group + "_train"] = camera_opt_params_train + param_groups[self.config.camera_optimizer.param_group + "_eval"] = camera_opt_params_eval else: - assert len(camera_opt_params) == 0 + assert len(camera_opt_params_train) == 0 + assert len(camera_opt_params_eval) == 0 return param_groups diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 1066155c7..761485572 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -22,7 +22,7 @@ from typing import Dict import torch -def collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): +def collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, sample_all_pixels: bool = False): """ Operates on a batch of images and samples pixels to use for generating rays. Returns a collated batch which is input to the Graph. @@ -31,38 +31,46 @@ def collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, keep_full_ Args: batch: batch of images to sample from 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 + sample_all_pixels: if True, will sample all pixels in the image """ device = batch["image"].device num_images, image_height, image_width, _ = batch["image"].shape - # only sample within the mask, if the mask is in the batch - if "mask" in batch: - nonzero_indices = torch.nonzero(batch["mask"][..., 0].to(device), as_tuple=False) - chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) - indices = nonzero_indices[chosen_indices] + if sample_all_pixels: + 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), + torch.arange(image_height, device=device), + torch.arange(image_width, device=device), + ), + dim=-1, + ) else: - indices = torch.floor( - torch.rand((num_rays_per_batch, 3), device=device) - * torch.tensor([num_images, image_height, image_width], device=device) - ).long() + # only sample within the mask, if the mask is in the batch + if "mask" in batch: + nonzero_indices = torch.nonzero(batch["mask"][..., 0].to(device), as_tuple=False) + chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) + indices = nonzero_indices[chosen_indices] + else: + indices = torch.floor( + torch.rand((num_rays_per_batch, 3), device=device) + * torch.tensor([num_images, image_height, image_width], device=device) + ).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} + 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} - assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape + assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape - # Needed to correct the random indices to their actual camera idx locations. - indices[:, 0] = batch["image_idx"][c] - collated_batch["indices"] = indices # with the abs camera indices - - if keep_full_image: - collated_batch["full_image"] = batch["image"] + # Needed to correct the random indices to their actual camera idx locations. + indices[:, 0] = batch["image_idx"][c] + collated_batch["indices"] = indices # with the abs camera indices return collated_batch -def collate_image_dataset_batch_list(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): +def collate_image_dataset_batch_list(batch: Dict, num_rays_per_batch: int, sample_all_pixels: bool = False): """ Does the same as collate_image_dataset_batch, except it will operate over a list of images / masks inside a list. @@ -74,7 +82,7 @@ def collate_image_dataset_batch_list(batch: Dict, num_rays_per_batch: int, keep_ Args: batch: batch of images to sample from 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 + sample_all_pixels: if True, will sample all the pixels in the image """ device = batch["image"][0].device @@ -84,31 +92,36 @@ def collate_image_dataset_batch_list(batch: Dict, num_rays_per_batch: int, keep_ all_indices = [] all_images = [] - if "mask" in batch: - num_rays_in_batch = num_rays_per_batch // num_images - for i in range(num_images): - if i == num_images - 1: - num_rays_in_batch = num_rays_per_batch - (num_images - 1) * num_rays_in_batch - nonzero_indices = torch.nonzero(batch["mask"][i][..., 0], as_tuple=False) - chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_in_batch) - indices = nonzero_indices[chosen_indices] - indices = torch.cat([torch.full((num_rays_in_batch, 1), i, device=device), indices], dim=-1) - all_indices.append(indices) - all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]]) - + if sample_all_pixels: + # sample all the pixels + raise NotImplementedError("sample_all_pixels is not implemented for collate_image_dataset_batch_list") else: - num_rays_in_batch = num_rays_per_batch // num_images - for i in range(num_images): - 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 = torch.floor( - torch.rand((num_rays_in_batch, 3), device=device) - * torch.tensor([1, image_height, image_width], device=device) - ).long() - indices[:, 0] = i - all_indices.append(indices) - all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]]) + # sample a subset of the pixels + if "mask" in batch: + num_rays_in_batch = num_rays_per_batch // num_images + for i in range(num_images): + if i == num_images - 1: + num_rays_in_batch = num_rays_per_batch - (num_images - 1) * num_rays_in_batch + nonzero_indices = torch.nonzero(batch["mask"][i][..., 0], as_tuple=False) + chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_in_batch) + indices = nonzero_indices[chosen_indices] + indices = torch.cat([torch.full((num_rays_in_batch, 1), i, device=device), indices], dim=-1) + all_indices.append(indices) + all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]]) + + else: + num_rays_in_batch = num_rays_per_batch // num_images + for i in range(num_images): + 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 = torch.floor( + torch.rand((num_rays_in_batch, 3), device=device) + * torch.tensor([1, image_height, image_width], device=device) + ).long() + indices[:, 0] = i + all_indices.append(indices) + all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]]) indices = torch.cat(all_indices, dim=0) @@ -127,9 +140,6 @@ def collate_image_dataset_batch_list(batch: Dict, num_rays_per_batch: int, keep_ indices[:, 0] = batch["image_idx"][c] collated_batch["indices"] = indices # with the abs camera indices - if keep_full_image: - collated_batch["full_image"] = batch["image"] - return collated_batch @@ -138,12 +148,10 @@ class PixelSampler: # pylint: disable=too-few-public-methods 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 """ - def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False) -> None: + def __init__(self, num_rays_per_batch: int) -> None: self.num_rays_per_batch = num_rays_per_batch - self.keep_full_image = keep_full_image def set_num_rays_per_batch(self, num_rays_per_batch: int): """Set the number of rays to sample per batch. @@ -153,27 +161,24 @@ class PixelSampler: # pylint: disable=too-few-public-methods """ self.num_rays_per_batch = num_rays_per_batch - def sample(self, image_batch: Dict): + def sample(self, image_batch: Dict, sample_all_pixels: bool = False) -> Dict: """Sample an image batch and return a pixel batch. Args: image_batch: batch of images to sample from + sample_all_pixels: if True, will sample all pixels in the image """ 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, keep_full_image=self.keep_full_image - ) + 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, keep_full_image=self.keep_full_image - ) + 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 -def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): +def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch: int): """ Operates on a batch of equirectangular images and samples pixels to use for generating rays. Rays will be generated uniformly on the sphere. @@ -183,7 +188,6 @@ def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch: Args: batch: batch of images to sample from 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 """ # TODO(kevinddchen): make more DRY device = batch["image"].device @@ -215,9 +219,6 @@ def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch: indices[:, 0] = batch["image_idx"][c] collated_batch["indices"] = indices # with the abs camera indices - if keep_full_image: - collated_batch["full_image"] = batch["image"] - return collated_batch @@ -227,13 +228,10 @@ class EquirectangularPixelSampler(PixelSampler): # pylint: disable=too-few-publ 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(self, image_batch: Dict): - pixel_batch = collate_image_dataset_batch_equirectangular( - image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image - ) + pixel_batch = collate_image_dataset_batch_equirectangular(image_batch, self.num_rays_per_batch) return pixel_batch diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index 23d712d7e..765f93793 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -33,6 +33,7 @@ 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 CONSOLE = Console(width=120) @@ -136,119 +137,3 @@ class CacheDataloader(DataLoader): collated_batch = self.cached_collated_batch self.num_repeated += 1 yield collated_batch - - -class EvalDataloader(DataLoader): - """Evaluation dataloader base class - - Args: - input_dataset: InputDataset to load data from - device: Device to load data to - """ - - def __init__( - self, - input_dataset: InputDataset, - device: Union[torch.device, str] = "cpu", - **kwargs, - ): - self.input_dataset = input_dataset - self.cameras = input_dataset.cameras.to(device) - self.device = device - self.kwargs = kwargs - super().__init__(dataset=input_dataset) - - @abstractmethod - def __iter__(self): - """Iterates over the dataset""" - return self - - @abstractmethod - def __next__(self) -> Tuple[RayBundle, Dict]: - """Returns the next batch of data""" - - def get_camera(self, image_idx: int = 0) -> Cameras: - """Get camera for the given image index - - Args: - image_idx: Camera image index - """ - return self.cameras[image_idx] - - def get_data_from_image_idx(self, image_idx: int) -> Tuple[RayBundle, Dict]: - """Returns the data for a specific image index. - - Args: - image_idx: Camera image index - """ - ray_bundle = self.cameras.generate_rays(camera_indices=image_idx, keep_shape=True) - batch = self.input_dataset[image_idx] - batch = get_dict_to_torch(batch, device=self.device, exclude=["image"]) - return ray_bundle, batch - - -class FixedIndicesEvalDataloader(EvalDataloader): - """Dataloader that returns a fixed set of indices. - - Args: - input_dataset: InputDataset to load data from - image_indices: List of image indices to load data from. If None, then use all images. - device: Device to load data to - """ - - def __init__( - self, - input_dataset: InputDataset, - image_indices: Optional[Tuple[int]] = None, - device: Union[torch.device, str] = "cpu", - **kwargs, - ): - super().__init__(input_dataset, device, **kwargs) - if image_indices is None: - self.image_indices = list(range(len(input_dataset))) - else: - self.image_indices = image_indices - self.count = 0 - - def __iter__(self): - self.count = 0 - return self - - def __next__(self): - if self.count < len(self.image_indices): - image_idx = self.image_indices[self.count] - ray_bundle, batch = self.get_data_from_image_idx(image_idx) - self.count += 1 - return ray_bundle, batch - raise StopIteration - - -class RandIndicesEvalDataloader(EvalDataloader): - """Dataloader that returns random images. - - Args: - input_dataset: InputDataset to load data from - device: Device to load data to - """ - - def __init__( - self, - input_dataset: InputDataset, - device: Union[torch.device, str] = "cpu", - **kwargs, - ): - super().__init__(input_dataset, device, **kwargs) - self.count = 0 - - def __iter__(self): - self.count = 0 - return self - - def __next__(self): - if self.count < 1: - image_indices = range(self.cameras.size) - image_idx = random.choice(image_indices) - ray_bundle, batch = self.get_data_from_image_idx(image_idx) - self.count += 1 - return ray_bundle, batch - raise StopIteration diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index 5bae109d0..8f9ac51a0 100644 --- a/nerfstudio/engine/optimizers.py +++ b/nerfstudio/engine/optimizers.py @@ -21,12 +21,15 @@ from dataclasses import dataclass from typing import Any, Dict, List, Type import torch +from rich.console import Console from torch.cuda.amp.grad_scaler import GradScaler from torch.nn.parameter import Parameter from nerfstudio.configs import base_config from nerfstudio.utils import writer +CONSOLE = Console(width=120) + # Optimizer related configs @dataclass @@ -77,7 +80,7 @@ def setup_optimizers(config: base_config.Config, param_groups: Dict[str, List[Pa camera_optimizer_config = config.pipeline.datamanager.camera_optimizer if camera_optimizer_config.mode != "off": assert camera_optimizer_config.param_group not in optimizer_config - optimizer_config[camera_optimizer_config.param_group] = { + optimizer_config[camera_optimizer_config.param_group + "_train"] = { "optimizer": config.pipeline.datamanager.camera_optimizer.optimizer, "scheduler": config.pipeline.datamanager.camera_optimizer.scheduler, } @@ -97,6 +100,11 @@ class Optimizers: self.optimizers = {} self.schedulers = {} for param_group_name, params in param_groups.items(): + if param_group_name not in config: + CONSOLE.print( + f"[bold yellow]Warning: Skipping param group '{param_group_name}' because optimizer and scheduler are not in the config." + ) + continue lr_init = config[param_group_name]["optimizer"].lr self.optimizers[param_group_name] = config[param_group_name]["optimizer"].setup(params=params) if config[param_group_name]["scheduler"]: diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index f30d2ad9a..2b9f2e071 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -214,7 +214,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"] + mask = outputs["alive_ray_mask"] #.squeeze() 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 d45cda16c..02d79d110 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -23,6 +23,7 @@ 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 ( @@ -183,11 +184,15 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): """Configuration for pipeline instantiation""" _target: Type = field(default_factory=lambda: VanillaPipeline) - """target class to instantiate""" + """Target class to instantiate.""" datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig() - """specifies the datamanager config""" + """Specifies the datamanager config.""" model: ModelConfig = ModelConfig() - """specifies the model config""" + """Specifies the model config.""" + eval_optimize_cameras: bool = False + """Whether to optimize the cameras during evaluation.""" + eval_optimize_appearance: bool = False + """Whether to optimize the appearance during evaluation.""" class VanillaPipeline(Pipeline): @@ -255,7 +260,7 @@ class VanillaPipeline(Pipeline): model_outputs = self.model(ray_bundle) metrics_dict = self.model.get_metrics_dict(model_outputs, batch) - camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group + camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group + "_train" if camera_opt_param_group in self.datamanager.get_param_groups(): # Report the camera optimization metrics metrics_dict["camera_opt_translation"] = ( @@ -320,7 +325,7 @@ class VanillaPipeline(Pipeline): """ self.eval() metrics_dict_list = [] - num_images = len(self.datamanager.fixed_indices_eval_dataloader) + num_images = len(self.datamanager.eval_dataset) with Progress( TextColumn("[progress.description]{task.description}"), BarColumn(), @@ -329,11 +334,52 @@ class VanillaPipeline(Pipeline): transient=True, ) as progress: task = progress.add_task("[green]Evaluating all eval images...", total=num_images) - for camera_ray_bundle, batch in self.datamanager.fixed_indices_eval_dataloader: + for idx in range(num_images): + camera_ray_bundle, batch = self.datamanager.get_eval_image(idx) # time this the following line inner_start = time() height, width = camera_ray_bundle.shape num_rays = height * width + + # optimize camera poses on full image + if self.config.eval_optimize_cameras: + + # get the eval camera optimizer's parameters + camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group + "_eval" + param_groups = self.datamanager.get_param_groups() + optimizer = torch.optim.Adam(param_groups[camera_opt_param_group], lr=0.1) + + # check that camera optimizer is on + if self.config.datamanager.camera_optimizer.mode == "off": + raise ValueError( + "You must set datamanager.camera_optimizer.mode to optimize to optimize the camera poses." + ) + num_pose_iters = 100 + for i in range(num_pose_iters): + print("pose iter: ", i) + # need to optimize the ray generator's camera optimizer + # TODO: need to add require_grad in various locations to make it work + optimizer.zero_grad() + camera_ray_bundle, batch = self.datamanager.get_eval_image(idx) + outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + metrics_dict = self.model.get_metrics_dict(outputs, batch) + loss_dict = self.model.get_loss_dict(outputs, batch, metrics_dict) + loss_dict["rgb_loss"].backward() + optimizer.step() + + # optimize appearance on half image + if self.config.eval_optimize_appearance: + num_appearance_iters = 100 + side = random.choice([0, 1]) + half_width = width // 2 + if side == 0: + half_camera_ray_bundle = camera_ray_bundle[:, :half_width] + else: + half_camera_ray_bundle = camera_ray_bundle[:, half_width:] + for i in range(num_appearance_iters): + # need to optimize the "embedding_appearance" in the nerfacto field + pass + outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) metrics_dict, _ = self.model.get_image_metrics_and_images(outputs, batch) assert "num_rays_per_sec" not in metrics_dict