Update app.py

app.py

@@ -5,27 +5,17 @@ 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 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
-import gradio as gr
-import os
-import albumentations as A
-from albumentations.pytorch import ToTensorV2
-from torchvision import transforms
-import cv2
-
-# Device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# FastAPI instance
+# FastAPI app instance
 app = FastAPI()
 
-# ================== MODEL DEFINITIONS ==================
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Load segmentation model
 m1 = smp.Unet(
     encoder_name="resnet34",
     encoder_weights="imagenet",
@@ -35,91 +25,7 @@ m1 = smp.Unet(
 m1.load_state_dict(torch.load("Segmentation_Model.pth", map_location=device))
 m1.eval()
 
-# 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)
-
-
-unet_transform = A.Compose([
-    A.Resize(256, 256),
-    A.Normalize(mean=0.5, std=0.5),
-    ToTensorV2()
-])
-
-classifier_transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=0.41, std=0.16)
-])
-
-# Inference Function
-def analyze(image):
-    try:
-        if image is None:
-            return "No image uploaded", "0.00%"
-            
-        # Grayscale image for UNet
-        image_gray = image.convert("L")
-        img_np = np.array(image_gray)
-    
-        # UNet input
-        augmented = unet_transform(image=img_np)
-        unet_input = augmented["image"].unsqueeze(0).to(device)
-    
-        # Segmentation
-        with torch.no_grad():
-            mask_pred = m1(unet_input)
-            mask_pred = torch.sigmoid(mask_pred)
-            mask = (mask_pred > 0.5).float().squeeze().cpu().numpy()
-    
-        # Resize to match classifier input
-        resized_gray = image_gray.resize((224, 224), Image.BILINEAR)
-        mask_img = Image.fromarray((mask * 255).astype(np.uint8))
-        mask_resized = transforms.functional.resize(mask_img, [224, 224])
-        image_np = np.array(resized_gray).astype(np.float32)
-        mask_np = (np.array(mask_resized) > 127).astype(np.float32)
-    
-        lung_image = image_np * mask_np
-        lung_image_3ch = np.stack([lung_image] * 3, axis=-1)
-        lung_image_3ch = np.clip(lung_image_3ch, 0, 255).astype(np.uint8)
-        lung_image_pil = Image.fromarray(lung_image_3ch)
-    
-        input_tensor = classifier_transform(lung_image_pil).unsqueeze(0).to(device)
-    
-        # Classification
-        with torch.no_grad():
-            logits = m2(input_tensor)
-            probs = torch.softmax(logits, dim=1)
-            confidence, pred_class = torch.max(probs, dim=1)
-    
-        classes = ["COVID", "Lung_Opacity", "Normal", "Viral Pneumonia"]
-        confidence_percent = f"{confidence.item() * 100:.2f}%"
-        
-        return classes[pred_class.item()], confidence_percent
-        
-    except Exception as e:
-        return f"[Error] {str(e)}", "0.00%"
-
-# ================== GRADIO INTERFACE ==================
-
-with gr.Blocks() as interface:
-    with gr.Row():
-        img = gr.Image(type="pil", label="Upload")
-        btn = gr.Button("Analyze")
-    with gr.Row():
-        pred = gr.Text(label="Prediction")
-        conf = gr.Text(label="Confidence")
-
-    btn.click(fn=analyze, inputs=img, outputs=[pred, conf])
-
-interface.queue()
-
-app = gr.mount_gradio_app(app, interface, path="/")
-
-# ================== TRAINING API ==================
-# Data prepare
+# Dataset for incremental training
 class LungXrayDataset(Dataset):
     def __init__(self, entries, transform):
         self.entries = entries
@@ -154,17 +60,16 @@ async def trigger_train(request: Request):
 
     if not samples or len(samples) < 100:
         return {"error": "Not enough validated samples (minimum 100 required)."}
-    
-    # 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)
+
+    # 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)
+
+    m2_old = load_model().eval()
+    m2_new = load_model()
 
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
@@ -176,15 +81,12 @@ async def trigger_train(request: Request):
         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_loader = DataLoader(train_ds, batch_size=16, shuffle=True)
-    val_loader = DataLoader(val_ds, batch_size=16)
+    train_loader = DataLoader(LungXrayDataset(train_entries, transform), batch_size=16, shuffle=True)
+    val_loader = DataLoader(LungXrayDataset(val_entries, transform), 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:
@@ -208,25 +110,17 @@ 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_old = evaluate(m2_old, val_loader)
     eval_new = evaluate(m2_new, val_loader)
-    
-    if eval_new["f1_macro"] > eval_old["f1_macro"]:
+
+    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")
-        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,
         "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)