from typing import Union, Optional, List, Tuple, Text, BinaryIO import pathlib import torch import math import warnings import numpy as np from PIL import Image, ImageDraw, ImageFont, ImageColor __all__ = ["make_grid", "save_image", "draw_bounding_boxes"] @torch.no_grad() def make_grid( tensor: Union[torch.Tensor, List[torch.Tensor]], nrow: int = 8, padding: int = 2, normalize: bool = False, value_range: Optional[Tuple[int, int]] = None, scale_each: bool = False, pad_value: int = 0, **kwargs ) -> torch.Tensor: """Make a grid of images. Args: tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W) or a list of images all of the same size. nrow (int, optional): Number of images displayed in each row of the grid. The final grid size is ``(B / nrow, nrow)``. Default: ``8``. padding (int, optional): amount of padding. Default: ``2``. normalize (bool, optional): If True, shift the image to the range (0, 1), by the min and max values specified by :attr:`range`. Default: ``False``. value_range (tuple, optional): tuple (min, max) where min and max are numbers, then these numbers are used to normalize the image. By default, min and max are computed from the tensor. scale_each (bool, optional): If ``True``, scale each image in the batch of images separately rather than the (min, max) over all images. Default: ``False``. pad_value (float, optional): Value for the padded pixels. Default: ``0``. Example: See this notebook `here `_ """ if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))): raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}') if "range" in kwargs.keys(): warning = "range will be deprecated, please use value_range instead." warnings.warn(warning) value_range = kwargs["range"] # if list of tensors, convert to a 4D mini-batch Tensor if isinstance(tensor, list): tensor = torch.stack(tensor, dim=0) if tensor.dim() == 2: # single image H x W tensor = tensor.unsqueeze(0) if tensor.dim() == 3: # single image if tensor.size(0) == 1: # if single-channel, convert to 3-channel tensor = torch.cat((tensor, tensor, tensor), 0) tensor = tensor.unsqueeze(0) if tensor.dim() == 4 and tensor.size(1) == 1: # single-channel images tensor = torch.cat((tensor, tensor, tensor), 1) if normalize is True: tensor = tensor.clone() # avoid modifying tensor in-place if value_range is not None: assert isinstance(value_range, tuple), \ "value_range has to be a tuple (min, max) if specified. min and max are numbers" def norm_ip(img, low, high): img.clamp_(min=low, max=high) img.sub_(low).div_(max(high - low, 1e-5)) def norm_range(t, value_range): if value_range is not None: norm_ip(t, value_range[0], value_range[1]) else: norm_ip(t, float(t.min()), float(t.max())) if scale_each is True: for t in tensor: # loop over mini-batch dimension norm_range(t, value_range) else: norm_range(tensor, value_range) if tensor.size(0) == 1: return tensor.squeeze(0) # make the mini-batch of images into a grid nmaps = tensor.size(0) xmaps = min(nrow, nmaps) ymaps = int(math.ceil(float(nmaps) / xmaps)) height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding) num_channels = tensor.size(1) grid = tensor.new_full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value) k = 0 for y in range(ymaps): for x in range(xmaps): if k >= nmaps: break # Tensor.copy_() is a valid method but seems to be missing from the stubs # https://pytorch.org/docs/stable/tensors.html#torch.Tensor.copy_ grid.narrow(1, y * height + padding, height - padding).narrow( # type: ignore[attr-defined] 2, x * width + padding, width - padding ).copy_(tensor[k]) k = k + 1 return grid @torch.no_grad() def save_image( tensor: Union[torch.Tensor, List[torch.Tensor]], fp: Union[Text, pathlib.Path, BinaryIO], format: Optional[str] = None, **kwargs ) -> None: """Save a given Tensor into an image file. Args: tensor (Tensor or list): Image to be saved. If given a mini-batch tensor, saves the tensor as a grid of images by calling ``make_grid``. fp (string or file object): A filename or a file object format(Optional): If omitted, the format to use is determined from the filename extension. If a file object was used instead of a filename, this parameter should always be used. **kwargs: Other arguments are documented in ``make_grid``. """ grid = make_grid(tensor, **kwargs) # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy() im = Image.fromarray(ndarr) im.save(fp, format=format) @torch.no_grad() def draw_bounding_boxes( image: torch.Tensor, boxes: torch.Tensor, labels: Optional[List[str]] = None, colors: Optional[List[Union[str, Tuple[int, int, int]]]] = None, fill: Optional[bool] = False, width: int = 1, font: Optional[str] = None, font_size: int = 10 ) -> torch.Tensor: """ Draws bounding boxes on given image. The values of the input image should be uint8 between 0 and 255. If filled, Resulting Tensor should be saved as PNG image. Args: image (Tensor): Tensor of shape (C x H x W) boxes (Tensor): Tensor of size (N, 4) containing bounding boxes in (xmin, ymin, xmax, ymax) format. Note that the boxes are absolute coordinates with respect to the image. In other words: `0 <= xmin < xmax < W` and `0 <= ymin < ymax < H`. labels (List[str]): List containing the labels of bounding boxes. colors (List[Union[str, Tuple[int, int, int]]]): List containing the colors of bounding boxes. The colors can be represented as `str` or `Tuple[int, int, int]`. fill (bool): If `True` fills the bounding box with specified color. width (int): Width of bounding box. font (str): A filename containing a TrueType font. If the file is not found in this filename, the loader may also search in other directories, such as the `fonts/` directory on Windows or `/Library/Fonts/`, `/System/Library/Fonts/` and `~/Library/Fonts/` on macOS. font_size (int): The requested font size in points. """ if not isinstance(image, torch.Tensor): raise TypeError(f"Tensor expected, got {type(image)}") elif image.dtype != torch.uint8: raise ValueError(f"Tensor uint8 expected, got {image.dtype}") elif image.dim() != 3: raise ValueError("Pass individual images, not batches") ndarr = image.permute(1, 2, 0).numpy() img_to_draw = Image.fromarray(ndarr) img_boxes = boxes.to(torch.int64).tolist() if fill: draw = ImageDraw.Draw(img_to_draw, "RGBA") else: draw = ImageDraw.Draw(img_to_draw) txt_font = ImageFont.load_default() if font is None else ImageFont.truetype(font=font, size=font_size) for i, bbox in enumerate(img_boxes): if colors is None: color = None else: color = colors[i] if fill: if color is None: fill_color = (255, 255, 255, 100) elif isinstance(color, str): # This will automatically raise Error if rgb cannot be parsed. fill_color = ImageColor.getrgb(color) + (100,) elif isinstance(color, tuple): fill_color = color + (100,) draw.rectangle(bbox, width=width, outline=color, fill=fill_color) else: draw.rectangle(bbox, width=width, outline=color) if labels is not None: draw.text((bbox[0], bbox[1]), labels[i], fill=color, font=txt_font) return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1)