Update app.py

app.py

@@ -7,53 +7,15 @@ from tensorflow.keras.models import Model
 import os
 from reportlab.lib.pagesizes import letter
 from reportlab.pdfgen import canvas
-import tempfile
-
-# Print TensorFlow version for debugging
-print(f"TensorFlow version: {tf.__version__}")
-
-# Model loading with error handling
-def load_models():
-    try:
-        # Check if model files exist
-        if not os.path.exists("modelDense.h5"):
-            print("Warning: modelDense.h5 not found")
-            return None, None, None
-            
-        if not os.path.exists("modelVGG16.h5"):
-            print("Warning: modelVGG16.h5 not found")
-            return None, None, None
-            
-        if not os.path.exists("modelCovid.h5"):
-            print("Warning: modelCovid.h5 not found")
-            return None, None, None
-            
-        print("Loading models...")
-        model_step1 = tf.keras.models.load_model("modelDense.h5")
-        model_step2 = tf.keras.models.load_model("modelVGG16.h5")
-        model_bin = tf.keras.models.load_model("modelCovid.h5")
-        print("Models loaded successfully!")
-        return model_step1, model_step2, model_bin
-    except Exception as e:
-        print(f"Error loading models: {e}")
-        return None, None, None
-
-# Load models when app starts
-try:
-    print("Attempting to load models...")
-    model_step1, model_step2, model_bin = load_models()
-except Exception as e:
-    print(f"Exception during model loading: {e}")
-    model_step1, model_step2, model_bin = None, None, None
+from reportlab.lib import utils
+
+# Load the trained models
+model_step1 = tf.keras.models.load_model("modelDense.h5")  # Model for 3-class classification (bacterial-viral, covid, normal)
+model_step2 = tf.keras.models.load_model("modelVGG16.h5")  # Model for 2-class classification (bacterial, viral)
+model_bin = tf.keras.models.load_model("modelCovid.h5")
 
 # Function to preprocess and predict
 def predict(img):
-    if img is None:
-        return "Please upload an image."
-    
-    if model_step1 is None or model_step2 is None:
-        return "Models could not be loaded. Please check the model files."
-        
     img_array = analyze_image(img)  # Pass the PIL image directly
     img_array_expanded = np.expand_dims(img_array, axis=0)  # Add batch dimension
     img_array_expanded /= 255.0  # Normalize
@@ -86,12 +48,6 @@ def analyze_image(img):
 
 # Function for Grad-CAM visualization with center focus
 def generate_gradcam_heatmap_center_focus(img):
-    if img is None:
-        return None
-    
-    if model_step1 is None or model_step2 is None:
-        return None
-        
     img_array = analyze_image(img)
     img_array_expanded = np.expand_dims(img_array, axis=0)  # Add batch dimension
     img_array_expanded /= 255.0  # Normalize image
@@ -99,7 +55,6 @@ def generate_gradcam_heatmap_center_focus(img):
     step1_prediction = model_step1.predict(img_array_expanded)
     class_idx2 = np.argmax(step1_prediction)
 
-    # Select appropriate model for Grad-CAM
     if class_idx2 == 1 or class_idx2 == 2:
         last_conv_layer_name = 'block3_conv1'
         last_conv_layer = model_step2.get_layer(last_conv_layer_name)
@@ -139,26 +94,24 @@ def generate_gradcam_heatmap_center_focus(img):
     heatmap_resized /= np.max(heatmap_resized)
 
     heatmap_resized = np.uint8(255 * heatmap_resized)
+    #heatmap_resized = 255 - heatmap_resized
     heatmap_colormap = cv2.applyColorMap(heatmap_resized, cv2.COLORMAP_JET)
 
     original_img = np.uint8(255 * img_array / np.max(img_array))
     superimposed_img = cv2.addWeighted(original_img, 0.6, heatmap_colormap, 0.4, 0)
-    
-    return superimposed_img  # Return the image directly, not a file path
+
+    output_path = "gradcam_output.png"
+    cv2.imwrite(output_path, superimposed_img)
+
+    return output_path  # Return the path to the image file
 
 # Function to generate a PDF report
 def generate_pdf_report(prediction, gradcam_img, patient_name, patient_id, notes):
-    if gradcam_img is None:
-        return None
+    pdf_path = f"patient_pneumonia_report_{patient_id}.pdf"
+    gradcam_image_path = "gradcam_output.png"
     
-    # Create a temporary file for the PDF
-    with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as tmp:
-        pdf_path = tmp.name
-    
-    # Create a temporary file for the Grad-CAM image
-    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as img_tmp:
-        gradcam_path = img_tmp.name
-        cv2.imwrite(gradcam_path, gradcam_img)
+    # Save the Grad-CAM image
+    cv2.imwrite(gradcam_image_path, gradcam_img)  # Saving the image file
     
     # Create the PDF
     c = canvas.Canvas(pdf_path, pagesize=letter)
