Spaces:
Sleeping
Sleeping
| """Chest pathology detection using TorchXRayVision DenseNet121. | |
| This model is pretrained on 700K+ clinical X-ray images and outputs | |
| probabilities for 18 different chest pathologies. | |
| """ | |
| from __future__ import annotations | |
| import numpy as np | |
| import torch | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| # Lazy-loaded model singleton | |
| _model = None | |
| _pathology_names: list[str] = [] | |
| def _get_model(): | |
| """Load the DenseNet121 model (lazy singleton).""" | |
| global _model, _pathology_names | |
| if _model is None: | |
| logger.info("Loading TorchXRayVision DenseNet121 model...") | |
| try: | |
| import torchxrayvision as xrv | |
| _model = xrv.models.DenseNet(weights="densenet121-res224-all") | |
| _model.eval() | |
| _pathology_names = list(_model.pathologies) | |
| logger.info(f"DenseNet121 loaded. Pathologies: {_pathology_names}") | |
| except Exception as e: | |
| logger.error(f"Failed to load DenseNet121: {e}") | |
| raise | |
| return _model | |
| def predict_chest_pathologies(preprocessed_image: np.ndarray, | |
| confidence_threshold: float = 0.40) -> list[dict]: | |
| """Run chest pathology inference. | |
| Args: | |
| preprocessed_image: numpy array of shape (1, 1, 224, 224) | |
| normalized to [-1024, 1024] range. | |
| confidence_threshold: minimum probability to include in results. | |
| Returns: | |
| List of findings, each with name, confidence, severity, model, region. | |
| """ | |
| model = _get_model() | |
| with torch.no_grad(): | |
| tensor_input = torch.from_numpy(preprocessed_image).float() | |
| outputs = model(tensor_input) | |
| probabilities = outputs.cpu().numpy()[0] | |
| findings = [] | |
| for i, (name, prob) in enumerate(zip(_pathology_names, probabilities)): | |
| confidence = float(prob) * 100 # Convert to percentage | |
| if confidence < confidence_threshold * 100: | |
| continue | |
| # Determine severity based on confidence | |
| if confidence >= 80: | |
| severity = "high" | |
| color = "destructive" | |
| elif confidence >= 60: | |
| severity = "moderate" | |
| color = "warning" | |
| else: | |
| severity = "low" | |
| color = "info" | |
| # Map pathology to anatomical region | |
| region = _get_region(name) | |
| icd_code = _get_icd_code(name) | |
| findings.append({ | |
| "name": name, | |
| "confidence": round(confidence, 1), | |
| "severity": severity, | |
| "model": "DenseNet121", | |
| "region": region, | |
| "icd_code": icd_code, | |
| "color": color, | |
| }) | |
| # Sort by confidence descending | |
| findings.sort(key=lambda x: x["confidence"], reverse=True) | |
| return findings | |
| def _get_region(pathology: str) -> str: | |
| """Map pathology name to anatomical region.""" | |
| region_map = { | |
| "Atelectasis": "Lung parenchyma", | |
| "Cardiomegaly": "Mediastinum", | |
| "Consolidation": "Lung parenchyma", | |
| "Edema": "Bilateral", | |
| "Effusion": "Pleural space", | |
| "Emphysema": "Lung parenchyma", | |
| "Fibrosis": "Lung parenchyma", | |
| "Hernia": "Diaphragm", | |
| "Infiltration": "Lung parenchyma", | |
| "Mass": "Lung parenchyma", | |
| "Nodule": "Lung parenchyma", | |
| "Pleural_Thickening": "Pleural space", | |
| "Pneumonia": "Lung parenchyma", | |
| "Pneumothorax": "Pleural space", | |
| "Enlarged Cardiomediastinum": "Mediastinum", | |
| "Lung Opacity": "Lung parenchyma", | |
| "Lung Lesion": "Lung parenchyma", | |
| "Fracture": "Rib cage", | |
| } | |
| return region_map.get(pathology, "Chest") | |
| def _get_icd_code(pathology: str) -> str: | |
| """Map pathology name to ICD-10 code.""" | |
| icd_map = { | |
| "Atelectasis": "J98.1", | |
| "Cardiomegaly": "I51.7", | |
| "Consolidation": "J18.9", | |
| "Edema": "J81", | |
| "Effusion": "J90", | |
| "Emphysema": "J43.9", | |
| "Fibrosis": "J84.1", | |
| "Hernia": "K44.9", | |
| "Infiltration": "R91.8", | |
| "Mass": "R91.1", | |
| "Nodule": "R91.1", | |
| "Pleural_Thickening": "J92.9", | |
| "Pneumonia": "J18.9", | |
| "Pneumothorax": "J93.9", | |
| "Enlarged Cardiomediastinum": "R93.1", | |
| "Lung Opacity": "R91.8", | |
| "Lung Lesion": "R91.1", | |
| "Fracture": "S22.9", | |
| } | |
| return icd_map.get(pathology, "R93.1") | |