前言本文介绍了用于轻量端到端颜色校正的倒残差注意力模块 IRA融合了 MobileNet 风格倒残差结构与通道、空间注意力机制。该模块通过深度可分离卷积高效提取局部特征并利用注意力自适应强化与颜色、纹理、边缘相关的关键信息在较低计算成本下提升特征选择能力与表达稳定性。我们将 IRA 成功集成进 YOLO11 的 C3k2 模块构建 C3k2_IRA用于增强主干与颈部网络的局部细节建模和有效特征筛选能力从而提升目标检测中的特征表达效果。文章目录 YOLO11改进大全卷积层、轻量化、注意力机制、损失函数、Backbone、SPPF、Neck、检测头全方位优化汇总专栏链接: YOLO11改进专栏介绍近年来快照式多光谱MS成像的发展使紧凑、低成本的光谱传感器能够应用于消费级和移动设备。相比传统 RGB 传感器这类系统能够捕获更丰富的光谱信息从而增强包括颜色校正在内的关键成像任务。然而大多数现有方法将颜色校正流程拆分为多个独立阶段并且往往在流程早期就丢弃 MS 数据。本文提出一种统一的、基于学习的框架用于执行端到端颜色校正并联合利用来自高分辨率 RGB 传感器和辅助低分辨率 MS 传感器的数据。该方法将完整颜色校正流程整合到单一模型中生成一致且颜色准确的输出。作者通过重构两种不同的先进图像到图像架构展示了该框架的灵活性和通用性。为支持训练和评估本文聚合并重新利用公开光谱数据集构建了专用数据集并在多种 RGB 相机光谱灵敏度下进行渲染。大量实验表明该方法提升了颜色准确性和稳定性相比仅 RGB 和 MS 驱动的基线误差最高降低 50%。代码、模型和数据集均已公开。文章链接论文地址论文地址代码地址代码地址基本原理1. 解决的关键问题IRA 所在的这篇论文主要解决移动端相机颜色校正中 RGB 信息不足、传统管线误差逐级传播以及多光谱信息利用不充分的问题。传统相机颜色校正通常被拆成自动白平衡AWB和颜色空间变换CST两个阶段前者估计并抵消场景光照颜色后者将白平衡后的相机 RGB 映射到 CIE XYZ 或 ProPhoto RGB 等设备无关颜色空间。这种模块化流程虽然高效但每个阶段相互独立前一步的照明估计误差会影响后续颜色变换复杂光照下容易出现偏色和不稳定。更根本的问题在于普通 RGB 传感器只有三个宽光谱通道难以充分区分物体反射率和场景光照。当不同光照、不同材质或不同相机光谱响应组合在一起时RGB 数值可能高度相似但真实颜色却不同这会造成颜色恒常性和颜色空间映射的歧义。快照式低分辨率多光谱传感器可以提供额外光谱线索但现有方法常常只把 MS 信息用于照明估计之后仍然回到传统 RGB 管线导致多光谱信息没有贯穿整个颜色校正过程。IRA 模块所在的 LPIENet 改造方案就是为了在轻量图像到图像网络中更有效地处理 RGB 与 MS 特征让网络端到端地联合完成照明估计、照明抵消和颜色空间变换。2. 整体架构论文提出的是一个双输入端到端颜色校正框架输入包括高分辨率 RGB 图像和辅助低分辨率多光谱图像输出为颜色校正后的 CIE XYZ 图像。作者没有从零设计全新网络而是重构了两类已有轻量图像到图像架构其中 LPIENet 分支的核心组件就是 IRA block。LPIENet 原本面向移动端实时图像增强采用类似 U-Net 的编码器-解码器结构适合需要空间一致输出的颜色校正任务。在改造后的 LPIENet 中RGB 图像首先经过 RGB encoder 提取多尺度特征网络包含三个编码块、两个解码块、跳跃连接以及最终 refinement block。为了引入多光谱模态论文额外加入一个 spectral encoder该分支镜像 RGB encoder 的结构由三个 IRA block 组成但不执行下采样以适配低分辨率 MS 输入并提取光谱辅助特征。每一层多光谱编码器提取的特征会通过 skip connection 与 RGB 编码器对应层的特征融合具体采用逐元素相加并在每个解码阶段之前注入。这样RGB 分支负责保留高分辨率空间结构MS 分支负责提供更丰富的光谱先验IRA 则作为两个分支中的轻量特征提取与注意力增强单元为后续颜色校正输出提供稳定表征。3. 技术原理IRA 的全称是 Inverted Residual Attention即倒残差注意力模块。它结合了 MobileNet 风格的轻量卷积结构和 CBAM 类似的通道/空间注意力思想。倒残差结构的价值在于以较低参数量和计算量完成有效特征变换通常先通过 1×1 卷积进行通道变换再使用深度可分离卷积提取局部空间信息最后再通过 1×1 卷积投影回目标通道维度并通过残差连接保持信息流动。这种设计非常适合移动端颜色校正因为颜色校正需要对整幅图像产生空间连续、细节一致的输出同时又不能引入过重计算。在 IRA 中注意力机制进一步增强了倒残差卷积的选择能力。通道注意力用于判断哪些特征通道对当前颜色校正更重要例如与光照颜色、材质反射、局部色偏相关的响应空间注意力则用于定位图像中更需要校正或更具参考价值的区域例如高亮区域、阴影区域、色彩边界和材质变化位置。论文图示中IRA block 包含 IRB 卷积分支并并行结合 channel attention 与 spatial attention最终通过逐元素乘法和残差/融合操作强化有效特征。相比纯卷积模块IRA 能更自适应地选择颜色相关信息相比重型 Transformer 注意力它又保持了移动端友好的轻量特性。在多光谱颜色校正场景中IRA 的作用不仅是增强 RGB 图像特征还承担了光谱特征编码的基础单元。对于 spectral encoder 来说三个 IRA block 从低分辨率 MS 图像中提取与照明、反射率和相机光谱响应相关的辅助特征这些特征通过跳跃连接注入 RGB 解码流程使网络不再只依赖 RGB 三通道猜测真实颜色而是在端到端映射中持续利用 MS 线索。实验结果也支持这一设计论文在消融中移除 spectral encoder 后性能明显下降加入 MS 信息后LPIENet、LPIENet-small 和 cmKAN-light 的平均 ΔE00 与 Reproduction Error 均显著改善整体误差最高可比 RGB-only 设置降低约 50%。因此IRA 可以被理解为该轻量多光谱颜色校正框架中的关键积木它用低成本卷积保证移动端可部署性用通道与空间注意力增强颜色相关表征再通过 RGB-MS 双分支融合提升复杂光照下的颜色准确性和稳定性。核心代码YOLO11引入代码在根目录下的ultralytics/nn/目录新建一个C3k2目录然后新建一个以C3k2_IRA为文件名的py文件 把代码拷贝进去。importtorchimporttorch.nnasnnimporttorch.nn.functionalasFfromultralytics.nn.modules.convimportRepConv,Convfromultralytics.nn.modules.blockimportC3k,C3k2classAttentionBlock(nn.Module):def__init__(self,dim:int):super(AttentionBlock,self).__init__()self._spatial_attention_convnn.Conv2d(2,dim,kernel_size3,padding1)# Channel attention MLPself._channel_attention_conv0nn.Conv2d(1,dim,kernel_size1,padding0)self._channel_attention_conv1nn.Conv2d(dim,dim,kernel_size1,padding0)self._out_convnn.Conv2d(2*dim,dim,kernel_size1,padding0)defforward(self,x:torch.Tensor):iflen(x.shape)!4:raiseValueError(fExpected [B, C, H, W] input, got{x.shape}.)# Spatial attentionmeantorch.mean(x,dim1,keepdimTrue)# Mean/Max on C axismax,_torch.max(x,dim1,keepdimTrue)spatial_attentiontorch.cat([mean,max],dim1)# [B, 2, H, W]spatial_attentionself._spatial_attention_conv(spatial_attention)spatial_attentiontorch.sigmoid(spatial_attention)*x# Channel attention. TODO: Correct that it only uses average pool contrary to CBAM?# NOTE/TODO: This differs from CBAM as it uses Channel pooling, not spatial pooling!# In a way, this is 2x spatial attentionchannel_attentiontorch.relu(self._channel_attention_conv0(mean))channel_attentionself._channel_attention_conv1(channel_attention)channel_attentiontorch.sigmoid(channel_attention)*x attentiontorch.cat([spatial_attention,channel_attention],dim1)# [B, 2*dim, H, W]attentionself._out_conv(attention)returnxattentionclassBaseBlock(nn.Module):def__init__(self,channels:int):super(BaseBlock,self).__init__()self._conv0nn.Conv2d(channels,channels,kernel_size1)self._dw_convnn.Conv2d(channels,channels,kernel_size3,padding1,groupschannels)self._conv1nn.Conv2d(channels,channels,kernel_size1)self._conv2nn.Conv2d(channels,channels,kernel_size1)self._conv3nn.Conv2d(channels,channels,kernel_size1)defforward(self,x:torch.Tensor):featuresself._conv0(x)featuresF.elu(self._dw_conv(features))# TODO: ELU or ReLU?featuresself._conv1(features)xxfeatures featuresF.elu(self._conv2(x))featuresself._conv3(features)returnxfeaturesclassIRA(nn.Module):def__init__(self,inc,ouc):super().__init__()self.projnn.Conv2d(inc,ouc,kernel_size3,padding1)self.irbnn.Sequential(nn.Conv2d(ouc,ouc,kernel_size1),nn.Conv2d(ouc,ouc,kernel_size3,padding1,groupsouc),nn.ELU(inplaceTrue),nn.Conv2d(ouc,ouc,kernel_size1),)self.ffnnn.Sequential(nn.Conv2d(ouc,ouc,kernel_size1),nn.ELU(inplaceTrue),nn.Conv2d(ouc,ouc,kernel_size1),)self.attnAttentionBlock(ouc)defforward(self,x):xself.proj(x)identityx xidentityself.irb(x)identityx xidentityself.ffn(x)xself.attn(x)returnxclassC3k_IRA(C3k):def__init__(self,c1,c2,n1,shortcutFalse,g1,e0.5,k3):super().__init__(c1,c2,n,shortcut,g,e,k)c_int(c2*e)# hidden channelsself.mnn.Sequential(*(IRA(c_,c_)for_inrange(n)))classC3k2_IRA(C3k2):def__init__(self,c1,c2,n1,c3kFalse,e0.5,g1,shortcutTrue):super().__init__(c1,c2,n,c3k,e,g,shortcut)self.mnn.ModuleList(C3k_IRA(self.c,self.c,2,shortcut,g)ifc3kelseIRA(self.c,self.c)for_inrange(n))注册在ultralytics/nn/tasks.py中进行如下操作步骤1:fromultralytics.nn.C3k2.C3k2_IRAimportC3k2_IRA步骤2修改def parse_model(d, ch, verboseTrue):C3k2_IRA配置yolo11-C3k2_IRA.yaml# Ultralytics YOLO , AGPL-3.0 license# YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect# Parametersnc:80# number of classesscales:# model compound scaling constants, i.e. modelyolo11n.yaml will call yolo11.yaml with scale n# [depth, width, max_channels]n:[0.50,0.25,1024]# summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPss:[0.50,0.50,1024]# summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPsm:[0.50,1.00,512]# summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPsl:[1.00,1.00,512]# summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPsx:[1.00,1.50,512]# summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs# YOLO11n backbonebackbone:# [from, repeats, module, args]-[-1,1,Conv,[64,3,2]]# 0-P1/2-[-1,1,Conv,[128,3,2]]# 1-P2/4-[-1,2,C3k2_IRA,[256,False,0.25]]-[-1,1,Conv,[256,3,2]]# 3-P3/8-[-1,2,C3k2_IRA,[512,False,0.25]]-[-1,1,Conv,[512,3,2]]# 5-P4/16-[-1,2,C3k2_IRA,[512,True]]-[-1,1,Conv,[1024,3,2]]# 7-P5/32-[-1,2,C3k2_IRA,[1024,True]]-[-1,1,SPPF,[1024,5]]# 9-[-1,2,C2PSA,[1024]]# 10# YOLO11n headhead:-[-1,1,nn.Upsample,[None,2,nearest]]-[[-1,6],1,Concat,[1]]# cat backbone P4-[-1,2,C3k2_IRA,[512,False]]# 13-[-1,1,nn.Upsample,[None,2,nearest]]-[[-1,4],1,Concat,[1]]# cat backbone P3-[-1,2,C3k2_IRA,[256,False]]# 16 (P3/8-small)-[-1,1,Conv,[256,3,2]]-[[-1,13],1,Concat,[1]]# cat head P4-[-1,2,C3k2_IRA,[512,False]]# 19 (P4/16-medium)-[-1,1,Conv,[512,3,2]]-[[-1,10],1,Concat,[1]]# cat head P5-[-1,2,C3k2_IRA,[1024,True]]# 22 (P5/32-large)-[[16,19,22],1,Detect,[nc]]# Detect(P3, P4, P5)实验脚本importwarnings warnings.filterwarnings(ignore)fromultralyticsimportYOLOif__name____main__:# 修改为自己的配置文件地址modelYOLO(./ultralytics/cfg/models/11/yolo11-C3k2_IRA.yaml)# 修改为自己的数据集地址model.train(data./ultralytics/cfg/datasets/coco8.yaml,cacheFalse,imgsz640,epochs10,single_clsFalse,# 是否是单类别检测batch8,close_mosaic10,workers0,optimizerSGD,ampTrue,projectruns/train,nameC3k2_IRA,)结果
