From 0e77c731e2a583de63dd880504b285eccd620010 Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Mon, 9 May 2022 11:16:04 -0700
Subject: [PATCH] missing file

---
 mattport/utils/visualization.py | 42 ++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/mattport/utils/visualization.py b/mattport/utils/visualization.py
index df9f8de05..a8fac0487 100644
--- a/mattport/utils/visualization.py
+++ b/mattport/utils/visualization.py
@@ -1,5 +1,6 @@
 """ Helper functions for visualizing outputs """
 
+from typing import Optional
 import torch
 from matplotlib import cm
 from torchtyping import TensorType
@@ -19,6 +20,45 @@ def apply_colormap(image: TensorType[..., 1], cmap="viridis") -> TensorType[...,
     colormap = cm.get_cmap(cmap)
     colormap = torch.tensor(colormap.colors).to(image.device)
 
+    image2 = (image * 255).long()
+
+    if torch.min(image2) < 0 or torch.max(image2) > 255:
+        print("hi")
+
     image = (image * 255).long()
 
-    return colormap[image]
+    return colormap[image[..., 0]]
+
+
+def apply_depth_colormap(
+    depth: TensorType[..., 1],
+    accumulation: Optional[TensorType[..., 1]] = None,
+    near_plane: Optional[float] = None,
+    far_plane: Optional[float] = None,
+    cmap="turbo",
+) -> TensorType[..., 3]:
+    """Converts a depth image to color for easier analysis.
+
+    Args:
+        depth (TensorType[..., 1]): Depth image.
+        accumulation (TensorType[..., 1], optional): Ray accumulation used for masking vis. Defaults to None.
+        near_plane (float, optional): Closest depth to consider. If None, use min image value. Defaults to None.
+        far_plane (float, optional): Furthest depth to consider. If None, use max image value. Defaults to None.
+        cmap (str, optional): Colormap to apply. Defaults to "turbo".
+
+    Returns:
+        TensorType[..., 3]: Colored depth image
+    """
+
+    near_plane = near_plane or torch.min(depth)
+    far_plane = far_plane or torch.max(depth)
+
+    depth = (depth - near_plane) / (far_plane - near_plane + 1e-10)
+    depth = torch.clip(depth, 0, 1)
+
+    colored_image = apply_colormap(depth, cmap=cmap)
+
+    if accumulation is not None:
+        colored_image = colored_image * accumulation + (1 - accumulation)
+
+    return colored_image