简介面向图像分类与工业视觉检测场景这份数据集将微铣削刀具磨损状态划分为好、坏两个类别所有图片已按训练集与测试集目录分别存放并附有类别字典json文件可直接用于yolov5分类任务或其他分类网络训练。包内共2000个文件包含1998张jpg样本、1个py可视化脚本和1个json标签文件压缩包约165.2MB训练集共3600张、测试集共900张数据划分清晰无需额外整理即可投入模型训练与验证。配套的py脚本支持一键可视化数据集方便快速检查类别分布与典型样本json文件则明确了类别映射关系便于后续评估指标计算。资源已有502人学习适合正在研究刀具磨损识别、工业质检或希望利用现成数据开展yolov5分类实战的研究者与开发者直接取用。1. 微铣削刀具磨损图像识别为什么这份整理好的数据比模型更值钱微铣削加工中刀具磨损的判断长期依赖老师傅拿显微镜看刀尖一个人一个标准夜班和白班的标准还不一样。把磨损状态改成图像分类问题后最大的工作量其实不在模型而在数据谁来决定某张图是“正常”还是“磨损”训练集和验证集怎么切才不漏信息类别字典写错一个键整个评估就作废。这个项目给的正是这部分按文件夹划分好的图像数据外加类别字典文件拿到就能用现成的图像分类模型训练。适合手里有现场图像、但不想在数据工程上重复踩坑的人也适合想快速验证深度学习图像识别方案能不能在刀具管理上落地的小团队。2. 文件夹结构与类别字典先搞清楚这套数据怎么被模型消费2.1 为什么用文件夹保存而不是一张 CSV 清单图像分类任务里最常见的组织方式是 torchvision 的 ImageFolder 约定数据集根目录下按类别建子文件夹每个子文件夹里放对应类别的图片。这个项目特意强调“划分好的数据【文件夹保存】”本质上就是把训练集、验证集、测试集各自建成一个 ImageFolder 根目录模型训练时直接用目录加载不需要额外解析 CSV 映射关系。我一般会在拿到这样的数据后先做一件事把整棵目录树打出来确认划分是否完整。一个合格的划分至少要有 train、val、test 三个根目录每个根目录下的类别文件夹名字完全一致且图片数量符合预期。如果交付物里只有 train 和 val测试要用 val 替代这会让你后面调参时没有真正“没见过”的数据可用模型选型的可信度会打折扣。tree data -L 2输出里应当能看到 train/normal、train/wear、val/normal、val/wear 这样的路径。如果发现某个类别文件夹缺失或者图片后缀混了 .jpg 和 .png先统一格式再进训练环节。图像分类模型对输入尺寸和通道数敏感后缀混乱虽然不影响加载但会影响后续数据增强管线的统一性。2.2 类别字典文件为什么它值得单独作为一个交付物类别字典文件的典型内容是一个 JSON把类别名映射成整数索引。比如二分类时是 {normal: 0, wear: 1}三分类时可能是 {normal: 0, initial_wear: 1, severe_wear: 2}。这个文件的价值在推理阶段体现模型输出的 logits 只有索引 0、1、2如果没有字典反向映射你根本不知道 1 对应的到底是“初期磨损”还是“严重磨损”。我见过不止一次因为字典和文件夹顺序对不上导致的翻车写字典的人手动数了一下文件夹名的字母顺序写了 {wear: 0, normal: 1}但代码用 sorted(listdir) 生成了 {normal: 0, wear: 1}训练和推理各用各的映射模型在训练集上精度不错一到现场全部预测错。避免这种事故的办法是直接写脚本从文件夹生成字典而不是手写。import os import json def build_class_dict(root_dir, out_path): classes sorted(os.listdir(root_dir)) class_to_idx {cls: idx for idx, cls in enumerate(classes)} idx_to_class {idx: cls for cls, idx in class_to_idx.items()} with open(out_path, w, encodingutf-8) as f: json.dump({class_to_idx: class_to_idx, idx_to_class: idx_to_class}, f, ensure_asciiFalse, indent2) print(class_to_idx) build_class_dict(data/train, class_dict.json)这段脚本的逻辑是以 train 目录下的子文件夹名字作为唯一事实来源按字母排序后再生成映射保证每次运行结果一致。参数方面root_dir 必须传 train 目录而不是 val因为训练集通常类别最全如果有某个类别在 val 里出现而 train 里没有这个划分本身就有问题脚本会在生成时直接暴露出来。out_path 建议保存为 JSON带 UTF-8 编码类别名如果是中文也不会乱码。推理阶段读取同一个 JSON 做反向索引训练和部署用同一份字典才能把映射错位这个风险关死。2.3 数据划分的隐藏陷阱按刀具编号切分而不是按文件切分这是整个数据工程里最容易踩的坑也是大多数翻车的真正来源。微铣削刀具磨损图像有个特点同一把刀在不同磨损阶段拍出来的照片不仅磨损程度不同刀体上的特征纹理、编号刻印、甚至光照反射都会形成一种“刀具个体指纹”。如果划分数据时把所有文件混在一起随机切分同一把刀的图像会同时出现在训练集和验证集里模型完全可以靠识别刀体本身的纹理来猜类别而不是靠识别磨损状态。用这种划分训出来的模型验证集精度能到 99%但换一把新刀现场测试精度可能直接掉到 70%。这不是玄学这是数据泄漏。正确做法是按刀具编号分组让一把刀的所有图像要么全进训练集要么全进验证集或测试集。这样验证集里全是模型没见过的刀评估结果才接近真实工况。import os import shutil import random from collections import defaultdict raw_root data/raw dst_root data/split tool_images defaultdict(list) for label in os.listdir(raw_root): label_dir os.path.join(raw_root, label) if not os.path.isdir(label_dir): continue for fname in os.listdir(label_dir): # 假设文件名里带刀号信息例如 T12_0001.jpg tool_id fname.split(_)[0] tool_images[tool_id].append((label, fname)) tool_ids list(tool_images.keys()) random.shuffle(tool_ids) n len(tool_ids) train_tools tool_ids[:int(n * 0.7)] val_tools tool_ids[int(n * 0.7):int(n * 0.85)] test_tools tool_ids[int(n * 0.85):]这段代码的核心是把“按文件划分”改成“按刀具划分”。train_tools、val_tools、test_tools 三个列表分别装不同刀号然后按刀号把对应图片复制到目标目录。分割比例 70/15/15 是图像分类的常见做法如果数据量很小可以放宽到 75/12.5/12.5但验证集至少要有两把完整刀具的图像否则评估结果波动会很大。如果原始文件名里没有刀号信息就得靠生产记录或人工打标签把刀号补上这一步省不得。注意如果项目只给了划分好的文件夹而没有提供刀号信息你无法确认划分是否避免了数据泄漏。这种情况下建议随机抽取几张训练集和验证集图像人工对比是否存在肉眼可辨的相同刀具特征。特征相似度太高就要重新找原始数据源。3. 从文件夹数据到训练闭环加载、增强与模型选型3.1 用 ImageFolder 加载这套数据和类别字典对齐文件夹结构已经就绪我用 torchvision 的 ImageFolder 加载它会自动扫描 train 根目录下的子文件夹把每个子文件夹名字解析成类别标签。加载之后第一件事是打印 class_to_idx和项目给的类别字典文件核对。ImageFolder 内部按字母顺序分配索引所以如果类别文件夹名是中文或带特殊字符排序结果可能和你的预期不一致必须核对。from torchvision import datasets, transforms train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_ds datasets.ImageFolder(rootdata/train, transformtrain_transform) print(train_ds.class_to_idx)这里用 Resize 到 256 再 RandomResizedCrop 到 224属于图像分类的标准套路。RandomResizedCrop 的作用是从图像中随机裁剪一块区域并缩放到固定尺寸对微铣削刀具磨损这种局部特征明显的图像特别有用磨损区域不一定总在图像正中心随机裁剪相当于在模拟不同拍摄角度下磨损带的位置偏移。ColorJitter 的亮度增强参数我故意调到了 0.3就是为了抵消车间光照变化对磨损纹理判别的干扰。注意 Normalize 的均值方差用的是 ImageNet 预训练统计量因为后面要加载预训练模型这个必须对齐否则预训练权重等于白加载。3.2 模型选型为什么默认先上 ResNet18 而不是最新结构图像分类模型现在选择很多从 ResNet 到 EfficientNet再到 transformer 图像分类模型如 ViT、Swin。但微铣削刀具磨损图像识别的场景有其特殊性数据集规模通常是几百到几千张属于典型的小样本细粒度分类。我在这种数据规模下优先用 ResNet18 或者 ResNet50 的 ImageNet 预训练权重加载后替换最后一层全连接层输出类别数改为你的分类数量。如果是二分类“正常/磨损”ResNet18 已经够用如果是三分类包含“初期磨损”磨损带的范围更小、内部纹理更细我建议上 ResNet50它的深层特征图能保留更多细节。transformer 系列的图像分类模型在数据量超过一万张时优势明显少于这个量级容易过拟合训练周期也更长。不是说越新的模型越好而是你的数据量决定模型容量的上限。import torch.nn as nn from torchvision import models def build_model(num_classes2, archresnet18): if arch resnet18: model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) in_features model.fc.in_features model.fc nn.Linear(in_features, num_classes) elif arch resnet50: model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) in_features model.fc.in_features model.fc nn.Linear(in_features, num_classes) return model model build_model(num_classes2, archresnet18)这段代码的关键在最后三行把预训练模型的最后一层全连接替换成输出为 num_classes 的新层其余卷积层权重全部保留。训练时新层的学习率可以设置得大一些预训练部分用小学习率微调。ResNet18 的 fc.in_features 是 512ResNet50 是 2048代码里动态读取是为了防止写死数字造成维度不匹配的错误。如果显存有限ResNet18 在 batch_size 32、输入 224 的情况下大约占用 2GB 显存普通 4GB 显卡就能跑。3.3 类别不均衡与训练参数磨损样本少的时候怎么办微铣削刀具磨损图像的现实情况是正常刀具的图像远多于磨损刀具因为刀具大部分时间处于健康状态磨损是渐进的且在后期才被重点关注。这会导致模型偏向预测“正常”漏报磨损。处理办法有两个方向一是用 WeightedRandomSampler 给磨损样本更高的采样权重二是在损失函数里给磨损类别更高的权重。我一般两个都用采样层面的均衡更直接损失函数的权重再兜底。from torch.utils.data import WeightedRandomSampler labels [s[1] for s in train_ds.samples] class_counts torch.bincount(torch.tensor(labels)) weights 1.0 / class_counts[labels].float() sampler WeightedRandomSampler(weights, num_sampleslen(labels), replacementTrue) train_loader DataLoader(train_ds, batch_size32, samplersampler, num_workers4)逻辑说明train_ds.samples 里每个元素是一个 (图片路径, 标签索引) 的元组labels 取出所有标签class_counts 统计每个类别的样本数然后每个样本的权重是它所属类别样本数的倒数。磨损样本少它的权重就大采样器会更多地抽到它。replacementTrue 表示允许重复采样同一个样本保证每个 epoch 都能采满 num_samples 个样本。训练结束后验证时不需要 sampler用普通的顺序加载即可否则验证集损失的计算会被扰动。训练参数方面我常用的配置是优化器用 AdamW基础学习率 1e-4预训练部分用 1e-5新全连接层用 1e-3batch_size 取决于显存ResNet18 下 32 是安全值训练 30 到 50 个 epoch配合余弦退火学习率调度。训练过程中每 2 个 epoch 记录一次验证集精度只保留验证集精度最高的模型权重。这个策略避免在训练后期因为过拟合而丢失最佳权重相当于给调参留了后悔药。4. 微铣削磨损模型训练的避坑清单五条血泪经验4.1 验证集精度虚高现场直接翻车现象训练阶段验证集精度 98%模型部署到车间后新刀的正常图像频繁被报警为磨损。排查后发现训练集里的正常刀图像几乎都拍摄于灯光充足、刀面洁白的场景磨损刀图像则大多在灯光偏暗、刀面发灰的场景下采集模型实际学到的是图像亮度差异而不是磨损带纹理。原因数据采集时没有控制光照变量导致类别和亮度形成了虚假的强关联。深度学习图像识别模型会优先利用最容易区分的特征亮度差异比磨损纹理更容易拟合于是模型走了捷径。解决在训练增强中强制加入亮度、对比度扰动打破亮度与类别的关联。我在上面的 ColorJitter 里把 brightness0.3就是针对这个场景。更彻底的做法是重新采集一批相同光照下的正常刀图像加入训练集让模型意识到“正常刀也有暗面”。验证模型是否依赖亮度可以把验证集所有图像统一调暗 30% 再测试如果精度大幅下降说明模型还在靠亮度判断。4.2 换一批刀具后精度断崖式下降现象同一台机床、同一型号铣刀模型在上一批刀具图像上精度 95%换新批次后降到 75%。看混淆矩阵误报集中在把新刀的正常状态识别为磨损。原因不同批次的刀具在涂层颜色、刀面纹理、甚至加工后残留的切削液痕迹上存在肉眼可见差异。上一批刀具的“个体特征”被模型当成了可判别的类别特征数据划分时没有按刀具隔开。解决重新按刀具编号划分数据保证验证集里的刀具在生产批号、涂层批次上与训练集不重叠。如果现有划分已经无法追溯至少要保证训练集和验证集里包含多种批次的刀具图像而不是单独某一批。对于已经部署的模型建议在换批次后先采集 50 张新图像做快速测试精度低于阈值就触发模型重训流程。4.3 类别字典的顺序不稳定导致推理映射错位现象训练脚本和推理脚本各自读字典文件训练时 class_to_idx 是 {normal: 0, wear: 1}推理时因为列表排序变化或字典文件被重新生成变成了 {wear: 0, normal: 1}现场永远预测反。原因字典文件由不同脚本共享但没有规定唯一生成入口。有人直接改 JSON 内容有人用 set 集合去重类别名导致顺序随机最终字典和文件夹顺序不一致。解决用前面的 build_class_dict 脚本从 train 目录自动生成字典训练和推理代码都只读同一个 JSON 文件杜绝手写。每次训练前打印 class_to_idx 并确认它和 JSON 内容一致把这个验证写进训练启动脚本里不一致就报错退出。别以为看一眼不会错现场凌晨三点换模型的时候人的判断力是不可靠的。4.4 磨损样本太少导致 loss 震荡模型不收敛现象训练 loss 忽高忽低验证精度始终在 60% 到 70% 之间徘徊磨损类别的召回率极低。观察样本分布整个数据集中磨损刀图像只有 120 张正常刀图像却有 800 张。原因类别不均衡加上随机采样每个 batch 里磨损图像很少出现模型对磨损类别的学习率极低。即便加权采样到位增强过度也会让本来就少的磨损样本变得失真模型学不到有效信息。解决先启用 WeightedRandomSampler把磨损样本的采样频率拉高。如果数据量不足考虑对磨损类别做针对性增强比如只对磨损图像做小幅旋转和缩放不要用翻转增强因为刀具左右刃的磨损形态不一定对称。仍不收敛就把损失函数换为 Focal Loss降低易分类样本的权重让模型把注意力放在磨损样本上。4.5 resize 压扁磨损纹理细粒度特征丢失现象模型在原始分辨率下测试精度尚可但训练时把 512x512 的原图 Resize 到 224x224 后精度下降尤其是“初期磨损”这种小区域缺陷磨损带在压缩后的图像里几乎看不见。原因直接等比例缩放把磨损区域的纹理缩小到几个像素宽卷积核已经提取不到有效特征。微铣削刀具的前刀面月牙洼宽度在原始图像里可能只占 20 个像素缩到 224 后只剩 8 个像素特征信息基本消失。解决如果原图分辨率有 1024 甚至更高先按长边等比缩放到 512再中心裁剪出 448x448 的区域用于训练模型输入分辨率从 224 提升到 448。代价是显存占用翻几倍ResNet18 在 batch_size 16 的情况下大约需要 3GB 显存。另一种做法是保留原图做推理训练时用多尺度裁剪模拟不同放大倍率下的磨损形态。5. 进阶验证用 CAM 和混淆矩阵确认模型到底在看哪里5.1 Grad-CAM 定位模型的关注区域训练精度高不代表模型在看磨损区域。如果模型把目光聚焦在刀柄上的编号刻印、夹具边缘的反光甚至图像角落的日期水印换一个新的拍摄视角就可能失效。我用 Grad-CAM 对验证集样本做可视化确认模型的注意力集中在刃口磨损带上。import torch from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget from pytorch_grad_cam.utils.image import show_cam_on_image model.eval() cam GradCAM(modelmodel, target_layers[model.layer4[-1]]) input_tensor preprocess(image).unsqueeze(0) targets [ClassifierOutputTarget(1)] # 1 代表 wear 类别 grayscale_cam cam(input_tensorinput_tensor, targetstargets) cam_image show_cam_on_image(normalized_image, grayscale_cam[0])这段代码对模型输出的 wear 类别计算梯度热力图把热力图叠加到原图上。target_layers 选 model.layer4[-1]也就是 ResNet 最后一个残差块的最后一层卷积它的空间分辨率是 7x7足以定位磨损区域的大致位置。如果热力图的中心偏离刃口跑到刀体中央或背景区域说明模型学到的是其他特征。我一般对每个类别抽 20 张图做 CAM把热力图和原图拼在一起人工检查。需要注意的是Grad-CAM 的分辨率很低只有 7x7 或 14x14无法精确到磨损带的像素级边界。它能回答的问题是“模型在看图的哪个区域”而不是“磨损宽度是多少像素”。如果要更精细的定位可以改用 Layer-CAM 或 Ablation-CAM但作为验证工具Grad-CAM 已经足够判断模型是否在合理的位置上做判断。5.2 混淆矩阵找出最容易混的类别对微铣削刀具磨损的三分类场景里最常出现的是“初期磨损”被预测成“正常”或“严重磨损”。混淆矩阵能直观暴露这种类别间的混淆程度比只看总体精度有用得多。from sklearn.metrics import confusion_matrix y_true [] y_pred [] with torch.no_grad(): for images, labels in val_loader: outputs model(images.cuda()) _, preds torch.max(outputs, 1) y_true.extend(labels.tolist()) y_pred.extend(preds.cpu().tolist()) cm confusion_matrix(y_true, y_pred) print(cm)重点看对角线之外的数字如果正常刀被误判为磨损说明模型的误报率高现场会频繁停机检查操作员很快就会不信任这套系统如果磨损刀被误判为正常那就是漏检后果是刀具继续带病工作影响加工表面质量。根据误报和漏损的方向调整类别权重或阈值。从混淆矩阵出发我还会回看被分错的图像找出是图像本身模糊、磨损形态特殊还是采集角度问题。这类定性分析往往比调参带来更大的精度提升。漏检的后果更严重所以我对“严重磨损”类别使用更高的误分类惩罚或者在后处理逻辑里对“严重磨损”类别的判定阈值放宽。具体做法是把验证集上预测概率大于 0.3 的样本都算作严重磨损宁可让初期磨损偶尔误报成严重磨损也不能让严重磨损漏报成正常。5.3 多尺度推理让模型适应不同拍摄距离和放大倍率刀具磨损图像的采集在车间里做不到严格标准化操作员拍摄时距离稍有变化磨损区域在图像中的占比就不同。单尺度推理对距离变化敏感多尺度推理的思路是对同一张图分别缩放到 224、320、448 三个尺寸分别输入模型取平均概率再决定最终类别。def multi_scale_predict(model, image, scales[224, 320, 448]): probs [] for scale in scales: resized transforms.Resize((scale, scale))(image) transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) input_tensor transform(resized).unsqueeze(0).cuda() with torch.no_grad(): output torch.softmax(model(input_tensor), dim1) probs.append(output) avg_prob torch.mean(torch.cat(probs), dim0) return torch.argmax(avg_prob).item()多尺度推理本质上是在模拟不同放大倍率下同一把刀的成像结果把三个尺度的预测概率做平均能有效降低单一尺度下因为磨损区域过小或者被裁剪掉带来的误判。代价是推理时间变长单张图从一次前向变成三次用 ResNet18 在三尺度下处理一张图大约耗时 15 毫秒到 30 毫秒对刀具检测这种非高速场景完全够用。如果推理设备是 CPU建议只保留最有代表性的两个尺度一个接近训练分辨率一个偏高分辨率。6. 最后一招把类别阈值校准成产线可用的报警规则模型在验证集上的最优精度对应的是默认阈值 0.5也就是预测概率大于 0.5 判为磨损小于 0.5 判为正常。但这个阈值在实际生产中不一定最合适。刀具磨损的漏报代价远高于误报漏报一次可能导致加工件报废甚至刀具断裂误报一次顶多让操作员多看一眼。所以我在部署前会做一次阈值校准让报警规则匹配产线的真实容错偏好。具体方法是取验证集全部图像模型输出每个样本的磨损概率然后从小到大遍历阈值计算每个阈值下的精确率和召回率画 PR 曲线找召回率不低于 95% 时的最大精确率点把这个点对应的概率作为报警阈值。import numpy as np probs [] labels [] with torch.no_grad(): for images, targets in val_loader: outputs torch.softmax(model(images.cuda()), dim1) probs.extend(outputs[:, 1].cpu().numpy()) labels.extend(targets.numpy()) probs np.array(probs) thresholds np.arange(0.3, 0.9, 0.01) best_th 0.5 for th in thresholds: preds (probs th).astype(int) recall np.mean(preds[labels 1] 1) precision np.mean(preds[preds 1] labels[preds 1]) if preds.sum() 0 else 0 if recall 0.95 and precision best_precision: best_precision precision best_th th这段代码的意图是在保证磨损召回率不低于 95% 的前提下选择误报率最低的阈值。recall 0.95 意味着 100 个真实磨损样本里至少检出 95 个只允许 5 个漏到正常类别里。我实测下来磨损图像得分通常集中在 0.7 到 0.95 之间正常图像在 0.1 到 0.4 之间阈值提到 0.6 不会损失太多召回率但能让误报数量明显下降。如果验证集 PR 曲线显示 0.95 召回率对应的精确率过低说明模型本身判别力不足不要硬调阈值回到数据增强和模型容量上想办法。调整完阈值后报警逻辑也可以做成三档概率大于报警阈值的直接报磨损概率在 0.5 到报警阈值之间的提示人工复核低于 0.5 的判为正常。这样既保留了自动检测的效率又给操作员留了判断空间。刀具磨损本来就是渐变的连续过程正常和磨损之间不该是一条硬边界概率本身就是在表达这种不确定性。我自己的习惯是每次换一批新刀具或者调整采集摄像头后不急着重训模型先在旧模型下采集 50 张新图看概率分布相比训练时是否偏移。偏移幅度大说明数据分布变了该重训偏移小随手做个阈值微调就行。这样做能省掉很多不必要的重训开销也让产线系统保持安静。希望这些经验能帮你在同样的项目里少走几步弯路希望帮到你。本文还有配套的精品资源点击获取