代码开源&即插即用注意力机制&24年新晋可见光-红外(visible-infrared)融合算法RCAFusion——自动驾驶语义信息感知方案

原创已于 2024-07-28 10:11:07 修改 · 2k 阅读

31 ·

CC 4.0 BY-SA版权

文章标签：

#自动驾驶 #机器学习 #开源 #python #YOLO #图像处理 #目标检测

于 2024-07-28 10:08:48 首次发布

多模态融合算法专栏收录该内容

1 篇文章

订阅专栏

一、写在前面

本博客论文取自2024年自动驾驶领域顶级会议IV，引用格式为：

A. Li et al., "RCAFusion: Cross Rubik Cube Attention Network for Multi-modal Image Fusion of Intelligent Vehicles," 2024 IEEE Intelligent Vehicles Symposium (IV), Jeju Island, Korea, Republic of, 2024, pp. 2848-2854, doi: 10.1109/IV55156.2024.10588756.

目前该文章的论文源码已开源在Github中，如需下载请访问：

https://github.com/vehicle-AngLi/RCAFusion

如下载过程出现疑问，请联系作者：vehicle_liang@163.com

如你需要快速跑通该算法，直接跳转阅读至第三部分

如你需要插入即插即用的注意力机制模块/修改损失函数，直接跳转阅读至第四部分

二、RCAFusion贡献

1. 网络架构的优化

本文网络架构如图所示：

图1 RCAFusion整体网络架构

引入了一种交叉的网络结构，让特征提取过程中的多模态信息实现交互，同时将残差模块进行优化，用以代替经典的卷积方法

2. 全新的注意力机制

该即插即用的注意力机制如图所示：

图2 魔方注意力机制

和CBAM注意力机制相似的是，该算法也同时考虑了图像空间维度、通道维度的注意力添加，但是不同的是这个方法采用了并联（张量积）的方式，让网络分别学习两种维度，使得两个注意力不至于过度互相影响。同时该注意力机制还加入了类似于Transformer的自注意力机制算法，用图像转置的方式让图像添加自注意力。

3. 不对称的SSIM损失函数

大多数情况下融合算法的损失函数由两部分组成，一个控制融合图像的信息量尽可能保留更多，另一个则尽可能让融合图像保留显著特征，公式如下：

$L_\textit{int} = \frac{1}{H \times W} \Vert I_f-max(I_\textit{vis}, I_\textit{inf}) \Vert_1$

$L_\textit{text} = \frac{1}{H \times W} \Vert \nabla I_f-max(\nabla I_\textit{vis}, \nabla I_\textit{inf}) \Vert_1$

在之前引入的SSIM损失函数通常把输入的两个模态的重要程度视作相同。但实际上，在不同的自动驾驶工况下，不同模态发挥的作用重要程度存在着较大的差异。因此，本文将信息量作为评判的标准，对不同工况下的损失进行了一定程度的区分，公式如下：

$L_\textit{asymmetric\_SSIM} = \frac{\textit{En}_\textit{vis}}{\textit{En}_\textit{vis} + \textit{En}_\textit{inf}} \textit{SSIM}(I_\textit{fus}, I_\textit{vis})+\frac{\textit{En}_\textit{inf}}{\textit{En}_\textit{vis}+textit{En}_\textit{inf}} \textit{SSIM}(I_\textit{fus}, I_\textit{inf})$

在这一部分的损失函数中，EN代表源图像所包含的信息量，SSIM(*)则代表计算两张图像中的结构相似性结果。

三、源码复现

本博客将会从零开始带大家复现该项目的源码！

1. 环境配置

torch==2.0.1
torchvision==0.15.2
pandas==2.0.3
pillow==9.5.0
numpy==1.24.4
ultralytics==8.0.188

上述为作者在github中声明的环境配置文件，但部分配置未作要求，这里建议大家直接配置Yolo系列的图像处理环境即可，借用Yolov5官方文档中的requirements文件（文章中实验包含目标检测验证实验，因此一并安装）：

