update model input to image_shape

Author: sot176

utils/models/solveig_model.py

@@ -6,12 +6,12 @@ class SolveigModel(nn.Module):
     """
         A Convolutional Neural Network model for classification.
 
-        Args:
+         Args
         ----
-        in_channels : int
-            Number of input channels (e.g., 3 for RGB images, 1 for grayscale).
+        image_shape : tuple(int, int, int)
+            Shape of the input image (C, H, W).
         num_classes : int
-            The number of output classes (e.g., 2 for binary classification).
+            Number of classes in the dataset.
 
         Attributes:
         -----------
@@ -25,12 +25,14 @@ class SolveigModel(nn.Module):
             Fully connected layer that outputs the final classification scores.
         """
 
-    def __init__(self, in_channels, num_classes):
+    def __init__(self, image_shape, num_classes):
         super().__init__()
 
+        C, *_ = image_shape
+
         # Define the first convolutional block (conv + relu + maxpool)
         self.conv_block1 = nn.Sequential(
-            nn.Conv2d(in_channels=in_channels, out_channels=25, kernel_size=3, padding=1),
+            nn.Conv2d(in_channels=C, out_channels=25, kernel_size=3, padding=1),
             nn.ReLU(),
             nn.MaxPool2d(kernel_size=2, stride=2)
         )
@@ -62,9 +64,11 @@ def forward(self, x):
 
 
 if __name__ == "__main__":
-    model = SolveigModel(3, 3)
 
-    x = torch.randn(1, 3, 16, 16)
+    x = torch.randn(1,3, 16, 16)
+
+    model = SolveigModel(x.shape[1:], 3)
+
     y = model(x)
 
     print(y)