Dumpers/ecg-image-classifier

1import os2import numpy as np3import torch4import torch.nn as nn5import torch.nn.functional as F6from torchvision import models, transforms7from PIL import Image8import gradio as gr9import matplotlib10matplotlib.use("Agg")11import matplotlib.pyplot as plt12import matplotlib.cm as cm13 14CLASS_NAMES  = ["Abnormal Heartbeat", "History of MI", "Normal Person"]15CLASS_COLORS = ["#ef4444", "#f97316", "#22c55e"]16IMG_SIZE     = 22417 18device = torch.device("cuda" if torch.cuda.is_available() else "cpu")19 20 21def load_model():22    model = models.mobilenet_v2(weights=None)23    model.classifier[1] = nn.Linear(model.last_channel, len(CLASS_NAMES))24 25    local_path = "mobilenetv2_best.pth"26    hf_repo    = os.environ.get("HF_MODEL_REPO", "")27 28    if os.path.exists(local_path):29        checkpoint = torch.load(local_path, map_location=device)30        state_dict = checkpoint.get("model_state_dict", checkpoint)31        model.load_state_dict(state_dict)32    elif hf_repo:33        from huggingface_hub import hf_hub_download34        path = hf_hub_download(repo_id=hf_repo, filename="mobilenetv2_best.pth")35        checkpoint = torch.load(path, map_location=device)36        state_dict = checkpoint.get("model_state_dict", checkpoint)37        model.load_state_dict(state_dict)38    else:39        print("No weights found. Run train_and_save.py first.")40 41    model = model.to(device)42    model.eval()43    return model44 45 46model = load_model()47 48preprocess = transforms.Compose([49    transforms.Resize((IMG_SIZE, IMG_SIZE)),50    transforms.ToTensor(),51    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),52])53 54inv_normalize = transforms.Normalize(55    mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225],56    std=[1/0.229, 1/0.224, 1/0.225],57)58 59 60class GradCAM:61    def __init__(self, model, target_layer):62        self.model       = model63        self.activations = None64        self.gradients   = None65        self._hooks = [66            target_layer.register_forward_hook(self._save_activation),67            target_layer.register_full_backward_hook(self._save_gradient),68        ]69 70    def _save_activation(self, _, __, output):71        self.activations = output.detach()72 73    def _save_gradient(self, _, __, grad_output):74        self.gradients = grad_output[0].detach()75 76    def __call__(self, input_tensor, class_idx):77        self.model.zero_grad()78        output = self.model(input_tensor)79        output[0, class_idx].backward()80 81        weights = self.gradients.mean(dim=(2, 3), keepdim=True)82        cam     = (weights * self.activations).sum(dim=1, keepdim=True)83        cam     = F.relu(cam)84        cam     = F.interpolate(cam, size=(IMG_SIZE, IMG_SIZE), mode="bilinear", align_corners=False)85        cam     = cam.squeeze().cpu().numpy()86        cam     = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)87        return cam88 89 90gradcam = GradCAM(model, model.features[-1])91 92 93def make_bar_chart(probs):94    fig, ax = plt.subplots(figsize=(5, 2.8))95    fig.patch.set_facecolor("#0f172a")96    ax.set_facecolor("#1e293b")97 98    bars = ax.barh(CLASS_NAMES, probs, color=CLASS_COLORS, edgecolor="none", height=0.5)99    ax.set_xlim(0, 1)100    ax.set_xlabel("Confidence", color="white", fontsize=10)101    ax.tick_params(colors="white", labelsize=9)102    ax.spines[:].set_visible(False)103 104    for bar, p in zip(bars, probs):105        ax.text(106            min(p + 0.02, 0.97), bar.get_y() + bar.get_height() / 2,107            f"{p * 100:.1f}%", va="center", color="white", fontsize=9, fontweight="bold",108        )109 110    plt.tight_layout(pad=0.5)111    return fig112 113 114def make_gradcam_overlay(pil_img, cam_map):115    img_np  = np.array(pil_img.resize((IMG_SIZE, IMG_SIZE))).astype(np.float32) / 255.0116    heatmap = cm.jet(cam_map)[:, :, :3]117    overlay = np.clip(0.55 * img_np + 0.45 * heatmap, 0, 1)118    return Image.fromarray((overlay * 255).astype(np.uint8))119 120 121def predict(pil_img):122    if pil_img is None:123        return None, None, "Upload an ECG image."124 125    img_rgb = pil_img.convert("RGB")126    tensor  = preprocess(img_rgb).unsqueeze(0).to(device)127 128    with torch.no_grad():129        logits = model(tensor)130        probs  = F.softmax(logits, dim=1).squeeze().cpu().numpy()131 132    pred_idx = int(probs.argmax())133 134    tensor2 = preprocess(img_rgb).unsqueeze(0).to(device)135    cam_map  = gradcam(tensor2, class_idx=pred_idx)136 137    label = CLASS_NAMES[pred_idx]138    conf  = probs[pred_idx] * 100139 140    status_html = f"""141    <div style="background:#1e293b;border-radius:12px;padding:16px;text-align:center;142                font-family:'Inter',sans-serif;">143      <div style="font-size:1.5rem;font-weight:700;color:{CLASS_COLORS[pred_idx]};">{label}</div>144      <div style="color:#94a3b8;margin-top:4px;font-size:0.9rem;">145        Confidence: <b style="color:white;">{conf:.1f}%</b>146      </div>147      <div style="color:#475569;font-size:0.75rem;margin-top:8px;">148        MobileNetV2 · 5-Fold CV · 99.07% test accuracy149      </div>150    </div>151    """152 153    return make_gradcam_overlay(img_rgb, cam_map), make_bar_chart(probs), status_html154 155 156example_paths = []157for cls in ["Abnormal heartbeat", "History of MI", "Normal Person"]:158    cls_dir = os.path.join("ECG Dataset 2", cls)159    if os.path.isdir(cls_dir):160        imgs = [f for f in os.listdir(cls_dir) if f.lower().endswith((".png", ".jpg", ".jpeg"))]161        if imgs:162            example_paths.append([os.path.join(cls_dir, imgs[0])])163 164 165css = """166body, .gradio-container { background: #0f172a !important; font-family: 'Inter', sans-serif; }167.gr-button { background: linear-gradient(135deg,#6366f1,#8b5cf6) !important;168             border: none !important; color: white !important; border-radius: 8px !important; }169.gr-button:hover { opacity: 0.85 !important; }170h1 { color: #f1f5f9 !important; }171"""172 173with gr.Blocks(css=css, title="ECG Arrhythmia Classifier") as demo:174 175    gr.Markdown("""176    # ECG Arrhythmia Classifier177 178    Upload an ECG image to get a classification and a Grad-CAM heatmap showing which regions of the image the model focused on.179 180    Model: MobileNetV2 fine-tuned with 5-fold stratified cross-validation. Test accuracy: 99.07%.181    """)182 183    with gr.Row():184        with gr.Column(scale=1):185            inp = gr.Image(type="pil", label="ECG Image", height=280)186            btn = gr.Button("Analyse", variant="primary")187            gr.Examples(examples=example_paths, inputs=inp, label="Sample images")188 189        with gr.Column(scale=1):190            status_out  = gr.HTML(label="Prediction")191            gradcam_out = gr.Image(type="pil", label="Grad-CAM Heatmap", height=280)192            chart_out   = gr.Plot(label="Class Probabilities")193 194    btn.click(fn=predict, inputs=inp, outputs=[gradcam_out, chart_out, status_out])195    inp.change(fn=predict, inputs=inp, outputs=[gradcam_out, chart_out, status_out])196 197    gr.Markdown("Classes: Abnormal Heartbeat · History of Myocardial Infarction · Normal Person")198 199if __name__ == "__main__":200    demo.launch()201