图像分类任务的可视化脚本,生成类别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文件,然后对照目录按照上面的方式输入也行

相关推荐
*Lisen13 小时前
从零手写 FlashAttention(PyTorch实现 + 原理推导)
人工智能·pytorch·python
字节跳动数据库13 小时前
数据孤岛难打通、权限怕失控?DBW 助“小龙虾”落地最后一公里
人工智能
俊哥V13 小时前
AI一周事件 · 2026-04-22 至 2026-04-28
人工智能·ai
Black蜡笔小新14 小时前
AI大模型训练工作站/私有化本地化AI模型训推工作站DLTM为农业生产装上AI“慧眼”
人工智能·ai大模型
小星AI14 小时前
Claude Code Agent SDK 从入门到精通,一步到位
人工智能·agent·cursor
端平入洛14 小时前
梯度是什么:PyTorch 自动求导详解
人工智能·深度学习
时序之心14 小时前
上海交大、东北大学:时序分类与感知领域的两项前沿突破
人工智能·分类·时间序列
小五兄弟14 小时前
YouTube 肖像检测扩展背后:短剧出海版权保护的技术实现与实战策略
大数据·人工智能
阿聪谈架构14 小时前
第08章:MCP 模型上下文协议(下)
人工智能·后端
大囚长14 小时前
AI意识思想实验
人工智能