Update app.py

app.py

@@ -3,19 +3,33 @@ import torch.nn as nn
 import gradio as gr
 import numpy as np
 
-# Define the Simple1DCNN model
 class Simple1DCNN(nn.Module):
-    def __init__(self, input_channels=1, num_classes=5, sequence_length=500):
+    def __init__(self, input_channels, num_classes, sequence_length):
         super(Simple1DCNN, self).__init__()
-        self.conv1 = nn.Conv1d(input_channels, 16, kernel_size=5, stride=1, padding=2)
+        self.conv1 = nn.Conv1d(in_channels=input_channels, out_channels=64, kernel_size=3, padding=1)
         self.relu = nn.ReLU()
-        self.pool = nn.MaxPool1d(kernel_size=2, stride=2)
-        self.fc = nn.Linear(16 * (sequence_length // 2), num_classes)
+        self.conv2 = nn.Conv1d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool1d(kernel_size=2)
+
+        # Compute the output size after convolutions and pooling
+        self._to_linear = self._compute_flattened_size(input_channels, sequence_length)
+
+        self.fc1 = nn.Linear(self._to_linear, 256)
+        self.fc2 = nn.Linear(256, num_classes)
+
+    def _compute_flattened_size(self, input_channels, sequence_length):
+        x = torch.randn(1, input_channels, sequence_length)  # Dummy tensor
+        x = self.pool(self.relu(self.conv1(x)))
+        x = self.pool(self.relu(self.conv2(x)))
+        return x.numel()  # Total number of features after conv and pooling
 
     def forward(self, x):
         x = self.pool(self.relu(self.conv1(x)))
-        x = x.view(x.size(0), -1)
-        return self.fc(x)
+        x = self.pool(self.relu(self.conv2(x)))
+        x = x.view(x.shape[0], -1)  # Flatten
+        x = self.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
 
 # Load model
 model_path = "ecg.pth"  # Adjust if necessary