Update app.py

app.py

@@ -1,3 +1,4 @@
+from fastapi import FastAPI, Request
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -6,19 +7,20 @@ from torch.utils.data import Dataset, DataLoader
 from PIL import Image
 import requests
 from io import BytesIO
-from pymongo import MongoClient
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import accuracy_score, f1_score
 import segmentation_models_pytorch as smp
-from torchvision.models import densenet121
+from torchvision.models import densenet121, DenseNet121_Weights
 import gradio as gr
-import json
+import os
 
+# Device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-client = MongoClient("mongodb+srv://xcovidinsight:FYP25S108@cluster0.vqesy.mongodb.net/")
-db = client["covid_system"]
-collection = db["validated_xrays"]
+# FastAPI instance
+app = FastAPI()
+
+# ================== MODEL DEFINITIONS ==================
 
 class UNet(nn.Module):
     def __init__(self):
@@ -36,7 +38,6 @@ class UNet(nn.Module):
 class MyClassifier(nn.Module):
     def __init__(self, num_classes=4):
         super().__init__()
-        from torchvision.models import DenseNet121_Weights
         base_model = densenet121(weights=DenseNet121_Weights.IMAGENET1K_V1)
         in_features = base_model.classifier.in_features
         base_model.classifier = nn.Linear(in_features, num_classes)
@@ -45,63 +46,101 @@ class MyClassifier(nn.Module):
     def forward(self, x):
         return self.model(x)
 
-def fetch_image(url):
-    try:
-        response = requests.get(url)
-        image = Image.open(BytesIO(response.content)).convert("L").resize((256, 256))
-        return image
-    except Exception as e:
-        print(f"Failed to fetch image from {url}: {e}")
-        return None
+# ================== LOAD MODELS ==================
+m1 = UNet()
+m1.model.load_state_dict(torch.load("weights/Segmentation_Model.pth", map_location=device))
+m1.to(device).eval()
 
-def trigger_incremental_train():
-    samples = list(collection.find())
-    samples = [s for s in samples if "image_path" in s and "true_label" in s]
+m2 = MyClassifier()
+m2.load_state_dict(torch.load("weights/Classification_Model.pth", map_location=device))
+m2.eval().to(device)
 
-    if not samples or len(samples) < 100:
-        return {"error": "Not enough validated samples (minimum 100 required)."}
+classes = ["COVID", "Lung_Opacity", "Normal", "Viral Pneumonia"]
 
-    m1 = UNet()
-    m1.model.load_state_dict(torch.load("Segmentation Model.pth"))
-    m1.to(device).eval()
+# ================== INFERENCE FUNCTION ==================
 
-    m2_old = MyClassifier()
-    m2_old.load_state_dict(torch.load("Classification Model.pth", map_location=device))
-    m2_old.to(device).eval()
+def analyze(image):
+    image_gray = image.convert("L")
+    image_tensor = transforms.ToTensor()(image_gray).unsqueeze(0).to(device)
 
-    m2 = MyClassifier()
-    m2.load_state_dict(torch.load("Classification Model.pth", map_location=device))
-    m2.to(device)
+    with torch.no_grad():
+        mask = m1(image_tensor).sigmoid()
+        mask = (mask > 0.5).float()
+        masked = image_tensor * mask
+
+    masked_rgb = masked.squeeze(0).repeat(3, 1, 1)
+
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.Normalize(mean=[0.41]*3, std=[0.16]*3)
+    ])
 
-    class LungXrayDataset(Dataset):
-        def __init__(self, entries, transform):
-            self.entries = entries
-            self.transform = transform
-            self.label_map = {label: i for i, label in enumerate(sorted(set(s["true_label"] for s in entries)))}
+    processed = transform(masked_rgb).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        logits = m2(processed)
+        probs = torch.softmax(logits, dim=1)
+        confidence, pred_class = torch.max(probs, dim=1)
+
+    return classes[pred_class.item()], f"{confidence.item() * 100:.2f}%"
+
+# ================== GRADIO INTERFACE ==================
+
+gradio_ui = gr.Interface(
+    fn=analyze,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Text(label="Prediction"),
+        gr.Text(label="Confidence")
+    ],
+    title="Chest X-Ray Analyzer",
+    description="Upload a chest X-ray image to predict the disease."
+)
+
+# Mount Gradio app
+app = gr.mount_gradio_app(app, gradio_ui, path="/")
+
+# ================== TRAINING API ==================
+class LungXrayDataset(Dataset):
+    def __init__(self, entries, transform, unet):
+        self.entries = entries
+        self.transform = transform
+        self.label_map = {label: i for i, label in enumerate(sorted(set(s["true_label"] for s in entries)))}
+        self.unet = unet
+
+    def __len__(self):
+        return len(self.entries)
+
+    def __getitem__(self, idx):
+        entry = self.entries[idx]
+        response = requests.get(entry["image_path"])
+        image = Image.open(BytesIO(response.content)).convert("L").resize((256, 256))
 
