Update app.py

app.py

@@ -4,8 +4,9 @@ import torch.nn as nn
 import torchvision.transforms as transforms
 from efficientnet_pytorch import EfficientNet
 from PIL import Image
+from datetime import datetime
 
-# Define the HardSwish activation function
+# Define HardSwish activation
 class HardSwish(nn.Module):
     def __init__(self):
         super(HardSwish, self).__init__()
@@ -13,7 +14,7 @@ class HardSwish(nn.Module):
     def forward(self, x):
         return x * (torch.clamp(x + 3, 0, 6) / 6)
 
-# Define the model class
+# Define custom EfficientNet model
 class CustomEfficientNet(nn.Module):
     def __init__(self, num_classes):
         super(CustomEfficientNet, self).__init__()
@@ -30,14 +31,16 @@ class CustomEfficientNet(nn.Module):
         return self.model(x)
 
 # Class names
-class_names = ['No Finding', 'Enlarged Cardiomediastinum', 'Cardiomegaly', 'Lung Opacity',
-               'Lung Lesion', 'Edema', 'Consolidation', 'Pneumonia', 'Atelectasis', 'Pneumothorax',
-               'Pleural Effusion', 'Pleural Other', 'Fracture', 'Support Devices']
+class_names = [
+    'No Finding', 'Enlarged Cardiomediastinum', 'Cardiomegaly', 'Lung Opacity',
+    'Lung Lesion', 'Edema', 'Consolidation', 'Pneumonia', 'Atelectasis', 'Pneumothorax',
+    'Pleural Effusion', 'Pleural Other', 'Fracture', 'Support Devices'
+]
 
 # Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Load the model
+# Load model
 model = CustomEfficientNet(num_classes=14)
 checkpoint = torch.load('Final_global_model.pth.tar', map_location=device)
 if 'state_dict' in checkpoint:
@@ -47,7 +50,7 @@ else:
 model = model.to(device)
 model.eval()
 
-# Transformations for input image
+# Transformations
 transform = transforms.Compose([
     transforms.Resize((300, 300)),
     transforms.ToTensor(),
@@ -55,42 +58,53 @@ transform = transforms.Compose([
 ])
 
 # Prediction function
-def predict(image):
+def predict(patient_name, scan_date, image):
     if image is None:
         raise ValueError("❌ Error: No image uploaded.")
 
-    # Ensure image is in RGB mode
+    # Ensure image is RGB
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image)
 
     if image.mode != 'RGB':
         image = image.convert('RGB')
 
-    # Preprocess the image
     img = transform(image).unsqueeze(0).to(device)
 
-    # Prediction
     with torch.no_grad():
         outputs = model(img)
         probs = torch.sigmoid(outputs).cpu().numpy()[0]
 
-    # Prepare results
     results = {class_names[i]: float(probs[i]) for i in range(len(class_names))}
     sorted_results = dict(sorted(results.items(), key=lambda item: item[1], reverse=True))
-    top5 = {k: v for k, v in list(sorted_results.items())[:5]}
-
-    return top5
-
-# Gradio Interface
-interface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=5),
-    title="Chest X-ray Disease Classification",
-    description="Upload a chest X-ray image (JPG, PNG, BMP, TIFF, etc.) to get disease predictions.\n\nTop 5 diseases are shown with their probability.",
-    theme="default",
-    allow_flagging="never"
-)
+    top5 = {k: f"{v*100:.2f}%" for k, v in list(sorted_results.items())[:5]}
+
+    summary = f"📋 **Patient Name**: {patient_name}\n📅 **Scan Date**: {scan_date.strftime('%Y-%m-%d')}\n\n### Top 5 Predictions"
+    return summary, top5
+
+# Gradio UI
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🩺 Chest X-ray Disease Classifier
+        Upload a chest X-ray and get the top 5 predicted diseases with probability scores.
+        """
+    )
+    with gr.Row():
+        with gr.Column():
+            patient_name = gr.Textbox(label="Patient Name", placeholder="Enter full name...")
+            scan_date = gr.Date(label="Scan Date", value=datetime.today)
+            image = gr.Image(label="Chest X-ray Image", type="pil")
+            predict_button = gr.Button("🔍 Predict")
+        with gr.Column():
+            summary = gr.Markdown()
+            output = gr.Label(num_top_classes=5)
+
+    predict_button.click(
+        predict,
+        inputs=[patient_name, scan_date, image],
+        outputs=[summary, output]
+    )
 
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()