Update app.py

app.py

@@ -4,8 +4,11 @@ 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 the HardSwish activation function
+# Define HardSwish activation
 class HardSwish(nn.Module):
     def __init__(self):
         super(HardSwish, self).__init__()
@@ -13,7 +16,7 @@ class HardSwish(nn.Module):
     def forward(self, x):
         return x * (torch.clamp(x + 3, 0, 6) / 6)
 
-# Define the model class
+# Model class
 class CustomEfficientNet(nn.Module):
     def __init__(self, num_classes):
         super(CustomEfficientNet, self).__init__()
@@ -34,7 +37,7 @@ class_names = ['No Finding', 'Enlarged Cardiomediastinum', 'Cardiomegaly', 'Lung
                'Lung Lesion', 'Edema', 'Consolidation', 'Pneumonia', 'Atelectasis', 'Pneumothorax',
                'Pleural Effusion', 'Pleural Other', 'Fracture', 'Support Devices']
 
-# Device configuration
+# Device config
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load the model
@@ -47,50 +50,118 @@ else:
 model = model.to(device)
 model.eval()
 
-# Transformations for input image
+# Transform
 transform = transforms.Compose([
     transforms.Resize((300, 300)),
     transforms.ToTensor(),
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-# Prediction function
-def predict(image):
+# Prediction and PDF generation
+def predict_and_generate_pdf(image, patient_name, xray_date):
     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:
+        raise ValueError("❌ Error: X-ray date is required.")
 
-    # 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 the image
+    # Preprocess
     img = transform(image).unsqueeze(0).to(device)
 
-    # Prediction
+    # Predict
     with torch.no_grad():
         outputs = model(img)
         probs = torch.sigmoid(outputs).cpu().numpy()[0]
 
-    # Prepare results
+    # Process 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
+    top5 = list(sorted_results.items())[:5]
+
+    # Comment Logic
+    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 low confidence."
+
+    # Save thumbnail
+    image_path = "xray_thumbnail.jpg"
+    image_copy = image.copy()
+    image_copy.thumbnail((100, 100))
+    image_copy.save(image_path)
+
+    # Create PDF
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", 'B', 18)
+    pdf.cell(0, 10, "Chest X-ray Analysis 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 Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=True)
+    pdf.ln(10)
+
+    # X-ray Image
+    pdf.image(image_path, x=80, w=50)
+    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)
+
+    pdf.set_fill_color(230, 230, 230)
+    for disease, prob in top5:
+        pdf.cell(100, 10, disease, 1, 0, 'L', fill=True)
+        pdf.cell(40, 10, f"{prob*100:.2f}%", 1, 1, 'C', fill=True)
+
+    pdf.ln(10)
+
+    # Comments
+    pdf.set_font("Arial", 'B', 14)
+    pdf.cell(0, 10, "Doctor's Comment:", ln=True)
+    pdf.set_font("Arial", '', 12)
+    pdf.set_fill_color(240, 248, 255)
+    pdf.multi_cell(0, 10, comment, fill=True)
+
+    # Save PDF
+    output_pdf_path = "chest_xray_report.pdf"
+    pdf.output(output_pdf_path)
+
+    return output_pdf_path
 
 # 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"
-)
+with gr.Blocks(theme="default") as demo:
+    gr.Markdown("# 🩺 Chest X-ray Disease Classification App")
+    gr.Markdown("Upload a chest X-ray, enter patient information, and generate a detailed PDF report.")
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="Upload Chest X-ray Image")
+            name_input = gr.Textbox(label="Patient Name")
+            date_input = gr.Date(label="Date of X-ray")
+            submit_button = gr.Button("Analyze & Generate PDF Report")
+        with gr.Column():
+            file_output = gr.File(label="Download Generated Report (PDF)")
+
+    submit_button.click(
+        fn=predict_and_generate_pdf,
+        inputs=[image_input, name_input, date_input],
+        outputs=file_output
+    )
 
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()