Upload model.py

model.py

@@ -0,0 +1,79 @@
+# model.py
+# ECG Condition Classifier - Model Loading and Inference
+
+from transformers import pipeline
+
+# ── Model configuration ───────────────────────────────────────────────────────
+# ViT model fine-tuned on ECG image classification
+# Covers 5 cardiac conditions from the PTB-XL dataset
+MODEL_NAME = "Deranged/vit-base-patch16-224-finetuned-ecg"
+
+# Label mapping from model output to readable condition names
+LABEL_MAP = {
+    "NORM": "Normal Sinus Rhythm",
+    "AFIB": "Atrial Fibrillation",
+    "MI":   "Myocardial Infarction",
+    "STTC": "ST/T-Wave Change",
+    "CD":   "Conduction Disturbance",
+}
+
+# Global classifier instance (loaded once at startup)
+_classifier = None
+
+
+def load_model():
+    """
+    Loads the pretrained ECG classification pipeline.
+    Uses CPU (device=-1) for compatibility with Hugging Face free tier.
+    Called once at app startup to avoid reloading on every request.
+    """
+    global _classifier
+    if _classifier is None:
+        print(f"Loading model: {MODEL_NAME}")
+        _classifier = pipeline(
+            "image-classification",
+            model=MODEL_NAME,
+            device=-1  # CPU — required for HF Spaces free tier
+        )
+        print("Model loaded successfully.")
+    return _classifier
+
+
+def run_inference(image, top_k=3):
+    """
+    Runs the image classification pipeline on a preprocessed PIL image.
+
+    Args:
+        image: PIL Image (preprocessed by utils.py)
+        top_k: number of top predictions to return
+
+    Returns:
+        List of dicts: [{"label": str, "score": float}, ...]
+        Sorted by confidence score descending.
+    """
+    classifier = load_model()
+    results = classifier(image, top_k=top_k)
+
+    # Normalize labels to uppercase for consistent mapping
+    for r in results:
+        r["label"] = r["label"].upper().strip()
+
+    return results
+
+
+def get_top_prediction(results):
+    """
+    Extracts the top prediction from inference results.
+
+    Returns:
+        (label: str, confidence: float, full_name: str)
+    """
+    if not results:
+        return "UNKNOWN", 0.0, "Unknown"
+
+    top = results[0]
+    label = top["label"]
+    confidence = round(top["score"] * 100, 1)
+    full_name = LABEL_MAP.get(label, label)
+
+    return label, confidence, full_name