【论文阅读 | TPAMI 2025 | RWKVFusion：利用统一语言与掩码引导的高效图像融合网络】

1
    def _forward_implem(self, inp, modal, mask=None, llm_feature=None):
2
        # inp: vi/over/far/SPECT/LRMS
3
        # modal: ir/under/near/MRI/PAN
4

5
        # forward features
6
        feature_out, patch_embd_x = self._forward_features(
7
            inp, modal, mask, llm_feature
8
        )
9

10
        # forward prior
11
        prior = self._foward_prior(inp, modal)
12

13
        # reconstruction
14
        outp = self._forward_recon(feature_out, prior, patch_embd_x)
15

16
        return outp

1
    def _forward_features(self, inp, modal, mask=None, llm_feature=None):
2
        bs, _, H, W = inp.shape
3

4
        cat_x_modal = torch.cat([inp, modal], dim=1)
5
        patch_embd_x = self.patch_embd(cat_x_modal)
6
        modal = self.modal_in(cat_x_modal)
7

8
        if self.if_abs_pos:
9
            x = patch_embd_x + self.resize_pos_embd(
10
                self.abs_pos, inp_size=(H, W)
11
            ).expand(bs, patch_embd_x.size(1), -1, -1)
12
        else:
13
            x = patch_embd_x
14

15
        if llm_feature is not None and self.llm_channel is not None:
16
            llm_feature = F.normalize(llm_feature, p=2, dim=1)
17
            llm_feature = self.llm_embd(llm_feature)
18
            llm_feature = (llm_feature + self.txt_sine_pe).float()
19
            llm_feature = self.embd_drop(llm_feature)
20

21
        if mask is not None and hasattr(self, "mask_embd"):
22
            mask = self.mask_embd(mask)
23

24
        # init condition
25
        u_cond = ConditionInput(modal, llm_feature, mask)
26

27
        ## encoder
28
        encs = []
29
        for encoder, down in zip(self.encoders, self.downs):
30
            u_cond.modalities = self.resize_conds(
31
                H, W, u_cond.modalities, resample_type="bilinear"
32
            )
33
            u_cond.mask_input = self.resize_conds(
34
                H, W, u_cond.mask_input, resample_type="nearest"
35
            )
36
            x, u_cond = encoder.enc_forward(x, u_cond, (H, W))
37
            encs.append(x)
38
            x = down(x)
39
            H = H // 2
40
            W = W // 2
41

42
        ## middle layer
43
        u_cond.modalities = self.resize_conds(
44
            H, W, u_cond.modalities, resample_type="bilinear"
45
        )
46
        u_cond.mask_input = self.resize_conds(
47
            H, W, u_cond.mask_input, resample_type="nearest"
48
        )
49
        x, _ = self.middle_blks.enc_forward(x, u_cond, (H, W))
50

51
        ## decoder
52
        for decoder, up, enc_skip, skip_scale in zip(
53
            self.decoders, self.ups, encs[::-1], self.skip_scales
54
        ):
55
            x = up(x)
56
            H = H * 2
57
            W = W * 2
58
            u_cond.modalities = self.resize_conds(
59
                H, W, u_cond.modalities, resample_type="bilinear"
60
            )
61
            u_cond.mask_input = self.resize_conds(
62
                H, W, u_cond.mask_input, resample_type="nearest"
63
            )
64
            x = torch.cat([x, enc_skip * skip_scale.view(1, -1, 1, 1)], dim=1)
65
            x = decoder.dec_forward(x, u_cond, (H, W))
66

67
        return x, patch_embd_x

1
    def BRWKV_img_forward(
2
        self,
3
        x: torch.Tensor,
4
        MIMF_modals: torch.Tensor = None,
5
        MIFM_mm: torch.Tensor = None,
6
        llm_feat: torch.Tensor = None,
7
        patch_resolution=None,
8
    ):
