| import torch | |
| import torch.nn as nn | |
| from .bayar_conv import BayarConv2d | |
| from .srm_conv import SRMConv2d | |
| class EarlyFusionPreFilter(nn.Module): | |
| def __init__(self, bayar_magnitude: float, srm_clip: float): | |
| super().__init__() | |
| self.bayar_filter = BayarConv2d( | |
| 3, 3, 5, stride=1, padding=2, magnitude=bayar_magnitude | |
| ) | |
| self.srm_filter = SRMConv2d(stride=1, padding=2, clip=srm_clip) | |
| self.rgb_filter = nn.Identity() | |
| self.map = nn.Conv2d(9, 3, 1, stride=1, padding=0) | |
| def forward(self, x): | |
| x_bayar = self.bayar_filter(x) | |
| x_srm = self.srm_filter(x) | |
| x_rgb = self.rgb_filter(x) | |
| x_concat = torch.cat([x_bayar, x_srm, x_rgb], dim=1) | |
| x_concat = self.map(x_concat) | |
| return x_concat | |