Skip to content

Latest commit

Β 

History

14 Commits

Folders and files

NameName
Last commit message
Last commit date
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

TD-Net 🫁

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.


πŸ“‹ Paper Information

The paper is available at https://link.springer.com/chapter/10.1007/978-3-032-13497-4_6


🎯 Key Highlights

  • 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.


πŸ—‚οΈ Datasets

Dataset #1: Kaggle TB Chest X-Ray Dataset

  • 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

Dataset #2: TBX11K 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

βš™οΈ Installation

Prerequisites

  • Python 3.8+
  • CUDA-compatible GPU (recommended: NVIDIA T4 or better)
  • 16GB+ RAM

Setup

  1. Clone the repository:
git clone https://github.com/debg48/TD-Net.git
cd TD-Net
  1. Create a virtual environment (optional but recommended):
python -m venv env
source env/bin/activate  # On Windows: env\Scripts\activate
  1. Install dependencies:
pip install -r requirements.txt

πŸš€ Usage

TD-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.

Basic Training Pipeline

To run the complete pipeline (Downloads the default Kaggle dataset -> Splits data -> Augments -> Trains -> Evaluates):

python main.py

CLI Arguments

You can customize the pipeline using the following flags:

  • --skip-download: Skips dataset downloading (use if your data is already in the data/ 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 within data/ (e.g., --dataset "TBX11K").
  • --gradcam IMAGE: Run Grad-CAM visualization on a specific image path.

Using Custom Datasets

To train TD-Net on your own dataset:

  1. Place your dataset directory inside the data/ folder (e.g., data/my-dataset).
  2. Ensure it has subdirectories for each class (e.g., data/my-dataset/Normal and data/my-dataset/Tuberculosis).
  3. Run the pipeline with your custom dataset name:
python main.py --skip-download --dataset "my-dataset"

Inference & Grad-CAM

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.


πŸ”¬ Technical Details

Data Preprocessing

  • 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

Training Configuration

  • 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)

Attention Mechanisms

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))))

πŸ“ˆ Visualizations

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

πŸ” Grad-CAM Analysis

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

πŸ₯ Clinical Significance

Advantages

  • 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

Limitations

  • May struggle with very noisy or low-quality images
  • Cannot replace confirmatory sputum tests
  • Performance depends on image quality and positioning

🀝 Contributing

Contributions are welcome! Please feel free to:

  1. Fork the repository
  2. Create a feature branch (git checkout -b feature/improvement)
  3. Commit your changes (git commit -m 'Add improvement')
  4. Push to the branch (git push origin feature/improvement)
  5. Open a Pull Request

πŸ“œ Citation

@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"
}

πŸ“„ License

This project is licensed under the MIT License - see the LICENSE file for details.


πŸ™ Acknowledgements

  • 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

πŸ“§ Contact

For questions, collaborations, or issues:


🌟 Star History

If you find this project useful, please consider giving it a ⭐!

Star History Chart


πŸ“š Related Work

Check out our other projects:


Developed with ❀️ at Jadavpur University, Kolkata

About

No description, website, or topics provided.

Resources

Stars

3 stars

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages