Update
@ -8,11 +8,13 @@ Codes for ***CDDFuse: Correlation-Driven Dual-Branch Feature Decomposition for M
|
|||||||
-[*[Supplementary Materials]*]()
|
-[*[Supplementary Materials]*]()
|
||||||
|
|
||||||
|
|
||||||
|
## Update
|
||||||
|
- [2023/5] Training codes and config files will be public available before June.
|
||||||
|
- [2023/4] Release inference code for infrared-visible image fusion and medical image fusion.
|
||||||
|
|
||||||
|
|
||||||
## Citation
|
## Citation
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
```
|
```
|
||||||
@article{DBLP:journals/corr/abs-2211-14461,
|
@article{DBLP:journals/corr/abs-2211-14461,
|
||||||
author = {Zixiang Zhao and Haowen Bai and Jiangshe Zhang and Yulun Zhang and Shuang Xu and Zudi Lin and Radu Timofte and Luc Van Gool},
|
author = {Zixiang Zhao and Haowen Bai and Jiangshe Zhang and Yulun Zhang and Shuang Xu and Zudi Lin and Radu Timofte and Luc Van Gool},
|
||||||
|
BIN
models/CDDFuse_IVF.pth
Normal file
BIN
models/CDDFuse_MIF.pth
Normal file
403
net.py
Normal file
@ -0,0 +1,403 @@
|
|||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import math
|
||||||
|
import torch.nn.functional as F
|
||||||
|
import torch.utils.checkpoint as checkpoint
|
||||||
|
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
|
||||||
|
from einops import rearrange
|
||||||
|
|
||||||
|
|
||||||
|
def drop_path(x, drop_prob: float = 0., training: bool = False):
|
||||||
|
"""
|
||||||
|
Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
|
||||||
|
This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
|
||||||
|
the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
|
||||||
|
See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
|
||||||
|
changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
|
||||||
|
'survival rate' as the argument.
|
||||||
|
"""
|
||||||
|
if drop_prob == 0. or not training:
|
||||||
|
return x
|
||||||
|
keep_prob = 1 - drop_prob
|
||||||
|
# work with diff dim tensors, not just 2D ConvNets
|
||||||
|
shape = (x.shape[0],) + (1,) * (x.ndim - 1)
|
||||||
|
random_tensor = keep_prob + \
|
||||||
|
torch.rand(shape, dtype=x.dtype, device=x.device)
|
||||||
|
random_tensor.floor_() # binarize
|
||||||
|
output = x.div(keep_prob) * random_tensor
|
||||||
|
return output
|
||||||
|
|
||||||
|
class DropPath(nn.Module):
|
||||||
|
"""
|
||||||
|
Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, drop_prob=None):
|
||||||
|
super(DropPath, self).__init__()
|
||||||
|
self.drop_prob = drop_prob
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return drop_path(x, self.drop_prob, self.training)
|
||||||
|
|
||||||
|
|
||||||
|
class AttentionBase(nn.Module):
|
||||||
|
def __init__(self,
|
||||||
|
dim,
|
||||||
|
num_heads=8,
|
||||||
|
qkv_bias=False,):
|
||||||
|
super(AttentionBase, self).__init__()
|
||||||
|
self.num_heads = num_heads
|
||||||
|
head_dim = dim // num_heads
|
||||||
|
self.scale = nn.Parameter(torch.ones(num_heads, 1, 1))
|
||||||
|
self.qkv1 = nn.Conv2d(dim, dim*3, kernel_size=1, bias=qkv_bias)
|
||||||
|
self.qkv2 = nn.Conv2d(dim*3, dim*3, kernel_size=3, padding=1, bias=qkv_bias)
|
||||||
|
self.proj = nn.Conv2d(dim, dim, kernel_size=1, bias=qkv_bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
# [batch_size, num_patches + 1, total_embed_dim]
|
||||||
|
b, c, h, w = x.shape
|
||||||
|
qkv = self.qkv2(self.qkv1(x))
|
||||||
|
q, k, v = qkv.chunk(3, dim=1)
|
||||||
|
q = rearrange(q, 'b (head c) h w -> b head c (h w)',
|
||||||
|
head=self.num_heads)
|
||||||
|
k = rearrange(k, 'b (head c) h w -> b head c (h w)',
|
||||||
|
head=self.num_heads)
|
||||||
|
v = rearrange(v, 'b (head c) h w -> b head c (h w)',
|
||||||
|
head=self.num_heads)
|
||||||
|
q = torch.nn.functional.normalize(q, dim=-1)
|
||||||
|
k = torch.nn.functional.normalize(k, dim=-1)
|
||||||
|
# transpose: -> [batch_size, num_heads, embed_dim_per_head, num_patches + 1]
|
||||||
|
# @: multiply -> [batch_size, num_heads, num_patches + 1, num_patches + 1]
|
||||||
|
attn = (q @ k.transpose(-2, -1)) * self.scale
|
||||||
|
attn = attn.softmax(dim=-1)
|
||||||
|
|
||||||
|
out = (attn @ v)
|
||||||
|
|
||||||
|
out = rearrange(out, 'b head c (h w) -> b (head c) h w',
|
||||||
|
head=self.num_heads, h=h, w=w)
|
||||||
|
|
||||||
|
out = self.proj(out)
|
||||||
|
return out
|
||||||
|
|
||||||
|
class Mlp(nn.Module):
|
||||||
|
"""
|
||||||
|
MLP as used in Vision Transformer, MLP-Mixer and related networks
|
||||||
|
"""
|
||||||
|
def __init__(self,
|
||||||
|
in_features,
|
||||||
|
hidden_features=None,
|
||||||
|
ffn_expansion_factor = 2,
|
||||||
|
bias = False):
|
||||||
|
super().__init__()
|
||||||
|
hidden_features = int(in_features*ffn_expansion_factor)
|
||||||
|
|
||||||
|
self.project_in = nn.Conv2d(
|
||||||
|
in_features, hidden_features*2, kernel_size=1, bias=bias)
|
||||||
|
|
||||||
|
self.dwconv = nn.Conv2d(hidden_features*2, hidden_features*2, kernel_size=3,
|
||||||
|
stride=1, padding=1, groups=hidden_features, bias=bias)
|
||||||
|
|
||||||
|
self.project_out = nn.Conv2d(
|
||||||
|
hidden_features, in_features, kernel_size=1, bias=bias)
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.project_in(x)
|
||||||
|
x1, x2 = self.dwconv(x).chunk(2, dim=1)
|
||||||
|
x = F.gelu(x1) * x2
|
||||||
|
x = self.project_out(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
class BaseFeatureExtraction(nn.Module):
|
||||||
|
def __init__(self,
|
||||||
|
dim,
|
||||||
|
num_heads,
|
||||||
|
ffn_expansion_factor=1.,
|
||||||
|
qkv_bias=False,):
|
||||||
|
super(BaseFeatureExtraction, self).__init__()
|
||||||
|
self.norm1 = LayerNorm(dim, 'WithBias')
|
||||||
|
self.attn = AttentionBase(dim, num_heads=num_heads, qkv_bias=qkv_bias,)
|
||||||
|
self.norm2 = LayerNorm(dim, 'WithBias')
|
||||||
|
self.mlp = Mlp(in_features=dim,
|
||||||
|
ffn_expansion_factor=ffn_expansion_factor,)
|
||||||
|
def forward(self, x):
|
||||||
|
x = x + self.attn(self.norm1(x))
|
||||||
|
x = x + self.mlp(self.norm2(x))
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class InvertedResidualBlock(nn.Module):
|
||||||
|
def __init__(self, inp, oup, expand_ratio):
|
||||||
|
super(InvertedResidualBlock, self).__init__()
|
||||||
|
hidden_dim = int(inp * expand_ratio)
|
||||||
|
self.bottleneckBlock = nn.Sequential(
|
||||||
|
# pw
|
||||||
|
nn.Conv2d(inp, hidden_dim, 1, bias=False),
|
||||||
|
# nn.BatchNorm2d(hidden_dim),
|
||||||
|
nn.ReLU6(inplace=True),
|
||||||
|
# dw
|
||||||
|
nn.ReflectionPad2d(1),
|
||||||
|
nn.Conv2d(hidden_dim, hidden_dim, 3, groups=hidden_dim, bias=False),
|
||||||
|
# nn.BatchNorm2d(hidden_dim),
|
||||||
|
nn.ReLU6(inplace=True),
|
||||||
|
# pw-linear
|
||||||
|
nn.Conv2d(hidden_dim, oup, 1, bias=False),
|
||||||
|
# nn.BatchNorm2d(oup),
|
||||||
|
)
|
||||||
|
def forward(self, x):
|
||||||
|
return self.bottleneckBlock(x)
|
||||||
|
|
||||||
|
class DetailNode(nn.Module):
|
||||||
|
def __init__(self):
|
||||||
|
super(DetailNode, self).__init__()
|
||||||
|
# Scale is Ax + b, i.e. affine transformation
|
||||||
|
self.theta_phi = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||||
|
self.theta_rho = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||||
|
self.theta_eta = InvertedResidualBlock(inp=32, oup=32, expand_ratio=2)
|
||||||
|
self.shffleconv = nn.Conv2d(64, 64, kernel_size=1,
|
||||||
|
stride=1, padding=0, bias=True)
|
||||||
|
def separateFeature(self, x):
|
||||||
|
z1, z2 = x[:, :x.shape[1]//2], x[:, x.shape[1]//2:x.shape[1]]
|
||||||
|
return z1, z2
|
||||||
|
def forward(self, z1, z2):
|
||||||
|
z1, z2 = self.separateFeature(
|
||||||
|
self.shffleconv(torch.cat((z1, z2), dim=1)))
|
||||||
|
z2 = z2 + self.theta_phi(z1)
|
||||||
|
z1 = z1 * torch.exp(self.theta_rho(z2)) + self.theta_eta(z2)
|
||||||
|
return z1, z2
|
||||||
|
|
||||||
|
class DetailFeatureExtraction(nn.Module):
|
||||||
|
def __init__(self, num_layers=3):
|
||||||
|
super(DetailFeatureExtraction, self).__init__()
|
||||||
|
INNmodules = [DetailNode() for _ in range(num_layers)]
|
||||||
|
self.net = nn.Sequential(*INNmodules)
|
||||||
|
def forward(self, x):
|
||||||
|
z1, z2 = x[:, :x.shape[1]//2], x[:, x.shape[1]//2:x.shape[1]]
|
||||||
|
for layer in self.net:
|
||||||
|
z1, z2 = layer(z1, z2)
|
||||||
|
return torch.cat((z1, z2), dim=1)
|
||||||
|
|
||||||
|
# =============================================================================
|
||||||
|
|
||||||
|
# =============================================================================
|
||||||
|
import numbers
|
||||||
|
##########################################################################
|
||||||
|
## Layer Norm
|
||||||
|
def to_3d(x):
|
||||||
|
return rearrange(x, 'b c h w -> b (h w) c')
|
||||||
|
|
||||||
|
|
||||||
|
def to_4d(x, h, w):
|
||||||
|
return rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
|
||||||
|
|
||||||
|
|
||||||
|
class BiasFree_LayerNorm(nn.Module):
|
||||||
|
def __init__(self, normalized_shape):
|
||||||
|
super(BiasFree_LayerNorm, self).__init__()
|
||||||
|
if isinstance(normalized_shape, numbers.Integral):
|
||||||
|
normalized_shape = (normalized_shape,)
|
||||||
|
normalized_shape = torch.Size(normalized_shape)
|
||||||
|
|
||||||
|
assert len(normalized_shape) == 1
|
||||||
|
|
||||||
|
self.weight = nn.Parameter(torch.ones(normalized_shape))
|
||||||
|
self.normalized_shape = normalized_shape
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
sigma = x.var(-1, keepdim=True, unbiased=False)
|
||||||
|
return x / torch.sqrt(sigma+1e-5) * self.weight
|
||||||
|
|
||||||
|
|
||||||
|
class WithBias_LayerNorm(nn.Module):
|
||||||
|
def __init__(self, normalized_shape):
|
||||||
|
super(WithBias_LayerNorm, self).__init__()
|
||||||
|
if isinstance(normalized_shape, numbers.Integral):
|
||||||
|
normalized_shape = (normalized_shape,)
|
||||||
|
normalized_shape = torch.Size(normalized_shape)
|
||||||
|
|
||||||
|
assert len(normalized_shape) == 1
|
||||||
|
|
||||||
|
self.weight = nn.Parameter(torch.ones(normalized_shape))
|
||||||
|
self.bias = nn.Parameter(torch.zeros(normalized_shape))
|
||||||
|
self.normalized_shape = normalized_shape
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
mu = x.mean(-1, keepdim=True)
|
||||||
|
sigma = x.var(-1, keepdim=True, unbiased=False)
|
||||||
|
return (x - mu) / torch.sqrt(sigma+1e-5) * self.weight + self.bias
|
||||||
|
|
||||||
|
class LayerNorm(nn.Module):
|
||||||
|
def __init__(self, dim, LayerNorm_type):
|
||||||
|
super(LayerNorm, self).__init__()
|
||||||
|
if LayerNorm_type == 'BiasFree':
|
||||||
|
self.body = BiasFree_LayerNorm(dim)
|
||||||
|
else:
|
||||||
|
self.body = WithBias_LayerNorm(dim)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
h, w = x.shape[-2:]
|
||||||
|
return to_4d(self.body(to_3d(x)), h, w)
|
||||||
|
|
||||||
|
##########################################################################
|
||||||
|
## Gated-Dconv Feed-Forward Network (GDFN)
|
||||||
|
class FeedForward(nn.Module):
|
||||||
|
def __init__(self, dim, ffn_expansion_factor, bias):
|
||||||
|
super(FeedForward, self).__init__()
|
||||||
|
|
||||||
|
hidden_features = int(dim*ffn_expansion_factor)
|
||||||
|
|
||||||
|
self.project_in = nn.Conv2d(
|
||||||
|
dim, hidden_features*2, kernel_size=1, bias=bias)
|
||||||
|
|
||||||
|
self.dwconv = nn.Conv2d(hidden_features*2, hidden_features*2, kernel_size=3,
|
||||||
|
stride=1, padding=1, groups=hidden_features*2, bias=bias)
|
||||||
|
|
||||||
|
self.project_out = nn.Conv2d(
|
||||||
|
hidden_features, dim, kernel_size=1, bias=bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.project_in(x)
|
||||||
|
x1, x2 = self.dwconv(x).chunk(2, dim=1)
|
||||||
|
x = F.gelu(x1) * x2
|
||||||
|
x = self.project_out(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
##########################################################################
|
||||||
|
## Multi-DConv Head Transposed Self-Attention (MDTA)
|
||||||
|
class Attention(nn.Module):
|
||||||
|
def __init__(self, dim, num_heads, bias):
|
||||||
|
super(Attention, self).__init__()
|
||||||
|
self.num_heads = num_heads
|
||||||
|
self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
|
||||||
|
|
||||||
|
self.qkv = nn.Conv2d(dim, dim*3, kernel_size=1, bias=bias)
|
||||||
|
self.qkv_dwconv = nn.Conv2d(
|
||||||
|
dim*3, dim*3, kernel_size=3, stride=1, padding=1, groups=dim*3, bias=bias)
|
||||||
|
self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
b, c, h, w = x.shape
|
||||||
|
|
||||||
|
qkv = self.qkv_dwconv(self.qkv(x))
|
||||||
|
q, k, v = qkv.chunk(3, dim=1)
|
||||||
|
|
||||||
|
q = rearrange(q, 'b (head c) h w -> b head c (h w)',
|
||||||
|
head=self.num_heads)
|
||||||
|
k = rearrange(k, 'b (head c) h w -> b head c (h w)',
|
||||||
|
head=self.num_heads)
|
||||||
|
v = rearrange(v, 'b (head c) h w -> b head c (h w)',
|
||||||
|
head=self.num_heads)
|
||||||
|
|
||||||
|
q = torch.nn.functional.normalize(q, dim=-1)
|
||||||
|
k = torch.nn.functional.normalize(k, dim=-1)
|
||||||
|
|
||||||
|
attn = (q @ k.transpose(-2, -1)) * self.temperature
|
||||||
|
attn = attn.softmax(dim=-1)
|
||||||
|
|
||||||
|
out = (attn @ v)
|
||||||
|
|
||||||
|
out = rearrange(out, 'b head c (h w) -> b (head c) h w',
|
||||||
|
head=self.num_heads, h=h, w=w)
|
||||||
|
|
||||||
|
out = self.project_out(out)
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
##########################################################################
|
||||||
|
class TransformerBlock(nn.Module):
|
||||||
|
def __init__(self, dim, num_heads, ffn_expansion_factor, bias, LayerNorm_type):
|
||||||
|
super(TransformerBlock, self).__init__()
|
||||||
|
|
||||||
|
self.norm1 = LayerNorm(dim, LayerNorm_type)
|
||||||
|
self.attn = Attention(dim, num_heads, bias)
|
||||||
|
self.norm2 = LayerNorm(dim, LayerNorm_type)
|
||||||
|
self.ffn = FeedForward(dim, ffn_expansion_factor, bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = x + self.attn(self.norm1(x))
|
||||||
|
x = x + self.ffn(self.norm2(x))
|
||||||
|
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
##########################################################################
|
||||||
|
## Overlapped image patch embedding with 3x3 Conv
|
||||||
|
class OverlapPatchEmbed(nn.Module):
|
||||||
|
def __init__(self, in_c=3, embed_dim=48, bias=False):
|
||||||
|
super(OverlapPatchEmbed, self).__init__()
|
||||||
|
|
||||||
|
self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=3,
|
||||||
|
stride=1, padding=1, bias=bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.proj(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class Restormer_Encoder(nn.Module):
|
||||||
|
def __init__(self,
|
||||||
|
inp_channels=1,
|
||||||
|
out_channels=1,
|
||||||
|
dim=64,
|
||||||
|
num_blocks=[4, 4],
|
||||||
|
heads=[8, 8, 8],
|
||||||
|
ffn_expansion_factor=2,
|
||||||
|
bias=False,
|
||||||
|
LayerNorm_type='WithBias',
|
||||||
|
):
|
||||||
|
|
||||||
|
super(Restormer_Encoder, self).__init__()
|
||||||
|
|
||||||
|
self.patch_embed = OverlapPatchEmbed(inp_channels, dim)
|
||||||
|
|
||||||
|
self.encoder_level1 = nn.Sequential(*[TransformerBlock(dim=dim, num_heads=heads[0], ffn_expansion_factor=ffn_expansion_factor,
|
||||||
|
bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[0])])
|
||||||
|
self.baseFeature = BaseFeatureExtraction(dim=dim, num_heads = heads[2])
|
||||||
|
self.detailFeature = DetailFeatureExtraction()
|
||||||
|
|
||||||
|
def forward(self, inp_img):
|
||||||
|
inp_enc_level1 = self.patch_embed(inp_img)
|
||||||
|
out_enc_level1 = self.encoder_level1(inp_enc_level1)
|
||||||
|
base_feature = self.baseFeature(out_enc_level1)
|
||||||
|
detail_feature = self.detailFeature(out_enc_level1)
|
||||||
|
return base_feature, detail_feature, out_enc_level1
|
||||||
|
|
||||||
|
class Restormer_Decoder(nn.Module):
|
||||||
|
def __init__(self,
|
||||||
|
inp_channels=1,
|
||||||
|
out_channels=1,
|
||||||
|
dim=64,
|
||||||
|
num_blocks=[4, 4],
|
||||||
|
heads=[8, 8, 8],
|
||||||
|
ffn_expansion_factor=2,
|
||||||
|
bias=False,
|
||||||
|
LayerNorm_type='WithBias',
|
||||||
|
):
|
||||||
|
|
||||||
|
super(Restormer_Decoder, self).__init__()
|
||||||
|
self.reduce_channel = nn.Conv2d(int(dim*2), int(dim), kernel_size=1, bias=bias)
|
||||||
|
self.encoder_level2 = nn.Sequential(*[TransformerBlock(dim=dim, num_heads=heads[1], ffn_expansion_factor=ffn_expansion_factor,
|
||||||
|
bias=bias, LayerNorm_type=LayerNorm_type) for i in range(num_blocks[1])])
|
||||||
|
self.output = nn.Sequential(
|
||||||
|
nn.Conv2d(int(dim), int(dim)//2, kernel_size=3,
|
||||||
|
stride=1, padding=1, bias=bias),
|
||||||
|
nn.LeakyReLU(),
|
||||||
|
nn.Conv2d(int(dim)//2, out_channels, kernel_size=3,
|
||||||
|
stride=1, padding=1, bias=bias),)
|
||||||
|
self.sigmoid = nn.Sigmoid()
|
||||||
|
def forward(self, inp_img, base_feature, detail_feature):
|
||||||
|
out_enc_level0 = torch.cat((base_feature, detail_feature), dim=1)
|
||||||
|
out_enc_level0 = self.reduce_channel(out_enc_level0)
|
||||||
|
out_enc_level1 = self.encoder_level2(out_enc_level0)
|
||||||
|
if inp_img is not None:
|
||||||
|
out_enc_level1 = self.output(out_enc_level1) + inp_img
|
||||||
|
else:
|
||||||
|
out_enc_level1 = self.output(out_enc_level1)
|
||||||
|
return self.sigmoid(out_enc_level1), out_enc_level0
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
height = 128
|
||||||
|
width = 128
|
||||||
|
window_size = 8
|
||||||
|
modelE = Restormer_Encoder().cuda()
|
||||||
|
modelD = Restormer_Decoder().cuda()
|
||||||
|
|
10
requirements.txt
Normal file
@ -0,0 +1,10 @@
|
|||||||
|
einops==0.4.1
|
||||||
|
kornia==0.2.0
|
||||||
|
numpy==1.21.5
|
||||||
|
opencv_python==4.5.3.56
|
||||||
|
scikit_image==0.19.2
|
||||||
|
scikit_learn==1.1.3
|
||||||
|
scipy==1.7.3
|
||||||
|
tensorboardX==2.5.1
|
||||||
|
timm==0.4.12
|
||||||
|
torch==1.8.1+cu111
|
80
test_IVF.py
Normal file
@ -0,0 +1,80 @@
|
|||||||
|
from net import Restormer_Encoder, Restormer_Decoder, BaseFeatureExtraction, DetailFeatureExtraction
|
||||||
|
import os
|
||||||
|
import numpy as np
|
||||||
|
from utils.Evaluator import Evaluator
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
from utils.img_read_save import img_save,image_read_cv2
|
||||||
|
import warnings
|
||||||
|
import logging
|
||||||
|
warnings.filterwarnings("ignore")
|
||||||
|
logging.basicConfig(level=logging.CRITICAL)
|
||||||
|
|
||||||
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
||||||
|
ckpt_path=r"models/CDDFuse_IVF.pth"
|
||||||
|
for dataset_name in ["TNO","RoadScene"]:
|
||||||
|
print("\n"*2+"="*80)
|
||||||
|
model_name="CDDFuse "
|
||||||
|
print("The test result of "+dataset_name+' :')
|
||||||
|
test_folder=os.path.join('test_img',dataset_name)
|
||||||
|
test_out_folder=os.path.join('test_result',dataset_name)
|
||||||
|
|
||||||
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||||
|
Encoder = nn.DataParallel(Restormer_Encoder()).to(device)
|
||||||
|
Decoder = nn.DataParallel(Restormer_Decoder()).to(device)
|
||||||
|
BaseFuseLayer = nn.DataParallel(BaseFeatureExtraction(dim=64, num_heads=8)).to(device)
|
||||||
|
DetailFuseLayer = nn.DataParallel(DetailFeatureExtraction(num_layers=1)).to(device)
|
||||||
|
|
||||||
|
Encoder.load_state_dict(torch.load(ckpt_path)['DIDF_Encoder'])
|
||||||
|
Decoder.load_state_dict(torch.load(ckpt_path)['DIDF_Decoder'])
|
||||||
|
BaseFuseLayer.load_state_dict(torch.load(ckpt_path)['BaseFuseLayer'])
|
||||||
|
DetailFuseLayer.load_state_dict(torch.load(ckpt_path)['DetailFuseLayer'])
|
||||||
|
Encoder.eval()
|
||||||
|
Decoder.eval()
|
||||||
|
BaseFuseLayer.eval()
|
||||||
|
DetailFuseLayer.eval()
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
for img_name in os.listdir(os.path.join(test_folder,"ir")):
|
||||||
|
|
||||||
|
data_IR=image_read_cv2(os.path.join(test_folder,"ir",img_name),mode='GRAY')[np.newaxis,np.newaxis, ...]/255.0
|
||||||
|
data_VIS = image_read_cv2(os.path.join(test_folder,"vi",img_name), mode='GRAY')[np.newaxis,np.newaxis, ...]/255.0
|
||||||
|
|
||||||
|
data_IR,data_VIS = torch.FloatTensor(data_IR),torch.FloatTensor(data_VIS)
|
||||||
|
data_VIS, data_IR = data_VIS.cuda(), data_IR.cuda()
|
||||||
|
|
||||||
|
feature_V_B, feature_V_D, feature_V = Encoder(data_VIS)
|
||||||
|
feature_I_B, feature_I_D, feature_I = Encoder(data_IR)
|
||||||
|
feature_F_B = BaseFuseLayer(feature_V_B + feature_I_B)
|
||||||
|
feature_F_D = DetailFuseLayer(feature_V_D + feature_I_D)
|
||||||
|
data_Fuse, _ = Decoder(data_VIS, feature_F_B, feature_F_D)
|
||||||
|
data_Fuse=(data_Fuse-torch.min(data_Fuse))/(torch.max(data_Fuse)-torch.min(data_Fuse))
|
||||||
|
fi = np.squeeze((data_Fuse * 255).cpu().numpy())
|
||||||
|
img_save(fi, img_name.split(sep='.')[0], test_out_folder)
|
||||||
|
|
||||||
|
|
||||||
|
eval_folder=test_out_folder
|
||||||
|
ori_img_folder=test_folder
|
||||||
|
|
||||||
|
metric_result = np.zeros((8))
|
||||||
|
for img_name in os.listdir(os.path.join(ori_img_folder,"ir")):
|
||||||
|
ir = image_read_cv2(os.path.join(ori_img_folder,"ir", img_name), 'GRAY')
|
||||||
|
vi = image_read_cv2(os.path.join(ori_img_folder,"vi", img_name), 'GRAY')
|
||||||
|
fi = image_read_cv2(os.path.join(eval_folder, img_name.split('.')[0]+".png"), 'GRAY')
|
||||||
|
metric_result += np.array([Evaluator.EN(fi), Evaluator.SD(fi)
|
||||||
|
, Evaluator.SF(fi), Evaluator.MI(fi, ir, vi)
|
||||||
|
, Evaluator.SCD(fi, ir, vi), Evaluator.VIFF(fi, ir, vi)
|
||||||
|
, Evaluator.Qabf(fi, ir, vi), Evaluator.SSIM(fi, ir, vi)])
|
||||||
|
|
||||||
|
metric_result /= len(os.listdir(eval_folder))
|
||||||
|
print("\t\t EN\t SD\t SF\t MI\tSCD\tVIF\tQabf\tSSIM")
|
||||||
|
print(model_name+'\t'+str(np.round(metric_result[0], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[1], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[2], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[3], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[4], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[5], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[6], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[7], 2))
|
||||||
|
)
|
||||||
|
print("="*80)
|
86
test_MIF.py
Normal file
@ -0,0 +1,86 @@
|
|||||||
|
from net import Restormer_Encoder, Restormer_Decoder, BaseFeatureExtraction, DetailFeatureExtraction
|
||||||
|
import os
|
||||||
|
import numpy as np
|
||||||
|
from utils.Evaluator import Evaluator
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
from utils.img_read_save import img_save,image_read_cv2
|
||||||
|
import warnings
|
||||||
|
import logging
|
||||||
|
warnings.filterwarnings("ignore")
|
||||||
|
logging.basicConfig(level=logging.CRITICAL)
|
||||||
|
import cv2
|
||||||
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
||||||
|
CDDFuse_path=r"models/CDDFuse_IVF.pth"
|
||||||
|
CDDFuse_MIF_path=r"models/CDDFuse_MIF.pth"
|
||||||
|
for dataset_name in ["MRI_CT","MRI_PET","MRI_SPECT"]:
|
||||||
|
print("\n"*2+"="*80)
|
||||||
|
print("The test result of "+dataset_name+" :")
|
||||||
|
print("\t\t EN\t SD\t SF\t MI\tSCD\tVIF\tQabf\tSSIM")
|
||||||
|
for ckpt_path in [CDDFuse_path,CDDFuse_MIF_path]:
|
||||||
|
model_name=ckpt_path.split('/')[-1].split('.')[0]
|
||||||
|
test_folder=os.path.join('test_img',dataset_name)
|
||||||
|
test_out_folder=os.path.join('test_result',dataset_name)
|
||||||
|
|
||||||
|
|
||||||
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||||
|
Encoder = nn.DataParallel(Restormer_Encoder()).to(device)
|
||||||
|
Decoder = nn.DataParallel(Restormer_Decoder()).to(device)
|
||||||
|
BaseFuseLayer = nn.DataParallel(BaseFeatureExtraction(dim=64, num_heads=8)).to(device)
|
||||||
|
DetailFuseLayer = nn.DataParallel(DetailFeatureExtraction(num_layers=1)).to(device)
|
||||||
|
|
||||||
|
Encoder.load_state_dict(torch.load(ckpt_path)['DIDF_Encoder'])
|
||||||
|
Decoder.load_state_dict(torch.load(ckpt_path)['DIDF_Decoder'])
|
||||||
|
BaseFuseLayer.load_state_dict(torch.load(ckpt_path)['BaseFuseLayer'])
|
||||||
|
DetailFuseLayer.load_state_dict(torch.load(ckpt_path)['DetailFuseLayer'])
|
||||||
|
Encoder.eval()
|
||||||
|
Decoder.eval()
|
||||||
|
BaseFuseLayer.eval()
|
||||||
|
DetailFuseLayer.eval()
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
for img_name in os.listdir(os.path.join(test_folder,dataset_name.split('_')[0])):
|
||||||
|
data_IR=image_read_cv2(os.path.join(test_folder,dataset_name.split('_')[1],img_name),mode='GRAY')[np.newaxis,np.newaxis, ...]/255.0
|
||||||
|
data_VIS = image_read_cv2(os.path.join(test_folder,dataset_name.split('_')[0],img_name), mode='GRAY')[np.newaxis,np.newaxis, ...]/255.0
|
||||||
|
|
||||||
|
data_IR,data_VIS = torch.FloatTensor(data_IR),torch.FloatTensor(data_VIS)
|
||||||
|
data_VIS, data_IR = data_VIS.cuda(), data_IR.cuda()
|
||||||
|
|
||||||
|
feature_V_B, feature_V_D, feature_V = Encoder(data_VIS)
|
||||||
|
feature_I_B, feature_I_D, feature_I = Encoder(data_IR)
|
||||||
|
feature_F_B = BaseFuseLayer(feature_V_B + feature_I_B)
|
||||||
|
feature_F_D = DetailFuseLayer(feature_V_D + feature_I_D)
|
||||||
|
if ckpt_path==CDDFuse_path:
|
||||||
|
data_Fuse, _ = Decoder(data_IR+data_VIS, feature_F_B, feature_F_D)
|
||||||
|
else:
|
||||||
|
data_Fuse, _ = Decoder(None, feature_F_B, feature_F_D)
|
||||||
|
data_Fuse=(data_Fuse-torch.min(data_Fuse))/(torch.max(data_Fuse)-torch.min(data_Fuse))
|
||||||
|
fi = np.squeeze((data_Fuse * 255).cpu().numpy())
|
||||||
|
img_save(fi, img_name.split(sep='.')[0], test_out_folder)
|
||||||
|
eval_folder=test_out_folder
|
||||||
|
ori_img_folder=test_folder
|
||||||
|
|
||||||
|
metric_result = np.zeros((8))
|
||||||
|
for img_name in os.listdir(os.path.join(ori_img_folder,dataset_name.split('_')[0])):
|
||||||
|
ir = image_read_cv2(os.path.join(ori_img_folder,dataset_name.split('_')[1], img_name), 'GRAY')
|
||||||
|
vi = image_read_cv2(os.path.join(ori_img_folder,dataset_name.split('_')[0], img_name), 'GRAY')
|
||||||
|
fi = image_read_cv2(os.path.join(eval_folder, img_name.split('.')[0]+".png"), 'GRAY')
|
||||||
|
metric_result += np.array([Evaluator.EN(fi), Evaluator.SD(fi)
|
||||||
|
, Evaluator.SF(fi), Evaluator.MI(fi, ir, vi)
|
||||||
|
, Evaluator.SCD(fi, ir, vi), Evaluator.VIFF(fi, ir, vi)
|
||||||
|
, Evaluator.Qabf(fi, ir, vi), Evaluator.SSIM(fi, ir, vi)])
|
||||||
|
|
||||||
|
metric_result /= len(os.listdir(eval_folder))
|
||||||
|
|
||||||
|
print(model_name+'\t'+str(np.round(metric_result[0], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[1], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[2], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[3], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[4], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[5], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[6], 2))+'\t'
|
||||||
|
+str(np.round(metric_result[7], 2))
|
||||||
|
)
|
||||||
|
print("="*80)
|
||||||
|
|
||||||
|
|
BIN
test_img/MRI_CT/CT/11.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_CT/CT/12.png
Normal file
After Width: | Height: | Size: 41 KiB |
BIN
test_img/MRI_CT/CT/13.png
Normal file
After Width: | Height: | Size: 41 KiB |
BIN
test_img/MRI_CT/CT/14.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_CT/CT/15.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_CT/CT/16.png
Normal file
After Width: | Height: | Size: 41 KiB |
BIN
test_img/MRI_CT/CT/17.png
Normal file
After Width: | Height: | Size: 39 KiB |
BIN
test_img/MRI_CT/CT/18.png
Normal file
After Width: | Height: | Size: 37 KiB |
BIN
test_img/MRI_CT/CT/19.png
Normal file
After Width: | Height: | Size: 34 KiB |
BIN
test_img/MRI_CT/CT/20.png
Normal file
After Width: | Height: | Size: 33 KiB |
BIN
test_img/MRI_CT/CT/21.png
Normal file
After Width: | Height: | Size: 38 KiB |
BIN
test_img/MRI_CT/CT/22.png
Normal file
After Width: | Height: | Size: 41 KiB |
BIN
test_img/MRI_CT/CT/23.png
Normal file
After Width: | Height: | Size: 44 KiB |
BIN
test_img/MRI_CT/CT/24.png
Normal file
After Width: | Height: | Size: 44 KiB |
BIN
test_img/MRI_CT/CT/25.png
Normal file
After Width: | Height: | Size: 43 KiB |
BIN
test_img/MRI_CT/CT/26.png
Normal file
After Width: | Height: | Size: 41 KiB |
BIN
test_img/MRI_CT/CT/27.png
Normal file
After Width: | Height: | Size: 40 KiB |
BIN
test_img/MRI_CT/CT/28.png
Normal file
After Width: | Height: | Size: 40 KiB |
BIN
test_img/MRI_CT/CT/29.png
Normal file
After Width: | Height: | Size: 40 KiB |
BIN
test_img/MRI_CT/CT/30.png
Normal file
After Width: | Height: | Size: 39 KiB |
BIN
test_img/MRI_CT/CT/31.png
Normal file
After Width: | Height: | Size: 37 KiB |
BIN
test_img/MRI_CT/MRI/11.png
Normal file
After Width: | Height: | Size: 59 KiB |
BIN
test_img/MRI_CT/MRI/12.png
Normal file
After Width: | Height: | Size: 58 KiB |
BIN
test_img/MRI_CT/MRI/13.png
Normal file
After Width: | Height: | Size: 60 KiB |
BIN
test_img/MRI_CT/MRI/14.png
Normal file
After Width: | Height: | Size: 66 KiB |
BIN
test_img/MRI_CT/MRI/15.png
Normal file
After Width: | Height: | Size: 63 KiB |
BIN
test_img/MRI_CT/MRI/16.png
Normal file
After Width: | Height: | Size: 61 KiB |
BIN
test_img/MRI_CT/MRI/17.png
Normal file
After Width: | Height: | Size: 60 KiB |
BIN
test_img/MRI_CT/MRI/18.png
Normal file
After Width: | Height: | Size: 57 KiB |
BIN
test_img/MRI_CT/MRI/19.png
Normal file
After Width: | Height: | Size: 54 KiB |
BIN
test_img/MRI_CT/MRI/20.png
Normal file
After Width: | Height: | Size: 52 KiB |
BIN
test_img/MRI_CT/MRI/21.png
Normal file
After Width: | Height: | Size: 51 KiB |
BIN
test_img/MRI_CT/MRI/22.png
Normal file
After Width: | Height: | Size: 53 KiB |
BIN
test_img/MRI_CT/MRI/23.png
Normal file
After Width: | Height: | Size: 54 KiB |
BIN
test_img/MRI_CT/MRI/24.png
Normal file
After Width: | Height: | Size: 56 KiB |
BIN
test_img/MRI_CT/MRI/25.png
Normal file
After Width: | Height: | Size: 55 KiB |
BIN
test_img/MRI_CT/MRI/26.png
Normal file
After Width: | Height: | Size: 52 KiB |
BIN
test_img/MRI_CT/MRI/27.png
Normal file
After Width: | Height: | Size: 50 KiB |
BIN
test_img/MRI_CT/MRI/28.png
Normal file
After Width: | Height: | Size: 50 KiB |
BIN
test_img/MRI_CT/MRI/29.png
Normal file
After Width: | Height: | Size: 49 KiB |
BIN
test_img/MRI_CT/MRI/30.png
Normal file
After Width: | Height: | Size: 48 KiB |
BIN
test_img/MRI_CT/MRI/31.png
Normal file
After Width: | Height: | Size: 46 KiB |
BIN
test_img/MRI_PET/MRI/11.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/12.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/13.png
Normal file
After Width: | Height: | Size: 43 KiB |
BIN
test_img/MRI_PET/MRI/14.png
Normal file
After Width: | Height: | Size: 43 KiB |
BIN
test_img/MRI_PET/MRI/15.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/16.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/17.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/18.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/19.png
Normal file
After Width: | Height: | Size: 43 KiB |
BIN
test_img/MRI_PET/MRI/20.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/21.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/22.png
Normal file
After Width: | Height: | Size: 41 KiB |
BIN
test_img/MRI_PET/MRI/23.png
Normal file
After Width: | Height: | Size: 39 KiB |
BIN
test_img/MRI_PET/MRI/24.png
Normal file
After Width: | Height: | Size: 38 KiB |
BIN
test_img/MRI_PET/MRI/25.png
Normal file
After Width: | Height: | Size: 37 KiB |
BIN
test_img/MRI_PET/MRI/26.png
Normal file
After Width: | Height: | Size: 36 KiB |
BIN
test_img/MRI_PET/MRI/27.png
Normal file
After Width: | Height: | Size: 34 KiB |
BIN
test_img/MRI_PET/MRI/28.png
Normal file
After Width: | Height: | Size: 32 KiB |
BIN
test_img/MRI_PET/MRI/29.png
Normal file
After Width: | Height: | Size: 29 KiB |
BIN
test_img/MRI_PET/MRI/30.png
Normal file
After Width: | Height: | Size: 26 KiB |
BIN
test_img/MRI_PET/MRI/31.png
Normal file
After Width: | Height: | Size: 21 KiB |
BIN
test_img/MRI_PET/MRI/32.png
Normal file
After Width: | Height: | Size: 39 KiB |
BIN
test_img/MRI_PET/MRI/33.png
Normal file
After Width: | Height: | Size: 40 KiB |
BIN
test_img/MRI_PET/MRI/34.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/35.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/36.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/37.png
Normal file
After Width: | Height: | Size: 42 KiB |
BIN
test_img/MRI_PET/MRI/38.png
Normal file
After Width: | Height: | Size: 44 KiB |
BIN
test_img/MRI_PET/MRI/39.png
Normal file
After Width: | Height: | Size: 46 KiB |
BIN
test_img/MRI_PET/MRI/40.png
Normal file
After Width: | Height: | Size: 44 KiB |
BIN
test_img/MRI_PET/MRI/41.png
Normal file
After Width: | Height: | Size: 44 KiB |
BIN
test_img/MRI_PET/MRI/42.png
Normal file
After Width: | Height: | Size: 43 KiB |
BIN
test_img/MRI_PET/MRI/43.png
Normal file
After Width: | Height: | Size: 41 KiB |
BIN
test_img/MRI_PET/MRI/44.png
Normal file
After Width: | Height: | Size: 39 KiB |
BIN
test_img/MRI_PET/MRI/45.png
Normal file
After Width: | Height: | Size: 39 KiB |
BIN
test_img/MRI_PET/MRI/46.png
Normal file
After Width: | Height: | Size: 39 KiB |
BIN
test_img/MRI_PET/MRI/47.png
Normal file
After Width: | Height: | Size: 38 KiB |
BIN
test_img/MRI_PET/MRI/48.png
Normal file
After Width: | Height: | Size: 36 KiB |
BIN
test_img/MRI_PET/MRI/49.png
Normal file
After Width: | Height: | Size: 34 KiB |
BIN
test_img/MRI_PET/MRI/50.png
Normal file
After Width: | Height: | Size: 30 KiB |
BIN
test_img/MRI_PET/MRI/51.png
Normal file
After Width: | Height: | Size: 26 KiB |
BIN
test_img/MRI_PET/MRI/52.png
Normal file
After Width: | Height: | Size: 20 KiB |
BIN
test_img/MRI_PET/PET/11.png
Normal file
After Width: | Height: | Size: 16 KiB |
BIN
test_img/MRI_PET/PET/12.png
Normal file
After Width: | Height: | Size: 18 KiB |
BIN
test_img/MRI_PET/PET/13.png
Normal file
After Width: | Height: | Size: 21 KiB |
BIN
test_img/MRI_PET/PET/14.png
Normal file
After Width: | Height: | Size: 27 KiB |
BIN
test_img/MRI_PET/PET/15.png
Normal file
After Width: | Height: | Size: 35 KiB |
BIN
test_img/MRI_PET/PET/16.png
Normal file
After Width: | Height: | Size: 40 KiB |
BIN
test_img/MRI_PET/PET/17.png
Normal file
After Width: | Height: | Size: 47 KiB |
BIN
test_img/MRI_PET/PET/18.png
Normal file
After Width: | Height: | Size: 49 KiB |
BIN
test_img/MRI_PET/PET/19.png
Normal file
After Width: | Height: | Size: 50 KiB |