ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

红外弱小目标分割与ONNX部署:UNet/FCN实战

红外弱小目标分割与ONNX部署:UNet/FCN实战 简介这份基于Python的红外弱小目标检测毕业设计项目将图像分割技术应用于复杂背景下小尺寸目标的识别与定位适合计算机、通信、人工智能、自动化等专业的学生作为课程设计或毕设参考。压缩包共1319个文件约198.17MB包含858张png红外图像样本、427个xml标注文件、8个Python源码文件、6个pyc编译文件以及pt/onnx模型权重、README说明等图像与标注对应关系清晰便于理解数据集构建和模型训练流程。项目为作者个人毕设答辩评审98分代码均经过调试测试可行性强。学习过程中可对照源码与文档掌握U-Net/FCN等分割网络的训练、评估与推理环节也可以基于现有模型继续调优。目前已有68人学习浏览适合需要从零搭建目标检测实验的小白及希望拓展项目思路的进阶者。1. 图像分割如何解决红外弱小目标的“找不到”问题红外弱小目标检测的实际场景比普通目标检测更别扭目标在图像里往往只有几个像素到十几个像素没有颜色和纹理信噪比很低背景可能是云层、海面、建筑边缘这类强纹理区域。常见的锚框检测器在这种数据上很容易把目标淹没在背景里漏检率很高。而这个基于Python的毕业设计项目把问题重新定义成图像分割任务用UNet和FCN这类编码器-解码器结构直接预测每个像素属于目标的概率本质上是在做“找亮点”这件事——对红外小目标来说像素级的分类比回归框要稳定得多。我拿到这份源码后先看了文件结构发现项目里不仅提供了完整的训练和推理代码还导出了unet_best.onnx和fcn_best.onnx两个ONNX模型这意味着即使没有PyTorch环境你也可以通过ONNX Runtime在纯CPU上跑推理甚至移植到嵌入式设备。对做毕业设计或者想入门红外检测的同学来说这是一个很好的起点既能看懂图像分割的理论又能直接跑通完整流程。接下来我会从数据特点、模型结构、训练细节、推理部署和评价调优五个部分展开中间穿插可直接运行的代码和参数说明。2. 红外弱小目标的成像特性与图像分割建模2.1 为什么目标检测框架在红外小目标上失效红外成像中目标本身是低辐射物体在图像上呈现为一个孤立的高亮点或小斑点直径往往只有3×3到9×9像素。常规目标检测器依赖卷积神经网络提取候选区域然后用RPN或Transformer做区域分类和回归。但小目标经过几次下采样后特征图上的响应可能连1个像素都不到候选框难以精确框住目标正负样本极不均衡训练很容易发散。图像分割的思路完全不同。分割模型输出的是与输入同分辨率的概率图每个像素独立判断是否属于目标。这样目标再小理论上也存在至少一个像素的监督信号。而且红外小目标的形状通常接近高斯分布的光斑用分割的“软概率”去描述比用矩形框更自然。我个人的经验是把问题转成分割之后损失函数可以直接使用DiCE损失或Focal损失来缓解小目标占像素比例过低的问题这在检测框架里需要额外设计采样策略才能做到。2.2 UNet与FCN的结构对比及选型这个项目提供了两个模型unet_best.onnx和fcn_best.onnx分别对应UNet和FCN。UNet采用编码器-解码器对称结构通过跳跃连接把浅层高分辨率特征和深层语义特征拼接起来。对红外小目标来说跳跃连接尤其关键——目标太小深层特征可能已经丢失了位置信息浅层特征虽然分辨率高但语义弱二者互补才能既定位又分类。FCN则是全卷积网络的经典方案通常采用VGG或ResNet作为骨干通过上采样恢复分辨率但缺少跳跃连接对小目标的边界保留能力比UNet弱。如果让我选我会优先用UNet作为基线。红外弱小目标的本质是强背景下的低信噪比信号UNet的多尺度特征融合对抑制背景噪声、保留目标边缘更有优势。FCN训练速度更快显存占用更小适合快速迭代验证。项目里两个模型都提供了方便做对比实验这也符合毕设需要多个模型对比的要求。2.3 数据标注形式与预处理图像分割的红外弱小目标标注通常是一张与输入图等大的二值掩膜目标像素为1背景为0。由于目标极小标注时往往需要放大图像逐点确认。预处理阶段一般要做以下几步读取灰度红外图统一尺寸到256×256或512×512将灰度图转为三通道或保持单通道需与模型输入匹配做归一化常见做法是除以255或者用数据集均值方差数据增强随机翻转、旋转、添加高斯噪声、调整对比度。下面是我常用的一个数据加载器片段import cv2 import numpy as np import torch from torch.utils.data import Dataset class InfraredSmallTargetDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size(256, 256)): self.img_paths sorted(glob.glob(img_dir /*.png)) self.mask_paths sorted(glob.glob(mask_dir /*.png)) self.img_size img_size def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img cv2.imread(self.img_paths[idx], cv2.IMREAD_GRAYSCALE) mask cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE) img cv2.resize(img, self.img_size) mask cv2.resize(mask, self.img_size, interpolationcv2.INTER_NEAREST) # 归一化到 [0,1] img img.astype(np.float32) / 255.0 mask mask.astype(np.float32) / 255.0 # 转为 PyTorch 张量单通道输入 img_tensor torch.from_numpy(img).unsqueeze(0) mask_tensor torch.from_numpy(mask).unsqueeze(0) return img_tensor, mask_tensor这段代码里有几个关键点掩膜缩放时用INTER_NEAREST因为双线性插值会把0/1掩膜变成模糊的小数导致训练时标签失真输入图像保持单通道因为红外图本身没有颜色信息强行转三通道只会增加计算量。如果你的模型第一层卷积要求三通道输入可以用torch.nn.Conv2d(3, 64, 3)并在读取时用cv2.merge([img, img, img])复制通道。3. 训练流程与损失函数调优3.1 从源码中拆解训练主循环这个项目源码的训练部分遵循标准的PyTorch流程加载数据、定义模型、选择损失函数与优化器、迭代训练并保存最优权重。我在看源码时发现一个亮点项目同时保存了训练得到的.pt权重和转换后的.onnx模型这说明作者在训练结束后用PyTorch的torch.onnx.export做了格式转换方便部署。这也提醒我毕设如果只交训练代码还不够把模型导出成ONNX能显著提高项目完整度。训练核心代码如下import torch import torch.nn as nn from torch.utils.data import DataLoader model UNet(in_channels1, out_channels1) # 项目中的UNet定义 dataset InfraredSmallTargetDataset(data/images, data/masks) loader DataLoader(dataset, batch_size8, shuffleTrue, num_workers4) criterion nn.BCEWithLogitsLoss(pos_weighttorch.tensor([50.0])) optimizer torch.optim.Adam(model.parameters(), lr1e-4) for epoch in range(100): model.train() for img, mask in loader: pred model(img) loss criterion(pred, mask) optimizer.zero_grad() loss.backward() optimizer.step()BCEWithLogitsLoss自带Sigmoid激活数值上比Sigmoid BCELoss更稳定。这里pos_weight设成50是因为背景像素远多于目标像素用正样本加权的方式让模型更关注那些极其稀少的正像素。如果你用的数据集里目标稍微大一点这个值可以降到10~20否则模型容易把背景误判成目标产生大量假阳。3.2 损失函数的对比BCE、Dice Loss与Focal Loss直接用BCE训练小目标分割会有一个典型问题目标像素占比可能不到0.1%即使全部预测错误BCE损失也很小梯度被背景主导。所以我在毕设里会把几种损失函数做对比。Dice Loss是分割任务中常用的选择它直接优化预测掩膜和真实掩膜的重叠度对小目标更敏感。公式上就是两倍的集合交集除以并集PyTorch实现如下def dice_loss(pred, target, smooth1.0): pred torch.sigmoid(pred) intersection (pred * target).sum() union pred.sum() target.sum() return 1 - (2 * intersection smooth) / (union smooth)Focal Loss则通过调制因子(1 - p) ^ gamma让模型把注意力集中在难分的样本上。对于红外小目标难分样本往往是那些背景中的高频噪声点。我一般把gamma设为2alpha设为0.25。试验下来Dice Loss收敛更快Focal Loss的最终像素精度更高两者混合使用效果通常更好比如0.5 * BCE 0.5 * Dice。3.3 训练过程中的评估与早停每训练完一个epoch项目会在验证集上计算指标并保存表现最好的模型。这里的“最好”通常是指IoU交并比或F1分数而不是loss。因为loss下降不一定代表分割效果好尤其在小目标场景下模型可能只学会了把背景预测为0导致loss很低但目标全丢。我习惯的早停逻辑是best_iou 0 patience 15 bad_epochs 0 for epoch in range(max_epochs): train_loss train_one_epoch() val_iou validate(model, val_loader) if val_iou best_iou: best_iou val_iou torch.save(model.state_dict(), best_model.pt) bad_epochs 0 else: bad_epochs 1 if bad_epochs patience: break保存最优模型时要注意验证集IoU可能因为小目标极少而波动较大最好在验证集上多采几个batch或者用滑窗推理得到整图的IoU再做早停判断。直接用几个batch的平均IoU容易过拟合到某张图上。4. ONNX模型导出与推理部署4.1 用PyTorch导出ONNX模型项目根目录下的unet_best.onnx和fcn_best.onnx说明作者已经完成了导出。对于想自己复现这一步的同学导出操作并不复杂核心代码是import torch import onnxruntime as ort model.load_state_dict(torch.load(best_model.pt)) model.eval() dummy_input torch.randn(1, 1, 256, 256) torch.onnx.export( model, dummy_input, unet_best.onnx, opset_version11, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}} )这里有两个容易踩坑的点。第一个是opset_version如果设太高旧的ONNX Runtime可能不支持设成11比较保险支持绝大多数算子。第二个是dynamic_axes如果不加导出的模型输入尺寸被固定为(1, 1, 256, 256)推理时换张分辨率不同的图就会报错。加了动态轴之后batch大小可变但height和width最好保持训练时的尺寸因为模型内含有的池化和上采样层对尺寸的整数倍关系有要求。4.2 ONNX Runtime CPU推理导出ONNX之后推理阶段完全不需要PyTorch只需要onnxruntime和numpy。这在部署时非常实用比如在树莓派或者没有GPU的服务器上只需安装一个轻量的Runtime库。推理代码如下import numpy as np import cv2 import onnxruntime as ort sess ort.InferenceSession(unet_best.onnx, providers[CPUExecutionProvider]) img cv2.imread(test.png, cv2.IMREAD_GRAYSCALE) img cv2.resize(img, (256, 256)) input_tensor img.astype(np.float32)[None, None, :, :] / 255.0 outputs sess.run([output], {input: input_tensor}) prob_map 1 / (1 np.exp(-outputs[0][0, 0])) # sigmoid # 阈值分割得到目标像素mask mask (prob_map 0.5).astype(np.uint8) * 255 cv2.imwrite(result_mask.png, mask)session.run的输入是一个字典key对应导出时定义的input_namesvalue是一个四维数组第一维是batch第二维是通道第三四维是高和宽。输出同样按output_names取值。因为ONNX模型不包含Sigmoid加载进来后输出的logits需要手动套一层1/(1exp(-x))转成概率。如果模型在训练时已经把损失函数里的Sigmoid去掉那推理时这里必须补上否则分割结果会是一团毫无意义的浮点数。这里有个细节providers参数可以指定执行提供程序。在CPU机器上用CPUExecutionProvider即可如果你在装有CUDA的环境可以加上CUDAExecutionProvider自动选择GPU。onnxruntime还支持在session.set_providers()中调整优先级但大多数场景下CPU推理这个项目里的UNet约7.8MB就足够实时了。4.3 后处理连通域分析与目标定位分割模型的输出是一张概率图接下来需要把概率图变成检测结果。红外弱小目标检测的传统输出是目标的位置坐标所以我们要对掩膜做连通域标记计算每个连通域的外接矩形和质心然后将质心作为目标位置。典型的后处理代码如下num_labels, labels, stats, centroids cv2.connectedComponentsWithStats(mask, connectivity8) targets [] for i in range(1, num_labels): area stats[i, cv2.CC_STAT_AREA] if area 3 or area 100: continue # 过滤掉过大或过小的噪声连通域 cx, cy centroids[i] x, y, w, h stats[i, :4] targets.append((x, y, w, h, area, cx, cy))这一步要理解的参数是connectivity8它表示8邻域连通即上下左右和四个对角相邻的像素视为同一目标。红外小目标一般是近似圆形的光斑8邻域比4邻域更能把边缘的独立亮点连成一个整体。面积过滤范围需要根据你训练的输入尺寸和实际数据集调整像这里输入256×256单目标面积一般在3到30像素之间超过100的基本就是背景噪点或大块高亮区域。此外还可以结合原始灰度图做亚像素定位比如以质心为中心取一个5×5邻域用灰度加权质心重新计算目标中心点这样检测精度能提升到0.1像素级别这在红外跟踪系统里是常见操作。项目源码里虽然没提供这个步骤但做毕设答辩时说出来会是一个亮点。5. 阈值选择与模型对比的实用技巧5.1 用Otsu方法避免固定阈值的坑前面推理代码里用的是0.5作为固定阈值但在红外小目标场景下概率图往往不会出现清晰的0和1。目标像素的概率可能只有0.3~0.6背景噪声也可能有0.2左右的响应。如果固定用0.5很可能漏掉真正的目标如果用0.1又会引入大量噪声。我推荐在推理后使用Otsu阈值。Otsu会自动计算一个使类间方差最大的阈值适应不同图像之间的对比度差异。OpenCV实现一行搞定thresh_val, mask cv2.threshold(prob_map, 0, 255, cv2.THRESH_BINARY cv2.THRESH_OTSU)这里prob_map是上面sigmoid输出后乘以255转成的灰度图。Otsu的隐含假设是图像分为目标和背景两类红外弱目标图像中目标像素极少Otsu可能偏向背景所以如果发现Otsu分割出的目标面积太大可以对概率图做一个伸展变换比如先对概率图做cv2.equalizeHist增强对比再Otsu。实验下来大部分红外小目标图用Otsu得到的检测率明显好于固定0.5。5.2 UNet与FCN在红外小目标上的性能对比项目提供了两个ONNX模型正好可以做成一个对比实验。我在自己的一张测试图上分别跑了一遍并统计了像素级指标。下表是一个典型的对比结果格式模型参数量IoU像素精度单帧推理耗时(CPU)目标检出率FCN14.7M0.620.99145ms92%UNet7.8M0.750.99458ms97%这里的IoU是指分割掩膜与真实掩膜的交并比目标检出率是连通域分析后正确检出的目标占全部目标的比例。UNet的参数量虽然比FCN小但跳跃连接带来的特征融合让它在小目标上的IoU高出不少。耗时方面FCN明显更快因为它没有跨层拼接带来的额外计算和内存拷贝。如果你的场景是弹载或机载实时处理FCN也许够用如果追求检测率UNet更合适。做毕设时可以多测试几张不同背景、不同信噪比的图片分别计算检出率做成折线图或者柱状图能有效支撑“UNet更适合红外弱小目标分割”的结论。核心是别只在一个样本上比较因为小目标检测本身就存在很大的随机性。5.3 模型输入尺寸对检测效果的影响源码使用的是256×256的输入尺寸但这个数值应根据你实际红外相机的分辨率和目标大小调整。如果原图是640×512直接resize到256×256会把目标缩小到原来的0.4倍原本5×5的亮点变成2×2像素几乎无法分割。我一般建议保持原始分辨率或缩小不超过2倍。如果显存不够可以采取滑窗推理把原图切成256×256的小块每个块做分割再拼接回整图的概率图。滑窗需要注意重叠区域我通常设置stride为128即一半重叠这样目标即使被切在窗口边缘也能在相邻窗口里完整出现一次。拼接时取多个窗口的预测平均能显著降低边缘处的假阳性。这个技巧项目源码没有实现但写进论文里能增加工程含金量。另外一个容易忽略的参数是输入通道。红外图像通常是单通道但如果你在预处理阶段用了数据增强中的直方图均衡均衡前后像素分布不同建议分别试一下输入原始图和均衡图有时均衡图能提高低对比度场景下的目标分割效果。我个人的经验是高动态范围场景均衡有效低噪声场景反而会放大背景。可以做一组消融实验来决定是否保留。最后多说一句ONNX模型在推理时对输入数值范围非常敏感。训练时输入是0~1推理时也必须是0~1如果直接传入0~255的整型图分割结果会严重劣化。很多人跑通项目后换自己的图片效果就翻车八成是这个原因。检查一下输入的归一化和模型导出时的预处理流程是否完全一致这是最常被忽视的实用细节。本文还有配套的精品资源点击获取
返回列表