import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import gradio as gr

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class_names = [
    "ECG Images of Myocardial Infarction Patients (240x12=2880)",
    "ECG Images of Patient that have History of MI (172x12=2064)",
    "ECG Images of Patient that have abnormal heartbeat (233x12=2796)",
    "Normal Person ECG Images (284x12=3408)"
]

def load_model(path):
    # Load the model with the default classifier (1000 classes)
    model = models.resnext50_32x4d(weights=None)  # No pretrained here
    model.load_state_dict(torch.load(path, map_location=device))  # This loads the ImageNet weights

    # Now replace the classifier for 4-class output
    in_features = model.fc.in_features
    model.fc = nn.Linear(in_features, len(class_names))

    return model.to(device).eval()

model_path = "resnext50_32x4d-1a0047aa.pth"
model = load_model(model_path)

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def predict(image):
    image = image.convert("RGB")
    image = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        logits = model(image)
        probs = torch.softmax(logits, dim=1)[0]
    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}

gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload ECG Image"),
    outputs=gr.Label(num_top_classes=4, label="Prediction Probabilities"),
    title="ECG Image Classification using ResNeXt50",
    description="Classify ECG images into: MI, History of MI, Abnormal Heartbeat, or Normal."
).launch()