ARTICLE DETAIL

资讯详情

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

虾图像分割数据集实战:二值分割与U-Net训练全解析

虾图像分割数据集实战:二值分割与U-Net训练全解析 简介海洋生物虾图像分割数据集面向深度学习二值分割任务包含可直接投入使用的训练集与测试集适合图像分割入门练习、模型训练与算法对比验证尤其契合需要规范标注数据快速搭建实验流程的开发者。数据按目录分层组织训练集含800张原始图像与800个对应GT掩膜测试集含200张原始图像与200个对应GT掩膜解压后即可直接用于模型输入无需额外预处理。压缩包共2000个文件以png图片为主共1999张另附1个Python可视化脚本可随机抽取样本并同时呈现原始图、GT掩膜与GT叠加蒙版效果便于快速检查标注质量和训练预测差异。资源包约320.3MB目前已有278人学习下载适合需要标准二值分割数据开展实验、复现分割流程或积累项目经验的开发者使用。1. 图像分割数据集这件事为什么虾比猫狗更难做拿到《深度学习之图像分割数据集海洋生物虾图像分割数据集二值图像分割任务》这个标题第一反应是它和常见的“猫狗分割”“车辆分割”数据集有本质区别这是一份二值图像分割任务专用数据集目标只有两类——前景虾和背景水体/池底输出是一张和原图同尺寸的 0/1 mask。这类数据集在养殖密度估计、虾体长度测量、投饵量控制、病害个体筛查里是刚需落地场景比公开的通用分割数据集更聚焦也因此更适合用来跑通图像分割的完整训练链路。但虾有一个让分割模型很头疼的特点虾体是半透明的边缘和背景的灰度差极小尤其在养殖池的深色底质或LED补光场景下虾须、虾足、腹部节段的轮廓几乎融进背景里。用 COCO 或 Cityscapes 上表现很好的模型直接搬过来往往会得到边缘毛糙、虾须断裂、小目标漏检的结果。这篇笔记就围绕这份数据集把二值分割从数据组织、训练配置、评估指标到踩坑记录完整讲一遍按我的实操习惯给参数和命令读者可以照着复现。2. 先摸清二值分割数据集的底细目录结构、标签格式与预处理管线2.1 一份二值分割数据集的标准长相二值分割数据集的组织方式无论发布方怎么命名最终都会落成一套“图像对”结构一张是原始 RGB 图一张是同尺寸、同文件名的单通道 mask。我的习惯是直接按 VOC 风格组织因为后续接 U-Net、DeepLab、PaddleSeg 都不用改代码shrimp_seg/ ├── JPEGImages/ # 原始图像jpg 或 png │ ├── shrimp_001.jpg │ └── shrimp_002.jpg ├── SegmentationClass/ # 标签图单通道 PNG │ ├── shrimp_001.png │ └── shrimp_002.png ├── ImageSets/ │ └── Segmentation/ │ ├── train.txt │ └── val.txt这里有个容易被新手忽略的硬规矩mask 必须存 PNG不要存 JPG。JPG 是有损压缩边缘一圈会出现 248、251、252 这类“零碎值”二值分割的 mask 一旦出现非 0 非 255 的像素交叉熵损失会把这些噪声当成独立类别去拟合轻则训练震荡重则边缘预测一团糊。用 PIL 读取时也要用Image.open(path)而不是cv2.imread(path, cv2.IMREAD_COLOR)否则单通道图会被强行复制成 3 通道通道数对不上就要 debug 半天。2.2 训练集与验证集拆分按视频帧去重是第一个坑图像分割数据集的拆分比分类数据集敏感。拍摄虾的视频抽帧得到的图像“同一只虾在相邻帧里位置几乎没变”如果随机打乱再切 train/val验证集会包含训练集的近邻帧IoU 虚高 35 个点。我的做法是按视频源分组再在组内抽帧。import os import glob from sklearn.model_selection import GroupShuffleSplit images sorted(glob.glob(JPEGImages/*.jpg)) groups [os.path.basename(p).split(_)[0] for p in images] # 假设文件名形如 video01_0001.jpg取 video01 作为视频分组依据 splitter GroupShuffleSplit(n_splits1, test_size0.2, random_state42) train_idx, val_idx next(splitter.split(images, groupsgroups)) with open(ImageSets/Segmentation/train.txt, w) as f: for i in train_idx: f.write(os.path.basename(images[i]).replace(.jpg, ) \n) with open(ImageSets/Segmentation/val.txt, w) as f: for i in val_idx: f.write(os.path.basename(images[i]).replace(.jpg, ) \n)这段代码的逻辑是GroupShuffleSplit把属于同一视频源的帧归为一组划分时保证 train 和 val 里出现的是不同的视频源。random_state42固定随机种子保证多次运行拆分结果一致。对一份虾分割数据集来说如果原始素材来自 3 段不同光照条件的视频按视频分组能逼着模型学到“虾本身长什么样”而不是“这段视频的背景长什么样”。2.3 预处理尺寸策略决定小目标存亡虾分割的预处理有两套路线取决于原图分辨率和显存。原图如果是 1920×1080 甚至更高直接整图丢进模型训练12GB 显存基本撑不住。常见做法是训练时 resize 到 512×512 或 768×768。这里有一个分水岭虾太小的样本resize 会把虾缩成十几个像素模型根本学不到结构信息。我的规则是目标长度虾体主轴低于图像短边的 1/10 时放弃整图 resize改做随机裁剪让裁剪窗口刚好覆盖单只虾。import cv2 import numpy as np from torch.utils.data import Dataset class ShrimpSegDataset(Dataset): def __init__(self, root, splittrain, crop_size512): self.root root self.split split self.crop_size crop_size with open(f{root}/ImageSets/Segmentation/{split}.txt) as f: self.names [line.strip() for line in f.readlines()] def __getitem__(self, idx): name self.names[idx] img cv2.imread(f{self.root}/JPEGImages/{name}.jpg) mask cv2.imread(f{self.root}/SegmentationClass/{name}.png, cv2.IMREAD_GRAYSCALE) mask (mask 127).astype(np.uint8) # 统一转成 0/1避免 0/255 混用 if self.split train: h, w img.shape[:2] y np.random.randint(0, h - self.crop_size 1) x np.random.randint(0, w - self.crop_size 1) img img[y:yself.crop_size, x:xself.crop_size] mask mask[y:yself.crop_size, x:xself.crop_size] img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img img.astype(np.float32) / 255.0 img (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] img np.transpose(img, (2, 0, 1)) # HWC - CHWPyTorch 默认输入格式 return img, mask.astype(np.int64)预处理里有几个参数和细节值得说透。mask 127做了阈值二值化把标注软件可能导出的 0/255 统一收敛成 0/1后端损失函数不用猜标签范围。归一化用 ImageNet 的均值标准差这虽然是迁移学习沿袭下来的习惯但对虾这类边缘纹理敏感的任务均值为零、方差统一的输入能让梯度更新更平稳实测比不归一化收敛快 30% 左右。np.transpose这一步不能省PyTorch 的卷积层默认接受 NCHW忘了转置会直接报维度错误。3. 用 U-Net 训练虾分割模型网络选型、损失函数与完整训练脚本3.1 为什么二值分割首选 U-Net 而不是 DeepLabV3单类别二值分割是 U-Net 的主场理由有三条。第一U-Net 的跳跃连接把 encoder 的浅层特征直接拼到 decoder而虾恰恰是“边缘信息比语义信息更重要”的目标——虾须、虾足的位置全靠浅层纹理特征这些特征在深层网络里早就被池化抹掉了DeepLabV3 虽然也有浅层融合但它的 ASPP 模块偏向多尺度语义对细长结构的敏感度不如 U-Net。第二U-Net 的参数量在 30MB 左右一份几百到几千张的虾分割数据集完全喂得饱不容易过拟合。第三推理时 U-Net 对输入尺寸不敏感patch 大小可以直接复用训练时的 crop省去 resize 带来的形变误差。当然如果数据集规模超过 5000 张且目标是正式交付我会换 DeepLabV3 用 ResNet50 做 backbone它在边界平滑度上略好但起步阶段 U-Net 是性价比最高的选择。3.2 损失函数BCE 与 Dice 的组合拳比例怎么定二值分割的损失函数有一个显著问题背景像素占比 90% 以上是常态。虾体在画面里占比通常只有 5%15%用纯 BCE二元交叉熵训练模型会学成“全预测背景”loss 照样很低但 IoU 为 0。解决办法是引入 Dice Loss它直接优化区域重叠度天然免疫类别不平衡。我的经验是两者按 0.5:0.5 加权前 10 个 epoch 用 BCE 主导帮助稳定收敛后面 Dice 的权重逐步提升来精修边界import torch import torch.nn as nn import torch.nn.functional as F class CombinedLoss(nn.Module): def __init__(self, bce_weight0.5, dice_weight0.5, smooth1.0): super().__init__() self.bce_weight bce_weight self.dice_weight dice_weight self.smooth smooth def forward(self, logits, targets): bce F.binary_cross_entropy_with_logits(logits, targets.float()) preds torch.sigmoid(logits) preds_flat preds.contiguous().view(-1) targets_flat targets.contiguous().view(-1) intersection (preds_flat * targets_flat).sum() dice 1.0 - (2.0 * intersection self.smooth) / ( preds_flat.sum() targets_flat.sum() self.smooth ) return self.bce_weight * bce self.dice_weight * dice代码里有两处需要说明。binary_cross_entropy_with_logits接收的是未过 sigmoid 的 logits内部自己做了数值稳定的 sigmoid 交叉熵直接喂 logits 就好不要提前在外面套 sigmoid否则梯度会变软。Dice 系数计算里的smooth参数是平滑项既防止分子分母都为零时除零崩溃也起到标签平滑的效果——我一般设 1.0这个值对结果不敏感别在这个参数上花时间调。3.3 完整训练脚本学习率、batch size 与早停import torch from torch import nn from torch.utils.data import DataLoader from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0.0 for images, masks in tqdm(loader, descTraining): images, masks images.to(device), masks.to(device) masks masks.unsqueeze(1).float() # (B, H, W) - (B, 1, H, W) optimizer.zero_grad() logits model(images) loss criterion(logits, masks) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) return total_loss / len(loader.dataset) model UNet(in_channels3, out_channels1).to(device) criterion CombinedLoss(bce_weight0.5, dice_weight0.5) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-5) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50) for epoch in range(50): train_loss train_one_epoch(model, train_loader, optimizer, criterion, device) val_iou, val_dice evaluate(model, val_loader, device) scheduler.step() print(fEpoch {epoch1:02d} | Loss: {train_loss:.4f} | Val IoU: {val_iou:.4f}) if val_iou best_iou: best_iou val_iou torch.save(model.state_dict(), best_shrimp_seg.pth)训练参数的选择直接决定能不能收敛到可用水平。学习率初始 1e-4 搭配 AdamW比 1e-3 的 Adam 稳定——虾的 mask 边界很细梯度在边界像素上波动大太大的学习率会让边缘预测在“整片黑”和“整片白”之间震荡weight_decay1e-5 是轻量 L2 正则用来压住背景区域可能出现的棋盘格伪影。CosineAnnealingLR 把学习率从 1e-4 按余弦曲线降到近零后半程的极低学习率负责“抛光”边界细节这是我能稳定拿到良好边界质量的关键设置。masks.unsqueeze(1)不能忘分割任务的 mask 是单通道但卷积层的输出是 (B, 1, H, W)不补维度会直接报错。4. 评估与可视化Dice、IoU 的计算姿势和预测 mask 的保存4.1 指标体系只有 loss 曲线不够要盯着 IoU 和 Dice 两个数图像分割评估只看 loss 是新手最容易犯的错。BCE 降到 0.01 以下不代表分割得好——背景占了 90%背景像素全部预测正确就能把 BCE 压到极低。二值分割的标准度量是 IoU交并比和 DiceF1 的像素级变体。def evaluate(model, loader, device): model.eval() iou_list, dice_list [], [] with torch.no_grad(): for images, masks in loader: images, masks images.to(device), masks.to(device) logits model(images) preds (torch.sigmoid(logits) 0.5).float() preds_flat preds.view(-1) masks_flat masks.view(-1) intersection (preds_flat * masks_flat).sum().item() union preds_flat.sum().item() masks_flat.sum().item() - intersection iou intersection / (union 1e-8) dice (2.0 * intersection) / ( preds_flat.sum().item() masks_flat.sum().item() 1e-8 ) iou_list.append(iou) dice_list.append(dice) return np.mean(iou_list), np.mean(dice_list)阈值取 0.5 是惯例但实际操作中我会额外试 0.3 和 0.7如果 0.3 的 IoU 明显高于 0.5说明模型预测整体偏保守、虾体区域概率值不够高可能是训练不够或者损失函数里 Dice 权重太大。这个“阈值敏感性测试”虽然不严谨但作为快速诊断手段很有效。4.2 预测 mask 可视化叠加图比单纯存 mask 更能发现问题训练完只保存 mask 数组是看不到问题的。我习惯把预测结果叠加到原图上用颜色区分真阳性、假阳性、假阴性这样一张图就能定位模型的系统性错误——是边缘整体外扩还是虾须断裂还是错把池底的饲料残渣当成了虾。import cv2 import numpy as np def overlay_prediction(image_path, mask_path, pred_path, output_path): img cv2.imread(image_path) gt cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) 127 pred cv2.imread(pred_path, cv2.IMREAD_GRAYSCALE) 127 overlay img.copy() overlay[pred gt] (0, 255, 0) # 真阳性 - 绿色 overlay[pred ~gt] (0, 0, 255) # 假阳性 - 红色 overlay[~pred gt] (255, 0, 0) # 假阴性 - 蓝色 cv2.imwrite(output_path, overlay)这份代码的三行布尔索引就是诊断的全部核心。绿色区域是预测对的红色是模型“多画”的蓝色是模型“漏画”的。如果蓝色集中在虾须末端说明模型对细长结构的感受野不够如果红色集中在水底纹理处说明模型并没有真正学到虾的形态而是拿背景纹理当特征。基于这个可视化结果再决定调损失权重还是换模型架构比盯着 loss 曲线猜靠谱得多。5. 二值分割避坑指南虾数据集上最容易翻车的 4 个问题5.1 mask 边缘出现 12 像素的“亮边”模型预测整体外扩现象训练出来的模型在虾体轮廓外圈有一圈稳定的误检IoU 卡在 0.85 左右上不去。原因标注工具LabelMe、EISeg导出的多边形在栅格化时有插值边缘像素的灰度介于 0 和 255 之间。虽然我做了 127的阈值化但二值化方向如果反了边缘会整体外扩一圈。解决做一次形态学腐蚀把 mask 边缘向内收缩 1 个像素import cv2 import numpy as np kernel np.ones((3, 3), np.uint8) mask cv2.imread(SegmentationClass/shrimp_001.png, cv2.IMREAD_GRAYSCALE) mask_bin (mask 127).astype(np.uint8) mask_eroded cv2.erode(mask_bin, kernel, iterations1) morph (mask_eroded * 255).astype(np.uint8) cv2.imwrite(SegmentationClass/shrimp_001.png, morph)cv2.erode对二值图的作用是“蚕食”前景边缘一圈iteration1 就够多了会丢失虾须这类细结构。这么做本质是把标注工具的系统性误差在预处理阶段抵消掉比让模型自己去学“真实边缘比 mask 小一圈”要省事得多。5.2 训练 loss 降不下去卡在 0.35 左右不再变化现象训练到第 20 个 epoch训练 loss 和验证 loss 都纹丝不动Dice 也不涨。原因最常见的是 mask 和原图没对齐——原图是 1920×1080mask 却是 960×540DataLoader里直接crop导致 mask 和图像内容错位。分割任务里这类“标签整体偏移几个像素”的错误loss 不会报错只会让网络永远学不到精确边界。解决写一段校验代码遍历所有 train.txt 里的样本打印 img.shape 与 mask.shape不等的直接列出文件名with open(ImageSets/Segmentation/train.txt) as f: names [line.strip() for line in f] for name in names: img cv2.imread(fJPEGImages/{name}.jpg) mask cv2.imread(fSegmentationClass/{name}.png, cv2.IMREAD_GRAYSCALE) if img.shape[:2] ! mask.shape: print(fMismatch: {name}, img{img.shape}, mask{mask.shape})如果确认是标注工具导出分辨率不一致解决的笨办法是把全部分辨率统一到原图尺寸——用cv2.resize插值到原图大小。但要注意resize 之后的 mask 必须重新做一次 127阈值因为插值会在边缘产生大量中间灰度值。5.3 服务器训练时 OOM卡在第一个 epoch 就爆显存现象batch size 设为 16输入 512×51212GB 显存直接 OOM。原因U-Net 的显存占用主要来自 decoder 阶段的特征图拼接512×512 输入、batch size 16、通道 64 起步的 U-Net显存占用逼近 14GB。解决显存不够时的正确操作不是调小 batch size——batch size 小于 4 时 BatchNorm 的统计量会很不稳定模型容易震荡。正确做法是减小输入尺寸。512→384 的显存占用能降到一半而虾分割对精度的损失远小于 batch size 减半带来的训练不稳定。# 显存紧张时使用384x384 输入batch size 8 train_loader DataLoader( dataset, batch_size8, shuffleTrue, num_workers6, pin_memoryTrue )另一招是打开cudnn.benchmark True它会为当前输入尺寸搜索最快的卷积算法一般能省 10%20% 的显存import torch.backends.cudnn as cudnn cudnn.benchmark True5.4 验证集 IoU 很高但看到新环境照片就崩现象在验证集上 IoU 0.93换一批不同光照、不同水质条件下的虾图IoU 掉到 0.6。原因数据集的“域偏移”。虾分割数据集如果全部来自同一段视频、同一种光照模型学到的其实是“这个池子的背景特征”而非“虾的通用形态”。这是分割任务里最常见的测试集性能陷阱。解决没有捷径必须在数据收集阶段就覆盖多场景。如果数据集已固定无法补充能做的只有增强训练时随机调整亮度、对比度、色相模拟不同水质和光照from torchvision import transforms train_transform transforms.Compose([ transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.2), ])注意不要加随机旋转 90°。虾有明确的头尾轴向旋转 90° 会让模型把“横着的虾”和“竖着的虾”当成两个类别去学浪费拟合能力。6. 从跑通到用好推导式后处理与小目标提升技巧模型训完、指标达标距离交付还差一步后处理。对虾分割来说最有价值的是用连通域分析过滤噪声和针对性提升小目标召回两个技巧。连通域过滤处理的是假阳性碎片模型偶尔会在背景水花、气泡或饲料残渣上产生小面积误检。这类误检的特点是面积小、形状不规则。用 OpenCV 的连通域分析把面积小于阈值的预测块直接删掉是最快见效的“最后一公里”优化import cv2 import numpy as np pred cv2.imread(pred.png, cv2.IMREAD_GRAYSCALE) pred_bin (pred 127).astype(np.uint8) num_labels, labels, stats, _ cv2.connectedComponentsWithStats(pred_bin, connectivity8) cleaned np.zeros_like(pred_bin) for i in range(1, num_labels): # 0 是背景跳过 area stats[i, cv2.CC_STAT_AREA] if area 200: # 面积阈值按数据集分辨率调 cleaned[labels i] 1 cv2.imwrite(pred_cleaned.png, cleaned * 255)connectedComponentsWithStats返回的stats是每个连通域的包围盒和面积信息connectivity8表示 8 邻域连通对虾须这种细长结构必须用 8 邻域4 邻域会把一条连续的虾须断成好几截。面积阈值 200 在 512×512 输入下相当于 0.07% 的像素占比低于这个值的预测块基本可以认定为噪声。小目标提升则要从模型和数据两侧同时做。模型侧把 U-Net 的深度降到 3 层而不是默认的 4 层浅层特征的感受野更小对虾须末端的响应更敏感数据侧训练时用 448×448 的随机裁剪代替固定尺寸相当于隐式地做了多尺度训练让模型在不同尺度下都见过虾的形态。这两个改动都不会让大目标精度下降但对小虾的召回通常能提升 35 个 IoU 点。还有一个我常用的验证技巧把测试图的预测 mask 转成 RLE 编码跑 COCO eval 那种格式用pycocotools跑一下整个测试集的mAP[0.5:0.95]。二值分割很多人不跑这个但 mAP 的 IoU 阈值从 0.5 到 0.95 取平均值比只看 IoU0.5 更能暴露边界精度的短板——IoU0.5 高不代表边缘好mAP[0.5:0.95] 才能反映预测边界和真实边界的重合程度。我自己的习惯是每次训完模型先跑一遍连通域清洗再用测试集里最暗的一张图做叠加可视化。如果暗光下虾须没有断裂这个模型才算真正能交付。虾分割这事说来也玄学指标好不等于效果好真到现场拍回来的图总会给你一两个没想到的惊喜。不过把数据组织、训练配置和这四类坑都趟过一遍之后再遇到新场景至少心里有底了。希望帮到你。本文还有配套的精品资源点击获取
返回列表