Update app.py

app.py

@@ -7,8 +7,8 @@ from PIL import Image
 from fpdf import FPDF
 import os
 from datetime import datetime
-print("This app is not accurate , use with caution, images uploaded with bad pixels may effect the result")
-# Define HardSwish activation
+
+# Define HardSwish activation function
 class HardSwish(nn.Module):
     def __init__(self):
         super(HardSwish, self).__init__()
@@ -16,7 +16,7 @@ class HardSwish(nn.Module):
     def forward(self, x):
         return x * (torch.clamp(x + 3, 0, 6) / 6)
 
-# Model class
+# Define Custom EfficientNet
 class CustomEfficientNet(nn.Module):
     def __init__(self, num_classes):
         super(CustomEfficientNet, self).__init__()
@@ -37,10 +37,10 @@ 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
+# Load model
 model = CustomEfficientNet(num_classes=14)
 checkpoint = torch.load('Final_global_model.pth.tar', map_location=device)
 if 'state_dict' in checkpoint:
@@ -50,22 +50,27 @@ else:
 model = model.to(device)
 model.eval()
 
-# Transform
+# Transformations
 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):
+# Helper to sanitize filename
+def clean_filename(name):
+    return "".join(c for c in name if c.isalnum() or c in (' ', '_', '-')).rstrip()
+
+# Prediction and PDF generation function
+def predict(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:
+    if not xray_date.strip():
         raise ValueError("❌ Error: X-ray date is required.")
 
+    # Ensure correct image mode
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image)
 
@@ -80,88 +85,70 @@ def predict_and_generate_pdf(image, patient_name, xray_date):
         outputs = model(img)
         probs = torch.sigmoid(outputs).cpu().numpy()[0]
 
-    # 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 = list(sorted_results.items())[:5]
 
-    # Comment Logic
+    # 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"🔎 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)
+        comment = f"🔎 Possible {top_label}, but confidence is low."
 
-    # Create PDF
+    # Generate 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.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"Date of X-ray: {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.cell(0, 10, f"{disease}: {prob*100:.2f}%", ln=True)
 
     pdf.ln(10)
 
-    # Comments
+    # Comment Section
     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)
+    pdf.multi_cell(0, 10, comment)
 
-    # Save PDF
-    output_pdf_path = "chest_xray_report.pdf"
-    pdf.output(output_pdf_path)
+    # Save PDF with clean filename
+    safe_name = clean_filename(patient_name.replace(" ", "_"))
+    report_filename = f"{safe_name}_{xray_date}_Report.pdf"
+    pdf.output(report_filename)
 
-    return output_pdf_path
+    return report_filename
 
 # Gradio Interface
-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.Textbox(label="Date of X-ray (YYYY-MM-DD)", placeholder="e.g. 2025-04-27")
-            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
-    )
+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.",
+    theme="default",
+    allow_flagging="never"
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    interface.launch()