From 7e8fc5c97f4a4de38ce374cf261869b8f5f68db0 Mon Sep 17 00:00:00 2001
From: Chung Min Kim <chungmin99@gmail.com>
Date: Tue, 16 May 2023 10:46:15 -0700
Subject: [PATCH] Create PCA colormap for images with channels > 3 (#1929)

* Init pca colormap for images with channels>3

* --amend

* formatting etc

* formatting etc

* docstring abt >3 channels

---------

Co-authored-by: Justin Kerr <justin.g.kerr@gmail.com>
Co-authored-by: Saurav Maheshkar <sauravvmaheshkar@gmail.com>
---
 nerfstudio/utils/colormaps.py             | 47 ++++++++++++++++++++++-
 nerfstudio/viewer/server/control_panel.py |  2 +
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/nerfstudio/utils/colormaps.py b/nerfstudio/utils/colormaps.py
index 0ff3b66a1..7346ca379 100644
--- a/nerfstudio/utils/colormaps.py
+++ b/nerfstudio/utils/colormaps.py
@@ -48,7 +48,10 @@ def apply_colormap(
     eps: float = 1e-9,
 ) -> TensorType["bs":..., "rgb":3]:
     """
-    Applies a colormap to a tensor image. Currently only supports 1 and 3 channel inputs
+    Applies a colormap to a tensor image.
+    If single channel, applies a colormap to the image.
+    If 3 channel, treats the channels as RGB.
+    If more than 3 channel, applies a PCA reduction on the dimensions to 3 channels
 
     Args:
         image: Input tensor image.
@@ -80,6 +83,9 @@ def apply_colormap(
     if image.dtype == torch.bool:
         return apply_boolean_colormap(image)
 
+    if image.shape[-1] > 3:
+        return apply_pca_colormap(image)
+
     raise NotImplementedError
 
 
@@ -164,3 +170,42 @@ def apply_boolean_colormap(
     colored_image[image[..., 0], :] = true_color
     colored_image[~image[..., 0], :] = false_color
     return colored_image
+
+
+def apply_pca_colormap(image: TensorType["bs":..., "dim"]) -> TensorType["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
+
+    Returns:
+        TensorType: 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)
+    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]
+
+    image[:, 0] -= rins.min()
+    image[:, 1] -= gins.min()
+    image[:, 2] -= bins.min()
+
+    image[:, 0] /= rins.max() - rins.min()
+    image[:, 1] /= gins.max() - gins.min()
+    image[:, 2] /= bins.max() - bins.min()
+
+    image = torch.clamp(image, 0, 1)
+    image_long = (image * 255).long()
+    image_long_min = torch.min(image_long)
+    image_long_max = torch.max(image_long)
+    assert image_long_min >= 0, f"the min value is {image_long_min}"
+    assert image_long_max <= 255, f"the max value is {image_long_max}"
+    return image.view(*original_shape[:-1], 3)
diff --git a/nerfstudio/viewer/server/control_panel.py b/nerfstudio/viewer/server/control_panel.py
index dc87c47b8..456e3598c 100644
--- a/nerfstudio/viewer/server/control_panel.py
+++ b/nerfstudio/viewer/server/control_panel.py
@@ -362,4 +362,6 @@ def _get_colormap_options(dimensions: int, dtype: type) -> List[Colormaps]:
         colormap_options = ["default"]
     if dimensions == 1 and dtype == torch.float:
         colormap_options = [c for c in list(get_args(Colormaps)) if c != "default"]
+    if dimensions > 3:
+        colormap_options = ["pca"]
     return colormap_options