9
        """
10
        ##! 1. Image-only attention with clipping windows
11
        w_img = window_partition(img, window_size)
12
        w_img = scan(img)
13
        w_img = only_img_attention(img)
14

15
        ##! 2. Multi-modal attention with txt feature
16
        # solution 1 (downsample image)
17
        img = window_reverse(w_img, window_size, H, W)
18
        ds_img = downsample_scanned_img(img)
19
        ds_img, txt = multi_modal_attention(ds_img, txt)
20
        upsample_img = upsample_scanned_img(ds_img)
21
        img = (upsample_img + img) / 2
22

23
        # solution 2 (repeat txt feature on bs-dim)
24
        rep_bs_txt = repeat(txt, 'bs d l -> (bs n_winds) d l')
25
        w_img, rep_bs_txt = multi_modal_attention(w_img, rep_bs_txt)
26
        img = window_reverse(w_img, window_size, H, W)
27
        txt = rearrange(rep_bs_txt, '(bs n_winds) d l -> bs n_winds d l', bs=B)
28
        txt = txt.mean(dim=1)
29

30
        ##! 3. Multi-modal FFN
31
        #! Flux.1 says we should have two FFNs, one for image and one for txt
32
        img = FFN(img, txt)
33

34
        """
35

36
        B, C, H, W = x.shape
37
        # has_llm = (llm_feat is not None) and self.has_llm
38

39
        # pre-fusion previous (non-downsampled) image and MIFM image
40
        # x = x + MIFM_img
41

42
        # ver 3
43
        # MIFM_mm = MIFM_mm * self.lerp_factor.view(1, -1, 1, 1)
44
        # x = self.mm_proj_in(torch.cat([x, MIFM_mm], dim=1))
45
        x = self.mm_proj_in(
46
            torch.cat(
47
                [x, MIFM_mm * self.lerp_factor.view(1, -1, 1, 1), MIMF_modals], dim=1
48
            )
49
        )
50

51
        # we perform scanning first and then window partition
52
        # different from v11 where we perform window partition first
53

54
        # window partition
55
        # (bs, c, h * w) -> (bs x n_winds, c, window_size, window_size)
56
        h, w, x = self.partition_img_by_window(H, W, x)
57

58
        # ESS
59
        x = (
60
            self.scan(x).view(x.size(0) * self.K, -1, h * w).transpose(1, 2)
61
        )  # [b*k, l, c]
62

63
        ######################## Image-only Attention ##########################
64

65
        # Image spatial mixing
66
        x = self.only_img_attn_ln(x)
67
        sc1, sh1, ga1 = self.modulator1(x)
68
        prenorm_x = (1 + sc1) * x + sh1  # modulated img
69
        x_attn = self.drop_path(self.att_img(prenorm_x, None, (h, w)))
70
        x = ga1 * x_attn + x
71

72
        # merge ESS
73
        # (bs, c * K, H * W) -> (bs, c, H, W)
74
        x = rearrange(x, "(b k) (h w) d -> b k d h w", k=self.K, h=h, w=w)
75
        x = self.merge(x) / self.K
76

77
        # window reverse
78
        # (bs x n_winds, c, wind x wind) -> (bs, c, h x w)
79
        x = self.reverse_img_by_window(H, W, x)
80
        x = x.view(B, C, H * W).transpose(1, 2)
81

82
        ########################################################################
83

84
        ######################## Image-modality FFN ############################
85

86
        # Channel mixing
87
        x = self.fusion_img_llm_ln(x)
88
        sc2, sh2, ga2 = self.modulator2(x)
89
        prenorm_x = (1 + sc2) * x + sh2
90

91
        # 2. pure image FFN
92
        x_ffn = self.drop_path(self.ffn_fusion(prenorm_x, None, (H, W)))
93
        x = ga2 * x_ffn + x
94

95
        x = x.transpose(1, 2).view(B, C, H, W)
96

97
        #########################################################################
98

99
        # txt reduced channel
100
        if (not self.last_enc_block) and self.has_llm and llm_feat is not None:
101
            llm_feat = self.txt_chan_reduce(llm_feat)
102

103
        return x, llm_feat

1
    def forward(
2
        self,
3
        x,
4
        txt=None,
5
        patch_resolution=None,
6
        mm_tokens: "tuple[torch.Tensor] | None" = None,
7
    ):
8
        def _inner_forward(x, txt, patch_resolution, mm_tokens):
9
            B, T, C = x.size()
10
            # sr, k, v, T = self.jit_func(x, txt, patch_resolution, mm_tokens)
11
            r, k, v, g, T = self.jit_func(x, txt, patch_resolution, mm_tokens)
12
            # x = RUN_CUDA_RWKV5(B, T, C, self.spatial_decay / T, self.spatial_first / T, k, v)
13
            x = RUN_CUDA_RWKV5_2(
14
                B, T, C, self.n_head, r, k, v, w=self.time_decay, u=self.time_faaaa
15
            )
