Create app.py

app.py

@@ -0,0 +1,157 @@
+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 segmentation_models_pytorch as smp
+from torchvision.models import densenet121
+from pymongo import MongoClient
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, f1_score
+import gradio as gr
+
+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__()
+        self.model = smp.Unet(
+            encoder_name="resnet34",
+            encoder_weights=None,
+            in_channels=1,
+            classes=1
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+# Classifier model
+class MyClassifier(nn.Module):
+    def __init__(self, num_classes=4):
+        super().__init__()
+        base_model = densenet121(weights=None)
+        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)
+
+# 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 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) < 100:
+        return "Not enough validated samples (need at least 100)."
+
+    m1 = UNet().to(device)
+    m1.load_state_dict(torch.load("Segmentation Model.pth", map_location=device))
+    m1.eval()
+
+    m2_old = MyClassifier().to(device)
+    m2_old.load_state_dict(torch.load("Classification Model.pth", map_location=device))
+    m2_old.eval()
+
+    m2 = MyClassifier().to(device)
+    m2.load_state_dict(torch.load("Classification Model.pth", map_location=device))
+
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        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_ds = LungXrayDataset(train_entries, transform, m1, device)
+    val_ds = LungXrayDataset(val_entries, transform, m1, device)
+
+    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()
+
+    for epoch in range(5):
+        m2.train()
+        for imgs, labels in train_loader:
+            imgs, labels = imgs.to(device), labels.to(device)
+            out = m2(imgs)
+            loss = criterion(out, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+    eval_old = evaluate(m2_old, val_loader)
+    eval_new = evaluate(m2, val_loader)
+
+    torch.save(m2.state_dict(), "New Classification Model.pth")
+
+    return {
+        "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()