Spaces:
Running on L40S
Running on L40S
File size: 5,726 Bytes
9f818c5 | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""image visualization utilities.
based on https://gitlab.com/qsh.zh/jam/-/blob/master/jamviz/img.py MIT License
"""
import os
from typing import Union
import numpy as np
import torch
from einops import rearrange
from PIL import Image
from torchvision.utils import make_grid
__all__ = [
"show_batch_img",
"save_batch_img",
]
def _reshape_viz_batch_img(img_data: torch.Tensor | np.ndarray, shape: int | str = 7) -> tuple:
"""
Reshapes a batch of images for visualization, organizing them into a grid format.
Args:
img_data (torch.Tensor | np.ndarray): The image data to be reshaped, can be either a PyTorch tensor or a NumPy array.
shape (int | str, optional): Defines the layout of the grid. If an integer is provided, it specifies both the number of rows and columns. If a string is provided in the format 'nrowxncol', it parses to individual row and column numbers. Defaults to 7.
Returns:
tuple: A tuple containing:
img (np.ndarray | torch.Tensor): The image data arranged in grid format.
nrow (int): Number of rows in the grid.
ncol (int): Number of columns in the grid.
Raises:
RuntimeError: If the shape parameter is neither an int nor a string, or if it's a string that doesn't contain 'x'.
Example:
>>> tensor_images = torch.rand(64, 3, 28, 28) # Example tensor of 64 images
>>> img_grid, rows, cols = _reshape_viz_batch_img(tensor_images, '8x8')
>>> img_grid.shape
(224, 224, 3)
"""
if isinstance(shape, int):
nrow, ncol = shape, shape
elif isinstance(shape, str):
if "x" not in shape:
nrow, ncol = int(shape), int(shape)
else:
shape = shape.split("x")
nrow, ncol = int(shape[0]), int(shape[1])
else:
raise RuntimeError(f"shape {shape} not support")
if isinstance(img_data, torch.Tensor):
assert img_data.shape[1] in [1, 3]
grid_img = make_grid(img_data[: nrow * ncol].detach().cpu(), ncol)
img = grid_img.permute(1, 2, 0)
elif isinstance(img_data, np.ndarray):
if img_data.shape[1] in [1, 3]:
img = rearrange(img_data[: nrow * ncol], "(b t) c h w -> (b h) (t w) c", b=nrow)
else:
img = rearrange(img_data[: nrow * ncol], "(b t) h w c -> (b h) (t w) c", b=nrow)
return img, nrow, ncol
def show_batch_img(
img_data: torch.Tensor | np.ndarray,
shape: int | str = 7,
grid: int = 3,
is_n1p1: bool = False,
auto_n1p1: bool = True,
) -> None:
"""
Displays a batch of images using matplotlib after arranging them into a specified grid layout.
Args:
img_data (torch.Tensor | np.ndarray): The image data to be displayed.
shape (int | str, optional): The grid shape to organize the images. Defaults to 7.
grid (int, optional): Scaling factor for each image in the grid, affecting the overall size of the displayed figure. Defaults to 3.
is_n1p1 (bool, optional): Whether to normalize the images from [-1, 1] to [0, 1] for visualization. Defaults to False.
auto_n1p1 (bool, optional): If true, automatically adjusts images from [-1, 1] to [0, 1] based on minimum pixel value detection. Defaults to True.
Returns:
None: This function does not return anything but displays the image grid using matplotlib.
Example:
>>> tensor_images = torch.rand(64, 3, 28, 28) # Example tensor of 64 images
>>> show_batch_img(tensor_images, '8x8')
"""
import matplotlib.pyplot as plt
if is_n1p1:
img_data = (img_data + 1) / 2
else:
if auto_n1p1:
if isinstance(img_data, torch.Tensor):
if img_data.min().item() < -0.5:
img_data = (img_data + 1) / 2
elif isinstance(img_data, np.ndarray):
if np.min(img_data) < -0.5:
img_data = (img_data + 1) / 2
img, nrow, ncol = _reshape_viz_batch_img(img_data, shape)
plt.figure(figsize=(ncol * grid, nrow * grid))
plt.axis("off")
plt.imshow(img)
def save_batch_img(fpath: str, img_data: Union[torch.Tensor, np.ndarray], shape: Union[int, str] = 7) -> None:
"""
Saves a batch of images to a file after arranging them into a grid format. Handles both PyTorch tensors and NumPy arrays as input.
Args:
fpath (str): File path where the image will be saved.
img_data (Union[torch.Tensor, np.ndarray]): The image data to be saved. Can be a PyTorch tensor or a NumPy array.
shape (Union[int, str], optional): The grid shape to organize the images. Can be an integer specifying equal number of rows and columns, or a string specifying 'nrowxncol'. Defaults to 7.
Returns:
None: This function does not return anything but saves the image to the specified file path.
Raises:
RuntimeError: If the input shape is neither an integer nor a string, or it does not include 'x' when provided as a string.
Example:
>>> tensor_images = torch.rand(64, 3, 28, 28) # Example tensor of 64 images
>>> save_batch_img('path/to/save/image.png', tensor_images, '8x8')
# This saves the image grid to 'path/to/save/image.png'
"""
img, _, _ = _reshape_viz_batch_img(img_data, shape)
if isinstance(img, np.ndarray):
img = torch.from_numpy(img)
ndarr = img.mul(255).add_(0.5).clamp_(0, 255).to("cpu", torch.uint8).numpy()
im = Image.fromarray(ndarr)
os.makedirs(os.path.dirname(fpath), exist_ok=True)
im.save(fpath)
|