Spaces:
Sleeping
Sleeping
File size: 4,502 Bytes
e7bcdd2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | # Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import imageio
import matplotlib
import numpy as np
import torch
from einops import rearrange
def visualize_depth(
depth: np.ndarray,
depth_min=None,
depth_max=None,
percentile=2,
ret_minmax=False,
ret_type=np.uint8,
cmap="Spectral",
):
"""
Visualize a depth map using a colormap.
Args:
depth: Input depth map array
depth_min: Minimum depth value for normalization. If None, uses percentile
depth_max: Maximum depth value for normalization. If None, uses percentile
percentile: Percentile for min/max computation if not provided
ret_minmax: Whether to return min/max depth values
ret_type: Return array type (uint8 or float)
cmap: Matplotlib colormap name to use
Returns:
Colored depth visualization as numpy array
If ret_minmax=True, also returns depth_min and depth_max
"""
depth = depth.copy()
depth.copy()
valid_mask = depth > 0
depth[valid_mask] = 1 / depth[valid_mask]
if depth_min is None:
if valid_mask.sum() <= 10:
depth_min = 0
else:
depth_min = np.percentile(depth[valid_mask], percentile)
if depth_max is None:
if valid_mask.sum() <= 10:
depth_max = 0
else:
depth_max = np.percentile(depth[valid_mask], 100 - percentile)
if depth_min == depth_max:
depth_min = depth_min - 1e-6
depth_max = depth_max + 1e-6
cm = matplotlib.colormaps[cmap]
depth = ((depth - depth_min) / (depth_max - depth_min)).clip(0, 1)
depth = 1 - depth
img_colored_np = cm(depth[None], bytes=False)[:, :, :, 0:3] # value from 0 to 1
if ret_type == np.uint8:
img_colored_np = (img_colored_np[0] * 255.0).astype(np.uint8)
elif ret_type == np.float32 or ret_type == np.float64:
img_colored_np = img_colored_np[0]
else:
raise ValueError(f"Invalid return type: {ret_type}")
if ret_minmax:
return img_colored_np, depth_min, depth_max
else:
return img_colored_np
# GS video rendering visulization function, since it operates in Tensor space...
def vis_depth_map_tensor(
result: torch.Tensor, # "*batch height width"
color_map: str = "Spectral",
) -> torch.Tensor: # "*batch 3 height with"
"""
Color-map the depth map.
"""
far = result.reshape(-1)[:16_000_000].float().quantile(0.99).log().to(result)
try:
near = result[result > 0][:16_000_000].float().quantile(0.01).log().to(result)
except (RuntimeError, ValueError) as e:
logger.error(f"No valid depth values found. Reason: {e}")
near = torch.zeros_like(far)
result = result.log()
result = (result - near) / (far - near)
return apply_color_map_to_image(result, color_map)
def apply_color_map(
x: torch.Tensor, # " *batch"
color_map: str = "inferno",
) -> torch.Tensor: # "*batch 3"
cmap = matplotlib.cm.get_cmap(color_map)
# Convert to NumPy so that Matplotlib color maps can be used.
mapped = cmap(x.float().detach().clip(min=0, max=1).cpu().numpy())[..., :3]
# Convert back to the original format.
return torch.tensor(mapped, device=x.device, dtype=torch.float32)
def apply_color_map_to_image(
image: torch.Tensor, # "*batch height width"
color_map: str = "inferno",
) -> torch.Tensor: # "*batch 3 height with"
image = apply_color_map(image, color_map)
return rearrange(image, "... h w c -> ... c h w")
def export_to_depth_vis(
depth,
export_dir,
):
os.makedirs(os.path.join(export_dir, "depth_vis"), exist_ok=True)
for idx in range(depth.shape[0]):
depth_vis = visualize_depth(depth[idx])
depth_vis = depth_vis.astype(np.uint8)
save_path = os.path.join(export_dir, "depth_vis", f"depth_vis_{idx:04d}.png")
imageio.imwrite(save_path, depth_vis, quality=95)
|