mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 18:44:21 +08:00
Merge 63c5e5842e into 50e0e3c70c
This commit is contained in:
commit
2d98d0381f
@ -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)
|
||||
|
||||
14
tests/utils/test_grayscale_colormap.py
Normal file
14
tests/utils/test_grayscale_colormap.py
Normal file
@ -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])
|
||||
Loading…
Reference in New Issue
Block a user