from .imports import import_torch
import os, numpy
import warnings
torch = import_torch()
[docs]
@torch.no_grad()
def gridify(
tensor: torch.Tensor,
max_width: int | None = None,
columns: int | None = None,
padding: int = 3,
pad_value: float = 0.0,
) -> torch.Tensor:
"""
Makes a grid of images/videos from a batch of images.
Like torchvision's make_grid, but more flexible.
Accepts (B,\\*,H,W)
Args:
tensor : (B,\\*,H,W) tensor
max_width : max width of the output grid. Resizes images to fit the width
columns : number of columns of the grid. If None, uses 8 or less
padding : padding to add to the images
pad_value : color of the padding
Returns:
(\\*,H',W') tensor, representing the grid of images/videos
"""
B, H, W = tensor.shape[0], tensor.shape[-2], tensor.shape[-1]
device = tensor.device
if columns is not None:
numCol = columns
else:
numCol = min(8, B)
black_cols = (-B) % numCol
tensor = torch.cat(
[tensor, torch.zeros(black_cols, *tensor.shape[1:], device=device)], dim=0
) # (B',*,H,W)
tensor = torch.nn.functional.pad(
tensor, (padding, padding, padding, padding), mode="constant", value=pad_value
) # (B',*,H+padding*2,W+padding*2)
B, H, W = tensor.shape[0], tensor.shape[-2], tensor.shape[-1]
rest_dim = tensor.shape[1:-2]
rest_dim_prod = 1
for dim in rest_dim:
rest_dim_prod *= dim
if max_width is not None:
resize_ratio = max_width / (W * numCol)
if resize_ratio < 1:
indiv_tens_size = int(H * resize_ratio), int(W * resize_ratio)
tensor = tensor.reshape((B, rest_dim_prod, H, W))
tensor = torch.nn.functional.interpolate(
tensor, size=indiv_tens_size, mode="bilinear", align_corners=False, antialias=True
) # (B',rest_dim_prod,H',W')
B, H, W = tensor.shape[0], tensor.shape[-2], tensor.shape[-1]
assert B % numCol == 0
numRows = B // numCol
tensor = tensor.reshape((numRows, numCol, rest_dim_prod, H, W)) # (numRows,numCol,rest_dim_prod,H',W')
tensor = torch.einsum("nmrhw->rnhmw", tensor) # (rest_prod,numRows,H',numCol,W')
tensor = tensor.reshape((rest_dim_prod, numRows * H, numCol * W)) # (rest_prod,numRows*H,numCol*W)
tensor = tensor.reshape((*rest_dim, numRows * H, numCol * W)) # (*,numRows*H,numCol*W)
return tensor
@torch.no_grad()
def _create_folder(folder: str, create_folder: bool = True):
if create_folder:
os.makedirs(folder, exist_ok=True)
else:
if not (os.path.exists(folder)):
raise FileNotFoundError(f"Folder {folder} does not exist !")
@torch.no_grad()
def _to_tensor(tensor: torch.Tensor | numpy.ndarray) -> torch.Tensor:
"""
Casts numpy arrays to torch tensors, leaving torch tensors untouched.
The dtype is preserved, so integer arrays stay integer.
Args:
tensor : (\\*) tensor or numpy array
Returns:
(\\*) tensor
"""
if not isinstance(tensor, numpy.ndarray):
return tensor
# torch cannot share memory with negatively-strided nor read-only arrays
array = numpy.ascontiguousarray(tensor)
if not array.flags.writeable:
array = array.copy()
return torch.from_numpy(array)
@torch.no_grad()
def _to_float(tensor: torch.Tensor) -> torch.Tensor:
"""
Casts a tensor to float, in the [0,1] range. uint8 tensors are taken to hold 0-255
values, and bool tensors to be masks. Float tensors are left untouched, use rescale
and clamp_range to control their range.
Any other dtype is rejected : an integer tensor has no unambiguous mapping to [0,1].
Args:
tensor : (\\*) tensor
Returns:
(\\*) float tensor
"""
if tensor.dtype == torch.uint8:
return tensor.float() / 255.0
if tensor.dtype == torch.bool:
return tensor.float()
if not tensor.is_floating_point():
raise TypeError(
f"Expected a float tensor with values in [0,1], a uint8 tensor or a bool tensor, "
f"but got dtype {tensor.dtype} ! Cast it yourself, or use rescale=True to min-max scale it."
)
return tensor
@torch.no_grad()
def _format_image(
tensor: torch.Tensor | numpy.ndarray,
columns: int = None,
rescale: bool = False,
clamp_range: tuple[float] = (0.0, 1.0),
max_width: int = None,
padding: int = 3,
pad_value: float = 0.0,
):
"""
Shows tensor as an image using pyplot.
Any extra dimensions **(\\*,C,H,W)** are treated as batch dimensions.
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)
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 = _to_tensor(tensor).detach().cpu()
tensor = _to_float(tensor)
tensor = _verify_chw(tensor, channels=(1, 3, 4)) # 4 channels for RGBA images
extra_params = dict(
columns=columns,
max_width=max_width,
pad_value=pad_value,
padding=padding,
rescale=rescale,
clamp_range=clamp_range,
)
if len(tensor.shape) == 2:
# Add batch and channel dimensions
return _format_image(tensor[None, :, :], **extra_params)
elif len(tensor.shape) == 3:
# Reached (C,H,W)
tensor = _rescale_images(tensor, rescale=rescale, clamp_range=clamp_range)
return tensor
elif len(tensor.shape) == 4:
# Gridify assuming (B,C,H,W)
B = tensor.shape[0]
if columns is not None:
numCol = columns
else:
numCol = min(8, B)
tensor = _rescale_images(tensor, rescale=rescale, clamp_range=clamp_range) # (B,C,H,W), rescaled
tensor = gridify(
tensor, columns=numCol, max_width=max_width, pad_value=pad_value, padding=padding
) # (C,H',W')
return tensor
elif len(tensor.shape) > 4:
# Collapse extra dimension to batch
tensor = tensor.reshape(
(-1, tensor.shape[-3], tensor.shape[-2], tensor.shape[-1])
) # assume all batch dimensions
print("Assuming extra dimension are all batch dimensions, newshape : ", tensor.shape)
return _format_image(tensor, **extra_params)
else:
raise Exception(
f"Tensor shape should be (H,W), (C,H,W), (*,C,H,W), (H,W,C) or (*,H,W,C), C=1 or 3, but got : {tensor.shape} !"
)
@torch.no_grad()
def _rescale_images(tensor: torch.Tensor, rescale: bool, clamp_range: tuple[float] = None):
"""
Clamp to clamp_range if not None, then rescale the images to 0-1 range.
The rescaling is dynamic with min and max values for each image channel.
Args:
tensor : (B,C,H,W) or (C,H,W) tensors to display
rescale : whether to rescale the images to 0-1 range
clamp_range : (min,max), tuple of values to clamp the images to, before rescaling
"""
if clamp_range is not None:
tensor = tensor.clamp(min=clamp_range[0], max=clamp_range[1])
wasthree = False
if len(tensor.shape) == 3:
tensor = tensor[None, :, :, :]
wasthree = True
B, C, H, W = tensor.shape
tensor = tensor.reshape(B, C, H * W)
if rescale:
maxes, _ = torch.max(tensor, dim=2, keepdim=True) # (B,C,1) Max value for each image channel
mins, _ = torch.min(tensor, dim=2, keepdim=True)
tensor = (tensor - mins) / (torch.maximum(maxes - mins, torch.tensor(1e-5))) # rescale
tensor = tensor.reshape(B, C, H, W)
if wasthree:
tensor = tensor[0]
return tensor
def _verify_chw(tensor: torch.Tensor, channels: tuple[int, ...] = (1, 3)) -> torch.Tensor:
"""
Changes a (\\*,H,W,C) tensor to (\\*,C,H,W). To detect the channel position,
it simply checks if the third-from-last dimension is a valid channel count. If it
is, leaves the tensor unchanged (i.e. giving (C,H,W) returns the same)
Args:
tensor : (\\*,C,H,W) or (\\*,H,W,C) tensor
channels : valid channel counts. Images also allow 4, for RGBA
Returns:
(\\*,C,H,W) tensor
"""
if len(tensor.shape) < 3:
return tensor
is_chw = tensor.shape[-3] in channels
is_hwc = tensor.shape[-1] in channels
if not is_chw and not is_hwc:
raise ValueError(
f"Tensor shape should be (H,W), (C,H,W), (*,C,H,W), (H,W,C) or (*,H,W,C) with C in {channels}, "
f"but got : {tuple(tensor.shape)} !"
)
if is_hwc:
if not is_chw:
return torch.einsum("...hwc->...chw", tensor)
else:
# (C,H,W) is compatible, we guess that.
warnings.warn(
f"Shape {tuple(tensor.shape)} is ambiguous, both (*,C,H,W) and (*,H,W,C) are plausible."
" Assuming (*,C,H,W) ; transpose the tensor yourself to remove the ambiguity.",
stacklevel=2,
)
return tensor