# YOLOv5 requirements
gitpython>=3.1.30
matplotlib>=3.3
numpy>=1.23.5
opencv-python>=4.1.1
pillow>=10.3.0
psutil # system resources
PyYAML>=5.3.1
requests>=2.32.0
scipy>=1.4.1
thop>=0.1.1 # FLOPs computation
torch>=1.8.0 # see https://pytorch.org/get-started/locally (recommended)
torchvision>=0.9.0
tqdm>=4.64.0
ultralytics>=8.2.34 # https://ultralytics.com
pandas>=1.1.4
seaborn>=0.11.0
setuptools>=70.0.0 # Snyk vulnerability fix

如有现成的requirements文件则直接运行下一步，若无，则新建txt文档将上述文本粘贴进去，并重命名为requirements.txt文件。之后在搭建的环境（搭建环境方法可参考其他博文）中运行：

pip install -r requirements.txt

2. 复现融合结果

开源的预训练模型存放于：解压后的文件夹/model/RCAFusion_pre.pth

之后将待融合的红外图像放置在：解压后的文件夹/test_imgs/ir/中

待融合的可见光图像放置在：解压后的文件夹/test_imgs/vi/中

在解压后的文件夹运行如下指令，之后可以在Save文件夹中找到融合结果：

python test.py

如果你保持原有的图像不变，应当可以获得如下的图像：

3. 训练自己的RCAFusion模型

训练的数据集存放的位置作者并没有提前准备，这里需要我们自己进行创建：

首先在解压后的目录中创建文件夹并命名为“train_data”

进入“train_data”文件夹，分别将可见光训练图像、红外训练图像存放在两个新文件夹中，分别命名为“vis”，“inf”。即训练图像分别存放在了 "./train_data/inf/"（红外），"./train_data/vis/"（可见光）中。

运行如下指令：

python train_rca.py --batch_size 2 --gpu 0 --epochs 50

在这里你可以自己调整训练的batch大小（指令中的2）和步数（指令中的50）。

训练结束后，你应当可以在“解压目录/model/Fusion/”中找到你的训练结果，fusion_model.pth

你可以将其重命名为：your_model.pth

之后，运行如下指令来让代码输出你的模型融合结果（测试用的数据集按三.2说明的位置存放）：

python test.py --model_path model/Fusion/your_model.pth --gpu 0

最后可以在Save文件夹中得到融合结果。

四、注意力机制/损失函数快速替换方法

1. 调用RCA注意力机制

RCA注意力机制源码：

class Self_Attention(nn.Module):
def __init__(self, channel):
super(Self_Attention, self).__init__()
self.conv1 = nn.Conv2d(channel, channel, kernel_size=1)
self.sig = nn.Sigmoid()
self.conv2 = nn.Conv2d(channel, channel, kernel_size=1)
self.conv3 = nn.Conv2d(channel, channel, kernel_size=1)

def forward(self,x):
x1 = self.conv1(x)
x1 = self.sig(x1)
x1 = torch.transpose(x1, dim0=3, dim1=2)
x2 = self.conv2(x)
x2 = self.sig(x2)
# print('\n', x1.size(), x2.size())
x3 = torch.matmul(x1,x2)
x4 = torch.matmul(x,x3)
out = self.conv3(x4)
return self.sig(out)

