Update app.py

app.py

@@ -1,103 +1,141 @@
-from fastapi import FastAPI, Request
 import torch
 import torch.nn as nn
 import torch.optim as optim
 from torchvision import transforms
 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, DenseNet121_Weights
-
-# FastAPI app instance
-app = FastAPI()
+from torchvision.models import densenet121
+import gradio as gr
+import json
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load segmentation model
-m1 = smp.Unet(
-    encoder_name="resnet34",
-    encoder_weights="imagenet",
-    in_channels=1,
-    classes=1
-).to(device)
-m1.load_state_dict(torch.load("Segmentation_Model.pth", map_location=device))
-m1.eval()
-
-# Dataset for incremental training
-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 = Image.open(entry["image_path"]).convert("L").resize((256, 256))
-        image_tensor = transforms.ToTensor()(image).unsqueeze(0).to(device)
+client = MongoClient("mongodb+srv://xcovidinsight:FYP25S108@cluster0.vqesy.mongodb.net/")
+db = client["covid_system"]
+collection = db["validated_xrays"]
+
+class UNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = smp.Unet(
+            encoder_name="resnet34",
+            encoder_weights="imagenet",
+            in_channels=1,
+            classes=1
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+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)
+        self.model = base_model
+
+    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
+
+def trigger_incremental_train():
+    samples = list(collection.find())
+    samples = [s for s in samples if "image_path" in s and "true_label" in s]
 
-        with torch.no_grad():
-            mask = m1(image_tensor).sigmoid()
-            mask = (mask > 0.5).float()
-            masked = image_tensor * mask
+    if not samples or len(samples) < 100:
+        return {"error": "Not enough validated samples (minimum 100 required)."}
 
-        masked_rgb = masked.squeeze(0).repeat(3, 1, 1).cpu()
+    m1 = UNet()
+    m1.model.load_state_dict(torch.load("Segmentation Model.pth"))
+    m1.to(device).eval()
 
-        if self.transform:
-            masked_rgb = self.transform(masked_rgb)
+    m2_old = MyClassifier()
+    m2_old.load_state_dict(torch.load("Classification Model.pth", map_location=device))
+    m2_old.to(device).eval()
 
-        label = self.label_map[entry["true_label"]]
-        return masked_rgb, label
+    m2 = MyClassifier()
+    m2.load_state_dict(torch.load("Classification Model.pth", map_location=device))
+    m2.to(device)
 
-@app.post("/trigger_incremental_train")
-async def trigger_train(request: Request):
-    data = await request.json()
-    samples = data.get("samples", [])
+    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)))}
 
-    if not samples or len(samples) < 100:
-        return {"error": "Not enough validated samples (minimum 100 required)."}
+        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.")
 
-    # Load models
-    def load_model():
-        model = densenet121(weights=DenseNet121_Weights.IMAGENET1K_V1)
-        model.classifier = nn.Linear(1024, 4)
-        model.load_state_dict(torch.load("Classification_Model.pth", map_location=device))
-        return model.to(device)
+            image_tensor = transforms.ToTensor()(image).unsqueeze(0).to(device)
 
-    m2_old = load_model().eval()
-    m2_new = load_model()
+            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)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.41]*3, std=[0.16]*3)
+        transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
     ])
 
     train_entries, val_entries = train_test_split(
         samples, test_size=0.2, stratify=[s["true_label"] for s in samples], random_state=42
     )
 
-    train_loader = DataLoader(LungXrayDataset(train_entries, transform), batch_size=16, shuffle=True)
-    val_loader = DataLoader(LungXrayDataset(val_entries, transform), batch_size=16)
+    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)
 
     criterion = nn.CrossEntropyLoss()
-    optimizer = optim.Adam(m2_new.parameters(), lr=1e-4)
+    optimizer = optim.Adam(m2.parameters(), lr=1e-4)
 
     for epoch in range(5):
-        m2_new.train()
+        m2.train()
+        total_loss, correct = 0, 0
         for imgs, labels in train_loader:
             imgs, labels = imgs.to(device), labels.to(device)
-            out = m2_new(imgs)
+            out = m2(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):
+    def evaluate(model, loader, version):
         model.eval()
         y_true, y_pred = [], []
         with torch.no_grad():
@@ -107,20 +145,26 @@ async def trigger_train(request: Request):
                 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)
-    eval_new = evaluate(m2_new, val_loader)
+    eval_old = evaluate(m2_old, val_loader, "Old")
+    eval_new = evaluate(m2, val_loader, "New")
 
-    model_used = "new" if eval_new["f1_macro"] > eval_old["f1_macro"] else "old"
-    if model_used == "new":
-        torch.save(m2_new.state_dict(), "Classification_Model.pth")
+    torch.save(m2.state_dict(), "New Classification Model.pth")
 
     return {
         "old_model": eval_old,
-        "new_model": eval_new,
-        "model_used": model_used,
-        "updated_model_path": "Classification_Model.pth" if model_used == "new" else "unchanged"
+        "new_model": eval_new
     }
+
+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)