Tile a list of images as a matrix (along h and w dimensions). Args: images: list of (*, h, w, c) ncols: number of images in a column. If -1, unlimited. background_color: If None, the image will be replaced by (h,w,c) * background
(
images: T.Union[T.List[torch.Tensor], T.List[np.ndarray]],
ncols: int = -1,
background_color: T.Union[float, T.List[float]] = 0., # [0, 1]
)
| 658 | |
| 659 | |
| 660 | def tile_images( |
| 661 | images: T.Union[T.List[torch.Tensor], T.List[np.ndarray]], |
| 662 | ncols: int = -1, |
| 663 | background_color: T.Union[float, T.List[float]] = 0., # [0, 1] |
| 664 | ) -> T.Union[torch.Tensor, np.ndarray]: |
| 665 | """ |
| 666 | Tile a list of images as a matrix (along h and w dimensions). |
| 667 | |
| 668 | Args: |
| 669 | images: |
| 670 | list of (*, h, w, c) |
| 671 | ncols: |
| 672 | number of images in a column. If -1, unlimited. |
| 673 | background_color: |
| 674 | If None, the image will be replaced by (h,w,c) * background |
| 675 | |
| 676 | Returns: |
| 677 | tiled image: |
| 678 | (*, h', w', c) |
| 679 | """ |
| 680 | total = len(images) |
| 681 | if ncols < 0: |
| 682 | ncols = total |
| 683 | |
| 684 | # find the first non-None image |
| 685 | img = None |
| 686 | for i in range(total): |
| 687 | if images[i] is not None: |
| 688 | img = images[i] |
| 689 | break |
| 690 | if img is None: |
| 691 | raise RuntimeError |
| 692 | |
| 693 | if isinstance(img, np.ndarray): |
| 694 | is_numpy = True |
| 695 | images = [ |
| 696 | torch.from_numpy(img) |
| 697 | if img is not None else None |
| 698 | for img in images |
| 699 | ] |
| 700 | else: |
| 701 | is_numpy = False |
| 702 | |
| 703 | *b_shape, h, w, c = img.shape |
| 704 | if isinstance(background_color, (int, float)): |
| 705 | background_color = [background_color] * c |
| 706 | |
| 707 | blank = torch.ones(*b_shape, h, w, c) |
| 708 | for ic in range(c): |
| 709 | blank[..., ic] = background_color[ic] |
| 710 | |
| 711 | nrows = math.ceil(total / ncols) |
| 712 | if nrows == 1: |
| 713 | ncols = total |
| 714 | rows = [] |
| 715 | for _ in range(nrows): |
| 716 | rows.append([blank] * ncols) |
| 717 |