From c3c151f3ae33c83c69348ee7a7cbdd561a48b0cb Mon Sep 17 00:00:00 2001
From: MrKarimiD <mr.karimi.d@gmail.com>
Date: Fri, 10 May 2024 20:26:34 -0400
Subject: [PATCH] ns-eval save images by name

---
 nerfstudio/lantern/config.py                  |  4 ++--
 nerfstudio/lantern/dataset.py                 |  2 ++
 nerfstudio/lantern/model_pano_hdr_nerfacto.py |  2 +-
 nerfstudio/pipelines/base_pipeline.py         | 19 ++++++++++++++++---
 4 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/nerfstudio/lantern/config.py b/nerfstudio/lantern/config.py
index 14e384fbc..d9356e353 100644
--- a/nerfstudio/lantern/config.py
+++ b/nerfstudio/lantern/config.py
@@ -64,8 +64,8 @@ def Get_PanoHDR_Nerfacto_config():
     pipeline=VanillaPipelineConfig(
             datamanager=HDRVanillaDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(),
-            train_num_rays_per_batch=8192,
-            eval_num_rays_per_batch=8192,
+            train_num_rays_per_batch=4096,
+            eval_num_rays_per_batch=4096,
             # camera_optimizer=CameraOptimizerConfig(
             #     mode="SO3xR3",
             #     optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2),
diff --git a/nerfstudio/lantern/dataset.py b/nerfstudio/lantern/dataset.py
index b6f9a1030..a0b5a7c80 100644
--- a/nerfstudio/lantern/dataset.py
+++ b/nerfstudio/lantern/dataset.py
@@ -108,6 +108,8 @@ class HDRInputDataset(InputDataset):
         """
         image = self.get_image(image_idx)
         data = {"image_idx": image_idx, "image": image}
+        image_filename = self._dataparser_outputs.image_filenames[image_idx]
+        data["image_filename"] = str(image_filename.name)
         if self._dataparser_outputs.mask_filenames is not None:
             mask_filepath = self._dataparser_outputs.mask_filenames[image_idx]
             data["mask"] = get_image_mask_tensor_from_path(filepath=mask_filepath, scale_factor=self.scale_factor)
diff --git a/nerfstudio/lantern/model_pano_hdr_nerfacto.py b/nerfstudio/lantern/model_pano_hdr_nerfacto.py
index 438554aa6..5c0ed8f90 100644
--- a/nerfstudio/lantern/model_pano_hdr_nerfacto.py
+++ b/nerfstudio/lantern/model_pano_hdr_nerfacto.py
@@ -222,7 +222,7 @@ class PanoHDRNeRFModel(NerfactoModel):
         gt_rgb = batch["mask"] * gt_rgb
         predicted_rgb = batch["mask"] * predicted_rgb
         
-        combined_rgb = torch.cat([gt_rgb, predicted_rgb], dim=1)
+        combined_rgb = predicted_rgb # torch.cat([gt_rgb, predicted_rgb], dim=1)
         combined_acc = torch.cat([acc], dim=1)
         combined_depth = torch.cat([depth], dim=1)
 
diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py
index 75fc493cf..d5c867a9f 100644
--- a/nerfstudio/pipelines/base_pipeline.py
+++ b/nerfstudio/pipelines/base_pipeline.py
@@ -395,9 +395,22 @@ class VanillaPipeline(Pipeline):
                         camera_indices = camera_ray_bundle.camera_indices
                         assert camera_indices is not None
                         for key, val in images_dict.items():
-                            Image.fromarray((val * 255).byte().cpu().numpy()).save(
-                                output_path / "{0:06d}-{1}.jpg".format(int(camera_indices[0, 0, 0]), key)
-                            )
+                            if 'image_filename' in batch.keys():
+                                output_name = batch['image_filename'].split('.')[0] + '-' + key
+                                if key == 'img':
+                                    output_name = output_name + '.' + batch['image_filename'].split('.')[-1]
+                                else:
+                                    output_name = output_name + '.jpg'
+                            else:
+                                output_name = "{0:06d}-{1}.jpg".format(int(camera_indices[0, 0, 0]), key)
+                            if output_name.endswith('.exr'):
+                                import cv2
+                                render_image = cv2.cvtColor(val.cpu().numpy(), cv2.COLOR_BGR2RGB)
+                                cv2.imwrite(str(output_path / output_name), render_image)
+                            else:
+                                Image.fromarray((val * 255).byte().cpu().numpy()).save(
+                                    output_path / output_name
+                                )
                     assert "num_rays_per_sec" not in metrics_dict
                     metrics_dict["num_rays_per_sec"] = num_rays / (time() - inner_start)
                     fps_str = "fps"