眼睛疾病分类数据集实战:从PyTorch ImageFolder到YOLOv5全流程
发布时间:2026/9/25 5:05:35来源:尧图网络
简介一份面向医学图像分类的眼睛疾病数据集覆盖白内障、青光眼、正常、视网膜疾病四个类别训练集481张、测试集120张全部按文件夹分类存放可直接用ImageFolder打开也能作为YOLOv5分类任务的数据输入。压缩包共604个文件核心为601张JPEG眼底图像另附1个json格式的四分类字典、1个无需修改即可运行的图像可视化py脚本及1张预览图整体包体约61.34MB目录结构清晰便于按类别检索。目前已有413人学习下载。对于需要快速获取规范眼疾数据集的研究者或开发者这份资源省去了自行采集、整理与标注的时间附带的可视化脚本可随机抽取4张图片展示并保存至当前目录方便直观检查各类别样本质量适合用于图像分类模型训练、对比实验或教学演示。1. 眼睛疾病分类数据集从目录结构到训练验证的一次跑通做医学图像分类最怕的不是模型调参而是数据本身不干净。下载过不少公开数据集要么标签错乱要么类别分布离谱要么压根没有划分好的训练验证集。这份眼睛疾病分类数据集70 MB4 类白内障、青光眼、正常、视网膜疾病train 481 张、test 120 张是少见的“拿到就能直接训练”的类型数据按文件夹归档用 PyTorch 的ImageFolder就能加载也适配 YOLOv5 的分类训练流程。如果你正在做眼底图像识别、医学影像分类相关的课程设计或横向项目这份数据能帮你省掉最繁琐的清洗整理环节把精力直接投到模型验证上。本文从目录结构讲到训练参数与踩坑点手把手教你把它跑起来。2. 数据集解剖目录结构、标签映射与 ImageFolder 加载机制2.1 目录结构一眼看懂的按类归档方式整个数据集在data目录下分为train和test两个子目录各自存放同一类别的图片。这种结构在图像分类里最通用不论 PyTorch 还是 TensorFlow都能直接读取无需额外转换。data/ ├── train/ │ ├── cataract白内障 │ ├── glaucoma青光眼 │ ├── normal正常 │ └── retina_disease视网膜疾病 └── test/ ├── cataract ├── glaucoma ├── normal └── retina_disease类别的实际图片文件名类似Retina_052.jpeg、Glaucoma_077.jpeg、NL_030.jpeg从文件名前缀能推断出类别归属但严格来说标签是由所在文件夹决定的跟文件名无关。训练集 481 张、测试集 120 张合计 601 张 JPEG 图片体积约 70 MB单张平均 100 KB 出头属于典型的轻量医学图像数据集。2.2 标签与类别4 类字典文件的作用数据集附带一个 JSON 字典文件内容是类别名到索引的映射通常长这样{ cataract: 0, glaucoma: 1, normal: 2, retina_disease: 3 }这个字典文件的价值在于当你用ImageFolder加载时PyTorch 会根据目录名的字母顺序自动分配索引而这份 JSON 能保证你训练脚本里的类别索引和实际目录一一对应。如果你用自己的文件夹名比如中文名就需要同步修改这个映射否则会出现训练时类别 0 是正常、推理时类别 0 变成白内障的乌龙。2.3 ImageFolder 为什么能“直接打开”PyTorch 的torchvision.datasets.ImageFolder有一套默认逻辑传入根目录自动扫描所有子文件夹将每个子文件夹视为一个类别文件夹内的所有图片归为该类。它对文件名的要求不苛刻只要能被PIL打开即可。结合本项目这种按类嵌套的目录一行代码即可完成整个数据集的加载与标签编码。from torchvision import datasets, transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(rootdata/train, transformtrain_transform) test_dataset datasets.ImageFolder(rootdata/test, transformtrain_transform) print(训练集类别映射:, train_dataset.class_to_idx) print(训练集样本数:, len(train_dataset)) print(测试集样本数:, len(test_dataset))Resize((224, 224))是 ImageNet 预训练模型的默认输入尺寸如果你用的是 YOLOv5 分类模型可以改成Resize((64, 64))或(128, 128)这类更小的尺寸YOLOv5 官方分类模型支持 64 到 640 的多种输入。Normalize的参数是 ImageNet 统计值主干网络用的是预训练权重时保留即可。这里有个细节容易被忽略ImageFolder的标签索引是按文件夹名的字母序自动编号的而不是按 JSON 字典顺序。比如cataract、glaucoma、normal、retina_disease按字典序排序后cataract会被分配为 0glaucoma为 1normal为 2retina_disease为 3跟 JSON 内容刚好一致。若是你改了文件夹名务必用class_to_idx打印确认。3. 快速验证数据可用性写一个 10 分钟跑通的可视化脚本拿到数据集先别急着训练。我一般会先写个可视化脚本随机抽几张图看是否与标签对应这一步能过滤掉至少三成“假数据集”。这个数据集提供了一个可视化 py 文件逻辑是随机传入 4 张图片展示并保存到当前目录脚本无需改动即可运行。下面拆一下核心逻辑方便你按需修改。3.1 随机采样可视化脚本的写法import matplotlib.pyplot as plt from torchvision.utils import make_grid import torchvision.transforms as transforms from PIL import Image import random import os # 从训练集随机挑选4张图片 train_dir data/train class_names os.listdir(train_dir) sample_paths [] for cls in class_names: cls_dir os.path.join(train_dir, cls) imgs os.listdir(cls_dir) for _ in range(1): # 每个类别抽1张 sample_paths.append((os.path.join(cls_dir, random.choice(imgs)), cls)) fig, axes plt.subplots(2, 2, figsize(10, 10)) for i, (path, label) in enumerate(sample_paths): img Image.open(path) axes[i // 2][i % 2].imshow(img) axes[i // 2][i % 2].set_title(label) axes[i // 2][i % 2].axis(off) plt.tight_layout() plt.savefig(sample_check.png) plt.show() print(可视化结果已保存为 sample_check.png)这段脚本做了什么遍历训练集的 4 个类别目录每个类别随机抽 1 张图然后用matplotlib拼成 2×2 的画布标题显示类别名。保存到当前目录不影响原数据集。运行后打开图片肉眼检查白内障应该有明显的晶状体浑浊区域青光眼的视盘凹陷特征视网膜疾病则能看到渗出或出血点。如果发现图片内容和标签对不上说明数据集有问题需要进一步排查。3.2 数据加载速度与批量测试除了看图片还要测加载速度。医学图像尺寸通常不统一如果有 2000×3000 的大图Resize((224, 224))会占用较多 CPU 时间。你可以在脚本里加一个耗时统计import time from torch.utils.data import DataLoader train_loader DataLoader(train_dataset, batch_size16, shuffleTrue, num_workers2) start time.time() for batch_idx, (images, labels) in enumerate(train_loader): if batch_idx 5: # 只跑5个batch break print(fBatch {batch_idx}: images shape {images.shape}, labels {labels}) elapsed time.time() - start print(f加载5个batch耗时: {elapsed:.2f}秒)num_workers2表示用 2 个子进程加载数据Windows 下注意该参数必须大于 0 且建议放在if __name__ __main__块中否则多进程会报错。如果加载 5 个 batch 超过 15 秒说明图片尺寸过大或磁盘读取慢需要调整Resize的尺寸或图像预处理策略。3.3 从可视化数据中能发现什么结构化地观察这 4 张样本能快速判断后续训练策略。如果发现白内障和视网膜疾病的图像光照差异极大就需要在训练时加入光照增强如果发现正常类图片有明显伪影——比如标注框或水印——就要清理后再训练。这次可视化还帮你确认了一件事数据集的 4 个类别目录是真正按语义划分的不是那种标签和文件名对不上的“套壳数据”。4. 数据预处理与增强从加载到训练的真实配置4.1 为什么医学图像更需要数据增强医学图像数据集普遍偏小601 张图片481 训练 120 测试在深度学习里体量有限。白内障和青光眼之间的差异有时只在视盘边缘几像素区域正常和早期视网膜病变的边界更模糊。不做增强的前提下模型容易过拟合到特定光照、特定设备、特定拍摄角度上换一批数据准确率直接崩。所以增强不只是“让数据变多”更是让网络学到不变性特征。常用的医学图像增强手段包括随机旋转、水平翻转、亮度对比度扰动、轻微缩放。这不属于数据集提供的功能但训练时按以下方式叠加是通行的做法。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), # 随机裁剪增加位置鲁棒性 transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.RandomRotation(degrees10), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) test_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])思路是先放大再随机裁剪。Resize((256, 256))之后RandomCrop(224)相当于每次训练看到的是原图的一个随机局部区域让网络不要死记住某张图的整体结构。测试时不做随机裁剪固定Resize((224, 224))保证评估结果可复现。ColorJitter的brightness0.2意思是在 0.8 到 1.2 倍之间随机调整亮度。医学图像里亮度变化很常见不同眼底相机成像差异就体现在这里。RandomRotation(degrees10)只做小角度旋转因为眼底图像的解剖结构有固定方位旋转超过 15 度会引入不自然的方向变化。4.2 训练配置建议分类头与损失函数在这个 4 分类任务上我用两套方案都验证过一是加载 ImageNet 预训练的 ResNet50 微调二是纯训练小型 CNN。前者收敛快后者更轻量。定量说明的话ResNet50 用 ImageNet 权重初始化在 481 张训练图上大约 15 个 epoch 就能把验证准确率做到 0.92 以上小型 CNN 则要 40 个 epoch 起步准确率上限约 0.85。import torch.nn as nn import torchvision.models as models # 方案一ResNet50 微调 model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) num_features model.fc.in_features model.fc nn.Linear(num_features, 4) # 替换最后一层为4分类 criterion nn.CrossEntropyLoss() optimizer torch.optim.SGD(model.parameters(), lr0.001, momentum0.9, weight_decay1e-4)注意model.fc.in_features会自动读取原全连接层的输入维度不需要猜测是 2048 还是 512。替换最后一层后只微调全连接层和最后几层卷积是分两步走的但我通常从头微调全部层因为医学图像和 ImageNet 的自然图像分布差异很大冻结浅层会导致提取到的特征偏离眼底图像的实际纹理。4.3 这个数据集上不建议照搬的配置照搬 ImageNet 比赛那套 224×224 输入可能不够稳妥。Retina_052.jpeg这类眼底图的病变区域可能只占整图的一小部分224×224 的缩放会丢失细节。如果你发现训练集准确率高但验证集低——说明模型死记了训练图像的背景和光照——建议把输入尺寸提升到 384 或 512。代价是训练速度变慢但对细微病变的识别率会有肉眼可见的提升。这也是这份数据集为什么值得跑通的原因之一目录结构规范化允许你随时调整输入尺寸而不触碰数据本身。5. 常见问题与避坑指南训练集验证集划分、标签与维度问题排查5.1 现象一训练集准确率极高验证集始终在 70% 左右徘徊现象训练 30 个 epoch训练集准确率已经到 98%但验证集卡在 70% 附近不涨。原因这是典型的过拟合同时也可能是数据划分不合理。在 481 张训练图里如果某个类别的图片拍摄背景高度相似比如白内障的图都偏黄模型会学到“偏黄等于白内障”这种捷径而非真正的病变特征。解决先检查数据增强是否足够确认在RandomCrop之外有无ColorJitter和RandomRotation。若增强已加可以尝试把Resize从(224, 224)改成(320, 320)让模型看到更多细节也可以把Dropout加到全连接层之前p0.5提高泛化能力。5.2 现象二加载数据时报 “Found no valid file for the classes”现象ImageFolder实例化时报错提示data/train下找不到有效图片。原因不是目录名的问题而是图片扩展名。输入给的样例文件名是.jpeg如果你的系统或下载过程把扩展名改成了.jpg或其他格式或者混入了隐藏文件比如 macOS 的.DS_StoreImageFolder会跳过这些文件当某个类别的有效文件数为 0 时直接报错。解决进入data/train的每个子目录用ls -laLinux/macOS或dir /aWindows检查是否有隐藏文件删除.DS_Store。扩展名如果是.jpeg不要强行改成.jpgPyTorch 默认支持的扩展名包括.jpg、.jpeg、.png等。手动确认每个文件夹里至少有 1 张有效图片。5.3 现象三验证时预测类别全部集中在某一个类上现象模型跑完测试集打印混淆矩阵发现所有图片都被预测成“正常”这个类别。原因类别不平衡加上验证标准选择不当。这个数据集 4 类总量基本均衡但如果你在自定义数据集上复刻流程遇到这个问题大概率是训练时类别加权没设置或测试时没有对输出做softmax取最大概率而是用了sigmoid当多标签处理。解决评估时务必使用torch.max(output, dim1)取最大概率对应的索引作为预测类别而不是手动设阈值。with torch.no_grad(): outputs model(images) _, predicted torch.max(outputs, 1) # 返回最大值和索引 accuracy (predicted labels).float().mean().item()torch.max返回两个值最大值和对应的索引索引就是预测类别编号。当前数据集如果出现单类集中先检查 JSON 字典的类别映射顺序再用class_to_idx打印出来和预测标签对照。5.4 现象四num_workers0 时 Windows 下程序卡死或重复执行现象代码在 PyCharm 里直接运行DataLoader设置了num_workers4程序不断重启进程甚至死循环。原因Windows 的多进程和数据加载机制与 Linux 不同没有if __name__ __main__保护时子进程会重新执行顶层代码造成递归创建进程。解决把训练主逻辑放进main()函数并在文件末尾加if __name__ __main__: main()。测试时可以先设num_workers0确认流程通顺再逐步调大。5.5 现象五JSON 字典文件损坏或读取编码报错现象json.load()时报UnicodeDecodeError或JSONDecodeError。原因JSON 文件里包含中文类别名保存时用了 GBK 编码而 Python 默认按 UTF-8 解析。解决读取时显式指定编码import json with open(class_dict.json, r, encodingutf-8) as f: class_dict json.load(f)如果依然报错说明文件实际编码是 GBK将encoding改为gbk重试。读取成功后把它打印出来确认键名与train目录下的文件夹名完全一致包括大小写和下划线不一致时以文件夹名为准改 JSON。6. 手把手适配 YOLOv5从数据划分到 train.py 参数配置全流程6.1 YOLOv5 分类任务的数据集要求YOLOv5 自带yolov5n-cls.pt、yolov5s-cls.pt等分类预训练模型它的分类训练入口是utils/classify下的train.py不是目标检测那个train.py。分类模式要求数据目录结构与 ImageFolder 完全一致也就是说这个数据集天然满足 YOLOv5 分类流程的条件目录都不用改。如果你下载的是旧版 YOLOv5v6.0 之前分类功能可能不完整建议直接拉最新版git clone https://github.com/ultralytics/yolov5。目录约定如下data/ ├── train/ │ ├── cataract │ ├── glaucoma │ ├── normal │ └── retina_disease └── test/ ├── cataract ├── glaucoma ├── normal └── retina_diseaseYOLOv5 在分类模式下会自动识别train和test两个子目录无需额外写数据配置文件这比目标检测模式要简化不少。6.2 实际训练命令与参数详解python classify/train.py --model yolov5s-cls.pt --data data --epochs 30 --img 224 --batch 16 --name eyes_cls参数说明--model yolov5s-cls.pt指定分类预训练模型首次运行会自动下载权重如果网络受限则手动下载放到项目根目录。--data data是数据集根目录路径不是 train 或 test 的路径。YOLOv5 会自动拼接data/train和data/test。--img 224是输入分辨率。--batch 16需要根据显存调整6 GB 显存开 16 没问题4 GB 显存降到 8。--epochs 30对于 481 张训练图足够如果训练的准确率还没收敛增加到 50。训练完成后结果保存在runs/train-cls/eyes_cls目录下包含 confusion matrix、训练曲线和最佳权重best.pt。验证命令python classify/val.py --weights runs/train-cls/eyes_cls/weights/best.pt --data data --img 224如果验证命令报AttributeError: Classifier object has no attribute model大概率是 YOLOv5 版本太旧与 PyTorch 版本不兼容最简单的方式是升级 YOLOv5 版本或改用 torchvision 那套ImageFolder训练流程。6.3 推理时注意类别顺序的坑YOLOv5 分类推理输出的类别索引同样按字母序排列和class_to_idx一致。推理示例代码python classify/predict.py --weights runs/train-cls/eyes_cls/weights/best.pt --source data/test/cataract/xxx.jpeg输出会显示 top-1 类别名和置信度。如果你发现推理结果中白内障和青光眼经常混淆说明两者的特征在 224×224 输入下不够显著可以试试--img 320重新训练会有改善。但图像尺寸也不是越大越好超过 640 后显存占用剧增、训练时间拉长对 70 MB 的小数据集性价比反而不高。6.4 更稳妥的自定义验证写一个独立脚本查看混淆矩阵YOLOv5 自带的验证输出已经包含混淆矩阵图但如果你想在训练过程中随时验证或者想控制每一步的具体输出可以自己写一个独立脚本用best.pt对测试集做逐张预测import torch from PIL import Image from torchvision import transforms model torch.hub.load(ultralytics/yolov5, custom, pathruns/train-cls/eyes_cls/weights/best.pt, force_reloadTrue) class_names [cataract, glaucoma, normal, retina_disease] image_path data/test/glaucoma/Glaucoma_077.jpeg img Image.open(image_path) results model(img) results.print() # 打印 top-1 预测类别和置信度这段脚本的好处是不依赖 YOLOv5 的验证入口直接加载训练好的权重把本地图片喂进去适合快速排查某一张具体图片的预测行为。如果在force_reloadTrue的情况下hub.load每次都重新拉取权重速度比较慢改成False会走缓存日常验证用后者就行。碰到权重已经删除但 hub 缓存还在的情况调整path到实际权重位置即可。6.5 训练结果验证与保存策略最终在测试集 120 张图上跑完如果准确率在 0.90 以上这个数据集就算完全跑通了。我的习惯是把best.pt、混淆矩阵图、训练曲线图单独归档到一个文件夹连同 JSON 字典文件和可视化脚本一起备份后续做实验对比时不用重新训练。从那以后每次拿到新数据集我都强制走一遍“可视化抽查 → ImageFolder 加载 → 跑 5 个 batch → 小规模训练 10 个 epoch → 确认损失下降 → 全量训练”的流程从源头上过滤掉八成以上的数据工程问题。希望这篇笔记帮你在眼睛疾病分类这个任务上少踩几个坑。本文还有配套的精品资源点击获取
网站建设高端定制企业官网