一、写在前面
本博客论文取自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损失函数
大多数情况下融合算法的损失函数由两部分组成,一个控制融合图像的信息量尽可能保留更多,另一个则尽可能让融合图像保留显著特征,公式如下:
在之前引入的SSIM损失函数通常把输入的两个模态的重要程度视作相同。但实际上,在不同的自动驾驶工况下,不同模态发挥的作用重要程度存在着较大的差异。因此,本文将信息量作为评判的标准,对不同工况下的损失进行了一定程度的区分,公式如下:
在这一部分的损失函数中,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 entropyen_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