Add detailed logging to diagnose startup failure

app.py

@@ -1,24 +1,42 @@
 """
 Chest X-Ray Classifier with Grad-CAM Explainability
-=====================================================
-A Hugging Face Space that classifies chest X-rays into 14 thoracic
-conditions and generates Grad-CAM heatmaps for explainability.
 """
 
-import torch
-import numpy as np
-import gradio as gr
-from PIL import Image
-from torchvision import transforms
-import timm
-from pytorch_grad_cam import GradCAM
-from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
-from pytorch_grad_cam.utils.image import show_cam_on_image
-
-# ─────────────────────────────────────────────────────────────
-# CONFIGURATION
-# ─────────────────────────────────────────────────────────────
-
+import traceback
+import sys
+
+print("=" * 50)
+print("STEP 1: Starting imports...")
+sys.stdout.flush()
+
+try:
+    import torch
+    print("  torch imported OK")
+    import numpy as np
+    print("  numpy imported OK")
+    import gradio as gr
+    print("  gradio imported OK")
+    from PIL import Image
+    print("  PIL imported OK")
+    from torchvision import transforms
+    print("  torchvision transforms imported OK")
+    import timm
+    print("  timm imported OK")
+    from pytorch_grad_cam import GradCAM
+    print("  GradCAM imported OK")
+    from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+    print("  ClassifierOutputTarget imported OK")
+    from pytorch_grad_cam.utils.image import show_cam_on_image
+    print("  show_cam_on_image imported OK")
+except Exception as e:
+    print(f"IMPORT ERROR: {e}")
+    traceback.print_exc()
+    sys.stdout.flush()
+    raise
+
+sys.stdout.flush()
+
+# CONFIG
 CHEXPERT_LABELS = [
     "No Finding",
     "Enlarged Cardiomediastinum",
@@ -36,239 +54,144 @@ CHEXPERT_LABELS = [
     "Support Devices",
 ]
 
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+DEVICE = torch.device("cpu")
 
+# PREPROCESS
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
 
-# ─────────────────────────────────────────────────────────────
-# MODEL LOADING
-# ─────────────────────────────────────────────────────────────
 
 def load_model():
-    """
-    Load a DenseNet-121 model (CheXNet architecture).
-    Uses ImageNet pretrained weights as a demo.
-    """
-    print("Loading DenseNet-121 model...")
+    """Load DenseNet-121."""
+    print("STEP 2: Loading model...")
+    sys.stdout.flush()
     model = timm.create_model("densenet121", pretrained=True, num_classes=14)
     model = model.to(DEVICE)
     model.eval()
-    print(f"Model loaded on {DEVICE}")
+    print("STEP 3: Model loaded OK")
+    sys.stdout.flush()
     return model
 
 
-# ─────────────────────────────────────────────────────────────
-# IMAGE PREPROCESSING
-# ─────────────────────────────────────────────────────────────
-
-preprocess = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],
-        std=[0.229, 0.224, 0.225],
-    ),
-])
-
-
-# ─────────────────────────────────────────────────────────────
-# GRAD-CAM GENERATION
-# ─────────────────────────────────────────────────────────────
-
-def generate_gradcam(model, input_tensor, target_class_idx, target_layers):
-    """Generate a Grad-CAM heatmap for the predicted class."""
-    cam = GradCAM(model=model, target_layers=target_layers)
-    targets = [ClassifierOutputTarget(target_class_idx)]
-    grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
-    return grayscale_cam[0]
-
-
-# ─────────────────────────────────────────────────────────────
-# MAIN INFERENCE FUNCTION
-# ─────────────────────────────────────────────────────────────
-
 def classify_and_explain(image):
