Update app.py

app.py

@@ -13,14 +13,9 @@ import segmentation_models_pytorch as smp
 from torchvision.models import densenet121, DenseNet121_Weights
 import gradio as gr
 import os
-import torch
-import gradio as gr
-from PIL import Image
-import numpy as np
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
 from torchvision import transforms
-import torch.nn as nn
 import cv2
 import matplotlib.cm as cm
 
@@ -47,7 +42,6 @@ m2.classifier = nn.Linear(1024, 4)
 m2.load_state_dict(torch.load("Classification_Model.pth", map_location=device))
 m2.eval().to(device)
 
-classes = ["COVID", "Lung_Opacity", "Normal", "Viral Pneumonia"]
 
 unet_transform = A.Compose([
     A.Resize(256, 256),
@@ -119,24 +113,23 @@ gradio_ui = gr.Interface(
 app = gr.mount_gradio_app(app, gradio_ui, path="/")
 
 # ================== TRAINING API ==================
+# Data prepare
 class LungXrayDataset(Dataset):
-    def __init__(self, entries, transform, unet):
+    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)))}
-        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))
-
+        image = Image.open(entry["image_path"]).convert("L").resize((256, 256))
         image_tensor = transforms.ToTensor()(image).unsqueeze(0).to(device)
+
         with torch.no_grad():
-            mask = self.unet(image_tensor).sigmoid()
+            mask = m1(image_tensor).sigmoid()
             mask = (mask > 0.5).float()
             masked = image_tensor * mask
 
@@ -155,28 +148,37 @@ async def trigger_train(request: Request):
 
     if not samples or len(samples) < 100:
         return {"error": "Not enough validated samples (minimum 100 required)."}
-
-    m2_new = MyClassifier()
-    m2_new.load_state_dict(torch.load("weights/Classification_Model.pth", map_location=device))
+    
+    # Classification Model
+    m2 = densenet121(weights=DenseNet121_Weights.IMAGENET1K_V1)
+    m2.classifier = nn.Linear(1024, 4)
+    m2.load_state_dict(torch.load("Classification_Model.pth", map_location=device))
+    m2.eval().to(device)
+
+    m2_new = densenet121(weights=DenseNet121_Weights.IMAGENET1K_V1)
+    m2_new.classifier = nn.Linear(1024, 4)
+    m2_new.load_state_dict(torch.load("Classification_Model.pth", map_location=device))
     m2_new.to(device)
 
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
-        transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.41]*3, std=[0.16]*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)
-    val_ds = LungXrayDataset(val_entries, transform, m1)
+    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)
 
+    print(f"Training triggered with {len(samples)} samples by API")
     for epoch in range(5):
         m2_new.train()
         for imgs, labels in train_loader:
@@ -200,14 +202,25 @@ async def trigger_train(request: Request):
             "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, val_loader)
     eval_new = evaluate(m2_new, val_loader)
-
-    torch.save(m2_new.state_dict(), "weights/Classification_Model.pth")
-
+    
+    if eval_new["f1_macro"] > eval_old["f1_macro"]:
+        torch.save(m2_new.state_dict(), "Classification_Model.pth")
+        model_used = "new"
+        print("[NOTE] New model used and saved based on higher F1 score.")
+    else:
+        model_used = "old"
+        print("[NOTE] Old model retained (new model did not improve F1 score).")
+    
     return {
         "old_model": eval_old,
         "new_model": eval_new,
-        "updated_model_path": "weights/Classification_Model.pth"
+        "model_used": model_used,
+        "updated_model_path": "Classification_Model.pth" if model_used == "new" else "unchanged"
     }
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)