图像分类任务的可视化脚本，生成类别json字典文件

news2024/11/19 5:56:29

1. 前言

之前的图像分类任务可视化，都是在train脚本里，用torch中dataloader将图片和类别加载，然后利用matplotlib库进行可视化。

如这篇文章中：CNN 卷积神经网络对染色血液细胞分类(blood-cells)

在分类任务中，必定经历过图像预处理，缩放啊、随即裁剪啊之类的，可视化效果不太明显

本章将从数据角度出发，直接根据数据目录将图像可视化，随机展示所有图片的四张图片，可视化后并且保存

目标检测的可视化可以参考：

关于目标检测任务中，YOLO(txt格式)标注文件的可视化

关于目标检测任务中，XML(voc格式)标注文件的可视化

2. 根据目录可视化（无需类别的json文件）

目录如下：代码应该data同一路径

2.1 代码介绍

root 传入的是文件夹路径，也就是多个类别文件夹的上一级目录

将所有图像保存，为了知道图片的类别，需要把图片的父目录保存。为了方便，这里生成一个列表文件，key 是目录类别，value 是相应的图像路径

展示的代码很简单，生成随机数，将列表的文件提取出来，然后展示四张就行了

2.2 可视化结果

可视化结果

代码会在当前目录生成刚刚可视化展示的图片

2.3 完整代码

如下：

import os
import matplotlib.pyplot as plt
import random
from PIL import Image


def main(path):
    classes = [i for i in os.listdir(path)]         # ['cat', 'dog']

    # 将所有图片按照 类别:路径 字典形式保存
    images_path = []  # [{'cat': './data/train\\cat\\Baidu_0000.jpeg'}, {'cat': './data/train\\cat\\Baidu_0002.jpeg'}]
    for cla in classes:
        for i in os.listdir(os.path.join(path,cla)):
            dic = {}  # 类别:图像路径
            img_path = os.path.join(path,cla,i)
            dic[cla] = img_path             # {'cat': './data/train\\cat\\Baidu_0000.jpeg'}
            images_path.append(dic)

    # 随机展示4张图像
    plt.figure(figsize=(12,8))
    for i in range(4):
        r = random.randint(0,len(images_path)-1)      # 生成随机数
        label,im_path= list(images_path[r].keys())[0],list(images_path[r].values())[0]
        #  cat , ./data/train\cat\Baidu_0049.jpeg

        im = Image.open(im_path)
        plt.subplot(2,2,i+1)
        plt.title(label)
        plt.imshow(im)

    plt.savefig('show.png')     # 保存图片
    plt.show()


if __name__ == '__main__':
    root = './data/train'       # 传入目录
    main(path=root)

3.生成类别json字典文件

图像分类任务，有的没有提供类别的字典文件，这里也记录一下如何生成json文件

可以通过下面代码生成

import os
import json


def main(path):
    classes = [i for i in os.listdir(path)]  # ['cat', 'dog']

    labels = {}         # 类别的字典文件
    for index,name in enumerate(classes):
        labels[index] = name
    labels = json.dumps(labels,indent=4)

    with open('./class_indices.json','w') as f:         # 保存成json文件
        f.write(labels)


if __name__ == '__main__':
    root = './data/train'  # 传入目录
    main(path=root)

结果如下：