Create app.py

app.py

@@ -0,0 +1,126 @@
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from torchvision import transforms, datasets
+from albumentations import Compose, HorizontalFlip, ShiftScaleRotate, Resize, Normalize
+from albumentations.pytorch import ToTensorV2
+import timm
+import gradio as gr
+import numpy as np
+from PIL import Image
+
+# Hyperparameters
+h = {
+    "num_epochs": 10,
+    "batch_size": 64,
+    "image_size": 224,
+    "lr": 0.001,
+    "model": "efficientnetv2",
+    "scheduler": "CosineAnnealingLR10",
+    "balance": True,
+    "early_stopping_patience": 5
+}
+
+# Custom Dataset and DataModule for PyTorch Lightning
+class CustomImageFolder(torch.utils.data.Dataset):
+    def __init__(self, root, transform=None):
+        self.dataset = datasets.ImageFolder(root)
+        self.transform = transform
+
+    def __getitem__(self, index):
+        image, label = self.dataset[index]
+        if self.transform:
+            image = self.transform(image=np.array(image))["image"]
+        return image, label
+
+    def __len__(self):
+        return len(self.dataset)
+
+class PneumoniaDataModule(pl.LightningDataModule):
+    def __init__(self, h, data_dir):
+        super().__init__()
+        self.h = h
+        self.data_dir = data_dir
+
+    def setup(self, stage=None):
+        train_transform = Compose([
+            ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=20),
+            HorizontalFlip(),
+            Resize(self.h["image_size"], self.h["image_size"]),
+            Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ToTensorV2()
+        ])
+        
+        self.train_dataset = CustomImageFolder(self.data_dir + "/train", transform=train_transform)
+        
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, batch_size=self.h["batch_size"], shuffle=True)
+
+# Model definition using LightningModule
+class PneumoniaModel(pl.LightningModule):
+    def __init__(self, h):
+        super().__init__()
+        self.h = h
+        self.model = timm.create_model("tf_efficientnetv2_b0", pretrained=True, num_classes=2)
+        self.criterion = nn.CrossEntropyLoss()
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        inputs, labels = batch
+        outputs = self(inputs)
+        loss = self.criterion(outputs, labels)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.h["lr"])
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.h["num_epochs"], eta_min=self.h["lr"] * 0.1)
+        return {"optimizer": optimizer, "lr_scheduler": scheduler}
+
+# Load model after training
+def load_model(h):
+    model = PneumoniaModel(h)
+    model.load_state_dict(torch.load("pneumonia_model.pth", map_location=torch.device('cpu')))
+    model.eval()
+    return model
+
+trained_model = load_model(h)
+
+# Gradio Prediction Function
+def predict_pneumonia(image):
+    transform = Compose([
+        Resize(h["image_size"], h["image_size"]),
+        Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ToTensorV2()
+    ])
+    
+    # Preprocess the image
+    image = np.array(image)
+    image = transform(image=image)["image"]
+    image = image.unsqueeze(0)  # Add batch dimension
+    
+    # Predict with the model
+    with torch.no_grad():
+        outputs = trained_model(image)
+        prediction = torch.argmax(outputs, dim=1).item()
+    
+    # Map prediction to label
+    label = "Pneumonia Detected" if prediction == 1 else "Normal"
+    return label
+
+# Gradio Interface
+input_image = gr.inputs.Image(type="pil", label="Upload Chest X-ray Image")
+output_label = gr.outputs.Label(label="Diagnosis")
+
+app = gr.Interface(
+    fn=predict_pneumonia, 
+    inputs=input_image, 
+    outputs=output_label,
+    title="Pneumonia Detection",
+    description="Upload a chest X-ray image to detect potential pneumonia using AI."
+)
+
+# Launch the app
+app.launch()