-    """Run chest X-ray classification and generate Grad-CAM visualization."""
-    if image is None:
-        return None, "Please upload a chest X-ray image."
-
-    # Preprocess
-    image_rgb = image.convert("RGB")
-    input_tensor = preprocess(image_rgb).unsqueeze(0).to(DEVICE)
-    rgb_display = np.array(image_rgb.resize((224, 224))) / 255.0
-
-    # Predict
-    with torch.no_grad():
-        logits = model(input_tensor)
-        probabilities = torch.sigmoid(logits).cpu().numpy()[0]
-
-    # Build results
-    top_idx = int(np.argmax(probabilities))
-    top_label = CHEXPERT_LABELS[top_idx]
-    top_conf = float(probabilities[top_idx])
-
-    # Sort all conditions by probability
-    sorted_indices = np.argsort(probabilities)[::-1]
-    sorted_labels = [
-        (CHEXPERT_LABELS[i], float(probabilities[i]))
-        for i in sorted_indices
-    ]
-
-    # Grad-CAM
-    grayscale_cam = generate_gradcam(
-        model, input_tensor, top_idx, target_layers
-    )
-    heatmap_overlay = show_cam_on_image(
-        rgb_display, grayscale_cam, use_rgb=True
+    """Run classification and Grad-CAM."""
+    try:
+        if image is None:
+            return None, "Please upload a chest X-ray image."
+
+        print("INFERENCE: Starting...")
+        sys.stdout.flush()
+
+        image_rgb = image.convert("RGB")
+        input_tensor = preprocess(image_rgb).unsqueeze(0).to(DEVICE)
+        rgb_display = np.array(image_rgb.resize((224, 224))).astype(np.float32) / 255.0
+
+        with torch.no_grad():
+            logits = model(input_tensor)
+            probabilities = torch.sigmoid(logits).cpu().numpy()[0]
+
+        top_idx = int(np.argmax(probabilities))
+        top_label = CHEXPERT_LABELS[top_idx]
+        top_conf = float(probabilities[top_idx])
+
+        print(f"INFERENCE: Top prediction = {top_label} ({top_conf:.1%})")
+        sys.stdout.flush()
+
+        # Grad-CAM
+        print("INFERENCE: Computing Grad-CAM...")
+        sys.stdout.flush()
+        cam = GradCAM(model=model, target_layers=target_layers)
+        targets = [ClassifierOutputTarget(top_idx)]
+        grayscale_cam = cam(input_tensor=input_tensor, targets=targets)[0]
+
+        print("INFERENCE: Creating overlay...")
+        sys.stdout.flush()
+        heatmap_overlay = show_cam_on_image(rgb_display, grayscale_cam, use_rgb=True)
+
+        print("INFERENCE: Done!")
+        sys.stdout.flush()
+
+        # Results
+        sorted_indices = np.argsort(probabilities)[::-1]
+        result_md = f"## Top Prediction: **{top_label}** ({top_conf:.1%})\n\n"
+        result_md += "### All Condition Probabilities\n\n"
+        result_md += "| Rank | Condition | Probability | Confidence |\n"
+        result_md += "|:----:|-----------|:-----------:|:----------:|\n"
+
+        for rank, idx in enumerate(sorted_indices, 1):
+            label = CHEXPERT_LABELS[idx]
+            prob = float(probabilities[idx])
+            bar_len = int(prob * 20)
+            bar = "X" * bar_len + "-" * (20 - bar_len)
+            emoji = "HIGH" if prob > 0.5 else "MED" if prob > 0.3 else "LOW"
+            result_md += f"| {rank} | {label} | {prob:.1%} | {emoji} `{bar}` |\n"
+
+        result_md += "\n---\n"
+        result_md += f"**Grad-CAM target:** {top_label} - red regions indicate areas the model focused on.\n"
+
+        return heatmap_overlay, result_md
+
+    except Exception as e:
+        error_msg = f"Error during inference: {e}\n{traceback.format_exc()}"
+        print(error_msg)
+        sys.stdout.flush()
+        return None, f"**Error:** {e}"
+
+
+# BUILD UI
+print("STEP 4: Building Gradio UI...")
+sys.stdout.flush()
+
+with gr.Blocks(
+    title="Chest X-Ray Classifier with Grad-CAM",
+    theme=gr.themes.Soft(),
+) as demo:
+
+    gr.Markdown(
+        """
+        # Chest X-Ray Classifier with Grad-CAM Explainability
+
+        Upload a chest X-ray to classify **14 thoracic conditions** and see
+        **where the model is looking** via Grad-CAM heatmaps.
+
+        > **Disclaimer:** This is an educational tool only. NOT a medical device.
+        """
     )
 
-    # Format markdown output
-    result_md = f"## Top Prediction: **{top_label}** ({top_conf:.1%})\n\n"
-    result_md += "### All Condition Probabilities\n\n"
-    result_md += "| Rank | Condition | Probability | Confidence |\n"
-    result_md += "|:----:|-----------|:-----------:|:----------:|\n"
-
-    for rank, (label, prob) in enumerate(sorted_labels, 1):
-        bar_len = int(prob * 20)
-        bar = "X" * bar_len + "-" * (20 - bar_len)
-        emoji = "HIGH" if prob > 0.5 else "MED" if prob > 0.3 else "LOW"
-        result_md += (
-            f"| {rank} | {label} | {prob:.1%} | {emoji} `{bar}` |\n"
-        )
-
-    result_md += "\n---\n"
-    result_md += f"**Grad-CAM target:** {top_label} - "
-    result_md += "red regions indicate areas the model focused on.\n"
-
-    return heatmap_overlay, result_md
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="pil", label="Upload Chest X-Ray", height=400)
+            classify_btn = gr.Button("Classify and Explain", variant="primary", size="lg")
+        with gr.Column(scale=1):
+            output_image = gr.Image(label="Grad-CAM Heatmap Overlay", height=400)
+            output_text = gr.Markdown(value="*Upload an image and click Classify to see results.*")
 
