RGB-IR融合网络IRFS理解+复现

IRFS简介其核心思想是在YCrCb颜色空间中仅对亮度通道Y进行融合利用红外图像提供的目标显著性与结构信息与可见光图像的纹理细节共同输入融合网络FSFNet生成增强后的亮度分量同时保留原始色度信息Cr、Cb从而获得既突出目标又保持自然色彩的融合图像此外IRFS通过引入显著性目标检测SOD对融合过程进行约束与引导使网络更加关注关键目标区域而非全局平均增强最终在低光或复杂环境下显著提升目标表达能力与下游检测性能。IRFS 主要由两个核心子网络一个颜色空间解耦机制组成1.颜色空间解耦YCrCb分离首先将 RGB 图像转换为 YCrCb 空间网络只对 Y通道进行学习与融合。融合后拼接CRCB2.显著性检测网络ΔY 通道去引导其是一个联合优化机制Fusion 生成图像SOD 判断目标是否突出反向指导 Fusion 学习更优结果创新点1.Y通道解耦融合机制2.FSFNet 端到端亮度融合3.SOD引导融合IRFS复现首先去下载源码原论文abstract下有链接:https://github.com/wdhudiekou/IRFS.这是环境配置根据自己的电脑有无显卡版本来安装这是最低的版本注意Kornia必须一致。我这里给一串命令conda create -n irfs python3.8 -y conda activate irfs pip install torch1.8.1 torchvision0.9.1 pip install opencv-python4.5.5.64 pip install kornia0.5.11 pip install numpy matplotlib tqdm pillow scipy环境配置好了后首先你需要下载ResNet-34的预训练模型这里写的很清楚去下载下来如果你要训练按它代码下的指示来就可以了如果你想跑自己的数据不进行二次开发直接用他的预训练模型就好了这一定是最好的这是你下载的预权重9是最好的效果所以我们用9这个就是融合的权重我们只是用它其它是sod校正用的我们这里给出代码跑自己的图片在TEST文件夹下新建.py文件将下列代码拷入。#!/usr/bin/python # -*- encoding: utf-8 -*- # Author: WangDi (修改版) # 功能: 用 FSFNet 模型对自定义 RGBIR 图像进行融合保持原始尺寸 import os import glob import cv2 import torch import numpy as np from models.fsfnet import FSFNet def test_fusion_custom( fusion_model_pathrE:\multfusio\IRFS-master\checkpoint\Fusion\fusion_model.pth, rgb_dirrE:\multfusio\IRFS-master\dataes\RGB, ir_dirrE:\multfusio\IRFS-master\dataes\IR, save_dirrE:\multfusio\IRFS-master\result ): # --- 1. 设置设备 --- device torch.device(cuda if torch.cuda.is_available() else cpu) print(Using device:, device) # --- 2. 加载模型 --- fusionNet FSFNet(64).to(device) fusionNet.load_state_dict(torch.load(fusion_model_path, map_locationdevice)) fusionNet.eval() print(Loaded Fusion model:, fusion_model_path) # --- 3. 准备保存路径 --- os.makedirs(save_dir, exist_okTrue) # --- 4. 获取图片列表 --- rgb_list sorted(glob.glob(os.path.join(rgb_dir, *.jpg))) if len(rgb_list) 0: print(No RGB images found in, rgb_dir) return # --- 5. 遍历融合 --- for rgb_path in rgb_list: name os.path.basename(rgb_path) ir_path os.path.join(ir_dir, name) if not os.path.exists(ir_path): print(IR image not found, skipping:, ir_path) continue # --- 5.1 读图 --- rgb cv2.imread(rgb_path) ir cv2.imread(ir_path, 0) # 灰度图 # --- 5.2 保持原始尺寸 --- h, w rgb.shape[:2] if ir.shape[:2] ! (h, w): ir cv2.resize(ir, (w, h)) print(fResized IR {name} to match RGB size.) # --- 5.3 转 YCrCb --- ycrcb cv2.cvtColor(rgb, cv2.COLOR_BGR2YCrCb) y, cr, cb cv2.split(ycrcb) # --- 5.4 转 tensor --- y torch.from_numpy(y / 255.).float().unsqueeze(0).unsqueeze(0).to(device) ir torch.from_numpy(ir / 255.).float().unsqueeze(0).unsqueeze(0).to(device) # --- 5.5 融合 --- with torch.no_grad(): fus_y fusionNet(y, ir) fus_y torch.clamp(fus_y, 0, 1).cpu().numpy().squeeze() fus_y (fus_y * 255).astype(np.uint8) # --- 5.6 合成回 BGR --- fusion_ycrcb cv2.merge([fus_y, cr, cb]) fusion_img cv2.cvtColor(fusion_ycrcb, cv2.COLOR_YCrCb2BGR) # --- 5.7 保存 --- save_path os.path.join(save_dir, name) cv2.imwrite(save_path, fusion_img) print(Saved:, save_path) print(All images fused. Done!) if __name__ __main__: test_fusion_custom( fusion_model_pathrE:\multfusio\IRFS-master\checkpoint\Fusion\fusion_model.pth, rgb_dirrE:\multfusio\IRFS-master\dataes\RGB, ir_dirrE:\multfusio\IRFS-master\dataes\IR, save_dirrE:\multfusio\IRFS-master\result )改成自己的路径运行就可以了注意你的RGB和IR图片大小名字要对应看看结果人明显变亮了