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

1. 前言

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

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

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

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

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

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

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

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

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

2.1 代码介绍

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

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

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

2.2 可视化结果

可视化结果

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

2.3 完整代码

如下:

python 复制代码
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文件

可以通过下面代码生成

python 复制代码
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)

结果如下:


或者直接新建json文件,然后对照目录按照上面的方式输入也行

相关推荐
懂AI的老郑1 天前
智能体演变路径预测:AI未来行为的核心解码
人工智能
wayz111 天前
Day 7:第一周复习与模型综合比较
人工智能·算法·机器学习·量化交易
小超同学你好1 天前
Transformer 27. Vision Transformer(ViT):把图像当作「词序列」的编码器
人工智能·深度学习·transformer
登山人在路上1 天前
RAGAS 中的已知偏差详解
人工智能·全文检索·ai编程
空白诗1 天前
ModelEngine Nexent:一句话创建智能体的全新范式
人工智能·nexent
元直数字电路验证1 天前
AI入门体系化知识框架
人工智能
zhangshuang-peta1 天前
MCP 与治理:当 Agent 成为组织决策的一部分
大数据·人工智能·ai agent·mcp·peta
SCBAiotAigc1 天前
2026.4.21:在做yolo26分类任务时出现save_dir一直是个固定的一个值的避坑技巧
人工智能·具身智能·yolo26
原我归来是少年1 天前
一 InsightMemory - 从 RAG 到 Memory:AI 应用缺的不是搜索,而是持续认知
人工智能
字节架构前端1 天前
Life of a Token:像调试代码一样看懂大模型如何生成 Token
人工智能·aigc·openai