diff --git a/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaInter/convert.py b/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaInter/convert.py
new file mode 100755
index 0000000000000000000000000000000000000000..ca12f8eb89a61ec573a1d0139c21032d3878c7f1
--- /dev/null
+++ b/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaInter/convert.py
@@ -0,0 +1,26 @@
+import torch
+import torch.nn as nn
+from net import ConditionalNet
+import numpy as np
+
+luma = np.ones((1, 1, 100, 100), dtype=np.float32)
+yuv = np.ones((1, 2, 50, 50), dtype=np.float32)
+pred = np.ones((1, 2, 50, 50), dtype=np.float32)
+bs = np.ones((1, 2, 50, 50), dtype=np.float32)
+qp = np.ones((1, 1, 50, 50), dtype=np.float32)
+
+# model = nn.DataParallel(ConditionalNet(96, 8)) # if model is trained on multiple GPUs
+model = ConditionalNet(96, 8) # if model is trained with single GPU
+state = torch.load('50.ckpt', map_location=torch.device('cpu'))
+model.load_state_dict(state)
+
+dummy_input = (torch.from_numpy(luma), torch.from_numpy(yuv), torch.from_numpy(pred), torch.from_numpy(bs), torch.from_numpy(qp))
+torch.onnx.export(model.module, dummy_input, "NnlfSet1_ChromaCNNFilter_InterSlice.onnx")
+
+
+
+
+
+
+
+
diff --git a/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaInter/net.py b/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaInter/net.py
new file mode 100644
index 0000000000000000000000000000000000000000..90e4af26576eb9155a69ccfc1c2e9b1df8979c93
--- /dev/null
+++ b/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaInter/net.py
@@ -0,0 +1,136 @@
+import torch
+import torch.nn as nn
+
+
+def conv3x3(in_channels, out_channels, stride=1, padding=1):
+ return nn.Conv2d(in_channels, out_channels, kernel_size=3,
+ stride=stride, padding=padding)
+
+
+def conv1x1(in_channels, out_channels, stride=1, padding=0):
+ return nn.Conv2d(in_channels, out_channels, kernel_size=1,
+ stride=stride, padding=padding)
+
+
+# Conv3x3 + PReLU
+class conv3x3_f(nn.Module):
+ def __init__(self, in_channels, out_channels, stride=1):
+ super(conv3x3_f, self).__init__()
+ self.conv = conv3x3(in_channels, out_channels, stride)
+ self.relu = nn.PReLU()
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.relu(x)
+ return x
+
+
+# Conv1x1 + PReLU
+class conv1x1_f(nn.Module):
+ def __init__(self, in_channels, out_channels, stride=1):
+ super(conv1x1_f, self).__init__()
+ self.conv = conv1x1(in_channels, out_channels, stride)
+ self.relu = nn.PReLU()
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.relu(x)
+ return x
+
+
+# Residual Block
+class ResidualBlock(nn.Module):
+ def __init__(self, in_channels, out_channels, stride=1):
+ super(ResidualBlock, self).__init__()
+ self.conv1 = conv3x3(in_channels, out_channels, stride)
+ self.relu = nn.PReLU()
+ self.conv2 = conv3x3(out_channels, out_channels)
+
+ def forward(self, x):
+ residual = x
+ out = self.conv1(x)
+ out = self.relu(out)
+ out = self.conv2(out)
+ out += residual
+ return out
+
+
+class Flatten(nn.Module):
+ def forward(self, x):
+ return x.view(x.size(0), -1)
+
+
+# Channel Gate
+class ChannelGate(nn.Module):
+ def __init__(self, channels):
+ super(ChannelGate, self).__init__()
+ self.mlp = nn.Sequential(
+ Flatten(),
+ nn.Linear(1, channels),
+ nn.PReLU(),
+ nn.Linear(channels, channels)
+ )
+
+ def forward(self, x):
+ out = self.mlp(x)
+ return out
+
+
+class SpatialGate(nn.Module):
+ def __init__(self, in_channels, num_features):
+ super(SpatialGate, self).__init__()
+ self.conv1 = conv3x3(in_channels, num_features, stride=2)
+ self.relu = nn.PReLU()
+ self.conv2 = conv3x3(num_features, 1)
+
+ def forward(self, x):
+ out = self.conv1(x)
+ out = self.relu(out)
+ out = self.conv2(out)
+ return out
+
+
+class ConditionalNet(nn.Module):
+ def __init__(self, f, rbn):
+ super(ConditionalNet, self).__init__()
+ self.rbn = rbn
+ self.convLuma = conv3x3_f(1, f, 2)
+ self.convRec = conv3x3_f(2, f)
+ self.convPred = conv3x3_f(2, f)
+ self.convBs = conv3x3_f(2, f)
+ self.convQp = conv3x3_f(1, f)
+ self.fuse = conv1x1_f(5 * f, f)
+ self.transitionH = conv3x3_f(f, f, 2)
+ self.backbone = nn.ModuleList([ResidualBlock(f, f)])
+ for _ in range(self.rbn - 1):
+ self.backbone.append(ResidualBlock(f, f))
+ self.last_layer = nn.Sequential(
+ nn.Conv2d(
+ in_channels=f,
+ out_channels=f,
+ kernel_size=3,
+ stride=1,
+ padding=1),
+ nn.PReLU(),
+ nn.Conv2d(
+ in_channels=f,
+ out_channels=8,
+ kernel_size=3,
+ stride=1,
+ padding=1),
+ )
+
+ def forward(self, luma, rec, pred, bs, qp):
+ luma_f = self.convLuma(luma)
+ rec_f = self.convRec(rec)
+ pred_f = self.convPred(pred)
+ bs_f = self.convBs(bs)
+ qp_f = self.convQp(qp)
+ xh = torch.cat((luma_f, rec_f, pred_f, bs_f, qp_f), 1)
+ xh = self.fuse(xh)
+ x = self.transitionH(xh)
+ for i in range(self.rbn):
+ x = self.backbone[i](x)
+ # output
+ x = self.last_layer(x)
+ return x
diff --git a/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaIntra/convert.py b/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaIntra/convert.py
new file mode 100755
index 0000000000000000000000000000000000000000..2a39ee7003b75788f9f52d080587568909676c8a
--- /dev/null
+++ b/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaIntra/convert.py
@@ -0,0 +1,27 @@
+import torch
+import torch.nn as nn
+from net import ConditionalNet
+import numpy as np
+
+luma = np.ones((1, 1, 100, 100), dtype=np.float32)
+yuv = np.ones((1, 2, 50, 50), dtype=np.float32)
+pred = np.ones((1, 2, 50, 50), dtype=np.float32)
+split = np.ones((1, 2, 50, 50), dtype=np.float32)
+bs = np.ones((1, 2, 50, 50), dtype=np.float32)
+qp = np.ones((1, 1, 50, 50), dtype=np.float32)
+
+# model = nn.DataParallel(ConditionalNet(96, 8)) # if model is trained on multiple GPUs
+model = ConditionalNet(96, 8) # if model is trained with single GPU
+state = torch.load('50.ckpt', map_location=torch.device('cpu'))
+model.load_state_dict(state)
+
+dummy_input = (torch.from_numpy(luma), torch.from_numpy(yuv), torch.from_numpy(pred), torch.from_numpy(split), torch.from_numpy(bs), torch.from_numpy(qp))
+torch.onnx.export(model.module, dummy_input, "NnlfSet1_ChromaCNNFilter_IntraSlice.onnx")
+
+
+
+
+
+
+
+
diff --git a/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaIntra/net.py b/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaIntra/net.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c0013f2cba255d5a7ba97914e0118fdb28b05aa
--- /dev/null
+++ b/training/training_scripts/Nn_Filtering_Set_1/Conversion/ChromaIntra/net.py
@@ -0,0 +1,138 @@
+import torch
+import torch.nn as nn
+
+
+def conv3x3(in_channels, out_channels, stride=1, padding=1):
+ return nn.Conv2d(in_channels, out_channels, kernel_size=3,
+ stride=stride, padding=padding)
+
+
+def conv1x1(in_channels, out_channels, stride=1, padding=0):
+ return nn.Conv2d(in_channels, out_channels, kernel_size=1,
+ stride=stride, padding=padding)
+
+
+# Conv3x3 + PReLU
+class conv3x3_f(nn.Module):
+ def __init__(self, in_channels, out_channels, stride=1):
+ super(conv3x3_f, self).__init__()
+ self.conv = conv3x3(in_channels, out_channels, stride)
+ self.relu = nn.PReLU()
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.relu(x)
+ return x
+
+
+# Conv1x1 + PReLU
+class conv1x1_f(nn.Module):
+ def __init__(self, in_channels, out_channels, stride=1):
+ super(conv1x1_f, self).__init__()
+ self.conv = conv1x1(in_channels, out_channels, stride)
+ self.relu = nn.PReLU()
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.relu(x)
+ return x
+
+
+# Residual Block
+class ResidualBlock(nn.Module):
+ def __init__(self, in_channels, out_channels, stride=1):
+ super(ResidualBlock, self).__init__()
+ self.conv1 = conv3x3(in_channels, out_channels, stride)
+ self.relu = nn.PReLU()
+ self.conv2 = conv3x3(out_channels, out_channels)
+
+ def forward(self, x):
+ residual = x
+ out = self.conv1(x)
+ out = self.relu(out)
+ out = self.conv2(out)
+ out += residual
+ return out
+
+
+class Flatten(nn.Module):
+ def forward(self, x):
+ return x.view(x.size(0), -1)
+
+
+# Channel Gate
+class ChannelGate(nn.Module):
+ def __init__(self, channels):
+ super(ChannelGate, self).__init__()
+ self.mlp = nn.Sequential(
+ Flatten(),
+ nn.Linear(1, channels),
+ nn.PReLU(),
+ nn.Linear(channels, channels)
+ )
+
+ def forward(self, x):
+ out = self.mlp(x)
+ return out
+
+
+class SpatialGate(nn.Module):
+ def __init__(self, in_channels, num_features):
+ super(SpatialGate, self).__init__()
+ self.conv1 = conv3x3(in_channels, num_features, stride=2)
+ self.relu = nn.PReLU()
+ self.conv2 = conv3x3(num_features, 1)
+
+ def forward(self, x):
+ out = self.conv1(x)
+ out = self.relu(out)
+ out = self.conv2(out)
+ return out
+
+
+class ConditionalNet(nn.Module):
+ def __init__(self, f, rbn):
+ super(ConditionalNet, self).__init__()
+ self.rbn = rbn
+ self.convLuma = conv3x3_f(1, f, 2)
+ self.convRec = conv3x3_f(2, f)
+ self.convPred = conv3x3_f(2, f)
+ self.convSplit = conv3x3_f(2, f)
+ self.convBs = conv3x3_f(2, f)
+ self.convQp = conv3x3_f(1, f)
+ self.fuse = conv1x1_f(6 * f, f)
+ self.transitionH = conv3x3_f(f, f, 2)
+ self.backbone = nn.ModuleList([ResidualBlock(f, f)])
+ for _ in range(self.rbn - 1):
+ self.backbone.append(ResidualBlock(f, f))
+ self.last_layer = nn.Sequential(
+ nn.Conv2d(
+ in_channels=f,
+ out_channels=f,
+ kernel_size=3,
+ stride=1,
+ padding=1),
+ nn.PReLU(),
+ nn.Conv2d(
+ in_channels=f,
+ out_channels=8,
+ kernel_size=3,
+ stride=1,
+ padding=1),
+ )
+
+ def forward(self, luma, rec, pred, split, bs, qp):
+ luma_f = self.convLuma(luma)
+ rec_f = self.convRec(rec)
+ pred_f = self.convPred(pred)
+ split_f = self.convSplit(split)
+ bs_f = self.convBs(bs)
+ qp_f = self.convQp(qp)
+ xh = torch.cat((luma_f, rec_f, pred_f, split_f, bs_f, qp_f), 1)
+ xh = self.fuse(xh)
+ x = self.transitionH(xh)
+ for i in range(self.rbn):
+ x = self.backbone[i](x)
+ # output
+ x = self.last_layer(x)
+ return x
diff --git a/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaInter/convert.py b/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaInter/convert.py
new file mode 100755
index 0000000000000000000000000000000000000000..b6525b350929bfb6b6cd1be34f868ca8fb5b3917
--- /dev/null
+++ b/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaInter/convert.py
@@ -0,0 +1,26 @@
+import torch
+import torch.nn as nn
+from net import ConditionalNet
+import numpy as np
+import os
+
+# input
+yuv = np.ones((1, 1, 32, 32), dtype=np.float32)
+pred = np.ones((1, 1, 32, 32), dtype=np.float32)
+bs = np.ones((1, 1, 32, 32), dtype=np.float32)
+qp = np.ones((1, 1, 32, 32), dtype=np.float32)
+
+# model
+# model = nn.DataParallel(ConditionalNet(96, 8)) # if model is trained on multiple GPUs
+model = ConditionalNet(96, 8) # if model is trained with single GPU
+state = torch.load('50.ckpt', map_location=torch.device('cpu'))
+model.load_state_dict(state)
+
+dummy_input = (torch.from_numpy(yuv), torch.from_numpy(pred), torch.from_numpy(bs), torch.from_numpy(qp))
+torch.onnx.export(model.module, dummy_input, "NnlfSet1_LumaCNNFilter_InterSlice.onnx")
+
+
+
+
+
+
diff --git a/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaInter/net.py b/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaInter/net.py
new file mode 100644
index 0000000000000000000000000000000000000000..1785466a8d64f65a97edbfdb4b77aee7c0de5e02
--- /dev/null
+++ b/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaInter/net.py
@@ -0,0 +1,137 @@
+import torch
+import torch.nn as nn
+
+
+def conv3x3(in_channels, out_channels, stride=1, padding=1):
+ return nn.Conv2d(in_channels, out_channels, kernel_size=3,
+ stride=stride, padding=padding)
+
+
+def conv1x1(in_channels, out_channels, stride=1, padding=0):
+ return nn.Conv2d(in_channels, out_channels, kernel_size=1,
+ stride=stride, padding=padding)
+
+
+# Conv3x3 + PReLU
+class conv3x3_f(nn.Module):
+ def __init__(self, in_channels, out_channels, stride=1):
+ super(conv3x3_f, self).__init__()
+ self.conv = conv3x3(in_channels, out_channels, stride)
+ self.relu = nn.PReLU()
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.relu(x)
+ return x
+
+
+# Conv1x1 + PReLU
+class conv1x1_f(nn.Module):
+ def __init__(self, in_channels, out_channels, stride=1):
+ super(conv1x1_f, self).__init__()
+ self.conv = conv1x1(in_channels, out_channels, stride)
+ self.relu = nn.PReLU()
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.relu(x)
+ return x
+
+
+# Residual Block
+class ResidualBlock(nn.Module):
+ def __init__(self, in_channels, out_channels):
+ super(ResidualBlock, self).__init__()
+ self.conv1 = conv3x3(in_channels, out_channels)
+ self.relu = nn.PReLU()
+ self.conv2 = conv3x3(out_channels, out_channels)
+
+ def forward(self, x):
+ out = self.conv1(x)
+ out = self.relu(out)
+ out = self.conv2(out)
+ return out
+
+
+class Flatten(nn.Module):
+ def forward(self, x):
+ return x.view(x.size(0), -1)
+
+
+# Channel Gate
+class ChannelGate(nn.Module):
+ def __init__(self, channels):
+ super(ChannelGate, self).__init__()
+ self.mlp = nn.Sequential(
+ Flatten(),
+ nn.Linear(1, channels),
+ nn.PReLU(),
+ nn.Linear(channels, channels)
+ )
+
+ def forward(self, x):
+ out = self.mlp(x)
+ return out
+
+
+class SpatialGate(nn.Module):
+ def __init__(self, in_channels, num_features):
+ super(SpatialGate, self).__init__()
+ self.conv1 = conv3x3(in_channels, num_features, stride=2)
+ self.relu = nn.PReLU()
+ self.conv2 = conv3x3(num_features, 1)
+
+ def forward(self, x):
+ out = self.conv1(x)
+ out = self.relu(out)
+ out = self.conv2(out)
+ return out
+
+
+class ConditionalNet(nn.Module):
+ def __init__(self, f, rbn):
+ super(ConditionalNet, self).__init__()
+ self.rbn = rbn
+ self.convRec = conv3x3_f(1, f)
+ self.convPred = conv3x3_f(1, f)
+ self.convBs = conv3x3_f(1, f)
+ self.convQp = conv3x3_f(1, f)
+ self.fuse = conv1x1_f(4 * f, f)
+ self.transitionH = conv3x3_f(f, f, 2)
+ self.backbone = nn.ModuleList([ResidualBlock(f, f)])
+ for _ in range(self.rbn - 1):
+ self.backbone.append(ResidualBlock(f, f))
+ self.mask = nn.ModuleList([SpatialGate(4, 32)])
+ for _ in range(self.rbn - 1):
+ self.mask.append(SpatialGate(4, 32))
+ self.last_layer = nn.Sequential(
+ nn.Conv2d(
+ in_channels=f,
+ out_channels=f,
+ kernel_size=3,
+ stride=1,
+ padding=1),
+ nn.PReLU(),
+ nn.Conv2d(
+ in_channels=f,
+ out_channels=4,
+ kernel_size=3,
+ stride=1,
+ padding=1),
+ )
+
+ def forward(self, rec, pred, bs, qp):
+ rec_f = self.convRec(rec)
+ pred_f = self.convPred(pred)
+ bs_f = self.convBs(bs)
+ qp_f = self.convQp(qp)
+ xh = torch.cat((rec_f, pred_f, bs_f, qp_f), 1)
+ xh = self.fuse(xh)
+ x = self.transitionH(xh)
+ for i in range(self.rbn):
+ x_resi = self.backbone[i](x)
+ attention = self.mask[i](torch.cat((rec, pred, bs, qp), 1))
+ x = attention.expand_as(x_resi) * x_resi + x_resi + x
+ # output
+ x = self.last_layer(x)
+ return x
diff --git a/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaIntra/convert.py b/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaIntra/convert.py
new file mode 100755
index 0000000000000000000000000000000000000000..21e3743bdd873fba980d5b2221003a5d412213b3
--- /dev/null
+++ b/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaIntra/convert.py
@@ -0,0 +1,26 @@
+import torch
+import torch.nn as nn
+from net import ConditionalNet
+import numpy as np
+import os
+
+# input
+yuv = np.ones((1, 1, 32, 32), dtype=np.float32)
+pred = np.ones((1, 1, 32, 32), dtype=np.float32)
+split = np.ones((1, 1, 32, 32), dtype=np.float32)
+bs = np.ones((1, 1, 32, 32), dtype=np.float32)
+qp = np.ones((1, 1, 32, 32), dtype=np.float32)
+
+# model = nn.DataParallel(ConditionalNet(96, 8)) # if model is trained on multiple GPUs
+model = ConditionalNet(96, 8) # if model is trained with single GPU
+state = torch.load('50.ckpt', map_location=torch.device('cpu'))
+model.load_state_dict(state)
+
+dummy_input = (torch.from_numpy(yuv), torch.from_numpy(pred), torch.from_numpy(split), torch.from_numpy(bs), torch.from_numpy(qp))
+torch.onnx.export(model.module, dummy_input, "NnlfSet1_LumaCNNFilter_IntraSlice.onnx")
+
+
+
+
+
+
diff --git a/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaIntra/net.py b/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaIntra/net.py
new file mode 100644
index 0000000000000000000000000000000000000000..de510fd7f659fbe7139e8dc4625648d776a44b95
--- /dev/null
+++ b/training/training_scripts/Nn_Filtering_Set_1/Conversion/LumaIntra/net.py
@@ -0,0 +1,139 @@
+import torch
+import torch.nn as nn
+
+
+def conv3x3(in_channels, out_channels, stride=1, padding=1):
+ return nn.Conv2d(in_channels, out_channels, kernel_size=3,
+ stride=stride, padding=padding)
+
+
+def conv1x1(in_channels, out_channels, stride=1, padding=0):
+ return nn.Conv2d(in_channels, out_channels, kernel_size=1,
+ stride=stride, padding=padding)
+
+
+# Conv3x3 + PReLU
+class conv3x3_f(nn.Module):
+ def __init__(self, in_channels, out_channels, stride=1):
+ super(conv3x3_f, self).__init__()
+ self.conv = conv3x3(in_channels, out_channels, stride)
+ self.relu = nn.PReLU()
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.relu(x)
+ return x
+
+
+# Conv1x1 + PReLU
+class conv1x1_f(nn.Module):
+ def __init__(self, in_channels, out_channels, stride=1):
+ super(conv1x1_f, self).__init__()
+ self.conv = conv1x1(in_channels, out_channels, stride)
+ self.relu = nn.PReLU()
+
+ def forward(self, x):
+ x = self.conv(x)
+ x = self.relu(x)
+ return x
+
+
+# Residual Block
+class ResidualBlock(nn.Module):
+ def __init__(self, in_channels, out_channels):
+ super(ResidualBlock, self).__init__()
+ self.conv1 = conv3x3(in_channels, out_channels)
+ self.relu = nn.PReLU()
+ self.conv2 = conv3x3(out_channels, out_channels)
+
+ def forward(self, x):
+ out = self.conv1(x)
+ out = self.relu(out)
+ out = self.conv2(out)
+ return out
+
+
+class Flatten(nn.Module):
+ def forward(self, x):
+ return x.view(x.size(0), -1)
+
+
+# Channel Gate
+class ChannelGate(nn.Module):
+ def __init__(self, channels):
+ super(ChannelGate, self).__init__()
+ self.mlp = nn.Sequential(
+ Flatten(),
+ nn.Linear(1, channels),
+ nn.PReLU(),
+ nn.Linear(channels, channels)
+ )
+
+ def forward(self, x):
+ out = self.mlp(x)
+ return out
+
+
+class SpatialGate(nn.Module):
+ def __init__(self, in_channels, num_features):
+ super(SpatialGate, self).__init__()
+ self.conv1 = conv3x3(in_channels, num_features, stride=2)
+ self.relu = nn.PReLU()
+ self.conv2 = conv3x3(num_features, 1)
+
+ def forward(self, x):
+ out = self.conv1(x)
+ out = self.relu(out)
+ out = self.conv2(out)
+ return out
+
+
+class ConditionalNet(nn.Module):
+ def __init__(self, f, rbn):
+ super(ConditionalNet, self).__init__()
+ self.rbn = rbn
+ self.convRec = conv3x3_f(1, f)
+ self.convPred = conv3x3_f(1, f)
+ self.convSplit = conv3x3_f(1, f)
+ self.convBs = conv3x3_f(1, f)
+ self.convQp = conv3x3_f(1, f)
+ self.fuse = conv1x1_f(5 * f, f)
+ self.transitionH = conv3x3_f(f, f, 2)
+ self.backbone = nn.ModuleList([ResidualBlock(f, f)])
+ for _ in range(self.rbn - 1):
+ self.backbone.append(ResidualBlock(f, f))
+ self.mask = nn.ModuleList([SpatialGate(5, 32)])
+ for _ in range(self.rbn - 1):
+ self.mask.append(SpatialGate(5, 32))
+ self.last_layer = nn.Sequential(
+ nn.Conv2d(
+ in_channels=f,
+ out_channels=f,
+ kernel_size=3,
+ stride=1,
+ padding=1),
+ nn.PReLU(),
+ nn.Conv2d(
+ in_channels=f,
+ out_channels=4,
+ kernel_size=3,
+ stride=1,
+ padding=1),
+ )
+
+ def forward(self, rec, pred, split, bs, qp):
+ rec_f = self.convRec(rec)
+ pred_f = self.convPred(pred)
+ split_f = self.convSplit(split)
+ bs_f = self.convBs(bs)
+ qp_f = self.convQp(qp)
+ xh = torch.cat((rec_f, pred_f, split_f, bs_f, qp_f), 1)
+ xh = self.fuse(xh)
+ x = self.transitionH(xh)
+ for i in range(self.rbn):
+ x_resi = self.backbone[i](x)
+ attention = self.mask[i](torch.cat((rec, pred, split, bs, qp), 1))
+ x = attention.expand_as(x_resi) * x_resi + x_resi + x
+ # output
+ x = self.last_layer(x)
+ return x
diff --git a/training/training_scripts/Nn_Filtering_Set_1/readme.pdf b/training/training_scripts/Nn_Filtering_Set_1/readme.pdf
index 08fe05a7952bdaf569b871546ab8c7dc6d68f609..d9da02c7a078a1551fb1fd15b6b9742dff64fe31 100644
Binary files a/training/training_scripts/Nn_Filtering_Set_1/readme.pdf and b/training/training_scripts/Nn_Filtering_Set_1/readme.pdf differ