File size: 2,423 Bytes
03aa71c
b282b59
 
 
 
 
c1de4f6
 
0856747
 
b282b59
77ac1d8
 
 
 
 
 
 
 
 
 
 
6105568
b282b59
 
a042204
c1de4f6
b282b59
 
c1de4f6
 
 
 
 
d7002fb
c1de4f6
 
 
 
 
b282b59
 
 
 
 
77ac1d8
 
 
b282b59
77ac1d8
b282b59
 
 
0856747
 
 
 
 
6105568
0856747
 
 
 
6105568
 
 
 
 
 
 
 
b282b59
0856747
 
 
 
b282b59
0856747
b282b59
0856747
 
b282b59
0856747
 
b282b59
0856747
 
 
77ac1d8
 
b282b59
03aa71c
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import streamlit as st
import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import numpy as np
import timm
from pathlib import Path
import traceback


# ----------------------
# Page config
# ----------------------
st.set_page_config(
    page_title="Pneumonia X-ray Classifier",
    layout="centered"
)

st.title("🫁 Pneumonia Detection from Chest X-ray")
st.write("Upload a chest X-ray image to classify it as **Normal** or **Pneumonia**.")

DEVICE = "cpu"
CLASS_NAMES = ["Normal", "Pneumonia"]

MODEL_PATH = Path(__file__).resolve().parent / "efficientnet_b3_best.pt"

@st.cache_resource
def load_model():
    # Recreate model architecture EXACTLY as training
    model = timm.create_model(
        "efficientnet_b3",
        pretrained=False,
        num_classes=2   # Normal vs Pneumonia
    )

    state_dict = torch.load(MODEL_PATH, map_location=DEVICE)
    model.load_state_dict(state_dict)

    model.to(DEVICE)
    model.eval()
    return model

model = load_model()

# ----------------------
# File uploader
# ----------------------
uploaded_file = st.file_uploader(
    "Upload Chest X-ray Image",
    type=["jpg", "jpeg", "png"]
)

if uploaded_file is not None:
    try:
        # --- Image loading (HF-safe) ---
        image = Image.open(uploaded_file)
        image = image.convert("RGB")
        image = image.resize((300, 300))  # resize

        st.image(image, caption="Uploaded X-ray", use_column_width=True)

        # --- Preprocess ---
        input_tensor = transforms.ToTensor()(image)
        input_tensor = transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )(input_tensor)
        
        input_tensor = input_tensor.unsqueeze(0)


        # --- Inference ---
        with torch.no_grad():
            logits = model(input_tensor)
            probs = F.softmax(logits, dim=1).cpu().numpy()[0]

        st.subheader("🔍 Prediction Results")

        for i, class_name in enumerate(CLASS_NAMES):
            st.write(f"**{class_name}**: {probs[i]*100:.2f}%")

        pred_idx = np.argmax(probs)
        st.success(f"🩺 Diagnosis: **{CLASS_NAMES[pred_idx]}**")

    except Exception as e:
        st.error("❌ Error while processing the image")
        st.code(traceback.format_exc())
    pred_idx = np.argmax(probs)
    st.success(f"🩺 Diagnosis: **{CLASS_NAMES[pred_idx]}**")