简介:一份覆盖语义分割、实例分割与全景分割的图像分割Python项目,适合计算机视觉与深度学习方向的开发者、学生及科研人员,可直接用于毕设或课程设计参考。压缩包共5个文件,含3个Markdown说明文档、1个Python脚本与1个工程文件,整体约11KB,结构紧凑,便于快速理解项目逻辑。目前已有206人学习下载,适用于相关专业在校生、教师及企业算法工程师。内容涵盖基于深度学习方法的分割思路整理、CCF遥感图像分割赛题工程参考(含相关预处理工程与配置),以及标签转单通道等数据处理脚本,可帮助读者打通从预处理、模型搭建到工程部署的完整链路,亦可作为项目答辩或论文写作的基础素材。
1. 图像分割三兄弟:语义、实例、全景,先分清你要哪个
同样是“分割”,三个词背后是三套完全不同的技术选型。语义分割做像素级分类,把猫和狗各涂一片色就叫完成;实例分割要在这个基础上把两只黏在一起的猫拆成猫1和猫2;全景分割更狠,要求画面里每一块像素都有归属——人、车这种可数的“东西”各算各的,天空、道路这种不可数的“材料”也要覆盖到位。这个区别直接决定你下载来的深度学习图像分割源码能不能用:拿语义分割的项目去接实例分割的需求,改代码的痛苦远大于重新训练一个模型。
图像分割最密集的应用场景集中在医学影像(病灶区域提取)、自动驾驶(车道线和可行驶区域)、遥感解译、电商抠图和工业质检。适合谁?适合已经跑通过图像分类或目标检测、现在需要输出像素级结果的开发者和学生。如果你刚接触Python和深度学习,本文所有命令从环境配置开始,按步骤来;如果你已经跑过UNet,重点看第4章的排错和第5章的部署加速,那些才是生产环境里真正卡人的地方。
2. 环境与数据:从零跑通图像分割最小工程的硬指标
2.1 先固化Python环境:四个包版本能救命
图像分割项目最怕环境不一致。同一个UNet代码,在我机器上loss正常下降,换台机器就报shape mismatch,多半是库版本漂移。我一般会用conda单独建一个虚拟环境,Python版本直接锁定3.10,不要用系统自带的Python,也别用最新版,有些分割库对Python 3.12的适配还不完善。
conda create -n seg python=3.10 -y conda activate seg pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install opencv-python albumentations tqdm tensorboard pip install pycocotools提示:
--index-url这里指定的是CUDA 11.8的预编译版本。如果你用的是NVIDIA 30系以后的显卡,CUDA 11.8是兼容性最好的选择;如果你不装GPU版,直接pip install torch torchvision走CPU版本也可以跑通小数据集,但把batch size调小一些。
装完立刻验证GPU是否真的可用,这一步能排查掉80%的后续报错:
import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU mode")如果cuda.is_available()返回False,优先检查你安装的torch是不是CPU版本,用pip list | grep torch看版本号里有没有+cu后缀;其次是显卡驱动版本太旧,在命令行用nvidia-smi看驱动支持的CUDA版本号,必须大于等于你在pip里指定的CUDA版本。这一步很多人翻车,最常见的原因就是torch版本和驱动对不上。
另外强烈建议装一个segmentation-models-pytorch库,它把UNet、DeepLabV3、FPN等主流分割模型封装成了几行代码能调用的接口,做对比实验时能省下大量改模型结构的时间。但注意,这个库更适合快速验证想法,生产环境部署时还是要把模型结构固定下来,避免库升级导致推理结果变化。
2.2 数据准备:标注格式与目录结构决定后续所有代码
图像分割的数据格式主要分两种:一种是小数据集常用的单通道PNG掩码图,每个像素的灰度值就是类别编号;另一种是COCO格式的JSON标注文件,多边形坐标需要运行时转换为掩码。两种格式差异很大,下代码前先看清楚项目用的是哪种。
推荐目录结构按以下方式组织,约定一致后,训练、验证、推断脚本都不用改路径参数:
dataset/ images/ train/ # 原图,jpg或png val/ masks/ train/ # 语义分割掩码,单通道png val/ annotations/ # 实例/全景分割的json标注(如用COCO格式)如果拿到的是LabelMe标注的JSON文件,需要先转换为掩码PNG。这里给一个转换脚本的核心逻辑,用opencv从多边形填充掩码:
import json import numpy as np import cv2 import os def labelme_json_to_mask(json_path, img_shape, label_map): """ 把labelme的polygon标注转成单通道掩码, 背景为0,每个标注对象按label_map映射到对应类别编号。 """ with open(json_path, 'r', encoding='utf-8') as f: data = json.load(f) mask = np.zeros(img_shape[:2], dtype=np.uint8) # 生成空掩码 for shape in data['shapes']: label = shape['label'] class_id = label_map.get(label, 0) # 不认识的标签归为背景 points = np.array(shape['points'], dtype=np.int32) cv2.fillPoly(mask, [points], color=class_id) # 多边形填充 return mask # 假设label_map = {'tumor': 1, 'organ': 2} # 调用时对images/train每张图找到同名json即可逻辑说明:label_map是标签名到类别编号的映射表,语义分割要求掩码值从0开始且连续,因为损失函数计算交叉熵时类别索引就是像素值。cv2.fillPoly用多边形顶点坐标填充掩码区域,如果标注里同一个类别有多个对象(比如多张桌子),多次调用fillPoly累加即可,相同类别编号会自然融合。
这个转换有四个高频坑:一是确认图像通道顺序是BGR还是RGB,labelme原始图像读取后是BGR,但训练框架里预处理通常转RGB,掩码转换用的是坐标几何关系,不受通道顺序影响,但别忘了检查;二是多边形顶点坐标的坐标系原点在左上角,与图像像素坐标一致,不需要额外偏移;三是某些标注工具导出的是归一化坐标,需要乘图像宽高还原;四是掩码必须保存为PNG格式,不要用JPEG保存掩码,有损压缩会篡改类别编号,损失函数对错误的类标一无所知却会学出奇怪的东西。
2.3 第一次训练:从基线到看清损失曲线
第一次跑模型不需要花哨的结构,用UNet配合一个简单Dataset类就能建立基线。Dataset的核心是做好三个同步:图像读进来是什么格式、掩码读进来是什么格式、数据增强对两者做同样变换。
import torch from torch.utils.data import Dataset import cv2 import albumentations as A class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size=(256, 256), augment=False): self.img_paths = sorted(os.listdir(img_dir)) self.mask_dir = mask_dir self.img_size = img_size self.augment = augment # 训练时用albumentations做随机翻转、缩放,必须保证图和掩码同变换 self.transform = A.Compose([ A.RandomResizedCrop(256, 256, scale=(0.5, 1.0)), A.HorizontalFlip(p=0.5), A.RandomBrightnessContrast(p=0.2), ]) def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img_path = self.img_paths[idx] img = cv2.imread(os.path.join(self.img_paths_dir, img_path)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask = cv2.imread(os.path.join(self.mask_dir, img_path.replace('.jpg', '.png')), cv2.IMREAD_GRAYSCALE) if self.augment: augmented = self.transform(image=img, mask=mask) img, mask = augmented['image'], augmented['mask'] # 转为tensor,像素归一化到[0,1],掩码保持long类型 img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0 mask = torch.from_numpy(mask).long() return img, mask逻辑说明:RandomResizedCrop同时作用于img和mask,保证了空间对齐;HorizontalFlip是分割任务最常用的增强,因为分割标签的语义不受水平翻转影响。permute(2,0,1)把HWC转为CHW,PyTorch卷积网络默认输入是通道在前。掩码转换时用long()是因为交叉熵损失的target必须是整型。
训练时使用交叉熵损失配合AdamW优化器,学习率初始值1e-4,不要用默认的1e-3,分割任务的像素级预测对学习率更敏感,梯度爆炸的可能性更大。训练脚本循环里每过一定epoch做一次验证,记录mIoU和loss曲线,看到loss降到一定程度后不再下降,再开始调整模型结构或数据增强策略,不要一上来就换模型。
3. 三类分割全都要会:语义、实例、全景的最小可靠实现
3.1 语义分割:UNet的跳跃连接为什么对小目标友好
UNet是语义分割里最稳健的基线模型,尤其适合医学图像分割这类样本量不大的场景。它的核心设计是编码器-解码器结构加跳跃连接:编码器逐层下采样提取高层语义,解码器逐步恢复空间分辨率,跳跃连接把编码器的高分辨率特征直接拼接到解码器同层,让网络在预测像素类别时同时看到细节和语义。
一个精简UNet的核心结构如下,完整版比这个多几层block,但原理一致:
import torch.nn as nn class DoubleConv(nn.Module): """两个卷积+BN+ReLU,UNet的基本构建块""" def __init__(self, in_ch, out_ch): super().__init__() self.conv = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_ch=3, num_classes=2): super().__init__() # 编码器 self.enc1 = DoubleConv(in_ch, 64) self.enc2 = DoubleConv(64, 128) self.pool = nn.MaxPool2d(2) # 解码器 self.up2 = nn.ConvTranspose2d(128, 64, 2, stride=2) self.dec2 = DoubleConv(128, 64) # 跳跃连接后通道数翻倍 # 输出层 self.out = nn.Conv2d(64, num_classes, 1) def forward(self, x): # 编码 e1 = self.enc1(x) # 256x256 e2 = self.enc2(self.pool(e1)) # 128x128 # 解码 d2 = self.up2(e2) # 恢复空间尺寸 d2 = torch.cat([d2, e1], dim=1) # 跳跃连接,拼通道 d2 = self.dec2(d2) return self.out(d2)逻辑说明:DoubleConv里每个卷积后接BatchNorm2d和ReLU,BN在batch size较小时容易不稳定,所以UNet训练batch size至少设在8以上。ConvTranspose2d是转置卷积上采样,也可以用双线性插值上采样再卷积,后者参数更少且不容易产生棋盘效应。torch.cat在通道维度拼接,这是跳跃连接的关键,跳跃连接让解码器直接访问编码器的精细边缘信息,对分割结果有质的提升。
参数设置上有三个要点:第一,输入尺寸按训练时的尺寸固定,推理时图像尺寸如果与训练尺寸不一致,最好做resize而非直接输入,因为UNet全卷积结构理论上不限制输入尺寸,但BN层统计数据是基于训练尺寸的;第二,num_classes要和你的掩码类别数一致,二分类问题输出为1个通道加sigmoid还是2个通道加softmax,两种写法都在用,但推荐后者,收敛更稳定;第三,UNet系列找backbone时常用ResNet34或EfficientNet做编码器,segmentation-models-pytorch里直接支持Unet('resnet34', encoder_weights='imagenet')这种方式加载预训练权重。
3.2 实例分割:Mask R-CNN在torchvision里的快速落地
实例分割相比语义分割多出一个目标检测分支:每个实例要输出边界框、类别和像素掩码。Mask R-CNN是经典方案,在Faster R-CNN基础上增加了一个并行的掩码分支,RoI Align操作替代了RoI Pooling,解决了特征图对齐的像素偏差。
torchvision直接提供了预训练模型,这是快速落地的最佳路径:
import torchvision from torchvision.models.detection import maskrcnn_resnet50_fpn # 第一次调用会自动下载预训练权重,约178MB model = maskrcnn_resnet50_fpn(pretrained=True, num_classes=91) # 如果要在自己的数据集上微调,把分类头换掉 from torchvision.models.detection.faster_rcnn import FastRCNNPredictor # 假设你的数据只有1个类(加上背景共2类) num_classes = 2 in_features = model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) # 掩码分支的输出通道数也要改 from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, 256, num_classes)逻辑说明:maskrcnn_resnet50_fpn的预训练权重是在COCO上训的,如果目标域差异很大(比如医学影像),建议用pretrained=True做初始化然后全量微调;如果目标域接近COCO(比如通用物体),可以冻结backbone只训练检测头。替换分类头时,cls_score.in_features和conv5_mask.in_channels分别读取原模型对应层的输入维度,替换后原有权重参数会被随机初始化,只有这些层的梯度会更新。
实例分割的数据集格式一般整理成字典列表:每张图对应一个字典,包含boxes(Nx4的tensor,格式是x1, y1, x2, y2)、labels(N维tensor,从1开始,0保留给背景)、masks(NxHxW的二值tensor)。训练时数据加载需要自定义collate函数,因为每张图的实例数不同,不能简单堆叠成一个batch。torchvision官方文档给出的collate函数就是把每个样本的字典原样打包成list,实测最简单可靠。
3.3 全景分割:把thing和stuff统一到一个id空间
全景分割的实现思路是把实例分割(thing类)和语义分割(stuff类)的结果融合:thing类用实例分割来预测,stuff类用语义分割来预测,然后将两类结果统一编码到同一个标签空间。一种常见的工程实现是使用detectron2的Panoptic FPN模型,它用一个共享的backbone网络同时输出语义分割和实例分割结果。
安装detectron2在Windows上有些麻烦,官方建议在Linux环境编译或用预编译wheel包。下面给出推理的核心流程:
from detectron2.config import get_cfg from detectron2.engine import DefaultPredictor from detectron2 import model_zoo cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file( "COCO-PanopticSegmentation/panoptic_fpn_R_50_3x.yaml" )) cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url( "COCO-PanopticSegmentation/panoptic_fpn_R_50_3x.yaml" ) cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # 实例置信度阈值 predictor = DefaultPredictor(cfg) # 推理一张图 outputs = predictor(img) # img是BGR格式的numpy数组 panoptic_seg, segments_info = outputs["panoptic_seg"] # panoptic_seg是HxW的uint8张量,每个像素的编码方式是: # category_id * segment_id + segment_id的复合编码,需要用segments_info解码 for seg in segments_info: print(seg["category_id"], seg["isthing"], seg["area"])逻辑说明:panoptic_seg里每个像素的值不是直接对应类别编号,而是一个编码后的id,必须配合segments_info才能转换。segments_info里包含每个分割区域的category_id、isthing标记和面积。实际项目中要把panoptic_seg可视化,需要遍历segments_info,根据category_id查类别名到颜色的映射,然后把掩码区域填充为对应颜色。这一点是新手最容易踩坑的地方,千万别直接把raw输出当作类别掩码用。
全景分割对算力要求最高,如果业务只关心前景物体,不需要做全景分割;如果确实需要全场像素覆盖,也要确认stuff类的标注质量,因为stuff类区域面积大但标注粗糙,一旦训练数据里马路和路沿的边界标得含糊,模型输出就会在边缘处出现大片混叠,看起来比语义分割结果更糟糕。
4. 图像分割避坑实录:五个高发问题的现象、原因与解决
4.1 训练时loss为Nan,常见于前几个epoch就崩掉
现象:训练刚开始或中途loss值突然变成NaN,tensorboard曲线断崖式消失,GPU显存占用率骤降。
原因:最常见的是学习率太大导致梯度爆炸;其次是数据里有NaN的掩码值或者图像全黑;还有一部分是混合精度训练和某些归一化层冲突。
解决:先把学习率降到1e-5跑10个epoch确认loss可以下降,再逐步恢复。同时检查数据加载时是否过滤了不可读的图片,用np.isfinite(img).all()验证输入数据。混合精度训练的GradScaler不要一开始就用,先在单精度下跑通再启用。
4.2 验证集mIoU很高,但部署到真实图片上一塌糊涂
现象:训练集和验证集都是同一分布的数据,指标好看;部署到用户拍摄的图片上,分割结果出现大片缺失或错切。
原因:数据增强太弱,网络对光照变化、遮挡、相机视角变化过拟合了。很多开源数据集拍摄条件统一,训练出的模型很容易把背景颜色当成判别特征。
解决:增加强数据增强,比如随机亮度对比度、高斯噪声、随机仿射变换和弹性形变。把训练时用的增强方式记录到配置文件里,部署时无效。另外在验证集里混入一些域外图片作为烘干测试,即使不计算指标,目视检查也能看出泛化问题。
4.3 一个小目标被漏检,或者多个小目标融合成一整块
现象:小物体分割不出来,或者相邻的多个同类小物体在掩码里连成一片,视觉上是一整块。
原因:连续下采样过程中小目标信息丢失,特征图上只有几个像素响应;或者损失函数按像素平均,小目标在交叉熵中的占比极低,网络优化时优先保证大目标。
解决:一是用DiceLoss或FocalLoss替代纯交叉熵,DiceLoss对小目标的固有面积占比敏感度更低;二是在数据加载时做随机裁剪,把小目标相对放大;三是选择更高分辨率的输入,比如把训练尺寸从256提升到512,显存不够就用torch.utils.checkpoint梯度检查点技术来换取显存。
4.4 CUDA out of memory 高频发生,改小batch size后卡死了
现象:训练或推理时显存不足报错,把batch size减到1还是不够。
原因:显存除了保存输入、输出和中间激活值,还把整个计算图存下来用于反向传播。全卷积网络即使batch size为1,如果输入分辨率为1024,每一层的激活值累积起来照样能吃满12GB显存。
解决:先用torch.cuda.max_memory_allocated()统计显存峰值,定位是哪个阶段消耗最大。实践中最有效的手段是开启torch.cuda.amp.autocast()混合精度训练,显存直接减半;输入尺寸过大的用CenterCrop或Resize限制到模型能承受的范围。推理阶段用torch.no_grad()包裹,同时用model.eval()关闭dropout和BN的统计更新,这两步不做显存占用差距可达30%以上。
4.5 全景分割结果里thing类边界和stuff类边界重叠或留下缝隙
现象:panoptic输出里,人形区域边缘出现了细小的背景碎片缝隙,或者车辆边缘被“吃进去”一圈。
原因:thing类和stuff类的预测是分开的,后处理融合时没有做冲突消解。一些像素同时被实例分割判定为“车辆”又被语义分割判定为“道路”,融合规则如果简单按类别优先级覆盖,就会出现边缘不整齐。
解决:对照官方实现的merge_semantic_and_instance函数,规则是:实例分割的置信度高于阈值时,直接覆盖语义分割结果;低于阈值时遵从语义分割。这个阈值的选取需要调参,一般设在0.3到0.5之间,日志记录里要输出冲突像素数量,方便判断阈值是否过严或过松。
5. 部署前必须做的最后两步:ONNX导出与预测可视化
模型训练完不等于项目结束,生产环境最常用的是ONNX Runtime或TensorRT推理。把PyTorch模型导出为ONNX格式,需要固定输入尺寸,否则动态维度会让导出的图变得复杂且性能下降。
import torch import onnxruntime as ort # 假设model是训练好的UNet,输入尺寸256x256 model.eval() dummy_input = torch.randn(1, 3, 256, 256) torch.onnx.export( model, dummy_input, "unet.onnx", opset_version=12, input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}} ) # ONNX Runtime推理 ort_session = ort.InferenceSession("unet.onnx", providers=["CUDAExecutionProvider"]) ort_inputs = {"input": img_numpy.astype(np.float32)} outputs = ort_session.run(None, ort_inputs)[0] # 拿到分割概率图提示:导出ONNX前务必用
torch.jit.trace先跟踪一遍模型,确认没有控制流依赖输入数据的分支,否则导出会报错或得到错误结果。UNet这类全卷积模型结构固定,导出一帆风顺;如果模型里用了基于数据检测的动态判断逻辑,就要改写模型结构再导出。
ONNX导出成功后,我的调试习惯是先写一个可视化脚本,把每张验证图片的预测掩码叠加到原图上,保存为一张带原始图、标签图、预测图三栏的拼图。这个脚本写好后,每次调参训练完直接跑一遍,用眼睛看图而不是只看mIoU数字。mIoU高并不能保证边缘细节正确,肉眼扫一遍比任何指标都能更快发现模型在局部结构上的问题。
关于ONNX的优化,如果显存充足,使用TensorRT把模型精度转成FP16,推理延迟能降到原来的三分之一左右。TensorRT的engine文件与GPU型号强绑定,换卡就要重新生成。如果你的项目只支持CPU推理,ONNX Runtime的CPU线程数默认会占满所有核,在并发场景下要设置intra_op_num_threads限制线程数,否则多个推理请求会互相抢CPU导致整体吞吐量不增反降。这个细节是我在部署时踩过的坑,单路推理延迟很好看,并行一压测就暴露出来。
图像分割项目沿用至今的个人习惯是先做数据质量检查再做模型调优。拿到任何开源代码,先跑一个batch的图片和掩码可视化出来人工核对,确认标签对齐了再开始训练。很多问题追根溯源,不是模型结构不行,而是数据从标注到掩码的转换那一环出了错。希望这些经验能帮你把项目的踩坑时间压缩下来,一次跑通。
本文还有配套的精品资源,点击获取