Preserve ray and image batch dimensions in the grayscale colormap

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)

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])