Remove editor module

2024-06-22 18:37:24 +01:00 · 2024-06-22 18:37:24 +01:00 · c966866abc
commit c966866abc
parent e5959268c1
37 changed files with 7752 additions and 7345 deletions
--- a/src/models/residual3.py
+++ b/src/models/residual3.py
@ -0,0 +1,139 @@
+import torch
+import torch.nn as nn
+
+
+class Residual3(nn.Module):
+    def __init__(
+        self,
+        elu_alpha: float = 1,
+        dropout_prob: float = 0.05,
+        features: list[int] = [16, 32, 64],
+        kernel_sizes: list[int] = [3, 3, 3],
+        use_instance_norm: bool = True,
+        use_elu: bool = True,
+        leaky_relu_alpha: float = 0.01,
+        **_
+    ):
+        super(Residual3, self).__init__()
+        self._elu_alpha = elu_alpha
+        self._dropout_prob = dropout_prob
+        self._features = features
+        self._kernel_sizes = kernel_sizes
+        self._use_instance_norm = use_instance_norm
+        self._use_elu = use_elu
+        self._leaky_relu_alpha = leaky_relu_alpha
+
+        self.conv1 = self._make_conv_layer(1, features[0], kernel_sizes[0])
+        self.res1 = self._make_resblock(features[0], kernel_sizes[0])
+        self.conv2 = self._make_conv_layer(features[0], features[1], kernel_sizes[1])
+        self.res2 = self._make_resblock(features[1], kernel_sizes[1])
+        self.conv3 = self._make_conv_layer(features[1], features[2], kernel_sizes[2])
+        self.res3 = self._make_resblock(features[2], kernel_sizes[2])
+
+        self.deconv1 = self._make_deconv_layer(
+            features[2], features[1], kernel_sizes[2]
+        )
+        self.deconv2 = self._make_deconv_layer(
+            features[1], features[0], kernel_sizes[1]
+        )
+        self.deconv3 = self._make_deconv_layer(features[0], 1, kernel_sizes[0])
+
+        self._initialize_weights()
+
+    def _make_conv_layer(
+        self, in_channels: int, out_channels: int, kernel_size: int
+    ) -> nn.Sequential:
+        return nn.Sequential(
+            nn.Conv3d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                padding=1,
+                bias=False,
+            ),
+            (
+                nn.ELU(self._elu_alpha)
+                if self._use_elu
+                else nn.LeakyReLU(self._leaky_relu_alpha)
+            ),
+            (nn.InstanceNorm3d if self._use_instance_norm else nn.BatchNorm3d)(
+                out_channels
+            ),
+            nn.Dropout(p=self._dropout_prob),
+        )
+
+    def _make_resblock(self, channels: int, kernel_size: int) -> nn.Sequential:
+        return nn.Sequential(
+            nn.Conv3d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=kernel_size,
+                padding=1,
+                bias=False,
+            ),
+            (
+                nn.ELU(self._elu_alpha)
+                if self._use_elu
+                else nn.LeakyReLU(self._leaky_relu_alpha)(
+                    nn.InstanceNorm3d if self._use_instance_norm else nn.BatchNorm3d
+                )(channels)
+            ),
+            nn.Conv3d(
+                in_channels=channels,
+                out_channels=channels,
+                kernel_size=kernel_size,
+                padding=1,
+                bias=False,
+            ),
+            (
+                nn.ELU(self._elu_alpha)
+                if self._use_elu
+                else nn.LeakyReLU(self._leaky_relu_alpha)
+            ),
+            (nn.InstanceNorm3d if self._use_instance_norm else nn.BatchNorm3d)(
+                channels
+            ),
+        )
+
+    def _make_deconv_layer(
+        self, in_channels: int, out_channels: int, kernel_size: int
+    ) -> nn.Sequential:
+        return nn.Sequential(
+            nn.ConvTranspose3d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                padding=1,
+            ),
+            (
+                nn.ELU(self._elu_alpha)
+                if self._use_elu
+                else nn.LeakyReLU(self._leaky_relu_alpha)
+            ),
+        )
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = out + self.res1(out)
+        out = self.conv2(out)
+        out = out + self.res2(out)
+        out = self.conv3(out)
+        out = out + self.res3(out)
+
+        out = self.deconv1(out)
+        out = self.deconv2(out)
+        out = self.deconv3(out)
+
+        return self._normalize(out)
+
+    def _normalize(self, x):
+        x_sum = torch.sum(x, dim=(2, 3, 4), keepdim=True)
+        return x / (x_sum + 1e-6)
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, (nn.Conv3d, nn.ConvTranspose3d)):
+                # Applying He normal initialization
+                nn.init.kaiming_normal_(m.weight, mode="fan_in", nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)