Update app.py

app.py

@@ -4,21 +4,22 @@ import torch.optim as optim
 from torchvision import transforms
 from torch.utils.data import Dataset, DataLoader
 from PIL import Image
-import segmentation_models_pytorch as smp
-from torchvision.models import densenet121
+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
 import gradio as gr
+import json
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# MongoDB connection
 client = MongoClient("mongodb+srv://xcovidinsight:FYP25S108@cluster0.vqesy.mongodb.net/")
 db = client["covid_system"]
 collection = db["validated_xrays"]
 
-# UNet model
 class UNet(nn.Module):
     def __init__(self):
         super().__init__()
@@ -32,7 +33,6 @@ class UNet(nn.Module):
     def forward(self, x):
         return self.model(x)
 
-# Classifier model
 class MyClassifier(nn.Module):
     def __init__(self, num_classes=4):
         super().__init__()
@@ -44,67 +44,63 @@ class MyClassifier(nn.Module):
     def forward(self, x):
         return self.model(x)
 
-# Dataset class
-class LungXrayDataset(Dataset):
-    def __init__(self, entries, transform, m1, device):
-        self.entries = entries
-        self.transform = transform
-        self.m1 = m1
-        self.device = device
-        self.label_map = {label: i for i, label in enumerate(sorted(set(s["true_label"] for s in entries)))}
-
-    def __len__(self):
-        return len(self.entries)
-
-    def __getitem__(self, idx):
-        entry = self.entries[idx]
-        image = Image.open(entry["image_path"]).convert("L").resize((256, 256))
-        image_tensor = transforms.ToTensor()(image).unsqueeze(0).to(self.device)
-
-        with torch.no_grad():
-            mask = self.m1(image_tensor).sigmoid()
-            mask = (mask > 0.5).float()
-            masked = image_tensor * mask
-
-        masked_rgb = masked.squeeze(0).repeat(3, 1, 1).cpu()
-
-        if self.transform:
-            masked_rgb = self.transform(masked_rgb)
-
-        label = self.label_map[entry["true_label"]]
-        return masked_rgb, label
-
-def evaluate(model, loader):
-    model.eval()
-    y_true, y_pred = [], []
-    with torch.no_grad():
-        for imgs, labels in loader:
-            imgs = imgs.to(device)
-            outputs = model(imgs)
-            y_pred.extend(outputs.argmax(1).cpu().numpy())
-            y_true.extend(labels.numpy())
-    return {
-        "accuracy": round(accuracy_score(y_true, y_pred), 4),
-        "f1_macro": round(f1_score(y_true, y_pred, average="macro"), 4),
-    }
+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
 
 def trigger_incremental_train():
     samples = list(collection.find())
     samples = [s for s in samples if "image_path" in s and "true_label" in s]
 
-    if len(samples) < 1:
-        return "Not enough validated samples (need at least 100)."
+    if not samples or len(samples) < 100:
+        return {"error": "Not enough validated samples (minimum 100 required)."}
 
-    m1 = UNet().to(device)
+    m1 = UNet()
     m1.load_state_dict(torch.load("Segmentation Model.pth", map_location=device))
-    m1.eval()
+    m1.to(device).eval()
 
-    m2_old = MyClassifier().to(device)
+    m2_old = MyClassifier()
     m2_old.load_state_dict(torch.load("Classification Model.pth", map_location=device))
-    m2_old.eval()
+    m2_old.to(device).eval()
 
-    m2 = MyClassifier().to(device)
+    m2 = MyClassifier()
     m2.load_state_dict(torch.load("Classification Model.pth", map_location=device))
+    m2.to(device)
+
+    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)))}
+
+        def __len__(self):
+            return len(self.entries)
+
+        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.")
+
+            image_tensor = transforms.ToTensor()(image).unsqueeze(0).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).cpu()
+
+            if self.transform:
+                masked_rgb = self.transform(masked_rgb)
+
+            label = self.label_map[entry["true_label"]]
+            return masked_rgb, label
 
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
@@ -115,17 +111,17 @@ 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, m1, device)
-    val_ds = LungXrayDataset(val_entries, transform, m1, device)
-
+    train_ds = LungXrayDataset(train_entries, transform)
+    val_ds = LungXrayDataset(val_entries, transform)
     train_loader = DataLoader(train_ds, batch_size=16, shuffle=True)
     val_loader = DataLoader(val_ds, batch_size=16)
 
-    optimizer = optim.Adam(m2.parameters(), lr=1e-4)
     criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(m2.parameters(), lr=1e-4)
 
     for epoch in range(5):
         m2.train()
+        total_loss, correct = 0, 0
         for imgs, labels in train_loader:
             imgs, labels = imgs.to(device), labels.to(device)
             out = m2(imgs)
@@ -133,26 +129,41 @@ def trigger_incremental_train():
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
-
-    eval_old = evaluate(m2_old, val_loader)
-    eval_new = evaluate(m2, val_loader)
+            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):
+        model.eval()
+        y_true, y_pred = [], []
+        with torch.no_grad():
+            for imgs, labels in loader:
+                imgs = imgs.to(device)
+                outputs = model(imgs)
+                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")
 
     torch.save(m2.state_dict(), "New Classification Model.pth")
 
     return {
-        "Old Model": eval_old,
-        "New Model": eval_new,
-        "status": "Training completed successfully"
+        "old_model": eval_old,
+        "new_model": eval_new
     }
 
-# Gradio interface
-demo = gr.Interface(
-    fn=trigger_incremental_train,
-    inputs=[],
-    outputs="json",
-    title="COVID X-ray Incremental Trainer",
-    description="Trigger training on new validated X-ray samples from MongoDB"
-)
-
-if __name__ == "__main__":
-    demo.launch()
+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)