1

app/model_architectures.py

@@ -6,9 +6,12 @@ class ResNet50(nn.Module):
     def __init__(self, num_classes=7, channels=3):
         super(ResNet50, self).__init__()
         self.resnet = models.resnet50(pretrained=True)
+        
         # Modify the first convolutional layer if channels != 3
         if channels != 3:
             self.resnet.conv1 = nn.Conv2d(channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        
+        # Replace the fully connected layer
         num_features = self.resnet.fc.in_features
         self.resnet.fc = nn.Linear(num_features, num_classes)
 
@@ -16,6 +19,7 @@ class ResNet50(nn.Module):
         return self.resnet(x)
 
     def extract_features(self, x):
+        # Feature extraction using layers up to layer4
         x = self.resnet.conv1(x)
         x = self.resnet.bn1(x)
         x = self.resnet.relu(x)
@@ -39,8 +43,13 @@ class LSTMPyTorch(nn.Module):
         self.fc = nn.Linear(hidden_size, num_classes)
 
     def forward(self, x):
+        # Initialize hidden state and cell state with zeros
         h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
         c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
-        out, _ = self.lstm(x, (h0, c0))
-        out = self.fc(out[:, -1, :])
-        return out
+        
+        # Forward propagate LSTM
+        out, _ = self.lstm(x, (h0, c0))  # out: tensor of shape (batch_size, seq_length, hidden_size)
+        
+        # Pass through the fully connected layer
+        out = self.fc(out[:, -1, :])  # Only the last time step output
+        return out