class Rubik_Cube_Attention(nn.Module):
def __init__(self, channel, ratio=16, kernel_size=7):
super(Rubik_Cube_Attention, self).__init__()
self.channel_avg_pool = nn.AdaptiveAvgPool2d(1)
self.line1 = nn.Conv2d(channel, channel // ratio, 1, bias=False)
self.relu = nn.ReLU()
self.line2 = nn.Conv2d(channel // ratio, channel, 1, bias=False)
self.sigmoid = nn.Sigmoid()
self.space = nn.Conv2d(1, 1, kernel_size, padding = 3, bias = False)
self.satt = Self_Attention(channel)
self.conv = nn.Conv2d(channel, channel, 1)

def forward(self,x):
##训练通道权重
##channel attention
x_ca = self.channel_avg_pool(x)
x_ca = self.relu(self.line1(x_ca))
x_ca = self.relu(self.line2(x_ca))
x_ca = self.sigmoid(x_ca)
##同时训练空间权重
##spatial attention
x_sa = torch.mean(x, dim=1, keepdim=True)
x_sa = self.space(x_sa)
x_sa = self.sigmoid(x_sa)
##直接相乘
##multiple
x_attention = x*x_sa*x_ca
##训练自注意力机制
##self-attention
x_self = self.satt(x)
##相加
##addition
out = 0.1*x_self+x_attention
out = self.sigmoid(self.conv(out))
return out

在网络架构文件中（本文github文档就是指FusionNet.py）粘贴上述代码。之后可以用如下方式进行调用，放置在需要的位置：

class your_code(nn.Module):
def __init__(self, channels):
super(your_code, self).__init__()
self.rubik = Rubik_Cube_Attention(channels)

def forward(self,x):
out = self.rubik(x)
return out

2. 非对称性结构相似性损失函数

非对称性结构损失函数源码：

def EN_function(image_array):
image_array = image_array.cpu()
# 计算图像的直方图
histogram, bins = np.histogram(image_array, bins=256, range=(0, 255))
# 将直方图归一化
histogram = histogram / float(np.sum(histogram))
# 计算熵
entropy = -np.sum(histogram * np.log2(histogram + 1e-7))
return entropy

en_vis = EN_function(image_y)
en_inf = EN_function(image_ir)

sim_vis = ssim(image_y,generate_img)
sim_inf = ssim(image_ir,generate_img)
loss_sim_vis = (en_vis/(en_inf+en_vis)) * (1 - sim_vis)
loss_sim_inf = (en_inf/(en_inf+en_vis)) * (1 - sim_inf)
loss_sim = loss_sim_inf + loss_sim_vis

在损失函数的设定文件中（本文的github中的loss.py）进行上述修改即可。将上述函数与计算方式代入你自行设定的损失公式中即可得到结果。本文的loss整体源码如下：

class Fusionloss(nn.Module):
def __init__(self):
super(Fusionloss, self).__init__()
self.sobelconv=Sobelxy()

def forward(self,image_vis,image_ir,generate_img):
# loss_in
image_y=image_vis[:,:1,:,:]
x_in_max=torch.max(image_y,image_ir)
loss_in=F.l1_loss(x_in_max,generate_img)
# loss_grad
y_grad=self.sobelconv(image_y)
ir_grad=self.sobelconv(image_ir)
generate_img_grad=self.sobelconv(generate_img)
x_grad_joint=torch.max(y_grad,ir_grad)
loss_grad=F.l1_loss(x_grad_joint,generate_img_grad)
# 在这里引入非对称的SSIM损失函数
# Introduction of asymmetric SSIM Loss function
en_vis = EN_function(image_y)
en_inf = EN_function(image_ir)
trans_vis = image_y.cpu()
trans_inf = image_ir.cpu()
trans_vis = trans_vis.numpy()
trans_inf = trans_inf.numpy()
trans_fus = generate_img.cpu()
trans_fus = trans_fus.detach().numpy()
sim_vis = ssim(image_y,generate_img)
sim_inf = ssim(image_ir,generate_img)
loss_sim_vis = (en_vis/(en_inf+en_vis)) * (1 - sim_vis)
loss_sim_inf = (en_inf/(en_inf+en_vis)) * (1 - sim_inf)
loss_sim = loss_sim_inf + loss_sim_vis
# loss_total
loss_total=loss_in+10*loss_grad+10*loss_sim
return loss_total,loss_in,loss_grad,loss_sim