import os, numpy
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg
from matplotlib.figure import Figure
import imageio.v3 as iio
from .imports import import_torch
from .util import _create_folder, _format_image
torch = import_torch()
[docs]
@torch.no_grad()
def show_image(
tensor: torch.Tensor | numpy.ndarray,
columns: int | None = None,
rescale: bool = False,
clamp_range: tuple[float] = (0.0, 1.0),
colorbar: bool = False,
max_width: int | None = None,
padding: int = 3,
pad_value: float = 0.0,
) -> None:
"""
Shows tensor of shape **(\\*,C,H,W)** as an image using pyplot.
Any extra dimensions are treated as batch dimensions, and displayed in a grid.
By default, the image is simply clipped to the 0-1 range. Use rescale and clamp_range to modify this behaviour.
Args:
tensor : (H,W) or (C,H,W) or (\\*,C,H,W) tensor or numpy array to display.
columns : number of columns to use for the grid of images (default 8 or less)
rescale : whether to rescale the images to 0-1 range. Uses min-max scaling.
clamp_range : (min,max), tuple of values to clamp the images to, before rescaling. Use None to disable clamping
colorbar : whether to add a colorbar to the image, only works for grayscale images (default False)
max_width : maximum width of the image
padding : number of pixels between images in the grid
pad_value : inter-padding value for the grid of images
"""
tensor = _format_image(
tensor,
columns=columns,
rescale=rescale,
clamp_range=clamp_range,
max_width=max_width,
padding=padding,
pad_value=pad_value,
) # (C,H',W') ready to show
C, H, W = tensor.shape
plt.figure(figsize=(W / 50, H / 50), dpi=50)
plt.subplots_adjust(left=0, right=1, bottom=0, top=1) # Remove margins
show_colorbar = colorbar and C == 1
# 1-channel images are grayscale, unless a colorbar gives the colormap a meaning
plt.imshow(tensor.permute((1, 2, 0)), cmap=None if show_colorbar else "gray")
plt.axis("off")
if show_colorbar:
plt.colorbar()
plt.show()
[docs]
@torch.no_grad()
def save_image(
tensor: torch.Tensor | numpy.ndarray,
folder: str,
name: str = "imagetensor",
columns: int | None = None,
rescale: bool = False,
clamp_range: tuple[float] = (0.0, 1.0),
colorbar: bool = False,
max_width: int | None = None,
padding: int = 3,
pad_value: float = 0.0,
create_folder: bool = True,
) -> None:
"""
Saves tensor of shape **(\\*,C,H,W)** as a png image.
Any extra dimensions are treated as batch dimensions, and displayed in a grid.
By default, the image is simply clipped to the 0-1 range. Use rescale and clamp_range to modify this behaviour.
Args:
tensor : (H,W) or (C,H,W) or (\\*,C,H,W) tensor or numpy array to display
folder : relative path of folder where to save the image
name : name of the image (do not include extension)
columns : number of columns to use for the grid of images (default 8 or less)
rescale : whether to rescale the images to 0-1 range
clamp_range : (min,max), tuple of values to clamp the images to, before rescaling. Use None to disable clamping
colorbar : whether to add a colorbar to the image, only works for grayscale images (default False)
max_width : maximum width of the image
padding : number of pixels between images in the grid
pad_value : inter-padding value for the grid of images
create_folder : whether to create the folder if it does not exist (default True)
"""
_create_folder(folder, create_folder)
tensor = _format_image(
tensor,
columns=columns,
rescale=rescale,
clamp_range=clamp_range,
max_width=max_width,
padding=padding,
pad_value=pad_value,
) # (C,H',W') ready to save
save_path = os.path.join(folder, f"{name}.png")
if not (colorbar and tensor.shape[0] == 1):
# Nothing to draw around the image, so write the pixels straight out
image = tensor.clamp(0.0, 1.0).mul(255).round().to(torch.uint8) # (C,H',W')
image = image.permute((1, 2, 0)).numpy() # (H',W',C)
if image.shape[-1] == 1:
image = image[..., 0] # (H',W'), grayscale pngs carry no channel dimension
iio.imwrite(save_path, image)
return
H, W = tensor.shape[1], tensor.shape[2]
# Figure rather than pyplot : never instantiates a GUI backend, so this works headless
figure = Figure(figsize=(W / 50, H / 50), dpi=50)
FigureCanvasAgg(figure)
figure.subplots_adjust(left=0, right=1, bottom=0, top=1) # Remove margins
axes = figure.subplots()
shown = axes.imshow(tensor[0], extent=[0, W, 0, H])
axes.axis("off")
figure.colorbar(shown, ax=axes)
figure.savefig(save_path, dpi=50, bbox_inches="tight", pad_inches=0)