From 0de536c2687bde8a411c16600c4d1b2de471821b Mon Sep 17 00:00:00 2001 From: Matthias Blaicher Date: Sun, 18 Dec 2022 09:40:35 +0000 Subject: [PATCH] Support grayscale-only images in dataset --- 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..a509acd89 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."