File size: 1,938 Bytes
8931de0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import matplotlib.pyplot as plt
import torch
from torch import Tensor
def showtensor(tensor: Tensor, figsize: int = 4, axis: bool = True, title=None) -> None:
""" Show a tensor as an image using matplotlib.
Parameters
----------
tensor : Tensor
The tensor to show.
figsize : int, optional
The size of the figure. The default is 4.
axis : bool, optional
Whether to show the axis. The default is True.
"""
tensor = ensure_tensor_has_four_dimensions(tensor)
images_or_revisits, channels, height, width = tensor.shape
nrows, ncols = 1, images_or_revisits
fig, ax = plt.subplots(
nrows=nrows,
ncols=ncols,
figsize=(ncols * figsize, nrows * figsize),
tight_layout=True,
squeeze=False,
)
if title is not None:
plt.title(title)
for image in range(images_or_revisits):
x = tensor[image]
# Convert images to float
if x.dtype == torch.uint8:
x = x / 255.0
if x.is_floating_point():
if x.ndim == 3 and x.shape[0] >= 3:
# Channels first to channels last [1,2,3] -> [3,2,1]
x = x.permute(1, 2, 0)
ax[0, image].imshow(x, vmin=0, vmax=1, interpolation="nearest")
else:
x = x[0]
ax[0, image].imshow(x)
ax[0, image].axis(axis if image == 0 else False)
def ensure_tensor_has_four_dimensions(tensor):
""" Ensure that the tensor has four dimensions.
If it doesn't pad the missing dimensions with empty dimensions.
Parameters
----------
tensor : Tensor
The tensor to pad.
Returns
-------
Tensor
A tensor with four dimensions.
"""
if tensor.ndim == 4:
pass
elif tensor.ndim == 3:
tensor = tensor[None]
elif tensor.ndim == 2:
tensor = tensor[None, None]
return tensor
|