Spaces:
Sleeping
Sleeping
| """ | |
| Chest X-Ray Classifier with Grad-CAM Explainability | |
| """ | |
| 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", | |
| "Cardiomegaly", | |
| "Lung Opacity", | |
| "Lung Lesion", | |
| "Edema", | |
| "Consolidation", | |
| "Pneumonia", | |
| "Atelectasis", | |
| "Pneumothorax", | |
| "Pleural Effusion", | |
| "Pleural Other", | |
| "Fracture", | |
| "Support Devices", | |
| ] | |
| 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]), | |
| ]) | |
| def load_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("STEP 3: Model loaded OK") | |
| sys.stdout.flush() | |
| return model | |
| def classify_and_explain(image): | |
| """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. | |
| """ | |
| ) | |
| 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.*") | |
| with gr.Accordion("How Grad-CAM Works", open=False): | |
| gr.Markdown( | |
| """ | |
| **Grad-CAM** highlights which parts of the image the model focused on. | |
| - **Red** = High attention | |
| - **Yellow** = Moderate attention | |
| - **Green/Blue** = Low attention | |
| Reference: Selvaraju et al., "Grad-CAM", ICCV 2017. | |
| """ | |
| ) | |
| classify_btn.click(fn=classify_and_explain, inputs=[input_image], outputs=[output_image, output_text]) | |
| # LOAD MODEL AND LAUNCH | |
| print("STEP 5: Loading model at startup...") | |
| sys.stdout.flush() | |
| model = load_model() | |
| target_layers = [model.features[-1]] | |
| print("STEP 6: Launching server...") | |
| sys.stdout.flush() | |
| demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False) | |