Merge pull request #51 from probcomp/debug_with_eric

WTF

Author: esli999

bayes3d/viz/viz.py

@@ -45,13 +45,12 @@ def preprocess_for_viz(img):
     return depth_np
 
 
-cmap = copy.copy(plt.get_cmap("turbo"))
+cmap  = copy.copy(plt.get_cmap('turbo'))
 cmap.set_bad(color=(1.0, 1.0, 1.0, 1.0))
 
-
-def get_depth_image(image, min_val=None, max_val=None, remove_max=True):
+def get_depth_image(image, max=None):
     """Convert a depth image to a PIL image.
-
+    
     Args:
         image (np.ndarray): Depth image. Shape (H, W).
         min (float): Minimum depth value for colormap.
@@ -60,28 +59,22 @@ def get_depth_image(image, min_val=None, max_val=None, remove_max=True):
     Returns:
         PIL.Image: Depth image visualized as a PIL image.
     """
-    if len(image.shape) > 2:
-        depth = np.array(image[:, :, -1])
+    depth = np.array(image)
+    if max is None:
+        maxim = depth.max()
     else:
-        depth = np.array(image)
-
-    if max_val is None:
-        max_val = depth.max()
-    if not remove_max:
-        max_val += 1
-    if min_val is None:
-        min_val = depth.min()
-
-    mask = (depth < max_val) * (depth > min_val)
+        maxim = max
+    mask = depth < maxim
     depth[np.logical_not(mask)] = np.nan
-    depth = (depth - min_val) / (max_val - min_val + 1e-10)
+    vmin = depth[mask].min()
+    vmax = depth[mask].max()
+    depth = (depth - vmin) / (vmax - vmin)
 
     img = Image.fromarray(
         np.rint(cmap(depth) * 255.0).astype(np.int8), mode="RGBA"
     ).convert("RGB")
     return img
 
-
 def get_rgb_image(image, max=255.0):
     """Convert an RGB image to a PIL image.