tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-12 10:55:57 +08:00
Code Activity

updates to apply_pca_colormap (#3086)

Browse Source 
* improvements to pca_colormap: allow input pca matrix, optional ignore_zeros arg

* typo
This commit is contained in:
Justin Kerr authored and GitHub committed 2024-04-17 23:54:43 -07:00
1 parent babf57718c
commit 45d8bb79ea
1 file changed
+25 -14
Download Patch File Download Diff File Expand all files Collapse all files
nerfstudio/utils/colormaps.py
+25 -14
View File
@@ -171,35 +171,46 @@ def apply_boolean_colormap(
return colored_image
def apply_pca_colormap(image: Float[Tensor, "*bs dim"]) -> Float[Tensor, "*bs rgb=3"]:
def apply_pca_colormap(
image: Float[Tensor, "*bs dim"], pca_mat: Optional[Float[Tensor, "dim rgb=3"]] = None, ignore_zeros=True
) -> Float[Tensor, "*bs rgb=3"]:
"""Convert feature image to 3-channel RGB via PCA. The first three principle
components are used for the color channels, with outlier rejection per-channel
Args:
image: image of arbitrary vectors
pca_mat: an optional argument of the PCA matrix, shape (dim, 3)
ignore_zeros: whether to ignore zero values in the input image (they won't affect the PCA computation)
Returns:
Tensor: Colored image
"""
original_shape = image.shape
image = image.view(-1, image.shape[-1])
_, _, v = torch.pca_lowrank(image)
image = torch.matmul(image, v[..., :3])
d = torch.abs(image - torch.median(image, dim=0).values)
if ignore_zeros:
valids = (image.abs().amax(dim=-1)) > 0
else:
valids = torch.ones(image.shape[0], dtype=torch.bool)
if pca_mat is None:
_, _, pca_mat = torch.pca_lowrank(image[valids, :], q=3, niter=20)
assert pca_mat is not None
image = torch.matmul(image, pca_mat[..., :3])
d = torch.abs(image[valids, :] - torch.median(image[valids, :], dim=0).values)
mdev = torch.median(d, dim=0).values
s = d / mdev
m = 3.0 # this is a hyperparam controlling how many std dev outside for outliers
rins = image[s[:, 0] < m, 0]
gins = image[s[:, 1] < m, 1]
bins = image[s[:, 2] < m, 2]
m = 2.0 # this is a hyperparam controlling how many std dev outside for outliers
rins = image[valids, :][s[:, 0] < m, 0]
gins = image[valids, :][s[:, 1] < m, 1]
bins = image[valids, :][s[:, 2] < m, 2]
image[:, 0] -= rins.min()
image[:, 1] -= gins.min()
image[:, 2] -= bins.min()
image[valids, 0] -= rins.min()
image[valids, 1] -= gins.min()
image[valids, 2] -= bins.min()
image[:, 0] /= rins.max() - rins.min()
image[:, 1] /= gins.max() - gins.min()
image[:, 2] /= bins.max() - bins.min()
image[valids, 0] /= rins.max() - rins.min()
image[valids, 1] /= gins.max() - gins.min()
image[valids, 2] /= bins.max() - bins.min()
image = torch.clamp(image, 0, 1)
image_long = (image * 255).long()