(
x: Float[Tensor, " *batch"],
color_map: str = "inferno",
)
| 6 | |
| 7 | |
| 8 | def apply_color_map( |
| 9 | x: Float[Tensor, " *batch"], |
| 10 | color_map: str = "inferno", |
| 11 | ) -> Float[Tensor, "*batch 3"]: |
| 12 | cmap = cm.get_cmap(color_map) |
| 13 | |
| 14 | # Convert to NumPy so that Matplotlib color maps can be used. |
| 15 | mapped = cmap(x.detach().clip(min=0, max=1).cpu().numpy())[..., :3] |
| 16 | |
| 17 | # Convert back to the original format. |
| 18 | return torch.tensor(mapped, device=x.device, dtype=torch.float32) |
| 19 | |
| 20 | |
| 21 | def apply_color_map_to_image( |
no outgoing calls
no test coverage detected