| 216 | |
| 217 | # ref: https://github.com/hbb1/2d-gaussian-splatting/blob/main/utils/general_utils.py#L163 |
| 218 | def colormap(img, cmap="jet"): |
| 219 | W, H = img.shape[:2] |
| 220 | dpi = 300 |
| 221 | fig, ax = plt.subplots(1, figsize=(H / dpi, W / dpi), dpi=dpi) |
| 222 | im = ax.imshow(img, cmap=cmap) |
| 223 | ax.set_axis_off() |
| 224 | fig.colorbar(im, ax=ax) |
| 225 | fig.tight_layout() |
| 226 | fig.canvas.draw() |
| 227 | data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) |
| 228 | data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) |
| 229 | img = torch.from_numpy(data).float().permute(2, 0, 1) |
| 230 | plt.close() |
| 231 | return img |
| 232 | |
| 233 | |
| 234 | def apply_float_colormap(img: torch.Tensor, colormap: str = "turbo") -> torch.Tensor: |