-
-# ─────────────────────────────────────────────────────────────
-# GRADIO UI
-# ─────────────────────────────────────────────────────────────
-
-def build_ui():
-    """Construct the Gradio Blocks interface."""
-
-    with gr.Blocks(
-        title="Chest X-Ray Classifier with Grad-CAM",
-        theme=gr.themes.Soft(),
-    ) as demo:
-
-        # Header
+    with gr.Accordion("How Grad-CAM Works", open=False):
         gr.Markdown(
             """
-            # Chest X-Ray Classifier with Grad-CAM Explainability
+            **Grad-CAM** highlights which parts of the image the model focused on.
 
-            Upload a chest X-ray to classify **14 thoracic conditions** and see
-            **where the model is looking** via Grad-CAM heatmaps.
+            - **Red** = High attention
+            - **Yellow** = Moderate attention
+            - **Green/Blue** = Low attention
 
-            > **Disclaimer:** This is an educational tool only.
-            > It is NOT a medical device and must NOT be used for clinical diagnosis.
-            > Always consult a qualified physician.
+            Reference: Selvaraju et al., "Grad-CAM", ICCV 2017.
             """
         )
 
-        with gr.Row():
-
-            # Left Column: Input
-            with gr.Column(scale=1):
-                input_image = gr.Image(
-                    type="pil",
-                    label="Upload Chest X-Ray",
-                    height=400,
-                )
-                classify_btn = gr.Button(
-                    "Classify and Explain",
-                    variant="primary",
-                    size="lg",
-                )
-
-                gr.Markdown("### Example Images")
-                gr.Markdown(
-                    "Upload any PA/AP chest X-ray. "
-                    "Best results with standard radiograph formats (JPEG, PNG)."
-                )
-
-            # Right Column: Output
-            with gr.Column(scale=1):
-                output_image = gr.Image(
-                    label="Grad-CAM Heatmap Overlay",
-                    height=400,
-                )
-                output_text = gr.Markdown(
-                    label="Classification Results",
-                    value="*Upload an image and click Classify to see results.*",
-                )
-
-        # How It Works Section
-        with gr.Accordion("How Grad-CAM Works", open=False):
-            gr.Markdown(
-                """
-                **Grad-CAM** (Gradient-weighted Class Activation Mapping) produces
-                a coarse localization map highlighting important regions in the image
-                for predicting a particular class.
-
-                ### The Process:
-                1. **Forward pass** - the image passes through the network
-                2. **Backward pass** - gradients of the target class score flow back
-                   to the final convolutional layer
-                3. **Weighting** - feature maps are weighted by the average gradient
-                4. **ReLU** - only positive contributions are kept
-                5. **Overlay** - the heatmap is superimposed on the original image
-
-                ### Interpreting the Heatmap:
-                - **Red** = High attention (model focused heavily here)
-                - **Yellow** = Moderate attention
-                - **Green/Blue** = Low attention (model largely ignored)
-
-                ### Reference:
-                Selvaraju et al., "Grad-CAM: Visual Explanations from Deep Networks
-                via Gradient-based Localization", ICCV 2017.
-                """
-            )
-
-        # Model Info Section
-        with gr.Accordion("Model Architecture", open=False):
-            gr.Markdown(
-                """
-                | Component | Detail |
-                |-----------|--------|
-                | **Architecture** | DenseNet-121 (CheXNet-style) |
-                | **Parameters** | ~8M |
-                | **Input Size** | 224 x 224 x 3 |
-                | **Output** | 14 condition probabilities (sigmoid) |
-                | **Grad-CAM Layer** | model.features[-1] (final conv block) |
-                | **Pretraining** | ImageNet-1K (demo) |
-
-                ### For Production Use:
-                Fine-tune on a medical imaging dataset:
-                - NIH ChestX-ray14
-                - CheXpert
-                - MIMIC-CXR
-                """
-            )
-
-        # Wire up the button
-        classify_btn.click(
-            fn=classify_and_explain,
-            inputs=[input_image],
-            outputs=[output_image, output_text],
-        )
-
-    return demo
-
+    classify_btn.click(fn=classify_and_explain, inputs=[input_image], outputs=[output_image, output_text])
 
-# ─────────────────────────────────────────────────────────────
-# ENTRY POINT
-# ─────────────────────────────────────────────────────────────
+# LOAD MODEL AND LAUNCH
+print("STEP 5: Loading model at startup...")
+sys.stdout.flush()
 
-if __name__ == "__main__":
-    # Load model at startup
-    model = load_model()
+model = load_model()
+target_layers = [model.features[-1]]
 
-    # Target layer for Grad-CAM
-    target_layers = [model.features[-1]]
+print("STEP 6: Launching server...")
+sys.stdout.flush()
 
-    # Build and launch
-    demo = build_ui()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch(server_name="0.0.0.0", server_port=7860)