model.py

import torch
import torch.nn as nn
from loss import LossFunction
from fuse_block import TransformerBlock

class SemanticFusionUnit(nn.Module):
    def __init__(self, channels):
        super(SemanticFusionUnit, self).__init__()
        
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=channels+3, out_channels=channels, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(channels),
            nn.ReLU()
        )
        
        #TODO: WEIGHT initial 【x】
        #TODO: 模块嵌入【预处理的数据】->dataloader 【x】
        #TODO：TEST的流程【SAM】
        #TODO：数据归一化，因为涉及加法 ->dataloader【x】
        
    def forward(self, fea, sem):
        cat = torch.cat((fea, sem), dim = 1) # (b, c, h, w)
        fusion = self.conv(cat)
        return fusion
        
        

class EnhanceNetwork(nn.Module):
    def __init__(self, layers, channels):
        super(EnhanceNetwork, self).__init__()

        kernel_size = 3
        dilation = 1
        padding = int((kernel_size - 1) / 2) * dilation

        self.in_conv = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=channels, kernel_size=kernel_size, stride=1, padding=padding),
            nn.ReLU()
        )
        
        # self.fusion = SemanticFusionUnit(channels)
        self.fusion = TransformerBlock(channels, channels, num_heads=3)

        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=kernel_size, stride=1, padding=padding),
            nn.BatchNorm2d(channels),
            nn.ReLU()
        )

        self.blocks = nn.ModuleList()
        for i in range(layers):
            self.blocks.append(self.conv)

        self.out_conv = nn.Sequential(
            nn.Conv2d(in_channels=channels, out_channels=3, kernel_size=3, stride=1, padding=1),
            nn.Sigmoid()
        )

    def forward(self, input, sem):
        fea = self.in_conv(input)
        
        fea = fea + self.fusion(fea, sem)
        for conv in self.blocks:
            fea = fea + conv(fea)
        fea = self.out_conv(fea)

        illu = fea + input
        illu = torch.clamp(illu, 0.0001, 1)

        return illu


class CalibrateNetwork(nn.Module):
    def __init__(self, layers, channels):
        super(CalibrateNetwork, self).__init__()
        kernel_size = 3
        dilation = 1
        padding = int((kernel_size - 1) / 2) * dilation
        self.layers = layers

        self.in_conv = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=channels, kernel_size=kernel_size, stride=1, padding=padding),
            nn.BatchNorm2d(channels),
            nn.ReLU()
        )
        
        # self.fusion = SemanticFusionUnit(channels)

        self.convs = nn.Sequential(
            nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=kernel_size, stride=1, padding=padding),
            nn.BatchNorm2d(channels),
            nn.ReLU(),
            nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=kernel_size, stride=1, padding=padding),
            nn.BatchNorm2d(channels),
            nn.ReLU()
        )
        self.blocks = nn.ModuleList()
        for i in range(layers):
            self.blocks.append(self.convs)

        self.out_conv = nn.Sequential(
            nn.Conv2d(in_channels=channels, out_channels=3, kernel_size=3, stride=1, padding=1),
            nn.Sigmoid()
        )

    def forward(self, input, sem):
        fea = self.in_conv(input)
        
        # fea = fea + self.fusion(fea, sem)
        for conv in self.blocks:
            fea = fea + conv(fea)

        fea = self.out_conv(fea)
        delta = input - fea

        return delta



class Network(nn.Module):

    def __init__(self, stage=3):
        super(Network, self).__init__()
        self.stage = stage
        self.enhance = EnhanceNetwork(layers=1, channels=3)
        self.calibrate = CalibrateNetwork(layers=3, channels=16)
        self._criterion = LossFunction()

    def weights_init(self, m):
        if isinstance(m, nn.Conv2d):
            m.weight.data.normal_(0, 0.02)
            m.bias.data.zero_()
            print("111")

        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.normal_(1., 0.02)
            print("222")
        
        else:
            print("333")
            

    def forward(self, input, sem):

        ilist, rlist, inlist, attlist = [], [], [], []
        input_op = input
        for i in range(self.stage):
            inlist.append(input_op)
            i = self.enhance(input_op, sem)
            r = input / i
            r = torch.clamp(r, 0, 1)
            att = self.calibrate(r, sem)
            input_op = input + att
            ilist.append(i)
            rlist.append(r)
            attlist.append(torch.abs(att))

        return ilist, rlist, inlist, attlist

    def _loss(self, input, sem):
        i_list, en_list, in_list, _ = self(input, sem)
        loss = 0
        for i in range(self.stage):
            loss += self._criterion(in_list[i], i_list[i])
        return loss


class Network_woCalibrate(nn.Module):

    def __init__(self):
        super().__init__()
        self.enhance = EnhanceNetwork(layers=1, channels=3)
        self._criterion = LossFunction()

    def weights_init(self, m):
        if isinstance(m, nn.Conv2d):
            m.weight.data.normal_(0, 0.02)
            m.bias.data.zero_()

        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.normal_(1., 0.02)

    def forward(self, input, sem):
        i = self.enhance(input, sem)
        r = input / i
        r = torch.clamp(r, 0, 1)
        return i, r


    def _loss(self, input, sem):
        i, r = self(input, sem)
        loss = self._criterion(input, i)
        return loss