Update app.py

app.py

@@ -4,12 +4,8 @@ import torch.nn as nn
 import torchvision.transforms as transforms
 from efficientnet_pytorch import EfficientNet
 from PIL import Image
-import numpy as np
-import os
-from fpdf import FPDF
-from datetime import datetime
 
-# Define HardSwish activation
+# Define the HardSwish activation function
 class HardSwish(nn.Module):
     def __init__(self):
         super(HardSwish, self).__init__()
@@ -17,7 +13,7 @@ class HardSwish(nn.Module):
     def forward(self, x):
         return x * (torch.clamp(x + 3, 0, 6) / 6)
 
-# Define model class
+# Define the model class
 class CustomEfficientNet(nn.Module):
     def __init__(self, num_classes):
         super(CustomEfficientNet, self).__init__()
@@ -38,7 +34,7 @@ class_names = ['No Finding', 'Enlarged Cardiomediastinum', 'Cardiomegaly', 'Lung
                'Lung Lesion', 'Edema', 'Consolidation', 'Pneumonia', 'Atelectasis', 'Pneumothorax',
                'Pleural Effusion', 'Pleural Other', 'Fracture', 'Support Devices']
 
-# Device config
+# Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load the model
@@ -51,104 +47,50 @@ else:
 model = model.to(device)
 model.eval()
 
-# Transform
+# 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
-def predict_and_generate_pdf(image, patient_name, xray_date):
+# Prediction function
+def predict(image):
     if image is None:
         raise ValueError("❌ Error: No image uploaded.")
 
+    # 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]
 
-    # Process results
+    # 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]
-
-    # Generate 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"Possibility of {top_label}, but with low confidence."
-
-    # Save thumbnail
-    image_path = "xray_temp.jpg"
-    image.thumbnail((200, 200))
-    image.save(image_path)
-
-    # Create PDF
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", 'B', 16)
-    pdf.cell(0, 10, "Chest X-ray Report", ln=True, align='C')
-    pdf.ln(10)
-
-    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 Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=True)
-    pdf.ln(10)
-
-    pdf.image(image_path, w=80)
-    pdf.ln(10)
-
-    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)
-    pdf.set_font("Arial", 'B', 14)
-    pdf.cell(0, 10, "Comment:", ln=True)
-    pdf.set_font("Arial", '', 12)
-    pdf.multi_cell(0, 10, comment)
-
-    # Save PDF
-    pdf_output_path = "report.pdf"
-    pdf.output(pdf_output_path)
-
-    return pdf_output_path
+    top5 = {k: v for k, v in list(sorted_results.items())[:5]}
+
+    return top5
 
 # Gradio Interface
-with gr.Blocks(theme="default") as demo:
-    gr.Markdown("# 🩺 Chest X-ray Disease Classification")
-    gr.Markdown("Upload a chest X-ray, enter patient's name and date, and download a PDF report.")
-
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(type="pil", label="Upload Chest X-ray")
-            name_input = gr.Textbox(label="Patient Name")
-            date_input = gr.Textbox(label="X-ray Date (YYYY-MM-DD)", placeholder="2025-04-27")
-            submit_btn = gr.Button("Analyze & Generate Report")
-        with gr.Column():
-            file_output = gr.File(label="Download Report (PDF)")
-
-    submit_btn.click(
-        fn=predict_and_generate_pdf,
-        inputs=[image_input, name_input, date_input],
-        outputs=file_output
-    )
+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"
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    interface.launch()