@@ -181,69 +134,38 @@ def generate_pdf_report(prediction, gradcam_img, patient_name, patient_id, notes
     c.setFont("Helvetica-Bold", 12)
     c.drawString(100, 640, "Grad-CAM Visualization:")
 
-    c.drawImage(gradcam_path, 100, 350, width=224, height=224)
+    c.drawImage(gradcam_image_path, 100, 350, width=224, height=224)  # Adjust image size slightly smaller
+
+
     c.save()
 
-    # Clean up the temporary image file
-    os.unlink(gradcam_path)
-    
     return pdf_path
 
+
+# Create Gradio Blocks for the prediction function and Grad-CAM functionality
+with gr.Blocks() as demo:
     
-    # Define what happens when models aren't loaded
-    if model_step1 is None or model_step2 is None or model_bin is None:
-        gr.Markdown("# ⚠️ Error: Models could not be loaded")
-        gr.Markdown("""
-        This application requires the following model files to be present:
-        - modelDense.h5
-        - modelVGG16.h5
-        - modelCovid.h5
-        
-        Please make sure these files are uploaded to the Hugging Face Space.
-        """)
-    else:
-        # Create Gradio Blocks for the app
-        with gr.Blocks(theme=gr.themes.Default()) as demo:
-            
-            gr.Markdown("# Pneumonia Detection Model")
-            gr.Markdown("Upload a Chest X-ray to detect Pneumonia and classify bacterial/viral vs normal/covid.")
-            
-            with gr.Row():
-                with gr.Column(scale=1):
-                    img_input = gr.Image(type="pil", label="Upload X-ray Image")
-                    predict_btn = gr.Button("Predict", variant="primary")
-                    gradcam_btn = gr.Button("Generate Grad-CAM Heatmap")
-                
-                with gr.Column(scale=1):
-                    name_input = gr.Textbox(label="Patient Name")
-                    id_input = gr.Textbox(label="Patient ID")
-                    notes_input = gr.Textbox(label="Additional Notes", lines=3)
-                    report_btn = gr.Button("Generate Report")
-            
-            with gr.Row():
-                with gr.Column(scale=1):
-                    output_text = gr.Textbox(label="Prediction Result")
-                    gradcam_output = gr.Image(label="Grad-CAM Heatmap")
-                
-                with gr.Column(scale=1):
-                    report_output = gr.File(label="Download Report")
-                    
-            # Set up event handlers
-            predict_btn.click(predict, inputs=img_input, outputs=output_text)
-            gradcam_btn.click(generate_gradcam_heatmap_center_focus, inputs=img_input, outputs=gradcam_output)
-            report_btn.click(
-                generate_pdf_report, 
-                inputs=[output_text, gradcam_output, name_input, id_input, notes_input], 
-                outputs=report_output
-            )
-
-            # Add example images for demo purposes
-            # gr.Examples(
-            #     examples=["example1.jpg", "example2.jpg"],
-            #     inputs=img_input
-            # )
-
-# Launch the app - for Hugging Face, we don't need to specify share=True
-# Use queue=False to avoid the validation error
-if model_step1 is not None and model_step2 is not None and model_bin is not None:
-    demo.launch(debug=True, show_error=True, queue=False)
+    gr.Markdown("# Pneumonia Detection Model (2-step classification with Grad-CAM)")
+    gr.Markdown("Upload a Chest X-ray to detect Pneumonia and classify bacterial/viral vs normal/covid.")
+    
+    with gr.Row():
+        img_input = gr.Image(type="pil", label="Upload X-ray Image")
+        name_input = gr.Textbox(label="Patient Name")
+        id_input = gr.Textbox(label="Patient ID")
+        notes_input = gr.Textbox(label="Additional Notes", lines=5)
+    
+    with gr.Row():
+        predict_btn = gr.Button("Predict")
+        gradcam_btn = gr.Button("Generate Grad-CAM Heatmap")
+        report_btn = gr.Button("Generate Report")
+    
+    output_text = gr.Textbox(label="Prediction Result")
+    gradcam_output = gr.Image(label="Grad-CAM Heatmap")
+    report_output = gr.File(label="Download Report")
+
+    predict_btn.click(predict, inputs=img_input, outputs=output_text)
+    gradcam_btn.click(generate_gradcam_heatmap_center_focus, inputs=img_input, outputs=gradcam_output)
+    report_btn.click(generate_pdf_report, inputs=[output_text, gradcam_output, name_input, id_input, notes_input], outputs=report_output)
+
+# Launch the app
+demo.launch(share=True)  # Makes it accessible online