16
            # print(f'Spatial mix - x max: {x.abs().max()}, x norm: {x.norm()}')
17
            x = self.key_norm(x)
18
            # x = sr * x
19
            # x = self.output(x)
20
            x = self.jit_func_2(x, g)
21

22
            return x
23

24
        if self.with_cp and x.requires_grad:
25
            x = cp.checkpoint(
26
                _inner_forward, x, txt, patch_resolution, mm_tokens, use_reentrant=False
27
            )
28
        else:
29
            x = _inner_forward(x, txt, patch_resolution, mm_tokens)
30
        return x

1
class CrossScan(torch.autograd.Function):
2
    # ZSJ 这里是把图像按照特定方向展平的地方，改变扫描方向可以在这里修改
3
    @staticmethod
4
    def forward(ctx, x: torch.Tensor):
5
        B, C, H, W = x.shape
6
        ctx.shape = (B, C, H, W)
7
        # xs = x.new_empty((B, 4, C, H * W))
8
        xs = x.new_empty((B, 8, C, H * W))
9
        # 添加横向和竖向的扫描
10
        xs[:, 0] = x.flatten(2, 3)
11
        xs[:, 1] = x.transpose(dim0=2, dim1=3).flatten(2, 3)
12
        xs[:, 2:4] = torch.flip(xs[:, 0:2], dims=[-1])
13

14
        # 提供斜向和反斜向的扫描
15
        xs[:, 4] = diagonal_gather(x)
16
        xs[:, 5] = antidiagonal_gather(x)
17
        xs[:, 6:8] = torch.flip(xs[:, 4:6], dims=[-1])
18

19
        return xs
20

21
    @staticmethod
22
    def backward(ctx, ys: torch.Tensor):
23
        # out: (b, k, d, l)
24
        B, C, H, W = ctx.shape
25
        L = H * W
26
        # 把横向和竖向的反向部分再反向回来，并和原来的横向和竖向相加
27
        # ys = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, -1, L)
28
        y_rb = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, -1, L)
29
        # 把竖向的部分转成横向，然后再相加,再转回最初是的矩阵形式
30
        # y = ys[:, 0] + ys[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).contiguous().view(B, -1, L)
31
        y_rb = y_rb[:, 0] + y_rb[:, 1].view(B, -1, W, H).transpose(
32
            dim0=2, dim1=3
33
        ).contiguous().view(B, -1, L)
34
        y_rb = y_rb.view(B, -1, H, W)
35

36
        # 把斜向和反斜向的反向部分再反向回来，并和原来的斜向和反斜向相加
37
        y_da = ys[:, 4:6] + ys[:, 6:8].flip(dims=[-1]).view(B, 2, -1, L)
38
        # 把斜向和反斜向的部分都转成原来的最初的矩阵形式，再相加
39
        y_da = diagonal_scatter(y_da[:, 0], (B, C, H, W)) + antidiagonal_scatter(
40
            y_da[:, 1], (B, C, H, W)
41
        )
42

43
        y_res = y_rb + y_da
44
        # return y.view(B, -1, H, W)
45
        return y_res

1
class CrossMerge(torch.autograd.Function):
2
    @staticmethod
3
    def forward(ctx, ys: torch.Tensor):
4
        B, K, D, H, W = ys.shape
5
        ctx.shape = (H, W)
6
        ys = ys.view(B, K, D, -1)
7
        # ys = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
8
        # y = ys[:, 0] + ys[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).contiguous().view(B, D, -1)
9

10
        y_rb = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
11
        # 把竖向的部分转成横向，然后再相加,再转回最初是的矩阵形式
12
        y_rb = y_rb[:, 0] + y_rb[:, 1].view(B, -1, W, H).transpose(
13
            dim0=2, dim1=3
14
        ).contiguous().view(B, D, -1)
15
        y_rb = y_rb.view(B, -1, H, W)
16

17
        # 把斜向和反斜向的反向部分再反向回来，并和原来的斜向和反斜向相加
18
        y_da = ys[:, 4:6] + ys[:, 6:8].flip(dims=[-1]).view(B, 2, D, -1)
19
        # 把斜向和反斜向的部分都转成原来的最初的矩阵形式，再相加
