Upload 5 files

README.md

@@ -1,12 +1,8 @@
----
-title: Ecg Ai Classifier
-emoji: 🐨
-colorFrom: indigo
-colorTo: green
-sdk: gradio
-sdk_version: 5.49.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ECG Classification AI
+
+Upload an ECG image and receive:
+- Predicted class
+- Probability scores
+- Automatically generated PDF report
+
+Model: Custom CNN (3-class)

app.py

@@ -0,0 +1,94 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+from fpdf import FPDF
+from datetime import datetime
+import gradio as gr
+from model import get_model
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MODEL_PATH = "ecg_multiclass.pth"
+IMG_SIZE = (224, 224)
+
+DESCRIPTIONS = {
+    "Normal": "This ECG indicates a normal heart rhythm with no abnormalities detected.",
+    "Myocardial_Infarction": "This ECG suggests a myocardial infarction (heart attack). Immediate medical attention is recommended.",
+    "Abnormal_heartbeat": "This ECG shows an abnormal heartbeat pattern indicating possible arrhythmia or heart irregularities."
+}
+
+# Load model checkpoint
+checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
+classes = checkpoint["classes"]
+num_classes = len(classes)
+
+model = get_model(num_classes=num_classes, weights=None).to(DEVICE)
+model.load_state_dict(checkpoint["model_state"])
+model.eval()
+
+transform = transforms.Compose([
+    transforms.Grayscale(),
+    transforms.Resize(IMG_SIZE),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5], [0.5])
+])
+
+
+def generate_report(image):
+    if image is None:
+        return None, "Please upload an image."
+
+    # Convert
+    img = Image.fromarray(image).convert("L")
+    input_tensor = transform(img).unsqueeze(0).to(DEVICE)
+
+    # Predict
+    with torch.no_grad():
+        output = model(input_tensor)
+        probabilities = torch.softmax(output, dim=1)[0]
+        pred_index = torch.argmax(probabilities).item()
+        pred_class = classes[pred_index]
+
+    # Prepare probability text
+    prob_text = ""
+    for i, cls in enumerate(classes):
+        prob_text += f"{cls}: {probabilities[i] * 100:.2f}%\n"
+
+    # Generate PDF
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    pdf_path = f"ecg_report_{timestamp}.pdf"
+
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", "B", 16)
+    pdf.cell(0, 10, "ECG Prediction Report", ln=True, align="C")
+    pdf.ln(10)
+
+    pdf.set_font("Arial", "", 12)
+    pdf.multi_cell(0, 10, f"Predicted Class: {pred_class}\n")
+    pdf.multi_cell(0, 8, f"Probabilities:\n{prob_text}")
+    pdf.multi_cell(0, 8, f"Description:\n{DESCRIPTIONS.get(pred_class, 'No description.')}")
+    pdf.ln(10)
+
+    # Save uploaded image temporarily for embedding
+    img_path = f"temp_img_{timestamp}.png"
+    img.save(img_path)
+
+    pdf.image(img_path, x=30, w=150)
+    pdf.output(pdf_path)
+
+    return pdf_path, f"Prediction complete: {pred_class}"
+
+
+# Gradio UI
+interface = gr.Interface(
+    fn=generate_report,
+    inputs=gr.Image(type="numpy", label="Upload ECG Image"),
+    outputs=[
+        gr.File(label="Download PDF Report"),
+        gr.Textbox(label="Status")
+    ],
+    title="ECG Classification & PDF Report Generator",
+    description="Upload an ECG image to get an AI-generated PDF report."
+)
+
+interface.launch()

ecg_multiclass.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9bf90aa76d576a335b5c6c3c3f49c374d6d5ba6775c5db80678e63612a868739
+size 44765643

model.py

@@ -0,0 +1,17 @@
+# model.py
+import torch.nn as nn
+from torchvision.models import resnet18
+
+def get_model(num_classes, pretrained=True):
+    """
+    Returns a CNN model adapted for grayscale ECG images
+    """
+    model = resnet18(pretrained=pretrained)
+
+    # Change first layer to accept 1-channel input (grayscale)
+    model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
+
+    # Change the output layer for our number of classes
+    model.fc = nn.Linear(model.fc.in_features, num_classes)
+
+    return model

requirements.txt

@@ -0,0 +1,5 @@
+torch
+torchvision
+pillow
+fpdf
+gradio