Maisamilens/tb-detection-mamba-vit

TB classification and segmentation using novel hybridization of Vision Tranformer and Mamba based state space Model # RetinexFormer-Enhanced Mamba-ViT for Pulmonary Tuberculosis Detection, classification and segmentation A hybrid deep learning framework combining RetinexFormer image enhancement with Mamba-ViT dual-encoder architecture for automated tuberculosis classification and lung segmentation from chest X-rays. --- ## Overview This repository implements a multi-task deep learning model that simultaneously performs: - **Binary Classification**: TB-positive vs. Normal chest radiographs - **Lung Segmentation**: Precise delineation of lung fields ### Key Features - **RetinexFormer Enhancement**: Learnable illumination normalization for robust preprocessing across varying image quality - **Dual-Encoder Architecture**: Mamba-style CNN for local features + Vision Transformer for global context - **Cross-Attention Fusion**: Effective integration of complementary feature representations - **Uncertainty-Weighted Loss**: Automatic task balancing without manual hyperparameter tuning - **Comprehensive Visualization**: Training curves, GradCAM, predictions, and evaluation metrics --- ## Performance | Task | Metric | Value | |------|--------|-------| | Classification | Accuracy | 99.64% | | Classification | AUC-ROC | 0.9999 | | Classification | Sensitivity | 97.86% | | Classification | Specificity | 100% | | Segmentation | Dice Coefficient | 0.962 | | Segmentation | IoU | 0.928 | --- ## Installation ```bash # Clone repository git clone github.com cd tb-detection-mamba-vit # Create virtual environment python -m venv venv source venv/bin/activate # Linux/Mac # or venv\Scripts\activate # Windows # Install dependencies pip install -r requirements.txt ``` ### Requirements ``` torch>=2.0.0 torchvision>=0.15.0 numpy>=1.21.0 pandas>=1.3.0 scikit-learn>=1.0.0 scikit-image>=0.19.0 Pillow>=9.0.0 matplotlib>=3.5.0 seaborn>=0.11.0 tqdm>=4.62.0 scipy>=1.7.0 ``` --- ## Dataset Structure ``` project_root/ ├── classification/ │ └── T …