This repository contains the official implementation of TD-Net, a novel deep learning architecture that combines MobileNetV2 with spatial and channel attention mechanisms for efficient and accurate tuberculosis (TB) detection from chest X-ray images.
The methodology achieves state-of-the-art performance with 98% accuracy and 0.99 AUC-ROC on benchmark datasets while maintaining a lightweight architecture suitable for resource-constrained environments.
The paper is available at https://link.springer.com/chapter/10.1007/978-3-032-13497-4_6
- 98% accuracy and 0.99 AUC-ROC on Kaggle TB dataset
- 97% accuracy and 0.98 AUC-ROC on TBX11K dataset
- Lightweight architecture based on MobileNetV2 (only 12.41 MB)
- Fast inference time averaging 1.76 seconds per image
- Attention mechanisms for improved feature extraction and interpretability
- Grad-CAM visualization for model explainability
Note
Since the implementation was revamped, some experimental variance was introduced. Results obtained using random seed 42 are the closest to the performance metrics reported in the paper.
- Source: Curated by Tawsifur Rahman et al.
- Total Images: 4,200 chest X-rays
- Classes:
- Normal: 3,500 images
- TB Positive: 700 images
- Image Format: PNG (512Γ512 pixels)
- Download: Kaggle Dataset
- Total Images: 4,600 chest X-rays (for binary classification)
- Classes:
- Healthy: 3,800 images
- TB Positive: 800 images
- Image Format: PNG (512Γ512 pixels)
- Download: Kaggle TBX11
- Python 3.8+
- CUDA-compatible GPU (recommended: NVIDIA T4 or better)
- 16GB+ RAM
- Clone the repository:
git clone https://github.com/debg48/TD-Net.git
cd TD-Net- Create a virtual environment (optional but recommended):
python -m venv env
source env/bin/activate # On Windows: env\Scripts\activate- Install dependencies:
pip install -r requirements.txtTD-Net has been completely refactored into a modular, production-ready Python package. You can run the entire pipelineβfrom downloading the data to training and evaluationβusing the main.py CLI.
To run the complete pipeline (Downloads the default Kaggle dataset -> Splits data -> Augments -> Trains -> Evaluates):
python main.pyYou can customize the pipeline using the following flags:
--skip-download: Skips dataset downloading (use if your data is already in thedata/folder).--skip-augment: Skips the offline dataset augmentation step.--no-online-augment: Disables heavy on-the-fly image augmentation during training.--evaluate-only: Skips training and evaluates an existing saved model.--dataset DATASET: Specify a custom dataset directory name withindata/(e.g.,--dataset "TBX11K").--gradcam IMAGE: Run Grad-CAM visualization on a specific image path.
To train TD-Net on your own dataset:
- Place your dataset directory inside the
data/folder (e.g.,data/my-dataset). - Ensure it has subdirectories for each class (e.g.,
data/my-dataset/Normalanddata/my-dataset/Tuberculosis). - Run the pipeline with your custom dataset name:
python main.py --skip-download --dataset "my-dataset"To visualize what the model is focusing on for a specific X-ray image:
python main.py --gradcam "path/to/chest_xray.png"The resulting heatmap will be saved in the results/ folder, showing the attention mechanisms at work.
- Image resizing to 224Γ224Γ3
- Normalization (rescaling to [0, 1])
- Train/Val/Test split: 60%/10%/30%
- Image augmentation: rotation, flip, shift, zoom, shear
- Optimizer: RMSprop (learning rate: 0.0001)
- Loss Function: Weighted Binary Cross-Entropy
- Class Weights: {Normal: 0.94, TB: 1.07}
- Regularization: L1 (strength: 0.01), Dropout (0.4)
- Early Stopping: Patience of 7 epochs on validation AUC-ROC
- Hardware: NVIDIA T4 GPU (15GB VRAM)
Channel Attention:
CA(F) = F Β· Ο(W_avg(AvgPool(F)) + W_max(MaxPool(F)))
Spatial Attention:
SA(F) = F Β· Ο(Conv_7Γ7(Concat(AvgPool(F), MaxPool(F))))
The repository includes:
- Training curves: Accuracy, loss, F1-score, AUC-ROC over epochs
- Confusion matrices: For both datasets
- Grad-CAM heatmaps: Showing model attention on X-ray images
- Misclassified samples: Analysis of failure cases
TD-Net uses Gradient-weighted Class Activation Mapping (Grad-CAM) for model interpretability:
- Before attention modules: Model fails to focus on relevant lung regions
- After attention modules: Model correctly highlights TB-affected areas
- Clinical relevance: Enables doctors to verify model decisions
- Non-invasive: Uses standard chest X-rays
- Fast screening: Real-time inference for quick diagnosis
- Resource-efficient: Deployable in low-resource settings
- Explainable: Grad-CAM provides visual explanations
- High sensitivity: 99% recall on TB-positive cases
- May struggle with very noisy or low-quality images
- Cannot replace confirmatory sputum tests
- Performance depends on image quality and positioning
Contributions are welcome! Please feel free to:
- Fork the repository
- Create a feature branch (
git checkout -b feature/improvement) - Commit your changes (
git commit -m 'Add improvement') - Push to the branch (
git push origin feature/improvement) - Open a Pull Request
@Inbook{Ghosh2026,
author="Ghosh, Debgandhar
and Singh, Pawan Kumar",
editor="Rivera, Gilberto
and Pedrycz, Witold
and Labella, {\'A}lvaro
and Florencia, Rogelio",
title="TD-Net: A Deep Learning Technique with Spatial and Channel Attention for Efficient Tuberculosis Detection from Chest X-Ray Scans",
bookTitle="Recent Advances in Artificial Neural Networks: Case Studies and Applications",
year="2026",
publisher="Springer Nature Switzerland",
address="Cham",
pages="161--192",
abstract="The disease, Tuberculosis (TB) is a bacterial infection that results from Mycobacterium tuberculosis and can be fatal if not recognized early. Chest X-rays are used to identify TB, thus requiring close supervision from radiologists and doctors. Tuberculosis (TB) can be discovered and recognized through chest X-rays employing deep learning methodologies, leveraging the capabilities of Convolutional Neural Networks (CNNs) to differentiate between X-rays of individuals with TB and those that are normal. In this work, we developed a new deep learning methodology, named as TD-Net, incorporating pre-trained MobileNetV2 architecture with spatial and channel attention to elevate the accuracy, F1 Score, Recall and AUC-ROC of MobileNetV2. The experiment shows that the proposed TD-Net methodology gains a test accuracy of 98 and 97{\%} respectively and an AUC-ROC score of 0.99 and 0.98 respectively on two benchmark TB datasets namely, a publicly available dataset curated by Twasifur Rahman et al. and TBX11 dataset.",
isbn="978-3-032-13497-4",
doi="10.1007/978-3-032-13497-4_6",
url="https://doi.org/10.1007/978-3-032-13497-4_6"
}
This project is licensed under the MIT License - see the LICENSE file for details.
- Ms. Ongira Badyopadhyay (Medical College Kolkata) for medical validation and guidance
- Tawsifur Rahman et al. for the Kaggle TB dataset
- TBX11K dataset creators for the benchmark dataset
For questions, collaborations, or issues:
- Debgandhar Ghosh: debgandhar4000@gmail.com
- Pawan Kumar Singh: pawansingh.ju@gmail.com
If you find this project useful, please consider giving it a β!
Check out our other projects:
Developed with β€οΈ at Jadavpur University, Kolkata