Create app.py

app.py

@@ -0,0 +1,160 @@
+import gradio as gr
+import torch
+from torchvision import models, transforms
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
+from PIL import Image
+import numpy as np
+from skimage.transform import resize
+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
+
+# Constants
+REPO_ID = "itsomk/chexpert-densenet121"
+FILENAME = "pytorch_model.safetensors"
+
+# Model Definition
+class DenseNet121_CheXpert(torch.nn.Module):
+    def __init__(self, num_labels=14, pretrained=None):
+        super().__init__()
+        self.densenet = models.densenet121(weights=pretrained)
+        num_features = self.densenet.classifier.in_features
+        self.densenet.classifier = torch.nn.Linear(num_features, num_labels)
+    
+    def forward(self, x):
+        return self.densenet(x)
+
+# Labels
+LABELS = [
+    "No Finding", "Enlarged Cardiomediastinum", "Cardiomegaly", "Lung Opacity",
+    "Lung Lesion", "Edema", "Consolidation", "Pneumonia", "Atelectasis",
+    "Pneumothorax", "Pleural Effusion", "Pleural Other", "Fracture", "Support Devices"
+]
+
+# 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])
+])
+
+# Load model
+print("Loading model...")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+local_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+state = load_file(local_path)
+model = DenseNet121_CheXpert(num_labels=14, pretrained=None)
+model.load_state_dict(state, strict=False)
+model.to(device)
+model.eval()
+print(f"Model loaded successfully on {device}")
+
+def predict(image, threshold):
+    """Generate predictions and Grad-CAM visualizations"""
+    if image is None:
+        return None, None, "Please upload an X-ray image"
+    
+    try:
+        # Convert to PIL Image
+        if isinstance(image, np.ndarray):
+            img = Image.fromarray(image).convert("RGB")
+        else:
+            img = image.convert("RGB")
+        
+        # Preprocess
+        img_tensor = preprocess(img).unsqueeze(0).to(device)
+        rgb_img = np.array(img.resize((224, 224)), dtype=np.float32) / 255.0
+        
+        # Get predictions
+        with torch.no_grad():
+            logits = model(img_tensor)
+            probs = torch.sigmoid(logits).squeeze().cpu().numpy()
+        
+        # Setup Grad-CAM
+        target_layer = model.densenet.features.denseblock4
+        cam = GradCAM(model=model, target_layers=[target_layer])
+        
+        # Generate visualizations for conditions above threshold
+        gradcam_images = []
+        detected_conditions = []
+        
+        for i, prob in enumerate(probs):
+            if prob > threshold:
+                label = LABELS[i]
+                targets = [ClassifierOutputTarget(i)]
+                grayscale_cam = cam(input_tensor=img_tensor, targets=targets)
+                grayscale_cam = grayscale_cam[0, :]
+                
+                resized_rgb_img = resize(rgb_img, grayscale_cam.shape, anti_aliasing=True)
+                cam_image = show_cam_on_image(resized_rgb_img, grayscale_cam, use_rgb=True)
+                
+                gradcam_images.append(cam_image)
+                detected_conditions.append(f"**{label}**: {prob:.4f}")
+        
+        # Create summary text
+        all_predictions = "\n".join([f"{LABELS[i]}: {prob:.4f}" for i, prob in enumerate(probs)])
+        
+        if detected_conditions:
+            summary = f"## Detected Conditions (>{threshold}):\n" + "\n".join(detected_conditions)
+            summary += f"\n\n## All Predictions:\n{all_predictions}"
+            # Return first Grad-CAM image and original image
+            return gradcam_images[0], img, summary
+        else:
+            summary = f"No conditions detected above threshold {threshold}\n\n## All Predictions:\n{all_predictions}"
+            return None, img, summary
+            
+    except Exception as e:
+        return None, None, f"Error: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🩻 X-Ray Grad-CAM Visualization
+        
+        Upload a chest X-ray image to analyze potential conditions using DenseNet121 with Grad-CAM visualization.
+        
+        **Model**: [itsomk/chexpert-densenet121](https://huggingface.co/itsomk/chexpert-densenet121)
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="Upload X-Ray Image", type="pil")
+            threshold = gr.Slider(
+                minimum=0.0, 
+                maximum=1.0, 
+                value=0.5, 
+                step=0.05, 
+                label="Prediction Threshold"
+            )
+            analyze_btn = gr.Button("🔍 Analyze X-Ray", variant="primary", size="lg")
+        
+        with gr.Column():
+            output_gradcam = gr.Image(label="Grad-CAM Visualization")
+            output_image = gr.Image(label="Original Image")
+    
+    with gr.Row():
+        output_text = gr.Markdown(label="Analysis Results")
+    
+    # Examples
+    gr.Markdown("### 📋 Instructions:")
+    gr.Markdown(
+        """
+        1. Upload a chest X-ray image (JPG, PNG)
+        2. Adjust the prediction threshold if needed (default: 0.5)
+        3. Click 'Analyze X-Ray' to see results
+        4. View detected conditions with Grad-CAM heatmaps
+        """
+    )
+    
+    # Connect components
+    analyze_btn.click(
+        fn=predict,
+        inputs=[input_image, threshold],
+        outputs=[output_gradcam, output_image, output_text]
+    )
+
+if __name__ == "__main__":
+    demo.launch()