20
        y_da = diagonal_scatter(y_da[:, 0], (B, D, H, W)) + antidiagonal_scatter(
21
            y_da[:, 1], (B, D, H, W)
22
        )
23

24
        y_res = y_rb + y_da
25
        return y_res.view(B, D, -1)
26
        # return y
27

28
    @staticmethod
29
    def backward(ctx, x: torch.Tensor):
30
        # B, D, L = x.shape
31
        # out: (b, k, d, l)
32
        H, W = ctx.shape
33
        B, C, L = x.shape
34
        # xs = x.new_empty((B, 4, C, L))
35
        xs = x.new_empty((B, 8, C, L))
36

37
        # 横向和竖向扫描
38
        xs[:, 0] = x
39
        xs[:, 1] = x.view(B, C, H, W).transpose(dim0=2, dim1=3).flatten(2, 3)
40
        xs[:, 2:4] = torch.flip(xs[:, 0:2], dims=[-1])
41
        # xs = xs.view(B, 4, C, H, W)
42

43
        # 提供斜向和反斜向的扫描
44
        xs[:, 4] = diagonal_gather(x.view(B, C, H, W))
45
        xs[:, 5] = antidiagonal_gather(x.view(B, C, H, W))
46
        xs[:, 6:8] = torch.flip(xs[:, 4:6], dims=[-1])
47

48
        # return xs
49
        return xs.view(B, 8, C, H, W)
50

51

52
############### K = 4 ################

1
    def forward(self, img_feat: torch.Tensor, cond_input: ConditionInput):
2
        modals = cond_input.modalities
3
        mask = cond_input.mask_input
4
        llm_feat = cond_input.llm_feature
5

6
        # modalities encoder
7
        modals = self.modal_convs(modals)
8
        modals = self.modal_norm(modals)
9

10
        # feature encoder
11
        img_feat = self.feat_drop(self.feat_convs(img_feat))
12

13
        # feature gate
14
        img_gate = self.feat_norm(img_feat + modals)
15
        feat_pool = self.adap_pool(img_gate)
16
        feat_mul = self.act_conv(feat_pool)  # to gate
17
        # m1, m2 = feat_mul.chunk(2, dim=1)
18

19
        # gate
20
        img_feat_gate = img_feat * feat_mul  # + img_feat
21
        modals_gate = modals * feat_mul  # + modals
22
        img_feat_gate = img_feat_gate + modals_gate
23
        # img_feat_gate = img_feat * m1 + modals * m2
24

25
        # fusion mask in feat and modals
26
        # then needs to be normalized
27
        if mask is not None and self.with_mask:
28
            mask = self.mask_convs(mask)
29
            img_feat_mask = img_feat * mask  # + img_feat
30
            modals_mask = modals * mask  # + modals
31
            img_feat = img_feat_gate + img_feat_mask + modals_mask
32
            # img_feat = img_feat_gate + img_feat_mask * m1 + modals_mask * m2
33
        else:
34
            img_feat = img_feat_gate
35

36
        # fusion feat and modals
37
        img_feat = self.fusion_proj(img_feat)
38
        # img_feat_modals = self.modals_out_ln(img_feat)
39
        img_feat_modals = img_feat
40

41
        ################################# LLM text and image attention ###########################
42
        # fusion image and llm features
43
        if llm_feat is not None and self.with_llm_feat:
44
            llm_feat = self.llm_drop(self.llm_dense(llm_feat))
45
            llm_feat = self.add_llm_pe(llm_feat)
46

47
            # downsample image
48
            img_feat, full_size = self.downsample_img_by_llm(img_feat)
49

50
            # ESS
51
            bs, C, H, W = img_feat.shape
52
            img_feat = self.scan(img_feat).view(bs * self.K, C, H * W).transpose(1, 2)
53
            img_feat = self.add_img_pe(img_feat, H, W)
54

55
            # multi-modal RWKV attention
56
            img_feat = self.multi_modal_img_ln(img_feat)
57
            llm_feat = self.multi_modal_txt_ln(llm_feat)
58
            # llm_feat: [bs, llm_len, c], img_feat: [bs, img_len, c]
59

60
            if self.add_mm_tokens:
61
                multi_modal_attn_out = self.multi_modal_attn(
62
                    img_feat,
63
                    llm_feat.repeat(self.K, 1, 1),
64
                    (H, W),
65
                    (self.img_token, self.txt_token),
66
                )
67
            else:
