Plot multiple images by concatenate them in space.
(
imgs: T.Union[torch.Tensor, np.ndarray],
dpi=150,
mode='tile', # 'horizontal', 'vertical', 'tile'
fig=None,
ax=None,
colorbar=True,
valrange=None, # (min, max)
ncols: int = 6,
background_color: float = 0.,
)
| 1759 | |
| 1760 | |
| 1761 | def plot_multiple_images( |
| 1762 | imgs: T.Union[torch.Tensor, np.ndarray], |
| 1763 | dpi=150, |
| 1764 | mode='tile', # 'horizontal', 'vertical', 'tile' |
| 1765 | fig=None, |
| 1766 | ax=None, |
| 1767 | colorbar=True, |
| 1768 | valrange=None, # (min, max) |
| 1769 | ncols: int = 6, |
| 1770 | background_color: float = 0., |
| 1771 | ): |
| 1772 | """Plot multiple images by concatenate them in space. """ |
| 1773 | # imgs: (b, h, w, *) |
| 1774 | |
| 1775 | if isinstance(imgs, np.ndarray): |
| 1776 | imgs = torch.from_numpy(imgs) |
| 1777 | |
| 1778 | mask = torch.logical_not(imgs.isfinite()) |
| 1779 | imgs = imgs.masked_fill(mask, 0) |
| 1780 | |
| 1781 | imgs_list = torch.chunk(imgs, chunks=imgs.shape[0], dim=0) |
| 1782 | imgs_list = [img[0] for img in imgs_list] # list of (h, w, *) |
| 1783 | |
| 1784 | if mode == 'horizontal': |
| 1785 | imgs = torch.cat(imgs_list, dim=1) |
| 1786 | elif mode == 'vertical': |
| 1787 | imgs = torch.cat(imgs_list, dim=0) |
| 1788 | elif mode == 'tile': |
| 1789 | assert len(imgs_list) > 0 |
| 1790 | if imgs_list[0].ndim == 2: |
| 1791 | squeeze_last_dim = True |
| 1792 | imgs_list = [img.unsqueeze(-1) for img in imgs_list] # list of (h, w, c) |
| 1793 | else: |
| 1794 | squeeze_last_dim = False |
| 1795 | imgs = render.tile_images( |
| 1796 | images=imgs_list, |
| 1797 | ncols=ncols, |
| 1798 | background_color=background_color, |
| 1799 | ) |
| 1800 | if squeeze_last_dim: |
| 1801 | imgs = imgs.squeeze(-1) |
| 1802 | else: |
| 1803 | raise NotImplementedError |
| 1804 | |
| 1805 | if valrange is not None: |
| 1806 | imgs = (imgs - valrange[0]) / (valrange[1] - valrange[0]) |
| 1807 | |
| 1808 | imgs = imgs.detach().cpu().numpy() |
| 1809 | |
| 1810 | # plot |
| 1811 | fig, axes = imagesc( |
| 1812 | arr=imgs, |
| 1813 | fig=fig, |
| 1814 | axes=ax, |
| 1815 | dpi=dpi, |
| 1816 | colorbar=colorbar, |
| 1817 | ) |
| 1818 | plt.axis('scaled') |
nothing calls this directly
no test coverage detected