Add models

src/editor/models/simple_cnn.py

@@ -0,0 +1,37 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SimpleCNN(nn.Module):
+    def __init__(self, bin_count):
+        super(SimpleCNN, self).__init__()
+        self.bin_count = bin_count
+
+        # Define the convolutional layers
+        self.conv1 = nn.Conv3d(
+            1, 16, kernel_size=3, padding=1
+        )  # input channels = 1, output channels = 16
+        self.conv2 = nn.Conv3d(
+            16, 32, kernel_size=3, padding=1
+        )  # input channels = 16, output channels = 32
+        self.conv3 = nn.Conv3d(
+            32, 64, kernel_size=3, padding=1
+        )  # input channels = 32, output channels = 64
+        self.conv4 = nn.Conv3d(
+            64, 32, kernel_size=3, padding=1
+        )  # input channels = 64, output channels = 32
+        self.conv5 = nn.Conv3d(
+            32, 16, kernel_size=3, padding=1
+        )  # input channels = 32, output channels = 16
+        self.conv6 = nn.Conv3d(
+            16, 1, kernel_size=3, padding=1
+        )  # input channels = 16, output channels = 1
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = F.relu(self.conv4(x))
+        x = F.relu(self.conv5(x))
+        x = self.conv6(x)
+        return x