简介：本资源为苹果缺陷图像语义分割数据集，面向从事图像分割算法研究、深度学习模型训练与农业视觉检测方向的开发者及学生，可用于训练和评估UNet、SwinUnet、TransUNet等分割网络。数据集已按训练集与测试集划分完毕，训练集约3000张图片及对应mask，测试集约1200张，涵盖健康、病害区域等5类分割标签，具体类别可参考classes文件。压缩包共2000个文件，以1998个png图像与掩膜为主，另含1个txt类别说明和1个py可视化脚本，整体约73.28MB。该脚本可随机抽取一张图片，展示原始图像、GT图像及GT在原图上的蒙板效果并保存至当前目录，便于快速核验标注质量。目前已有180人学习下载，适合需要现成数据快速验证分割模型、开展缺陷检测实验的读者。
1. 苹果缺陷图像语义分割数据集：5 类标注、约 4000 张，能直接喂给 U-Net 吗
拿到一个「苹果缺陷图像语义分割数据集（5 类分割，约 4000 张数据和标签）」的标题，很多人第一反应是把它当成又一个普通图像分割数据集，下载、解压、丢进 U-Net 训练脚本，然后等着看 mIoU。但真正做过工业质检落地的都知道，苹果表面缺陷分割和自然场景语义分割完全是两回事：缺陷区域边界模糊、类别间高度相似、光照和果面反光干扰大，5 类标签里往往还藏着一类「背景」需要单独处理。这个数据集的价值不在于它有多少张图，而在于它把「苹果表面缺陷」这个细分场景的像素级标注做成了可复用的语义分割样本，适合做水果分选、产后质检、农业机器人视觉方向的人拿来当基线数据。约 4000 张的规模不算大，但足够跑通一个 U-Net 或 DeepLabV3+ 的完整训练流程，也足够暴露你在数据预处理、类别不平衡、标签格式转换上的真实问题。下面按「先搞清楚数据长什么样、再动手转格式和训练、最后讲坑」的顺序展开，中间会给出可直接抄的脚本和参数。
2. 先看清 5 类标签到底标了什么：苹果缺陷语义分割的数据结构拆解
2.1 语义分割标签的两种常见存储形态
拿到一个语义分割数据集，第一件事不是写模型，而是确认标签的存储形态。苹果缺陷数据集常见的标签有两种：一种是单通道灰度掩膜图（mask），每个像素值直接对应类别 ID，比如 0 背景、1 斑点、2 腐烂、3 疤痕、4 虫眼；另一种是彩色掩膜图，用不同 RGB 颜色区分类别，需要额外一张颜色映射表才能转成类别 ID。前者对训练最友好，后者多见于标注工具直接导出的结果。你拿到数据后先用一条命令看标签图的通道数和唯一像素值，就能判断属于哪种。
# 查看标签目录下第一张图的通道数和像素值分布 python -c " import cv2, numpy as np, glob f = sorted(glob.glob('labels/*.png'))[0] img = cv2.imread(f, cv2.IMREAD_UNCHANGED) print('shape:', img.shape, 'dtype:', img.dtype) print('unique values:', np.unique(img)[:20]) "这段代码用IMREAD_UNCHANGED读取，避免 OpenCV 默认把单通道图转成三通道。如果输出 shape 是(H, W)且 unique values 是 0 到 4 的整数，说明是标准灰度掩膜，可以直接用；如果 shape 是(H, W, 3)，就需要按颜色映射表转换。参数上唯一要注意的是np.unique只取前 20 个值，防止类别多时刷屏。
2.2 5 类缺陷的类别定义与样本分布检查
5 类分割通常包含背景加 4 类缺陷，或者 5 类全是缺陷、背景用 0 表示。不管哪种，你都要先统计每类的像素占比，因为苹果缺陷分割最典型的问题就是极端类别不平衡——背景可能占 90% 以上，虫眼这种小目标可能不到 0.5%。不先做这个统计，训练时 loss 会被背景主导，模型学出来全是背景，mIoU 看着不低但缺陷一个都分不出来。
import numpy as np, glob, cv2 from collections import Counter counter = Counter() for f in glob.glob('labels/*.png'): mask = cv2.imread(f, cv2.IMREAD_UNCHANGED) counter.update(mask.flatten().tolist()) total = sum(counter.values()) for cls_id in sorted(counter): print(f'class {cls_id}: {counter[cls_id]/total*100:.3f}%')这段统计脚本遍历所有标签图，累计每个像素值的出现次数。跑完后你会得到一张类别占比表。如果某一类低于 1%，后面就要考虑加权损失或者过采样。这里没有用np.bincount是因为掩膜里可能混入非连续 ID，用 Counter 更稳。
2.3 图像与标签的命名对齐规则
约 4000 张数据里，图像和标签必须一一对应，但命名规则经常不统一：有的数据集图像是apple_001.jpg，标签是apple_001.png；有的标签多了_mask后缀；还有的按文件夹分 train/val，但两边文件名对不上。训练前必须写一个对齐检查，把没有对应标签的图像和没有对应图像的标签都列出来，否则训练时读到空标签会直接报错或者静默跳过，你以为是模型不收敛，其实是数据没对上。
import os img_dir, lbl_dir = 'images', 'labels' imgs = {os.path.splitext(f)[0] for f in os.listdir(img_dir)} lbls = {os.path.splitext(f)[0] for f in os.listdir(lbl_dir)} print('图像无标签:', imgs - lbls) print('标签无图像:', lbls - imgs)用集合差集找出不匹配项。常见做法是把标签后缀统一去掉再比对，如果数据集本身用_mask区分，就在生成集合时用f.replace('_mask','')处理。这一步花两分钟，能省掉后面几小时的排查。
3. 把苹果缺陷数据转成训练可用的格式：掩膜、划分与增强
3.1 灰度掩膜转类别 ID 与忽略区域处理
如果标签是彩色掩膜，第一步是转成类别 ID。转换的核心是建立颜色到 ID 的映射字典，逐像素查表。这里有个容易翻车的地方：标注工具导出的彩色掩膜边缘常有抗锯齿产生的过渡色，这些颜色不在映射表里，直接查表会变成未定义值。稳妥做法是把未匹配的像素统一归到背景或忽略类（通常设 255，训练时用ignore_index=255跳过）。
import cv2, numpy as np color_map = { (0, 0, 0): 0, # 背景 (255, 0, 0): 1, # 斑点 (0, 255, 0): 2, # 腐烂 (0, 0, 255): 3, # 疤痕 (255, 255, 0): 4, # 虫眼 } def color_to_id(mask_bgr): h, w, _ = mask_bgr.shape out = np.full((h, w), 255, dtype=np.uint8) for color, cid in color_map.items(): bgr = color[::-1] # OpenCV 读入是 BGR match = np.all(mask_bgr == bgr, axis=-1) out[match] = cid return outcolor_map的键是 RGB，代码里反转成 BGR 再比对，因为 OpenCV 默认 BGR。未匹配像素初始化为 255，作为忽略类。参数上，如果你的数据集背景就是黑色且没有过渡色，可以把初始值改成 0，省掉忽略逻辑。
3.2 按类别分层划分训练集与验证集
随机划分在类别不平衡时会让验证集里某些缺陷类几乎消失，导致验证指标剧烈波动。正确做法是按图像中出现的缺陷类别做分层抽样，保证每个类在训练和验证里都有足够样本。常见做法是先给每张图打一个「包含哪些类」的标签，再按这个多标签做分层。
import numpy as np, cv2, glob from sklearn.model_selection import train_test_split files = sorted(glob.glob('labels/*.png')) labels_per_img = [] for f in files: m = cv2.imread(f, cv2.IMREAD_UNCHANGED) present = set(np.unique(m).tolist()) - {0, 255} labels_per_img.append(sorted(present)) # 用第一个出现的缺陷类做粗分层，简单有效 strata = [l[0] if l else -1 for l in labels_per_img] train_f, val_f = train_test_split(files, test_size=0.2, stratify=strata, random_state=42) print(len(train_f), len(val_f))这里用每张图里最小的缺陷类 ID 做分层键，虽然粗糙，但比纯随机稳。test_size=0.2对 4000 张来说验证集约 800 张，够用。random_state固定住，保证可复现。如果你的数据里多类共现很普遍，可以改用sklearn的MultiLabelBinarizer加iterstrat，但那个需要额外装包，上面这个够大多数场景。
3.3 针对果面反光的增强策略与参数
苹果表面反光强，同一缺陷在不同光照下像素差异很大，增强不能只用随机翻转。我一般会加亮度对比度扰动、高斯噪声和轻微弹性形变，但弹性形变对缺陷边界破坏大，参数要保守。下面是一个基于 Albumentations 的增强配置，注意ElasticTransform的 alpha 和 sigma 都调小。
import albumentations as A train_tf = A.Compose([ A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.3), A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5), A.GaussNoise(var_limit=(10.0, 30.0), p=0.3), A.ElasticTransform(alpha=30, sigma=5, p=0.2), A.Resize(512, 512), ]) val_tf = A.Compose([A.Resize(512, 512)])brightness_limit和contrast_limit控制在 0.2 以内，太大模拟不出真实光照反而引入噪声。ElasticTransform的 alpha=30、sigma=5 是保守值，再大缺陷形状就失真了。Resize 到 512 是 U-Net 常用输入，显存不够就降到 384。增强只对训练集做，验证集只 Resize，这点别搞反。
4. 用 U-Net 跑通苹果缺陷分割：训练脚本与关键参数
4.1 数据集类与 DataLoader 的最小实现
PyTorch 下写一个语义分割 Dataset，核心是读图、读掩膜、同步增强、返回 tensor。掩膜要转成 long 类型，因为交叉熵损失要求类别索引是 int64。图像归一化用 ImageNet 均值方差即可，苹果图像和自然图像分布接近。
import torch, cv2, numpy as np from torch.utils.data import Dataset, DataLoader class AppleDefectDataset(Dataset): def __init__(self, img_files, lbl_files, transform=None): self.img_files = img_files self.lbl_files = lbl_files self.transform = transform def __len__(self): return len(self.img_files) def __getitem__(self, idx): img = cv2.imread(self.img_files[idx]) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask = cv2.imread(self.lbl_files[idx], cv2.IMREAD_UNCHANGED) if self.transform: aug = self.transform(image=img, mask=mask) img, mask = aug['image'], aug['mask'] img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0 img = (img - 0.485) / 0.229 # 简化归一化，实际按三通道分别做 mask = torch.from_numpy(mask).long() return img, maskpermute(2,0,1)把 HWC 转 CHW，/255.0归一化到 0-1。这里归一化只写了单通道示例，实际要对 RGB 三通道分别减均值除标准差。掩膜用long()转 int64。ignore_index对应的 255 在损失函数里处理，Dataset 不用管。
4.2 损失函数选择：交叉熵、Dice 与类别权重
苹果缺陷分割的类别不平衡决定了不能只用普通交叉熵。常见做法是交叉熵加类别权重，或者交叉熵和 Dice 损失按比例相加。类别权重可以按像素频率的倒数来设，但不要设得太极端，否则小类过拟合。我一般用CrossEntropyLoss(weight=w, ignore_index=255)加一个 0.5 权重的 Dice。
import torch.nn as nn class DiceLoss(nn.Module): def __init__(self, ignore_index=255): super().__init__() self.ignore_index = ignore_index def forward(self, logits, targets): probs = torch.softmax(logits, dim=1) valid = targets != self.ignore_index targets = targets.clone() targets[~valid] = 0 dice = 0 for c in range(1, logits.shape[1]): # 跳过背景 p = probs[:, c][valid] t = (targets == c)[valid].float() inter = (p * t).sum() dice += (2 * inter + 1e-6) / (p.sum() + t.sum() + 1e-6) return 1 - dice / (logits.shape[1] - 1)Dice 只对前景类计算，背景不参与，避免背景主导。1e-6防止除零。类别权重w可以先用[1.0, 5.0, 5.0, 8.0, 10.0]这种经验值，再根据验证集表现微调。总损失loss = ce_loss + 0.5 * dice_loss，0.5 这个系数不是固定的，小目标多就调到 1.0。
4.3 训练循环与验证指标 mIoU 的计算
训练循环本身不复杂，关键是验证时 mIoU 的计算要排除忽略类，并且按类累计混淆矩阵再求平均，不能逐 batch 平均。下面给出核心片段。
def train_one_epoch(model, loader, optimizer, criterion, device): model.train() for img, mask in loader: img, mask = img.to(device), mask.to(device) optimizer.zero_grad() out = model(img) loss = criterion(out, mask) loss.backward() optimizer.step() @torch.no_grad() def evaluate(model, loader, num_classes, device): model.eval() conf = np.zeros((num_classes, num_classes), dtype=np.int64) for img, mask in loader: img = img.to(device) pred = model(img).argmax(1).cpu().numpy() mask = mask.numpy() valid = mask != 255 for t, p in zip(mask[valid], pred[valid]): conf[t, p] += 1 iou = np.diag(conf) / (conf.sum(1) + conf.sum(0) - np.diag(conf) + 1e-6) return np.nanmean(iou)混淆矩阵按像素累计，最后一次性算 IoU。np.nanmean跳过没有样本的类。验证时一定要model.eval()和torch.no_grad()，否则显存爆得莫名其妙。学习率用 1e-3 配 Adam，或者 1e-2 配 SGD，U-Net 上 Adam 收敛更稳。
5. 苹果缺陷分割避坑：5 个真实踩过的坑
5.1 标签像素值不是从 0 连续编号
现象：训练 loss 一直不降，打印标签唯一值发现是 0、1、2、3、5，缺了 4。原因：标注工具或导出脚本跳号，模型输出 5 类但标签最大是 5，交叉熵的类别数对不上。解决：训练前强制重映射，把出现的类别 ID 映射到 0 到 N-1 连续空间，并同步更新类别数和权重数组。
5.2 图像和掩膜增强不同步导致错位
现象：训练几个 epoch 后 mIoU 卡在 0.2 上不去，可视化发现掩膜和图像旋转角度不一致。原因：图像和掩膜分别用了两套增强，或者用了只支持图像的增强库。解决：统一用 Albumentations 的image=, mask=同步接口，或者自己写增强时对两者用同一组随机参数。这个坑血泪经验，查了两天才发现。
5.3 验证集里小类样本为零导致 mIoU 虚高
现象：验证 mIoU 0.75 看着不错，但虫眼类一个都分不出来。原因：随机划分让验证集里虫眼样本为零，nanmean直接跳过该类，指标虚高。解决：按类别分层划分，并在评估时单独打印每类 IoU，不要只看平均值。如果某类验证样本少于 10 张，指标不可信。
5.4 忽略类 255 被当成真实类别参与训练
现象：模型预测结果里出现 255 这个类别，可视化一片白。原因：损失函数没设ignore_index=255，或者 Dataset 里把 255 转成了 long 但没在 loss 里忽略。解决：CrossEntropyLoss(ignore_index=255)，Dice 里也要手动排除 255。检查方法是打印标签最大值，如果大于类别数减一，就是没处理忽略类。
5.5 输入尺寸和显存不匹配导致训练中途 OOM
现象：前几个 batch 正常，跑到一半 CUDA out of memory。原因：图像尺寸不统一，Resize 没做或者做了但某些图特别大，batch 内 padding 到最大尺寸。解决：在 Dataset 里统一 Resize 到固定尺寸，或者用batch_size=1加梯度累积。显存 8G 以下建议 384×384 配 batch 4，别硬上 512。
6. 把 4000 张用到极致：小数据下的迁移学习与难例挖掘技巧
约 4000 张对语义分割来说属于小数据，从零训练 U-Net 很容易过拟合。我一般会先用 ImageNet 预训练的 ResNet34 或 EfficientNet 做编码器，解码器随机初始化，学习率编码器设小一点（1e-4），解码器设大一点（1e-3）。如果数据集里缺陷形态和自然图像差异大，预训练权重可以只用来初始化浅层，深层照样从头学。另一个技巧是难例挖掘：每个 epoch 结束后，把验证集里 IoU 最低的 50 张图挑出来，人工检查是标注问题还是模型问题，标注错的直接修，模型错的下一轮加进训练集重点学。这个循环跑三轮，mIoU 通常能涨 5 到 10 个点。
验证方法上，不要只看 mIoU。苹果缺陷分割最终要落到产线，你得看单类召回和误检率。比如腐烂类召回 0.9 但误检 0.3，产线上会把好苹果踢掉，这个指标比 mIoU 重要得多。我习惯在验证脚本里额外输出每类的 precision、recall 和混淆矩阵，混淆矩阵能直接看出哪两类在互相混。常见的是斑点和疤痕混，因为颜色接近，这时候要么加数据，要么在损失里给这两类的混淆对加惩罚项。
最后一个具体技巧：把 5 类里的背景单独拿出来看。如果背景像素占比超过 95%，考虑先做一个二分类的前景分割，再在前景区域里做 4 类细分。两阶段比直接 5 类分割在小数据上更稳，因为第一阶段任务简单，第二阶段只关注缺陷区域，类别不平衡的影响小很多。这个思路我在多个工业质检项目里用过，比硬调 5 类分割省事。希望帮到你。
本文还有配套的精品资源，点击获取
