Colorize a scalar map of. Args: scalar_map: Map of with format BHW. val_min: Minimu value to display. val_max: Maximum value to display. color_map: Which color map to use. Will be passed to matplotlob. Returns: A colorized image with format BHWC.
(
scalar_map: torch.Tensor, val_min=0.0, val_max=1.0, color_map: str = "jet"
)
| 42 | |
| 43 | |
| 44 | def colorize_scalar_map( |
| 45 | scalar_map: torch.Tensor, val_min=0.0, val_max=1.0, color_map: str = "jet" |
| 46 | ) -> torch.Tensor: |
| 47 | """Colorize a scalar map of. |
| 48 | |
| 49 | Args: |
| 50 | scalar_map: Map of with format BHW. |
| 51 | val_min: Minimu value to display. |
| 52 | val_max: Maximum value to display. |
| 53 | color_map: Which color map to use. Will be passed to matplotlob. |
| 54 | |
| 55 | Returns: |
| 56 | A colorized image with format BHWC. |
| 57 | """ |
| 58 | if scalar_map.ndim not in (2, 3, 4): |
| 59 | raise ValueError("Only scalar maps of 2 or 3 or 4 dimensions supported.") |
| 60 | |
| 61 | cmap = plt.get_cmap(color_map) |
| 62 | |
| 63 | scalar_map_np = scalar_map.detach().cpu().float().numpy() |
| 64 | scalar_map_np = (scalar_map_np - val_min) / (val_max - val_min) |
| 65 | scalar_map_np = np.clip(scalar_map_np, a_min=0.0, a_max=1.0) |
| 66 | |
| 67 | color_map_np = cmap(scalar_map_np)[..., :3] |
| 68 | tensor = torch.as_tensor(color_map_np * 255.0, dtype=torch.uint8) |
| 69 | |
| 70 | if tensor.ndim == 3: |
| 71 | return tensor.permute(2, 0, 1) |
| 72 | elif tensor.ndim == 4: |
| 73 | return tensor.permute(0, 3, 1, 2) |
| 74 | elif tensor.ndim == 5: |
| 75 | return tensor.permute(0, 1, 4, 2, 3) |
| 76 | else: |
| 77 | assert False, "Invalid tensor shape encountered." |
no outgoing calls
no test coverage detected