Update app.py

app.py

@@ -1,14 +1,10 @@
-import gradio as gr
 import torch
 import torch.nn as nn
 import torchvision.transforms as transforms
 from efficientnet_pytorch import EfficientNet
 from PIL import Image
-from fpdf import FPDF
-import os
-from datetime import datetime
 
-# Define HardSwish activation function
+# Define the HardSwish activation function
 class HardSwish(nn.Module):
     def __init__(self):
         super(HardSwish, self).__init__()
@@ -16,7 +12,7 @@ class HardSwish(nn.Module):
     def forward(self, x):
         return x * (torch.clamp(x + 3, 0, 6) / 6)
 
-# Define Custom EfficientNet
+# Define the model class
 class CustomEfficientNet(nn.Module):
     def __init__(self, num_classes):
         super(CustomEfficientNet, self).__init__()
@@ -40,7 +36,7 @@ class_names = ['No Finding', 'Enlarged Cardiomediastinum', 'Cardiomegaly', 'Lung
 # Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Load model
+# Load the model
 model = CustomEfficientNet(num_classes=14)
 checkpoint = torch.load('Final_global_model.pth.tar', map_location=device)
 if 'state_dict' in checkpoint:
@@ -50,97 +46,47 @@ else:
 model = model.to(device)
 model.eval()
 
-# Transformations
+# Transformations for input image
 transform = transforms.Compose([
     transforms.Resize((300, 300)),
     transforms.ToTensor(),
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-# Prediction and PDF generation function
-def predict(image, patient_name, xray_date):
+# Prediction function
+def predict(image):
     if image is None:
         raise ValueError("❌ Error: No image uploaded.")
-    if not patient_name.strip():
-        raise ValueError("❌ Error: Patient name is required.")
-    if not xray_date.strip():
-        raise ValueError("❌ Error: X-ray date is required.")
 
-    # Ensure correct image mode
+    # Ensure image is in RGB mode
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image)
 
     if image.mode != 'RGB':
         image = image.convert('RGB')
 
-    # Preprocess
+    # Preprocess the image
     img = transform(image).unsqueeze(0).to(device)
 
-    # Predict
+    # 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 = list(sorted_results.items())[:5]
-
-    # Doctor's comment
-    top_label, top_prob = top5[0]
-    if top_label == "No Finding" and top_prob > 0.5:
-        comment = "✅ No major abnormal findings detected."
-    elif top_prob > 0.5:
-        comment = f"⚠️ High likelihood of {top_label}."
-    else:
-        comment = f"🔎 Possible {top_label}, but confidence is low."
-
-    # Generate PDF report
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", 'B', 16)
-    pdf.cell(0, 10, "Chest X-ray Disease Report", ln=True, align='C')
-    pdf.ln(10)
-
-    # Patient Details
-    pdf.set_font("Arial", '', 12)
-    pdf.cell(0, 10, f"Patient Name: {patient_name}", ln=True)
-    pdf.cell(0, 10, f"X-ray Date: {xray_date}", ln=True)
-    pdf.cell(0, 10, f"Report Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=True)
-    pdf.ln(10)
-
-    # Top 5 Predictions
-    pdf.set_font("Arial", 'B', 14)
-    pdf.cell(0, 10, "Top 5 Predictions:", ln=True)
-    pdf.set_font("Arial", '', 12)
-
-    for disease, prob in top5:
-        pdf.cell(0, 10, f"{disease}: {prob*100:.2f}%", ln=True)
-
-    pdf.ln(10)
-
-    # Comment Section
-    pdf.set_font("Arial", 'B', 14)
-    pdf.cell(0, 10, "Doctor's Comment:", ln=True)
-    pdf.set_font("Arial", '', 12)
-    pdf.multi_cell(0, 10, comment)
-
-    # Save PDF
-    output_pdf_path = "chest_xray_report.pdf"
-    pdf.output(output_pdf_path)
-
-    return output_pdf_path
+    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", label="Upload Chest X-ray Image"),
-        gr.Textbox(label="Patient Name"),
-        gr.Textbox(label="Date of X-ray (YYYY-MM-DD)", placeholder="e.g. 2025-04-27")
-    ],
-    outputs=gr.File(label="Download PDF Report"),
-    title="Chest X-ray Disease Classification with Report",
-    description="Upload an X-ray, enter patient details, and download a detailed PDF report.",
+    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"
 )