Update app.py

app.py

@@ -1,91 +1,83 @@
-import streamlit as st
-from datasets import load_dataset
-from torchvision import models, transforms
+import os
 import torch
-import torch.nn as nn
+import gradio as gr
+from torchvision import models, transforms
+from datasets import load_dataset
 from torch.utils.data import DataLoader
-from torchvision.datasets.folder import default_loader
+from torch import nn, optim
 from PIL import Image
-import os
-
-st.set_page_config(page_title="Chest X-ray Classifier", layout="centered")
-MODEL_PATH = "resnet_chest_xray.pt"
 
 # Define transformation
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485], [0.229])
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-# Define model
-def get_model():
-    model = models.resnet34(pretrained=True)
-    model.fc = nn.Linear(model.fc.in_features, 15)
-    return model
+# Load and train model if not already saved
+model_path = "model.pth"
 
-# Train the model once if not already saved
-@st.cache_resource
-def train_model_once():
-    if os.path.exists(MODEL_PATH):
-        model = get_model()
-        model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
-        model.eval()
-        return model
+if not os.path.exists(model_path):
+    dataset = load_dataset("hongrui/mimic_chest_xray_v_1", split="train")
 
-    st.write("🔄 Training model (first time only)...")
-    dataset = load_dataset("BahaaEldin0/NIH-Chest-Xray-14", split="train[:1%]")  # small subset
+    class XRayDataset(torch.utils.data.Dataset):
+        def __init__(self, dataset):
+            self.dataset = dataset
 
-    # Preprocess dataset
-    images, labels = [], []
-    label_names = dataset.features["labels"].feature.names
-    for sample in dataset:
-        img = default_loader(sample["image"].filename)
-        images.append(transform(img))
-        label_vector = torch.zeros(15)
-        for l in sample["labels"]:
-            label_vector[l] = 1
-        labels.append(label_vector)
+        def __len__(self):
+            return len(self.dataset)
 
-    dataloader = DataLoader(list(zip(images, labels)), batch_size=8, shuffle=True)
+        def __getitem__(self, idx):
+            img = Image.open(self.dataset[idx]['image_path']).convert("RGB")
+            label = int(self.dataset[idx]['label'])
+            return transform(img), label
 
-    model = get_model()
-    criterion = nn.BCEWithLogitsLoss()
-    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
+    train_dataset = XRayDataset(dataset)
+    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+
+    model = models.resnet18(pretrained=True)
+    model.fc = nn.Linear(model.fc.in_features, 2)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
 
     model.train()
-    for epoch in range(2):  # few epochs
-        for x, y in dataloader:
+    for epoch in range(3):  # Few epochs for speed
+        for inputs, labels in train_loader:
+            inputs, labels = inputs.to(device), labels.to(device)
             optimizer.zero_grad()
-            output = model(x)
-            loss = criterion(output, y)
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
             loss.backward()
             optimizer.step()
 
-    torch.save(model.state_dict(), MODEL_PATH)
-    model.eval()
-    return model
-
-# App interface
-st.title("🩻 Chest X-ray Classifier")
+    torch.save(model.state_dict(), model_path)
 
-model = train_model_once()
-
-uploaded_file = st.file_uploader("Upload a Chest X-ray Image", type=["jpg", "jpeg", "png"])
-
-if uploaded_file:
-    image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-    input_tensor = transform(image).unsqueeze(0)
+# Load model for inference
+model = models.resnet18(pretrained=False)
+model.fc = nn.Linear(model.fc.in_features, 2)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.load_state_dict(torch.load(model_path, map_location=device))
+model = model.to(device)
+model.eval()
 
+# Inference function
+def predict(img):
+    img = transform(img).unsqueeze(0).to(device)
     with torch.no_grad():
-        output = model(input_tensor)
-        pred = torch.sigmoid(output).squeeze()
-        label_names = ["Atelectasis", "Cardiomegaly", "Effusion", "Infiltration", "Mass",
-                       "Nodule", "Pneumonia", "Pneumothorax", "Consolidation", "Edema",
-                       "Emphysema", "Fibrosis", "Pleural_Thickening", "Hernia", "No Finding"]
-
-        st.subheader("Prediction Results:")
-        for idx, val in enumerate(pred):
-            if val.item() > 0.5:
-                st.markdown(f"- **{label_names[idx]}**: {val.item():.2f}")
+        outputs = model(img)
+        _, predicted = torch.max(outputs, 1)
+    return "Abnormal" if predicted.item() == 1 else "Normal"
+
+# Gradio app
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="Chest X-ray Classification",
+    description="Upload a chest X-ray image to classify it as Normal or Abnormal."
+)
+
+interface.launch()