Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import gradio as gr | |
| import numpy as np | |
| class Simple1DCNN(nn.Module): | |
| def __init__(self, input_channels, num_classes, sequence_length): | |
| super(Simple1DCNN, self).__init__() | |
| self.conv1 = nn.Conv1d(in_channels=input_channels, out_channels=64, kernel_size=3, padding=1) | |
| self.relu = nn.ReLU() | |
| 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 = 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 | |
| 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) | |
| model.load_state_dict(torch.load(model_path, map_location="cpu")) | |
| model.eval() | |
| # Preprocessing function | |
| def preprocess_ecg(data): | |
| """Convert input ECG data to PyTorch tensor and prepare for inference.""" | |
| ecg = np.array(data.split(','), dtype=np.float32) | |
| ecg = torch.from_numpy(ecg).unsqueeze(0).unsqueeze(0) # 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() | |