JeeKay/Malaria-classification

1import os2import streamlit as st3import torch4from PIL import Image5import numpy as np6import warnings7import torch.nn.functional as F8# Add root to PYTHONPATH9import sys10from pathlib import Path11 12# Add root directory to Python path13sys.path.append(str(Path(__file__).parent.parent))14 15# Avoid OMP error from PyTorch/OpenCV16os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'17 18# Suppress FutureWarning from Matplotlib19warnings.filterwarnings("ignore", category=UserWarning)20 21# Import custom modules22from models.resnet_model import MalariaResNet50 23from gradcam.gradcam import visualize_gradcam24 25 26# -----------------------------27# Streamlit Page Setup28# -----------------------------29st.set_page_config(page_title="🧬 Malaria Cell Classifier", layout="wide")30st.title("🧬 Malaria Cell Classifier with Grad-CAM")31st.write("Upload a blood smear image and the model will classify it as infected or uninfected, and highlight key regions using Grad-CAM.")32 33 34# -----------------------------35# Load Model36# -----------------------------37@st.cache_resource38def load_model():39    # Ensure model class doesn't wrap backbone40    model = MalariaResNet50(num_classes=2)41    model.load_state_dict(torch.load("models/malaria_model.pth", map_location='cpu'))42    model.eval()43    return model44 45model = load_model()46 47 48# -----------------------------49# Upload Image50# -----------------------------51uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])52 53if uploaded_file is not None:54    # Save uploaded image temporarily55    temp_image_path = f"temp_{uploaded_file.name}"56    with open(temp_image_path, "wb") as f:57        f.write(uploaded_file.getbuffer())58 59    # Display original image (resize if needed)60    image = Image.open(uploaded_file).convert("RGB")61    max_size = (400, 400)  # Max width and height62    image.thumbnail(max_size)63    st.image(image, caption="Uploaded Image", use_container_width=False)64 65    # Predict button66    if st.button("Predict"):67        with st.spinner("Classifying..."):68            # Run prediction69            pred_label, confidence = model.predict(temp_image_path, device='cpu', show_image=False)70            st.success(f"✅ Prediction: **{pred_label}** | Confidence: **{confidence:.2%}**")71 72            # Show Grad-CAM73            st.subheader("🔍 Grad-CAM Visualization")74            with st.expander("ℹ️ What is Grad-CAM?"):75                st.markdown("""76                **Grad-CAM (Gradient-weighted Class Activation Mapping)** is an interpretability method that shows which parts of an image are most important for a CNN's prediction.77            78                How it works:79                1. Gradients flow from the output neuron back to the last convolutional layer.80                2. These gradients are global average pooled to get importance weights.81                3. A weighted combination creates a coarse heatmap.82                4. Final heatmap is overlaid on the original image.83            84                🔬 In this app:85                - Helps understand *why* the model thinks a blood smear cell is infected86                - Makes predictions more transparent and reliable87                """)88            visualize_gradcam(model, temp_image_path)