M20VJ/tb_classification_files

Clinical-Grade TB Classifier Finetuning

This repo provides a turnkey PyTorch/Transformers pipeline for finetuning facebook/dinov3-convnext-large-pretrain-lvd1689m on chest X-ray DICOM files from indian_tb_xray_scan_dataset. The script automates clinically-motivated best practices: advanced augmentations (MixUp, CutMix, CLAHE, Gaussian noise), focal loss + class weighting, SMOTE-Tomek rebalancing, shortcut auditing, domain adversarial regularization, temperature scaling, MC-dropout uncertainty, TTA, and ensemble evaluation.

Prerequisites

1. Python 3.10+
2. GPU with ≥24 GB VRAM recommended (ConvNeXt-Large is heavy); script auto-falls back to CPU.
3. Install dependencies:

pip install -r requirements.txt

Dataset Layout

Place DICOM studies under indian_tb_xray_scan_dataset using the provided folder split:

indian_tb_xray_scan_dataset/
    abnormal-tuberculosis/
        *.dicom|*.dcm
    normal/
        *.dicom|*.dcm

Optionally provide metadata.csv with filepath,domain columns to expose acquisition-site domains for adversarial alignment.

Quickstart Training

python train_tb_classifier.py \
  --data_root indian_tb_xray_scan_dataset \
  --output_dir runs/baseline \
  --batch_size 8 \
  --epochs 10 \
  --amp

Key switches:

• --mixup_alpha/--cutmix_alpha: enable MixUp/CutMix (set to 0 to disable).
• --smote_tomek: rebalance training split using SMOTE-Tomek on radiograph statistics.
• --audit_shortcuts: dump shortcut_report.json with contrast/brightness deltas per label.
• --use_weighted_sampler --class_weights 1.0 2.0: increase recall on TB-positive cases.
• --tta and --mc_dropout_samples 16: activate test-time augmentation + Monte-Carlo dropout metrics.

Multi-Seed Ensembles

python train_tb_classifier.py \
  --num_models 3 \
  --seed 7 \
  --output_dir runs/ensemble \
  --tta --mc_dropout_samples 16

Each seed trains independently, calibrates via temperature scaling, and the script reports both per-model and ensemble metrics.

MaxViT Variant

Prefer the Hugging Face timm MaxViT checkpoint timm/maxvit_base_tf_512.in21k_ft_in1k? Use the companion script:

python train_tb_classifier_maxvit.py \
  --data_root indian_tb_xray_scan_dataset \
  --output_dir runs/maxvit \
  --batch_size 4 \
  --epochs 10 \
  --amp

It mirrors the Dinov3 pipeline (CLAHE, MixUp/CutMix, focal loss, SMOTE-Tomek, calibration, MC-dropout/TTA) but defaults to 512×512 inputs and pulls normalization metadata from the timm data_config.

Evaluation-Only Mode

python train_tb_classifier.py \
  --eval_only \
  --checkpoint runs/baseline/tb_dinov3_run0.pt \
  --tta --mc_dropout_samples 16

Pass extra checkpoints via --ensemble_checkpoints ckpt1.pt ckpt2.pt to fuse heterogeneous backbones.

Optimized Inference Script

Once training finishes, run calibrated inference with the dedicated helper (automatically mirrors CLAHE, normalization, and optional TTA/MC dropout):

python inference_tb_classifier.py \
  --checkpoint runs/baseline/tb_dinov3_run0.pt \
  --temperature_path runs/baseline/temperature_run1.pt \
  --data_root indian_tb_xray_scan_dataset \
  --output_dir runs/baseline/inference \
  --batch_size 64 \
  --amp --tta

Set --mc_dropout_samples 8 to retrieve epistemic uncertainty estimates, and add --save_predictions --save_logits to dump CSV/tensor artifacts for downstream audit.

Outputs

• metrics.json: full validation/test stats (accuracy, F1, AUC, ECE, MC-dropout uncertainty, etc.).
• temperature_run*.pt: calibrated temperature parameters per model.
• history_run*.json: per-epoch logs.
• Optional logits.pt (set --save_logits).

Implementation Notes

• Backbone: Hugging Face AutoModel loads Dinov3 ConvNeXt; AutoTokenizer is initialized for completeness even though this is a vision backbone.
• Loss: focal loss (γ configurable) with optional class weights; MixUp/CutMix batches fall back to a soft-label CE.
• Balancing: SMOTE-Tomek operates on lightweight intensity/texture embeddings and duplicates nearest neighbours to approximate synthetic samples, helping mitigate shortcut reliance.
• Domain Adaptation: when ≥2 domains are observed, a gradient-reversal domain discriminator jointly trains with the classifier.
• Calibration & Uncertainty: temperature scaling is fitted on the validation logits; TTA+MC dropout track calibration/epistemic uncertainty for expert triage.

Run python train_tb_classifier.py --help to inspect every flag.