ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

苹果缺陷语义分割:数据集解析与Unet训练实战

苹果缺陷语义分割:数据集解析与Unet训练实战 简介苹果缺陷图像语义分割数据集面向图像分割算法开发与缺陷检测场景包含健康、病害区域等5类标签约4000张已划分好的图片与掩膜训练集约3000张、测试集约1200张可直接用于训练和评估语义分割模型。整个压缩包约73.28MB共2000个文件以png格式图像为主涵盖原始图片和对应的mask模板另附带1个类别说明txt和1个Python可视化脚本便于快速了解数据分布与标注效果。此外压缩包内还提供可视化脚本可随机抽取一张样本将原图、GT掩膜及GT叠加在原图上的效果并列展示并保存方便直观检查标注质量与模型预测差异。整体来看数据集适合入门至进阶的图像分割学习者、缺陷检测研究者及需要标准benchmark的算法工程师使用目前已有180人学习。1. 苹果缺陷只给检测框不够用语义分割才能按面积分选在苹果分选产线上真正的问题不是“有没有烂点”而是“烂点面积占多少、属于病斑还是果梗”。检测框会把果梗、花萼和早期病斑圈进同一个矩形按框统计面积时产生大量虚假缺陷改成像素级语义分割后图像里每个像素都被归入健康表皮、病害区域或背景才能按类别统计缺陷占比分选逻辑才站得住。这份苹果缺陷图像语义分割数据集就是冲着这个场景整理的Canon相机俯拍的苹果图像约4000张配套同名mask划分为5个语义类别训练集3000张左右、测试集1200张左右切分已经完成不需要再手工整理标签。适合做农产品质检、工业视觉缺陷定位以及想在医学图像分割或Unet系模型上练手的人。这套数据里有两个值得留意的设计mask与图像同名、以PNG索引图存储没有用COCO那种JSON标注文件名里带着拍摄时间和光照条件all_on / side_only这给后续做光照变化下的泛化验证留了操作空间。下面会先解析目录结构与标签编码再复现自带的GT叠加可视化脚本然后接入PyTorch跑Unet训练与mIoU评估最后把类别不平衡和mask读取这类容易翻车的细节单独拿出来讲。2. 目录结构、文件名语义与5类标签的编码方式2.1 训练集与测试集目录先核对数量解压后典型的目录结构是images与masks两套平级目录原图和标签文件名一一对应训练集和测试集结构一致dataset/ ├── train/ │ ├── images/ │ │ ├── 21_01_03_19_15_00_Canon_top_all_on_msk_6.png │ │ └── 20_12_27_01_15_00_Canon_top_side_only_msk_6.png │ └── masks/ │ └── 与images同名的mask图 └── test/ ├── images/ └── masks/随包还有一个类别说明文件记录5个类别索引与名称的映射。拿到数据后我一般先做三件事核对train/test里图片和mask数量是否一致抽查是否有空图再用PIL读取一张mask确认索引范围。files目录在数据集里已经是images和masks分开存放这比标签混在文件名后缀里的格式省心训练时按路径拼接即可也不会把标签当原图喂给增强管线。先数数量再训练是个好习惯。如果images比masks多几张通常是导出时残留了缩略图或临时文件直接跑Dataset会在迭代到一半时才报错排查成本比预处理阶段高得多。2.2 文件名里藏着的拍摄条件与光照标记文件名不是随机字符串每个字段都能解析出拍摄信息。以21_01_03_19_15_00_Canon_top_all_on_msk_6.png为例字段示例值含义21_01_032021年1月3日采集日期19_15_0019:15:00采集时间CanonCanon拍摄相机标识toptop俯拍角度all_on / side_onlyall_on顶灯全开 / 侧光照明msk_6msk_6mask批次编号all_on和side_only代表两种打光模式这个细节值得重点利用。语义分割模型很容易对光照建立隐式依赖训练时如果只见过顶灯全开的图换到侧光环境mIoU会掉一截。利用文件名解析出光照标签后可以按光照模式划分训练集和测试集专门验证模型在打光变化下的稳定性。2.3 mask是索引PNG别用cv2直接读语义分割标签常见两种存储方式单通道灰度PNG像素值直接是类别索引或调色板PNGP模式。这份数据的mask更接近后者读取时有明确讲究用PIL打开不要用OpenCV的cv2.imread直接读。cv2.imread遇到P模式PNG会按调色板解码成三通道BGR类别索引在这个过程里被转成了RGB颜色之后再想恢复成04的标签基本不可能。判断mask格式用下面几行就够from PIL import Image mask Image.open(train/masks/21_01_03_19_15_00_Canon_top_all_on_msk_6.png) print(mode:, mask.mode) # 预期是 P 或 L print(value range:, mask.getextrema()) # 预期是 (0, 4) 或 (0, 5)mode输出P代表调色板索引图L代表灰度索引图二者都可以直接用PIL转为numpy数组得到二维索引矩阵。getextrema()返回像素最小值与最大值用来确认类别索引范围是否是04。如果出现255说明mask里有标注边界或ignore区域后续做损失计算时要单独处理不能直接当第6个类别。3. 复现GT叠加可视化脚本3.1 mask读取为什么要绕开OpenCV在可视化这一步最容易出问题的写法就是cv2.imread(xxxx_msk_6.png)然后直接resize。P模式PNG会被OpenCV按调色板解码成三通道BGR看起来是彩色图但颜色已经和类别索引绑定没法用阈值还原。更隐蔽的坑是后续用cv2.resize默认的双线性插值去缩放mask类别索引会变成小数比如像素值1和2中间插出1.5训练时torch.long()一转换就报错或者静默产生错误的类别。所以整套流程里读取mask只走PIL。PIL打开P模式PNG后转numpy得到的就是单通道索引矩阵数值干净处理逻辑清晰。这也是复现可视化脚本之前要先确认环境的原因。3.2 可视化脚本原图、GT与叠加图并存随包自带脚本的作用是随机抽一张图把原始图片、GT mask、GT在原图上的蒙版展示并保存到当前目录。我用等价方式实现一遍import os import random import numpy as np import cv2 from PIL import Image image_dir train/images mask_dir train/masks sample_id random.choice(os.listdir(image_dir)) # 原图用OpenCV读取BGR转RGB img cv2.imread(os.path.join(image_dir, sample_id)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # mask用PIL读取保留类别索引 mask np.array(Image.open(os.path.join(mask_dir, sample_id))) # 类别索引 - 可视化颜色实际以classes文件为准 color_map np.array([ [0, 0, 0], # 0 背景 [0, 255, 0], # 1 健康果皮 [0, 0, 255], # 2 病害区域 [255, 255, 0], # 3 果梗/花萼等 [255, 0, 255], # 4 其他缺陷 ], dtypenp.uint8) mask np.clip(mask, 0, len(color_map) - 1) mask_rgb color_map[mask] # 原图与彩色mask线性叠加 overlay cv2.addWeighted(img, 0.6, mask_rgb, 0.4, 0) result np.hstack([img, mask_rgb, overlay]) cv2.imwrite(visualize_sample.png, cv2.cvtColor(result, cv2.COLOR_RGB2BGR)) print(saved:, sample_id)逻辑说明原图用OpenCV读取后转成RGB是为了和后面的RGB顺序保持一致mask用PIL读取并直接转numpy数组拿到的就是H×W的整数索引图不经过任何颜色解码。color_map[mask]利用numpy高级索引把每个类别索引映射成RGB三通道。np.clip(mask, 0, len(color_map) - 1)防止标注里的255边界像素导致索引越界。cv2.addWeighted做alpha混合原图系数0.6、mask系数0.4缺陷边界能看清原图纹理也不会丢。最后保存时把RGB转回BGR否则用cv2.imwrite写出的图会红蓝互换。3.3 从叠加图能直接看出标签对齐问题叠加图是检查数据集质量最直接的窗口。看三个点背景类轮廓是否紧贴苹果边缘病害区域边界是否与图像上的颜色过渡吻合缺陷区里有没有大量细碎噪点。轮廓整体偏移多半是原图和mask尺寸不同步缺陷边界像水彩晕开说明mask被非最近邻插值污染过。这个可视化脚本也可以当环境自检工具新机器上PIL、numpy、OpenCV版本不兼容时大概率会在这个环节暴露。4. 把数据集改成PyTorch格式并用Unet跑通训练4.1 Dataset实现与三个关键约束这套数据对应的Unet、SwinUnet、TransUnet改进方案和医学图像分割使用的方法论是同一条线小目标、类别不平衡、像素级精度。要跑起来先写一个适配这份目录结构的Datasetimport os import numpy as np import torch from torch.utils.data import Dataset from PIL import Image class AppleDefectDataset(Dataset): def __init__(self, image_dir, mask_dir, size512): self.image_dir image_dir self.mask_dir mask_dir self.names sorted(os.listdir(image_dir)) self.size size def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img Image.open(os.path.join(self.image_dir, name)).convert(RGB) mask Image.open(os.path.join(self.mask_dir, name)) # 保持P/L模式 img img.resize((self.size, self.size), Image.Resampling.BILINEAR) mask mask.resize((self.size, self.size), Image.Resampling.NEAREST) img torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.0 mask torch.from_numpy(np.array(mask)).long() return img, mask参数与约束说明size512是把高分辨率原图统一缩放到512×512显存不足时降到256但苹果缺陷区域通常很小低于256容易丢失细节。原图用convert(RGB)是安全的mask不能加convert(RGB)否则类别索引会被替换成RGB颜色值。mask.resize必须用Image.Resampling.NEAREST最近邻插值不会产生1.5这样的伪类别。图像归一化到0~1mask转成torch.long()后作为CrossEntropy的target。4.2 训练参数、损失函数与优化器实际训练中苹果图像里背景像素占比很高直接CrossEntropy会把注意力全放在背景上病害区域学不干净。常见做法是CrossEntropy加Soft Dice组合我一般按1:1权重相加。模型直接用segmentation_models_pytorch的Unetbackbone选timm-efficientnet-b08G显存能跑起来。import torch.nn.functional as F from segmentation_models_pytorch.losses import DiceLoss import segmentation_models_pytorch as smp model smp.Unet( encoder_nametimm-efficientnet-b0, encoder_weightsimagenet, classes5, ) criterion_ce torch.nn.CrossEntropyLoss() criterion_dice DiceLoss(modemulticlass) for batch in loader: x, y batch[0].to(device), batch[1].to(device) logits model(x) loss criterion_ce(logits, y) criterion_dice( F.softmax(logits, dim1), y ) optimizer.zero_grad() loss.backward() optimizer.step()CrossEntropy直接吃logits不用手动softmaxDiceLoss需要先对logits做softmax得到概率图。Dice对类别不平衡更钝感但训练早期梯度噪声大所以和CE各占一半比较稳。想要更强backbone时把smp.Unet换成smp.DeepLabV3Plus或SwinUnet即可Dataset接口不用动。起步参数参考下表参数建议值说明输入尺寸512×512显存不足降到256batch size812G显存左右可调4~16优化器AdamWlr1e-4比纯Adam收敛更稳学习率策略poly或余弦衰减最后阶段对分割精度影响明显训练轮数60~80配合早停评估指标各类别Dice、mIoU按光照条件分组看4.3 mIoU评估与训练监控只看loss会掩盖小类目表现差的问题。语义分割一般用mIoU和平均Dice监控训练mIoU实现如下def compute_miou(pred_mask, gt_mask, num_classes5, ignore_index255): pred_mask pred_mask.argmax(dim1).cpu().numpy() gt_mask gt_mask.cpu().numpy() valid gt_mask ! ignore_index pred_mask np.where(valid, pred_mask, ignore_index) ious [] for c in range(num_classes): p (pred_mask c) valid g (gt_mask c) valid inter (p g).sum() union (p | g).sum() ious.append(inter / union if union 0 else 1.0) return np.nanmean(ious)传入的pred_mask是网络输出的logitsshape为[B, C, H, W]先argmax(dim1)得到每像素的预测类别。valid gt_mask ! ignore_index把255这类无效像素排除掉避免边界标注污染指标。当某个类别在当前样本中完全没有GT时union为0这时计1.0而不是跳过否则batch内多个样本的平均会不稳定。如果只是想对比目标检测路线把mask转成polygon后YOLOv8-seg也能训练。但苹果缺陷经常连成一片实例分割会把一个完整病斑切成多个实例按面积统计反而更麻烦。语义分割的像素级输出对这类质检场景更直接。5. 类别不平衡、mask读取细节与按光照分组的验证5.1 训练前先统计每类像素占比不要拍脑袋决定损失权重。抽500张mask统计像素分布成本很低import os import numpy as np from PIL import Image mask_dir train/masks files sorted(os.listdir(mask_dir))[:500] counts np.zeros(5, dtypenp.int64) for f in files: m np.array(Image.open(os.path.join(mask_dir, f))) for c in range(5): counts[c] (m c).sum() freq counts / counts.sum() print(class frequency:, freq)如果背景类占比超过60%而某个缺陷类不足5%说明类别不平衡已经很严重。此时CrossEntropy的默认权重基本会把前景淹没病害区域可能完全训不出来val mIoU看着还行实际缺陷Dice却很低。5.2 加权CE与同步增强策略处理类别不平衡有两种常用手段。一是给CrossEntropy加权权重按w_c 1 / log(1.02 freq_c)计算低频类别获得更高权重二是依赖Dice Loss自带的类别不敏感特性和CE按1:1叠加。两种损失搭配时加权CE对病斑这类小目标更友好。增强策略要注意同步问题翻转、旋转、缩放这类几何变换必须对原图和mask施加完全相同的操作颜色抖动、亮度变化只作用于原图。现在直接推albumentations它的Compose支持同时处理image和mask能省掉手写随机种子同步的麻烦。5.3 按all_on/side_only分组看mIoU这份数据的文件名里带光照标记验证时不要只看测试集整体mIoU。把测试集按all_on和side_only拆成两组分别计算mIoU再对比组间差距。如果侧光组明显低于顶灯组说明模型对光照条件建立了隐式依赖泛化性不达标。此时要加亮度对比度增强或者把两种光照按比例混合进训练集。这种按域分组的验证方式也是后续做域泛化实验时可以直接套用的经验。本文还有配套的精品资源点击获取
返回列表