JetX-GT/hades-hellix-tb-linear-probe

Hades Hellix TB Linear Probe v4

Calibrated Linear Classification Head for Tuberculosis Screening

A 2-layer MLP trained on MedSigLIP embeddings for WHO-compliant TB screening from chest X-rays.

Model Details

• Model Type: Linear Probe (2-layer MLP)
• Input: MedSigLIP-448 embeddings (1152-dim)
• Architecture: Linear(1152, 512) → ReLU → Dropout(0.3) → Linear(512, 1)
• Output: Calibrated TB probability (0-1)
• Base Model: google/medsiglip-448

Files

File	Size	Description
best_tb_model_v4.pth	~2.4 MB	Trained linear probe weights
platt_calibrator.pkl	~1 KB	Platt scaling probability calibrator
config.json	-	Model configuration

Usage

import torch
import torch.nn as nn
import pickle
from transformers import AutoModel, SiglipImageProcessor
from PIL import Image

# Load model config
class TBLinearProbe(nn.Module):
    def __init__(self):
        super().__init__()
        self.classifier = nn.Sequential(
            nn.Linear(1152, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 1)
        )
    def forward(self, x):
        return self.classifier(x).squeeze(-1)

# Load linear probe
probe = TBLinearProbe()
ckpt = torch.load("best_tb_model_v4.pth", map_location="cpu")
probe.load_state_dict(ckpt["model_state_dict"])
probe.eval()

# Load calibrator
with open("platt_calibrator.pkl", "rb") as f:
    calibrator = pickle.load(f)

# Load MedSigLIP (download from HuggingFace)
model = AutoModel.from_pretrained("google/medsiglip-448")
processor = SiglipImageProcessor.from_pretrained("google/medsiglip-448")

# Extract features
img = Image.open("chest_xray.png").convert("RGB")
inputs = processor(images=img, return_tensors="pt")
with torch.no_grad():
    embedding = model.get_image_features(**inputs)

# Predict
with torch.no_grad():
    logit = probe(embedding).item()
    calibrated_prob = calibrator.predict_proba([[logit]])[0, 1]

print(f"TB Probability: {calibrated_prob:.3f}")

Thresholds

Threshold	Category
< 0.15	Confirmed Normal
0.15 - 0.45	Gray Zone (review recommended)
> 0.45	High Probability TB

WHO Triage Mapping

Probability	Priority	Action
> 0.90	P1-RED	Immediate (< 24h)
0.70 - 0.90	P2-YELLOW	Urgent (24-48h)
0.40 - 0.70	P3-AMBER	Standard
< 0.40	P4-GREEN	Routine

Training Datasets

NOT INCLUDED - Download from original sources:

Dataset	Source
ICMR TB Portal	ICMR
TBX11K	GitHub
Kaggle TB Chest X-ray	Kaggle
NIH Montgomery	NIH LHNCBC

Preprocessing

• CLAHE: clipLimit=2.0, tileGridSize=(8, 8)
• Resize: 448 × 448 (INTER_AREA)
• Z-Score: Per-image normalization

Citation

@software{hades_hellix_linear_probe_2026,
  title={Hades Hellix TB Linear Probe v4},
  author={Hades Hellix Team},
  year={2026},
  note={Calibrated classification head for MedSigLIP-based TB screening}
}

License

MIT License

Disclaimer

FOR RESEARCH ONLY - Not approved for clinical use. Consult medical professionals for diagnosis.