68
                multi_modal_attn_out = self.multi_modal_attn(
69
                    img_feat, llm_feat.repeat(self.K, 1, 1), (H, W)
70
                )
71

72
            # extract img and llm feat
73
            img_feat_out, llm_feat_out = self.extract_img_and_llm_feat(
74
                H, W, multi_modal_attn_out, self.with_llm_feat
75
            )
76

77
            # residual
78
            img_feat = img_feat + img_feat_out
79
            llm_feat = llm_feat + llm_feat_out.view(self.K, bs, -1, C).mean(dim=0)
80

81
            # merge ESS
82
            img_feat = rearrange(
83
                img_feat, "(b k) (h w) d -> b k d h w", k=self.K, h=H, w=W
84
            )
85
            img_feat = self.merge(img_feat) / self.K
86

87
            # upsample image
88
            img_feat = self.upsample_img_by_llm(img_feat, full_size)
89

90
            # mlps
91
            img_feat = self.img_mlp(img_feat) + img_feat
92
            if not self.last_enc_block:
93
                llm_feat = self.txt_mlp(llm_feat) + llm_feat
94

95
            # img_feat_llm = self.llm_out_ln(img_feat)
96
            img_feat_llm = img_feat
97

98
            return img_feat_modals, img_feat_llm, modals, llm_feat
99
        ###########################################################################################
100

101
        return img_feat_modals, modals, llm_feat

1
    def check_multi_value_mask_to_one_hot(self, mask: torch.Tensor | None):
2
        if mask is not None and self.has_mask:
3
            if self.multi_value_mask_max_classes is not None and mask.ndim == 3:
4
                # * if mask is dtype float, must be careful about its values, must be in [0, multi_value_mask_max_classes]
5
                # * if float value in it, the code will raise CUDA error by `scatter_` without clear hint.
6

7
                mask[mask >= self.multi_value_mask_max_classes] = (
8
                    0.0  # larger than max classes are set to background
9
                )
10
                new_mask = torch.zeros(
11
                    mask.size(0),
12
                    self.multi_value_mask_max_classes,
13
                    *mask.shape[-2:],
14
                    device=mask.device,
15
                    dtype=torch.float32,
16
                )
17
                new_mask.scatter_(1, mask.long()[:, None], 1.0)
18
            elif mask.ndim == 4:
19
                new_mask = mask
20
            else:
21
                raise ValueError(
22
                    f"Invalid mask shape: {mask.shape} when `multi_value_mask_max_classes` is {self.multi_value_mask_max_classes}",
23
                    "or model has `has_mask` as True. Input mask should be 3D or 4D tensor",
24
                    "such as [Bs, 1, H, W] or [Bs, *, H, W] (to be one-hot encoded in this case).",
25
                )
26
        else:
27
            new_mask = None
28

29
        return new_mask

1
    def fusion_train_step(
2
        self,
3
        vi: "torch.Tensor",
4
        ir: "torch.Tensor",
5
        mask: "torch.Tensor | None" = None,
6
        gt: "torch.Tensor | None" = None,
7
        txt: "torch.Tensor | None" = None,
8
        fusion_criterion: "Callable | None" = None,
9
        to_rgb_fn: "Callable | None" = None,
10
        has_gt: bool = False,
11
        **_kwargs,
12
    ):
13
        vi, ir, mask, txt = self.check_multi_modal_inputs(
14
            vi, ir, mask, txt, vi.device, vi.dtype
15
        )
16
        txt = self.drop_txt(txt)
17

18
        mask = self.check_multi_value_mask_to_one_hot(mask)
19

20
        fused_outp = self.only_fusion_step(vi, ir, mask, txt)
21

22
        # mask for loss, should be detached
23
        if mask is not None:
24
            mask_for_loss = mask.clone().detach()
25
        else:
26
            mask_for_loss = None
27

28
        fused_for_loss = to_rgb_fn(fused_outp) if to_rgb_fn is not None else fused_outp
29
        if (
30
            has_gt or gt.size(1) == 3
31
        ):  # TODO: find more robust way to check if gt is available
32
            # if we have supervised GT, we use it to compute the supervised loss
33
            # sometimes, for MEF fusion task, we can access the GTs
34
            fusion_gt = gt
35
            # two different modalities for GT to compute the unsupervised loss
36
            boundary_gt = (vi, ir)
37
        else:
38
            fusion_gt = None
