diff --git a/nerfstudio/utils/colormaps.py b/nerfstudio/utils/colormaps.py index c771178c2..8b209ca11 100644 --- a/nerfstudio/utils/colormaps.py +++ b/nerfstudio/utils/colormaps.py @@ -105,7 +105,7 @@ def apply_float_colormap(image: Float[Tensor, "*bs 1"], colormap: Colormaps = "v image = torch.nan_to_num(image, 0) if colormap == "gray": - return image.repeat(1, 1, 3) + return image.repeat_interleave(3, dim=-1) image_long = (image * 255).long() image_long_min = torch.min(image_long) image_long_max = torch.max(image_long) diff --git a/tests/utils/test_grayscale_colormap.py b/tests/utils/test_grayscale_colormap.py new file mode 100644 index 000000000..0a800f8c7 --- /dev/null +++ b/tests/utils/test_grayscale_colormap.py @@ -0,0 +1,14 @@ +import pytest +import torch + +from nerfstudio.utils.colormaps import ColormapOptions, apply_colormap + + +@pytest.mark.parametrize("shape", [(5, 1), (2, 4, 1), (2, 3, 4, 1)]) +def test_grayscale_preserves_batch_dimensions(shape): + values = torch.linspace(0, 1, torch.Size(shape).numel()).reshape(shape) + rgb = apply_colormap(values, ColormapOptions(colormap="gray")) + + assert rgb.shape == (*shape[:-1], 3) + for channel in range(3): + torch.testing.assert_close(rgb[..., channel], values[..., 0])