Creating an image grid of tensors out of a batch using matplotlib/pytorch

user_19428268 0
Viewed 179

I am trying to create a grid of images (e.g. 3 by 3) from a batch of tensors that will be fed into a GAN through a data loader in the next step. With the below code I was able to transform the tensors into images that are displayed in a grid in the right position. The problem is, that they are all displayed in a separate grid as shown here: Figure 1 Figure 2 Figure 5. How can put them in one grid and just get one figure returned with all 9 images?? Maybe I am making it too complicated. :D In the end tensors out of the real_samples have to be transformed and put into a grid.

real_samples = next(iter(train_loader))
for i in range(9):
    plt.figure(figsize=(9, 9))
    plt.subplot(330 + i + 1)
    plt.imshow(np.transpose(vutils.make_grid(real_samples[i].to(device)
                      [:40], padding=1, normalize=True).cpu(),(1,2,0)))
    plt.show()
tensor pytorch matplotlib grid python
3 Answers

And here is how to display a variable number of wonderful CryptoPunks using matplotlib XD:

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import ImageGrid
import matplotlib.image as mpimg

row_count = 3
col_count = 3

cryptopunks = [
    mpimg.imread(f"cryptopunks/{i}.png") for i in range(row_count * col_count)
]

fig = plt.figure(figsize=(8., 8.))
grid = ImageGrid(fig, 111, nrows_ncols=(row_count, col_count), axes_pad=0.1)

for ax, im in zip(grid, cryptopunks):
    ax.imshow(im)

plt.show()

enter image description here

Please note that the code allows you to generate all the images you want, not only 3 times 3. I have a folder called cryptopunks with a lot of images called #.png (e.g., 1.png, ..., 34.png, ...). Just change the row_count and col_count variable values. For instance, for row_count=6 and col_count=8 you get: enter image description here

If your image files do not have that naming pattern above (i.e., just random names), just replace the first lines with the following ones:

import os
from pathlib import Path

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import ImageGrid
import matplotlib.image as mpimg

for root, _, filenames in os.walk("cryptopunks/"):
    cryptopunks = [
        mpimg.imread(Path(root, filename)) for filename in filenames
    ]

row_count = 3
col_count = 3

# Here same code as above.

(I have downloaded the CryptoPunks dataset from Kaggle.)

user_149528560

For a 9 image grid you probably want three rows and three columns. Probably the simplest way to do that is something like so:

import matplotlib.pyplot as plt

fig, axes = plt.subplots(nrows=3, ncols=3)
axes = axes.flatten() # Flattens the array so you can access individual axes

for ax in axes:
    # Do stuff with your individual axes here
    plt.show() # This call is here just for example, prob. better call outside of the loop

which outputs the following axes configuration with plt.tight_layout():

3x3 image grid

You might also be interested in the matplotlib mosaic functionality or gridspec one. Hope this helps.

EDIT: Here's a solution which annotates each plot with its number so you can see what goes where as well:

import matplotlib.pyplot as plt

fig, axes = plt.subplots(nrows=3, ncols=3)
axes = axes.flatten()

for idx, ax in enumerate(axes):
    ax.annotate(f"{idx+1}", xy=(0.5, 0.5), xytext=(0.5, 0.5))

with annotations

user_108477550

Matplotlib provides a function called subplot, I think this is what you are searching for!

plt.subplot(9,1) is the syntax I guess.

And then configure your plots

user_194168780
Related