39
            boundary_gt = (vi, ir)
40

41
        # compute supervised and unsupervised losses
42
        assert fusion_criterion is not None, "fusion_criterion should be provided"
43
        loss = list(
44
            fusion_criterion(
45
                fused_for_loss,
46
                boundary_gt=boundary_gt,
47
                fusion_gt=fusion_gt,
48
                mask=mask_for_loss,
49
            )
50
        )
51

52
        return fused_for_loss.clip(0, 1), loss
53

54
    @torch.no_grad()

1
    def forward(
2
        self,
3
        img_fusion: Tensor,
4
        boundary_gt: "Tensor | tuple",  # cat([vi, ir]) or tuple(vi, ir)
5
        fusion_gt: "Tensor" = None,  # ground truth provided in dataset
6
        mask: "Tensor | None" = None,
7
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
8
        if mask is not None and self.mask_loss:
9
            self.check_dtype_and_device(img_fusion, boundary_gt, mask)
10
            with torch.no_grad():
11
                if self.reduce_label:
12
                    mask2 = mask.detach().clone()
13
                    mask2[mask2 > 1.0] = 1.0
14
                else:
15
                    mask2 = mask
16

17
            ## TODO: consider this case
18
            # mask.size(1) == 2, two boundaries all have masks
19
            if mask2.size(1) == 2:
20
                assert False, "mask.size(1) == 2, two boundaries all have masks"
21
                mask_B, mask_A = mask2.chunk(2, dim=1)
22
            else:
23
                mask_B, mask_A = mask2, mask2
24
        elif not self.mask_loss:
25
            mask2 = None  # cast to None
26

27
        self.check_dtype_and_device(img_fusion, boundary_gt)
28

29
        wd = self.weight_dict
30
        loss_intensity = 0
31
        loss_color = 0
32
        loss_grad = 0
33
        loss_fusion = 0
34
        loss = {}
35

36
        # split boundary gt
37
        no_batch_ndim = img_fusion.ndim - 1
38
        broadcast_fn = lambda x: x.reshape(-1, *[1] * no_batch_ndim)  # noqa: py3.11 supported
39

40
        # img_A: ir/under/MRI, img_B: vi/over/spect
41
        if isinstance(boundary_gt, (tuple, list)):
42
            img_B, img_A = boundary_gt
43
        else:
44
            img_B, img_A = self.split_boundary_gt_tensor(boundary_gt)
45

46
        # if has gt (e.g., when we train model on multi-exposure image fusion task)
47
        if exist_gt := exists(fusion_gt):
48
            gt_loss = wd["fusion_gt"] * self.loss_func(img_fusion, fusion_gt).nanmean()
49
            loss["gt_loss"] = gt_loss
50
            loss_fusion += gt_loss
51

52
        if self.grad_norm:
53
            img_fusion = grad_norm(img_fusion)
54

55
        # ir, vi (detached in this function)
56
        (img_A, A_is_color), (img_B, B_is_color) = (
57
            self.check_rgb(img_A),
58
            self.check_rgb(img_B),
59
        )
60

61
        ## vgg latent weights
62
        # the latent weights comes from U2Fusion paper
63
        if self.latent_weighted:
64
            vi_w, ir_w = self.dynamic_weight(img_A, img_B)
65
            vi_w = broadcast_fn(vi_w)
66
            ir_w = broadcast_fn(ir_w)
67
        else:
68
            vi_w, ir_w = 1.0, 1.0
69

70
        # YCbCr decomposition
71
        # use for intensity, color, and gradient loss
72
        Y_fusion, Cb_fusion, Cr_fusion = kornia.color.rgb_to_ycbcr(img_fusion).chunk(
73
            3, dim=1
74
        )
75
        Y_A, Cb_A, Cr_A = kornia.color.rgb_to_ycbcr(img_A).chunk(3, dim=1)  # ir
76
        Y_B, Cb_B, Cr_B = kornia.color.rgb_to_ycbcr(img_B).chunk(3, dim=1)  # vi
77

78
        ## intensity and color loss
79
        # * compute boundary intensity and color loss when GT is not provided
80
        # * or we omit computing the loss
81
        # * in this case, we can say intensity and color are supervised by the GT
82
        # * we do not need to use the unsupervised loss
83
        # * `still_boundary_loss_when_gt` is to keep the boundary loss when GT is provided (optional)
84
        if (not exist_gt) or self.still_boundary_loss_when_gt:
85
            if self.prior == "max":
86
                Y_joint = torch.max(Y_A, Y_B)
87
            elif self.prior == "mean":
88
                Y_joint = (Y_A + Y_B) / 2
89

90
            if self.grad_norm:
91
                Y_fusion = grad_norm(Y_fusion)
92
                Cb_fusion = grad_norm(Cb_fusion)
93
                Cr_fusion = grad_norm(Cr_fusion)
94

95
            if mask2 is not None:
96
                # * intensity loss
97
                # 1. || fusion - max(vi, ir) ||                     max fusion
98
                # 2. || mask * fusion - mask * ir ||                ir loss (pedistrain)
99
                # 3. || (1 - mask) * fusion - (1 - mask) * vi ||    vi loss (background)
100

101
                loss_intensity = (
102
                    (
103
                        wd["inten_f_joint"] * self.loss_func(Y_fusion, Y_joint)
104
                        if self.use_prior
105
                        else 0.0
106
                    )
107
                    + (
108
                        wd["inten_f_ir"]
109
                        * ir_w
110
                        * self.loss_func(mask2 * Y_fusion, mask2 * Y_A)
111
                        if self.boundary
112
                        else 0.0
113
                    )
114
                    + (
115
                        wd["inten_f_vi"]
116
                        * vi_w
117
                        * self.loss_func(Y_fusion * (1 - mask2), Y_B * (1 - mask2))
118
                        if self.boundary
119
                        else 0.0
120
                    )
121
                )
122
                # * color loss
123
                # 1. || (1-mask) * fusion_Cb - (1-mask) * ir_Cb ||          ir_Cb loss
124
                # 2. || (1-mask) * fusion_Cr - (1-mask) * ir_Cr ||          ir_Cr loss
125
                if self.color_loss_bg_masked:
126
                    bg_mask = 1.0 - mask2
127
                else:
128
                    bg_mask = torch.ones_like(mask2)
129
                if self.color_loss:
130
                    loss_color = 0.0
131
                    if B_is_color:
132
                        loss_color += wd["color_f_cb"] * vi_w * self.loss_func(
133
                            Cb_fusion * bg_mask, Cb_B * bg_mask
134
                        ) + wd["color_f_cr"] * ir_w * self.loss_func(
135
                            Cr_fusion * bg_mask, Cr_B * bg_mask
136
                        )
137
                    if A_is_color:
138
                        loss_color += wd["color_f_cb"] * ir_w * self.loss_func(
139
                            Cb_fusion * bg_mask, Cb_A * bg_mask
140
                        ) + wd["color_f_cr"] * vi_w * self.loss_func(
141
                            Cr_fusion * bg_mask, Cr_A * bg_mask
142
                        )
143
            else:
144
                loss_intensity = (
145
                    (
146
                        wd["inten_f_joint"] * self.loss_func(Y_fusion, Y_joint)
147
                        if self.use_prior
148
                        else 0.0
149
                    )
150
                    + (
151
                        wd["inten_f_ir"] * ir_w * self.loss_func(Y_fusion, Y_A)
152
                        if self.boundary
153
                        else 0.0
154
                    )
155
                    + (
156
                        wd["inten_f_vi"] * vi_w * self.loss_func(Y_fusion, Y_B)
157
                        if self.boundary
158
                        else 0.0
159
                    )
160
                )
161

162
                if self.color_loss:
163
                    loss_color = 0.0
164
                    if B_is_color:
165
                        loss_color += wd["color_f_cb"] * self.loss_func(
166
                            Cb_fusion, Cb_B
167
                        ) + wd["color_f_cr"] * self.loss_func(Cr_fusion, Cr_B)
168
                    if A_is_color:
169
                        loss_color += wd["color_f_cb"] * self.loss_func(
170
                            Cb_fusion, Cb_A
171
                        ) + wd["color_f_cr"] * self.loss_func(Cr_fusion, Cr_A)
172

173
            loss_intensity = loss_intensity.nanmean()
174
            loss_fusion += loss_intensity
175
            loss["intensity_loss"] = loss_intensity
176
            if self.color_loss:
177
                loss_color = loss_color.nanmean()
178
                loss_fusion += loss_color
179
                loss["loss_color"] = loss_color
180

181
        ## lpips loss
182
        # * lpips loss is to enhance perceptual visuality
183
        if self.lpips:
184
            lpips_loss = 0.0
185
            if exist_gt:
186
                lpips_loss += self.lpips_loss(img_fusion, fusion_gt) * wd["lpips_f_gt"]
187

188
            # if self.still_boundary_loss_when_gt or (not exist_gt):
189
            #     lpips_loss += self.lpips_loss(img_fusion, img_A) * wd['lpips_f_ir'] + \
190
            #                   self.lpips_loss(img_fusion, img_B) * wd['lpips_f_vi']
191
            #     if self.use_prior:
192
            #         if self.prior == 'max':
193
            #             img_joint = torch.max(img_A, img_B)
194
            #         elif self.prior == 'mean':
195
            #             img_joint = (img_A + img_B) / 2
196
            #         lpips_loss += self.lpips_loss(img_fusion, img_joint) * wd['lpips_f_joint']
197

198
            loss_fusion += lpips_loss
199
            loss["lpips_loss"] = lpips_loss
200

201
        ## grad loss
202
        # * gradient loss is to enhance the fused image
203
        # * keep it although the GT is provided
204
        if self.grad:
205
            if self.grad_only_on_Y:
206
                grad_A = self.grad_op(Y_A)
207
                grad_B = self.grad_op(Y_B)
208
                grad_fusion = self.grad_op(Y_fusion)
209
            else:
210
                grad_A = self.grad_op(img_A)
211
                grad_B = self.grad_op(img_B)
212
                grad_fusion = self.grad_op(img_fusion)
213

214
            if self.grad_norm:
215
                grad_fusion = grad_norm(grad_fusion)
216

217
            grad_joint = torch.max(grad_A, grad_B)
218
            loss_grad += wd["grad_f_joint"] * self.loss_func(grad_fusion, grad_joint)
219
            # if mask is not None:
220
            #     mask_expand = mask
221
            #     if grad_fusion.ndim > 4:
222
            #         mask_expand = mask_expand.unsqueeze(1)
223
            #     loss_grad += wd['grad_f_ir'] * self.loss_func(grad_fusion * mask_expand, grad_A * mask_expand)
224

225
            # loss_grad += wd['grad_f_ir'] * self.loss_func(grad_fusion, grad_A) + \
226
            #              wd['grad_f_vi'] * self.loss_func(grad_fusion, grad_B)
227

228
            loss_grad = loss_grad.nanmean()
229

230
            loss_fusion += loss_grad
231
            loss.update({"loss_grad": loss_grad})
232

233
        ## ssim loss
234
        # * ssim loss is to enhance the fused image
235
        # * keep it although the GT is provided
236
        if self.ssim:
237
            w_A = grad_A.norm()
238
            w_B = grad_B.norm()
239
            Z = w_A + w_B
240
            w_A /= Z
241
            w_B /= Z
242

243
            if self.grad_only_on_Y:
244
                ssim_A = self.ssim_func(Y_fusion, Y_A)
245
                ssim_B = self.ssim_func(Y_fusion, Y_B)
246
            else:
247
                ssim_A = self.ssim_func(img_fusion, img_A)
248
                ssim_B = self.ssim_func(img_fusion, img_B)
249
            loss_ssim = wd["ssim_f_joint"] * (w_A * ssim_A + w_B * ssim_B)
250
            loss_fusion += loss_ssim
251
            loss.update({"loss_ssim": loss_ssim})
252

253
        ## tv loss
254
        if self.tv:
255
            img_fusion_tv = img_fusion
256
            if self.grad_norm:
257
                img_fusion_tv = grad_norm(img_fusion_tv)
258
            tv_loss = wd["tv_f"] * self.tv_loss(img_fusion_tv).nanmean()
259
            loss_fusion += tv_loss
260
            loss.update({"tv_loss": tv_loss})
261

262
        ## correlation loss
263
        if self.correlation:
264
            img_fusion_tv = img_fusion
265
            if self.grad_norm:
266
                img_fusion_tv = grad_norm(img_fusion_tv)
267
            loss_corr = wd["crr_f"] * self.cc_loss(img_A, img_B, img_fusion_tv)
268
            loss_fusion += loss_corr
269
            loss.update({"loss_corr": loss_corr})
270

271
        loss.update({"loss_fusion": loss_fusion})
272

273
        return loss_fusion, loss
274

275

276
## EMMA stage two fusion training loss