基于PyTorch在TensorBoard中设置图像颜色映射

Hey there! Since you're using TensorBoard to track your image reconstruction progress (even without ML/DL—totally valid use case!), here are two straightforward ways to ditch the default grayscale and use custom colormaps:

1. Preprocess Images with Matplotlib Colormaps Before Writing to TensorBoard #

TensorBoard renders single-channel images as grayscale by default, so we can convert the single-channel data to a 3-channel colored image using matplotlib's colormaps before logging it. This keeps things aligned with the add_image workflow you’re probably using.

Here’s a code snippet to implement this:

import matplotlib.cm as cm
import torch
from torch.utils.tensorboard import SummaryWriter

# Initialize your TensorBoard writer
writer = SummaryWriter(log_dir="./runs/recon_progress")

# Assume `recon_img` is your single-channel reconstruction tensor (shape: (1, H, W))
recon_img = ...  # Replace with your actual image tensor

# Step 1: Normalize the image to the 0-1 range required by colormaps
img_normalized = (recon_img - recon_img.min()) / (recon_img.max() - recon_img.min())

# Step 2: Apply your chosen colormap (e.g., viridis, plasma, jet, inferno)
# Squeeze to remove the channel dim, convert to numpy, apply colormap
img_colored = cm.viridis(img_normalized.squeeze().numpy())
# Colormaps return RGBA; strip the alpha channel to get RGB
img_colored = img_colored[:, :, :3]
# Convert back to tensor and rearrange to (3, H, W) for TensorBoard
img_colored_tensor = torch.tensor(img_colored).permute(2, 0, 1)

# Step 3: Log the colored image
writer.add_image("Reconstructed Image (Viridis)", img_colored_tensor, global_step=current_iteration)

You can swap cm.viridis with any matplotlib colormap name to get different color scales.

2. Log Matplotlib Figures Directly (With Optional Colorbar) #

If you want more control—like adding a colorbar to track pixel value ranges (super useful for optimization!)—you can create a matplotlib figure with your desired colormap and log it using add_figure.

Here’s how:

import matplotlib.pyplot as plt
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(log_dir="./runs/recon_progress")

recon_img = ...  # Your single-channel reconstruction tensor

# Create a figure with the colormap
fig, ax = plt.subplots(figsize=(8, 6))
# Display the image with your chosen colormap
im = ax.imshow(recon_img.squeeze().numpy(), cmap="plasma")
# Add a colorbar to show value-to-color mapping
plt.colorbar(im, ax=ax)
# Turn off axes for cleaner visualization
ax.axis("off")

# Log the figure to TensorBoard
writer.add_figure("Reconstructed Image with Colorbar", fig, global_step=current_iteration)
# Close the figure to free up memory
plt.close(fig)

Quick Notes: #

• If you’re working with RGB images (3-channel), TensorBoard will display them in color automatically—these methods are mainly for single-channel grayscale-like data from your reconstruction.
• Matplotlib has tons of built-in colormaps; feel free to experiment with cm.jet, cm.inferno, or even perceptually uniform ones like cm.viridis (great for avoiding bias in visual interpretation).

内容的提问来源于stack exchange，提问作者CommandoMan