Upload 3 files

network/discriminator.py

@@ -10,12 +10,12 @@ def conv_block(ndf, in_channels, out_channels, kernel_size, stride, padding):
     )
 
 class PatchDiscriminator(nn.Module):
-    def __init__(self, input_nc=1, ndf=16):
+    def __init__(self, input_nc=1, ndf=3, output_size=1):
         """Initializes the Patch Discriminator model.
         
         Args:
             input_nc (int): Number of input channels. Default is 1 (e.g., for grayscale images).
-            ndf (int): Number of filters in the first convolution layer. Default is 16.
+            ndf (int): Number of filters in the first convolution layer. Default is 64.
         """
         super(PatchDiscriminator, self).__init__()
         
@@ -25,27 +25,25 @@ class PatchDiscriminator(nn.Module):
         self.conv3 = conv_block(ndf * 2, ndf * 2, ndf * 4, kernel_size=4, stride=2, padding=1)
         self.conv4 = conv_block(ndf * 4, ndf * 4, ndf * 8, kernel_size=4, stride=2, padding=1)
 
-        # Final convolution layer to reduce to a single channel output
-        self.conv5 = nn.Conv3d(ndf * 8, 1, kernel_size=4, padding=1)
+        # Final convolution layer
+        self.conv5 = nn.Conv3d(ndf * 8, ndf * 8, kernel_size=4, padding=1)
 
         # Flatten layer
         self.flatten = nn.Flatten()
 
-        # Fully connected layer to adjust output size
-        self.fc = nn.Linear(539, 1)  # Adjust '539' based on the flattened output size
+        # Fully connected layer
+        self.fc = nn.Linear(ndf * 8 * 7 * 11 * 7, output_size)
 
-        # Sigmoid activation to obtain a probability
+        # Sigmoid activation
         self.sigmoid = nn.Sigmoid()
 
     def forward(self, x):
-        """Defines the forward pass of the discriminator."""
         x = self.conv1(x)
         x = self.conv2(x)
         x = self.conv3(x)
         x = self.conv4(x)
-        x = self.conv5(x)
+        x = self.conv5(x)        
         x = self.flatten(x)
         x = self.fc(x)
         x = self.sigmoid(x)
         return x
-

network/generator.py

@@ -27,9 +27,8 @@ class ResnetBlock(nn.Module):
 class DeUpBlock(nn.Module):
     def __init__(self, inf, onf):
         super(DeUpBlock, self).__init__()
-        # Upsampling only in the width dimension
         self.deupblock = nn.Sequential(
-            nn.ConvTranspose3d(inf, onf, kernel_size=(1, 6, 1), stride=(1, 6, 1), padding=(0, 0, 0)),
+            nn.ConvTranspose3d(inf, onf, kernel_size=(1, 3, 1), stride=(1, 3, 1), padding=(0, 0, 0)),
             nn.LeakyReLU(0.2)
         )
 
@@ -38,7 +37,7 @@ class DeUpBlock(nn.Module):
 
 # Resnet Generator
 class ResnetGenerator(nn.Module):
-    def __init__(self, input_nc=1, output_nc=1, ngf=16, n_residual_blocks=4):
+    def __init__(self, input_nc=1, output_nc=1, ngf=32, n_residual_blocks=2):
         super(ResnetGenerator, self).__init__()
         self.n_residual_blocks = n_residual_blocks
 
@@ -66,4 +65,3 @@ class ResnetGenerator(nn.Module):
         x = self.conv_block2(y) + x
         x = self.deup(x)
         return self.conv3(x)
-