-        def __len__(self):
-            return len(self.entries)
+        image_tensor = transforms.ToTensor()(image).unsqueeze(0).to(device)
+        with torch.no_grad():
+            mask = self.unet(image_tensor).sigmoid()
+            mask = (mask > 0.5).float()
+            masked = image_tensor * mask
 
-        def __getitem__(self, idx):
-            entry = self.entries[idx]
-            image = fetch_image(entry["image_path"])
-            if image is None:
-                raise RuntimeError(f"Image at {entry['image_path']} could not be loaded.")
+        masked_rgb = masked.squeeze(0).repeat(3, 1, 1).cpu()
 
-            image_tensor = transforms.ToTensor()(image).unsqueeze(0).to(device)
+        if self.transform:
+            masked_rgb = self.transform(masked_rgb)
 
-            with torch.no_grad():
-                mask = m1(image_tensor).sigmoid()
-                mask = (mask > 0.5).float()
-                masked = image_tensor * mask
+        label = self.label_map[entry["true_label"]]
+        return masked_rgb, label
 
-            masked_rgb = masked.squeeze(0).repeat(3, 1, 1).cpu()
+@app.post("/trigger_incremental_train")
+async def trigger_train(request: Request):
+    data = await request.json()
+    samples = data.get("samples", [])
 
-            if self.transform:
-                masked_rgb = self.transform(masked_rgb)
+    if not samples or len(samples) < 100:
+        return {"error": "Not enough validated samples (minimum 100 required)."}
 
-            label = self.label_map[entry["true_label"]]
-            return masked_rgb, label
+    m2_new = MyClassifier()
+    m2_new.load_state_dict(torch.load("weights/Classification_Model.pth", map_location=device))
+    m2_new.to(device)
 
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
@@ -112,30 +151,25 @@ def trigger_incremental_train():
         samples, test_size=0.2, stratify=[s["true_label"] for s in samples], random_state=42
     )
 
-    train_ds = LungXrayDataset(train_entries, transform)
-    val_ds = LungXrayDataset(val_entries, transform)
+    train_ds = LungXrayDataset(train_entries, transform, m1)
+    val_ds = LungXrayDataset(val_entries, transform, m1)
     train_loader = DataLoader(train_ds, batch_size=16, shuffle=True)
     val_loader = DataLoader(val_ds, batch_size=16)
 
     criterion = nn.CrossEntropyLoss()
-    optimizer = optim.Adam(m2.parameters(), lr=1e-4)
+    optimizer = optim.Adam(m2_new.parameters(), lr=1e-4)
 
     for epoch in range(5):
-        m2.train()
-        total_loss, correct = 0, 0
+        m2_new.train()
         for imgs, labels in train_loader:
             imgs, labels = imgs.to(device), labels.to(device)
-            out = m2(imgs)
+            out = m2_new(imgs)
             loss = criterion(out, labels)
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
-            total_loss += loss.item()
-            correct += (out.argmax(1) == labels).sum().item()
-        acc = correct / len(train_ds)
-        print(f"Epoch {epoch+1} ➤ Loss: {total_loss:.4f} | Accuracy: {acc:.4f}")
 
-    def evaluate(model, loader, version):
+    def evaluate(model, loader):
         model.eval()
         y_true, y_pred = [], []
         with torch.no_grad():
@@ -145,26 +179,17 @@ def trigger_incremental_train():
                 y_pred.extend(outputs.argmax(1).cpu().numpy())
                 y_true.extend(labels.numpy())
         return {
-            "version": version,
             "accuracy": round(accuracy_score(y_true, y_pred), 4),
             "f1_macro": round(f1_score(y_true, y_pred, average="macro"), 4),
         }
 
-    eval_old = evaluate(m2_old, val_loader, "Old")
-    eval_new = evaluate(m2, val_loader, "New")
+    eval_old = evaluate(m2, val_loader)
+    eval_new = evaluate(m2_new, val_loader)
 
-    torch.save(m2.state_dict(), "New Classification Model.pth")
+    torch.save(m2_new.state_dict(), "weights/Classification_Model.pth")
 
     return {
         "old_model": eval_old,
-        "new_model": eval_new
+        "new_model": eval_new,
+        "updated_model_path": "weights/Classification_Model.pth"
     }
-
-with gr.Blocks() as demo:
-    with gr.Row():
-        train_button = gr.Button("Start Incremental Training")
-        output_json = gr.JSON(label="Training Result")
-
-    train_button.click(fn=trigger_incremental_train, inputs=[], outputs=output_json)
-
-demo.launch(server_name="0.0.0.0", server_port=7860)