Upload app.py

app.py

@@ -1,63 +1,55 @@
-import torch
-import torch.nn as nn
-from torchvision import models, transforms
-from PIL import Image
-import gradio as gr
-
-# Load model
-def load_model(path="LR_model.pth"):
-    model = models.resnet50(weights=None)
-    model.fc = nn.Sequential(
-        nn.Linear(model.fc.in_features, 256),
-        nn.ReLU(),
-        nn.Dropout(0.4),
-        nn.Linear(256, 2)
-    )
-    checkpoint = torch.load(path, map_location="cpu")
-    model.load_state_dict(checkpoint["model_state_dict"])
-    model.eval()
-    return model
-
-model = load_model()
-
-# Preprocessing
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.4326, 0.4953, 0.3120], [0.2178, 0.2214, 0.2091])
-])
-
-# Prediction logic
-def predict(img):
-    img = img.convert("RGB")
-    tensor = transform(img).unsqueeze(0)
-    with torch.no_grad():
-        output = model(tensor)
-        probs = torch.nn.functional.softmax(output, dim=1)
-        idx = probs.argmax().item()
-        conf = probs[0][idx].item()
-    label = "🦠 Parasitized" if idx == 0 else "✅ Uninfected"
-    return {label: conf}
-
-# Custom theme and layout using Gradio Blocks (v3)
-with gr.Blocks(theme=gr.themes.Soft()) as interface:
-    gr.Markdown("# 🧬 Malaria Cell Detection App")
-    gr.Markdown(
-        "Upload a blood smear cell image to predict whether it contains malaria parasites or not. "
-        "The model is based on ResNet50 fine-tuned for binary classification (Parasitized vs Uninfected)."
-    )
-
-    with gr.Row():
-        image_input = gr.Image(type="pil", label="Upload Blood Smear Image", tool="editor")
-        output_label = gr.Label(label="Prediction Confidence")
-
-    analyze_btn = gr.Button("🔍 Analyze")
-    analyze_btn.click(fn=predict, inputs=image_input, outputs=output_label)
-
-    gr.Markdown("---")
-    gr.Markdown(
-        "📌 **Model Info**: Fine-tuned ResNet50 | Input size: 224x224 | Confidence shown as probability.\n"
-        "🧪 This is a demo model. For real medical use, consult healthcare professionals."
-    )
-
-interface.launch()
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
+import gradio as gr
+
+def load_model(path="LR_model.pth"):
+    model = models.resnet50(weights=None)
+
+    # Your saved model has a Sequential head, not just one linear layer
+    model.fc = nn.Sequential(
+        nn.Linear(model.fc.in_features, 256),
+        nn.ReLU(),
+        nn.Dropout(0.4),
+        nn.Linear(256, 2)
+    )
+
+    checkpoint = torch.load(path, map_location="cpu")
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+    return model
+
+
+
+# Image preprocessing
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.4326, 0.4953, 0.3120], [0.2178, 0.2214, 0.2091])
+])
+
+# Predict function
+def predict(img):
+    img = img.convert("RGB")
+    tensor = transform(img).unsqueeze(0)
+    with torch.no_grad():
+        output = model(tensor)
+        probs = torch.nn.functional.softmax(output, dim=1)
+        idx = probs.argmax().item()
+        conf = probs[0][idx].item()
+    return {"Parasitized" if idx == 0 else "Uninfected": conf}
+
+# Load model once
+model = load_model()
+
+# Gradio UI
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=2),
+    title=" Malaria Cell Detection",
+    description="Upload a blood smear cell image to check for malaria (parasitized or uninfected)."
+)
+
+interface.launch()