博客

背景

我们在做图像分类或者聚类的时候，通常需要将结果可视化。这时候的结果就是一个类别有很多张图片，像下图这样，本文就介绍如何使用matplotlib将聚类结果可视化。

如何安装matplotlib

使用如下pip命令

pip install matplotlib

如何加载并展示图片

可以使用PIL库将数据集中的图像进行加载，也可以使用pytorch中常用的datasets.ImageFolder函数加载图像。

使用PIL库

from PIL import image
img = Image.open(r"./jupyter/matplotlib/images/1.jpg")
fig = plt.figure(figsize=(8, 4))
ax1 = fig.add_subplot(121)
ax1.imshow(img)

使用datasets.ImageFolder库

from torchvision import datasets
train_dataset = datasets.ImageFolder(root='data2', transform=transform)

注意，这里的data2是一个文件夹，里边保存要进行聚类的照片文件夹，如data2/dog/1.jpg, data2/dog/2.jpg, data2/cat/1.jpg, data2/cat/2.jpg… 因此，使用ImageFolder一般读取的是一个文件夹中的所有图片，需要使用for循环等方式对图片展示，后文会有说明，请继续阅读。

如何使用matplotlib展示多个图片

接着上文，现在你使用datasets.ImageFolder的方式加载了一批图片，并对其使用k-means算法进行聚类，得到一个labels数组。labels数组中记录其中图片对应的标签，如labels = [1, 0, 1, 0, 1]，现在我们需要将结果可视化出来

index = 0
n = len(train_dataset)
lt = [i for i, x in enumerate(labels) if x == result]
rows = n // 5 + 1
fig = plt.figure(figsize=(20, 30), dpi=150)
    for i in lt:
        image_name = dataset[i][0]
        img = Image.open(image_name)
        fig.add_subplot(rows, 5, index)
        plt.imshow(img, interpolation='none')
        plt.tight_layout()
        plt.axis('off')
        index += 1
plt.show()

最后的效果如下: