From 36d5e358018647676182b0f5deffc570a084ecc5 Mon Sep 17 00:00:00 2001 From: Alex Jean Date: Wed, 21 Aug 2024 14:59:44 -0700 Subject: [PATCH] Addressed TODO in full_images_datamanager.py to return image resolution as train rays (#3312) * addressed a todo in full_images_datamanager.py that should return the resolution of the image being rendered * fixed docstring get_train_rays_per_batch --- nerfstudio/data/datamanagers/full_images_datamanager.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index a67d49299..9c37ce619 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -332,8 +332,13 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): return {} def get_train_rays_per_batch(self): - # TODO: fix this to be the resolution of the last image rendered - return 800 * 800 + """Returns resolution of the image returned from datamanager.""" + if len(self.cached_train) != 0: + h = self.cached_train[0]["image"].shape[0] + w = self.cached_train[0]["image"].shape[1] + return h * w + else: + return 800 * 800 def next_train(self, step: int) -> Tuple[Cameras, Dict]: """Returns the next training batch