Add heatmap to 2D tensor. Args: tensor (tensor): a 2D tensor. Tensor value must be in [0..1] range. Returns: heatmap (tensor): a 3D tensor. Result of applying heatmap to the 2D tensor.
(tensor)
| 409 | |
| 410 | |
| 411 | def add_heatmap(tensor): |
| 412 | """ |
| 413 | Add heatmap to 2D tensor. |
| 414 | Args: |
| 415 | tensor (tensor): a 2D tensor. Tensor value must be in [0..1] range. |
| 416 | Returns: |
| 417 | heatmap (tensor): a 3D tensor. Result of applying heatmap to the 2D tensor. |
| 418 | """ |
| 419 | assert tensor.ndim == 2, "Only support 2D tensors." |
| 420 | # Move tensor to cpu if necessary. |
| 421 | if tensor.device != torch.device("cpu"): |
| 422 | arr = tensor.cpu() |
| 423 | else: |
| 424 | arr = tensor |
| 425 | arr = arr.numpy() |
| 426 | # Get the color map by name. |
| 427 | cm = plt.get_cmap("viridis") |
| 428 | heatmap = cm(arr) |
| 429 | heatmap = heatmap[:, :, :3] |
| 430 | # Convert (H, W, C) to (C, H, W) |
| 431 | heatmap = torch.Tensor(heatmap).permute(2, 0, 1) |
| 432 | return heatmap |