Colorize a depth map using a matplotlib colormap. Args: depth: Depth tensor of shape (b, 1, h, w) or (b, h, w) with planar depth values ranging from 0 to inf. vmin: Minimum depth value to use for scaling the colormap. Can also be a percentile val
(
depth,
vmin=None,
vmax=None,
percentiles=False,
cmap="Spectral",
invalid_mask=None,
invalid_color=(0, 0, 0),
inverse=False
)
| 131 | |
| 132 | |
| 133 | def colorize_depth_map( |
| 134 | depth, |
| 135 | vmin=None, |
| 136 | vmax=None, |
| 137 | percentiles=False, |
| 138 | cmap="Spectral", |
| 139 | invalid_mask=None, |
| 140 | invalid_color=(0, 0, 0), |
| 141 | inverse=False |
| 142 | ): |
| 143 | """ |
| 144 | Colorize a depth map using a matplotlib colormap. |
| 145 | |
| 146 | Args: |
| 147 | depth: Depth tensor of shape (b, 1, h, w) or (b, h, w) with |
| 148 | planar depth values ranging from 0 to inf. |
| 149 | vmin: Minimum depth value to use for scaling the colormap. Can |
| 150 | also be a percentile value if percentiles is True. If None, |
| 151 | values in the batch are not min-clipped. |
| 152 | vmax: Maximum depth value to use for scaling the colormap. Can |
| 153 | also be a percentile value if percentiles is True. If None, |
| 154 | values in the batch are not max-clipped. |
| 155 | percentiles: If True, vmin and vmax are interpreted as percentiles |
| 156 | of the depth values in the batch (per sample!). |
| 157 | cmap: Name of the matplotlib colormap to use. |
| 158 | invalid_mask: Boolean mask of shape (b, h, w) that is True where |
| 159 | the depth values are invalid. |
| 160 | invalid_color: RGB color to use for invalid depth values. |
| 161 | inverse: If True, the depth values are inverted before colorization. |
| 162 | Returns: |
| 163 | Colorized depth map of shape (b, h, w, 3) with RGB values [0, 255]. |
| 164 | """ |
| 165 | if len(depth.shape) == 4: |
| 166 | assert depth.shape[1] == 1, "Depth must have 1 channel." |
| 167 | depth = depth.squeeze(1) |
| 168 | assert len(depth.shape) == 3, "Depth must have shape (b, h, w) or (b, 1, h, w)." |
| 169 | # assert depth.min() >= 0, "Depth must be non-negative." |
| 170 | |
| 171 | if isinstance(depth, torch.Tensor): |
| 172 | depth = depth.detach().cpu().numpy() |
| 173 | |
| 174 | # clip values with vmin and vmax |
| 175 | if vmin is not None and percentiles: |
| 176 | assert 0 <= vmin < 100, "vmin must be in [0, 100] if using percentiles" |
| 177 | vmin = percentile_per_sample(depth, vmin) |
| 178 | vmin = pad_vector_like_x(vmin, depth) |
| 179 | if vmax is not None and percentiles: |
| 180 | assert 0 < vmax <= 100, "vmax must be in [0, 100] if using percentiles" |
| 181 | vmax = percentile_per_sample(depth, vmax) |
| 182 | vmax = pad_vector_like_x(vmax, depth) |
| 183 | if exists(vmin) or exists(vmax): |
| 184 | # clip values between vmin and vmax |
| 185 | depth = np.clip(depth, vmin, vmax) |
| 186 | |
| 187 | # take inverse of depth |
| 188 | if inverse: |
| 189 | depth = 1.0 / depth |
| 190 |
nothing calls this directly
no test coverage detected