Update app.py

app.py

@@ -1,57 +1,57 @@
-import torch
-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):
-        super(Simple1DCNN, self).__init__()
-        self.conv1 = nn.Conv1d(input_channels, 16, kernel_size=5, stride=1, padding=2)
-        self.relu = nn.ReLU()
-        self.pool = nn.MaxPool1d(kernel_size=2, stride=2)
-        self.fc = nn.Linear(16 * (sequence_length // 2), num_classes)
-
-    def forward(self, x):
-        x = self.pool(self.relu(self.conv1(x)))
-        x = x.view(x.size(0), -1)
-        return self.fc(x)
-
-# Load model
-model_path = "ecg.pth"  # Adjust if necessary
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-sequence_length = 500  # Adjust based on training
-model = Simple1DCNN(input_channels=1, num_classes=5, sequence_length=sequence_length).to(device)
-model.load_state_dict(torch.load(model_path, map_location=device))
-model.eval()
-
-# Preprocessing function
-def preprocess_ecg(data):
-    """Convert input ECG data to PyTorch tensor and prepare for inference."""
-    ecg = np.array(data).astype(np.float32)  # Ensure NumPy array format
-    ecg = torch.from_numpy(ecg).unsqueeze(0).unsqueeze(0).to(device)  # Shape: (1, 1, seq_length)
-    return ecg
-
-# Prediction function
-def predict(ecg_data):
-    ecg_tensor = preprocess_ecg(ecg_data)
-    with torch.no_grad():
-        output = model(ecg_tensor)
-    predicted_class = int(output.argmax(dim=1).item())
-
-    # Define class labels
-    class_labels = {0: "Normal", 1: "AFib", 2: "PVC", 3: "ST", 4: "Other"}
-    return class_labels.get(predicted_class, "Unknown")
-
-# Create Gradio interface
-app = gr.Interface(
-    fn=predict,
-    inputs=gr.Textbox(lines=2, placeholder="Enter ECG values separated by commas"),
-    outputs="label",
-    title="ECG Classification",
-    description="Predicts ECG conditions based on input signal.",
-)
-
-# Launch app
-if __name__ == "__main__":
-    app.launch()
+import torch
+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):
+        super(Simple1DCNN, self).__init__()
+        self.conv1 = nn.Conv1d(input_channels, 16, kernel_size=5, stride=1, padding=2)
+        self.relu = nn.ReLU()
+        self.pool = nn.MaxPool1d(kernel_size=2, stride=2)
+        self.fc = nn.Linear(16 * (sequence_length // 2), num_classes)
+
+    def forward(self, x):
+        x = self.pool(self.relu(self.conv1(x)))
+        x = x.view(x.size(0), -1)
+        return self.fc(x)
+
+# Load model
+model_path = "ecg.pth"  # Adjust if necessary
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+sequence_length = 187  # Adjust based on training
+model = Simple1DCNN(input_channels=1, num_classes=5, sequence_length=sequence_length).to(device)
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.eval()
+
+# Preprocessing function
+def preprocess_ecg(data):
+    """Convert input ECG data to PyTorch tensor and prepare for inference."""
+    ecg = np.array(data).astype(np.float32)  # Ensure NumPy array format
+    ecg = torch.from_numpy(ecg).unsqueeze(0).unsqueeze(0).to(device)  # Shape: (1, 1, seq_length)
+    return ecg
+
+# Prediction function
+def predict(ecg_data):
+    ecg_tensor = preprocess_ecg(ecg_data)
+    with torch.no_grad():
+        output = model(ecg_tensor)
+    predicted_class = int(output.argmax(dim=1).item())
+
+    # Define class labels
+    class_labels = {0: "Normal", 1: "AFib", 2: "PVC", 3: "ST", 4: "Other"}
+    return class_labels.get(predicted_class, "Unknown")
+
+# Create Gradio interface
+app = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(lines=2, placeholder="Enter ECG values separated by commas"),
+    outputs="label",
+    title="ECG Classification",
+    description="Predicts ECG conditions based on input signal.",
+)
+
+# Launch app
+if __name__ == "__main__":
+    app.launch()