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" ] # Label Descriptions for Report label_descriptions = { "No Finding": "No significant cardiopulmonary abnormality is identified.", "Enlarged Cardiomediastinum": "The cardiomediastinal silhouette appears enlarged, which may reflect cardiac or mediastinal pathology.", "Cardiomegaly": "The cardiac silhouette is enlarged, which may be seen in a variety of cardiac conditions including cardiomyopathy or volume overload.", "Lung Opacity": "There are areas of increased lung opacity, which may represent infection, inflammation, or other parenchymal processes.", "Lung Lesion": "There is a focal abnormality in the lung that may represent an underlying lesion and may warrant further evaluation.", "Edema": "The pulmonary parenchyma demonstrates changes that may represent pulmonary edema.", "Consolidation": "There is focal or multifocal consolidation compatible with alveolar filling, such as infection or aspiration.", "Pneumonia": "The pattern of opacities is suspicious for pneumonia in the appropriate clinical context.", "Atelectasis": "There is volume loss with increased opacity, which may represent atelectasis.", "Pneumothorax": "There is suspicion for pneumothorax, which represents air within the pleural space and may be clinically significant.", "Pleural Effusion": "There is fluid in the pleural space, which may compress the adjacent lung parenchyma.", "Pleural Other": "There are pleural abnormalities that may represent pleural thickening, plaques, or other pleural processes.", "Fracture": "There is suspicion of osseous fracture, which may require correlation with dedicated imaging and clinical findings.", "Support Devices": "Support devices are present (e.g. lines, tubes, pacemaker leads) which should be correlated with position and clinical need.", } LABEL_THRESHOLDS = { "No Finding": 0.5, "Cardiomegaly": 0.6, "Pneumothorax": 0.6, "Pleural Effusion": 0.5, "Fracture": 0.6 } # 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() if device.type=='cuda': print(f"Model loaded successfully on GPU {torch.cuda.get_device_name(torch.cuda.current_device())}") else: print(f"Model loaded successfully on CPU") # Report Generation Functions def prob_to_phrase(p: float) -> str: if p >= 0.8: return "highly suggestive of" elif p >= 0.6: return "likely" else: return "may represent" def rule_based_labeling(probs, default_threshold: float = 0.5): if len(probs) != len(LABELS): raise ValueError(f"Expected {len(LABELS)} probabilities, got {len(probs)}") selected = [] for i, prob in enumerate(probs): label = LABELS[i] th = LABEL_THRESHOLDS.get(label, default_threshold) if prob >= th: selected.append((i, prob)) return selected def handle_no_finding(selected): label_names = [LABELS[i] for i, _ in selected] if "No Finding" in label_names and len(label_names) > 1: selected = [(i, p) for (i, p) in selected if LABELS[i] != "No Finding"] return selected def remove_redundant_labels(selected): name_to_prob = {LABELS[i]: p for i, p in selected} if "Pneumonia" in name_to_prob and "Lung Opacity" in name_to_prob: selected = [(i, p) for (i, p) in selected if LABELS[i] != "Lung Opacity"] name_to_prob = {LABELS[i]: p for i, p in selected} if "Consolidation" in name_to_prob and "Lung Opacity" in name_to_prob: selected = [(i, p) for (i, p) in selected if LABELS[i] != "Lung Opacity"] name_to_prob = {LABELS[i]: p for i, p in selected} if "Pleural Effusion" in name_to_prob and "Pleural Other" in name_to_prob: selected = [(i, p) for (i, p) in selected if LABELS[i] != "Pleural Other"] return selected def build_impression_from_labels(selected): name_to_prob = {LABELS[i]: p for i, p in selected} lines = [] has_edema = "Edema" in name_to_prob has_peff = "Pleural Effusion" in name_to_prob has_consolidation = "Consolidation" in name_to_prob has_pneumonia = "Pneumonia" in name_to_prob has_atelectasis = "Atelectasis" in name_to_prob if has_edema and has_peff: lines.append("Pattern consistent with pulmonary edema with associated pleural effusions.") elif has_edema: lines.append("Pattern consistent with pulmonary edema.") elif has_peff: lines.append("Pleural effusion is suspected, which may be clinically significant.") if has_pneumonia and has_atelectasis: lines.append("Focal pulmonary opacity suspicious for pneumonia, atelectasis remains a differential consideration.") elif has_pneumonia or has_consolidation: lines.append("Focal pulmonary opacity is suspicious for pneumonia in the appropriate clinical context.") elif has_atelectasis: lines.append("Areas of volume loss may represent atelectasis.") if "Cardiomegaly" in name_to_prob: lines.append("Cardiac silhouette appears enlarged, correlate clinically for cardiomegaly.") if "Support Devices" in name_to_prob: lines.append("Support devices/tubes are present, correlate with clinical indication and positioning.") if not lines: for i, p in selected: label = LABELS[i] phrase = prob_to_phrase(p) lines.append(f"{phrase} {label.lower()}.") return "Impression:\n- " + "\n- ".join(lines) def generate_textual_report(probs, default_threshold: float = 0.5, top_k: int = None) -> str: selected = rule_based_labeling(probs, default_threshold) if not selected: return ( "Findings:\n" "No significant cardiopulmonary abnormality is identified by the model.\n\n" "Impression:\n" "No acute cardiopulmonary process detected by the model." ) selected = handle_no_finding(selected) selected = remove_redundant_labels(selected) selected.sort(key=lambda x: x[1], reverse=True) if top_k is not None: selected = selected[:top_k] findings_lines = [] for idx, prob in selected: label = LABELS[idx] description = label_descriptions.get(label, "") phrase = prob_to_phrase(prob) prob_pct = int(round(prob * 100)) findings_lines.append(f"- {label}: {description}.") findings_text = "Findings:\n" + "\n".join(findings_lines) impression_text = build_impression_from_labels(selected) return findings_text + "\n\n" + impression_text def predict(image, threshold): """Generate predictions, Grad-CAM visualizations, and report""" 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)]) # Generate textual report report = generate_textual_report(probs, default_threshold=0.5, top_k=5) if detected_conditions: summary = f"## Detected Conditions (>{threshold}):\n" + "\n".join(detected_conditions) summary += f"\n\n## All Predictions:\n{all_predictions}" return gradcam_images[0], img, summary, report else: summary = f"No conditions detected above threshold {threshold}\n\n## All Predictions:\n{all_predictions}" return None, img, summary, report 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 with Report Generation 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") # Report Section with gr.Row(): with gr.Column(): gr.Markdown("## 📋 Generated Report") output_report = gr.Textbox( label="Clinical Report", lines=12, max_lines=20, show_copy_button=True ) download_btn = gr.DownloadButton( label="📥 Download Report", visible=True ) # Instructions 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 5. Review the generated clinical report 6. Download the report as a text file if needed """ ) # Connect components def analyze_and_prepare_download(image, threshold): gradcam, original, summary, report = predict(image, threshold) # Prepare file for download if report: report_file = "xray_report.txt" with open(report_file, "w") as f: f.write(report) return gradcam, original, summary, report, gr.DownloadButton(value=report_file, visible=True) else: return gradcam, original, summary, report, gr.DownloadButton(visible=False) analyze_btn.click( fn=analyze_and_prepare_download, inputs=[input_image, threshold], outputs=[output_gradcam, output_image, output_text, output_report, download_btn] ) if __name__ == "__main__": demo.launch()