From 06642a5871481ddda2d2594853d683a5be205bdd Mon Sep 17 00:00:00 2001
From: Matthias Blaicher <matthias@blaicher.com>
Date: Thu, 22 Dec 2022 18:16:03 +0000
Subject: [PATCH] Support grayscale-only images in dataset (#1142)

* Support grayscale-only images in dataset

* lint
---
 nerfstudio/data/datasets/base_dataset.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py
index 1d1a88a3f..44793fbb0 100644
--- a/nerfstudio/data/datasets/base_dataset.py
+++ b/nerfstudio/data/datasets/base_dataset.py
@@ -64,7 +64,9 @@ class InputDataset(Dataset):
             width, height = pil_image.size
             newsize = (int(width * self.scale_factor), int(height * self.scale_factor))
             pil_image = pil_image.resize(newsize, resample=Image.BILINEAR)
-        image = np.array(pil_image, dtype="uint8")  # shape is (h, w, 3 or 4)
+        image = np.array(pil_image, dtype="uint8")  # shape is (h, w) or (h, w, 3 or 4)
+        if len(image.shape) == 2:
+            image = image[:, :, None].repeat(3, axis=2)
         assert len(image.shape) == 3
         assert image.dtype == np.uint8
         assert image.shape[2] in [3, 4], f"Image shape of {image.shape} is in correct."