| 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] |
|
|
| |
| if x.dtype == torch.uint8: |
| x = x / 255.0 |
| if x.is_floating_point(): |
| if x.ndim == 3 and x.shape[0] >= 3: |
| |
| 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 |
|
|
