A 2-layer MLP trained on MedSigLIP embeddings for WHO-compliant TB screening from chest X-rays.
1import torch
2import torch.nn as nn
3import pickle
4from transformers import AutoModel, SiglipImageProcessor
5from PIL import Image
6
7# Load model config
8class TBLinearProbe(nn.Module):
9 def __init__(self):
10 super().__init__()
11 self.classifier = nn.Sequential(
12 nn.Linear(1152, 512),
13 nn.ReLU(),
14 nn.Dropout(0.3),
15 nn.Linear(512, 1)
16 )
17 def forward(self, x):
18 return self.classifier(x).squeeze(-1)
19
20# Load linear probe
21probe = TBLinearProbe()
22ckpt = torch.load("best_tb_model_v4.pth", map_location="cpu")
23probe.load_state_dict(ckpt["model_state_dict"])
24probe.eval()
25
26# Load calibrator
27with open("platt_calibrator.pkl", "rb") as f:
28 calibrator = pickle.load(f)
29
30# Load MedSigLIP (download from HuggingFace)
31model = AutoModel.from_pretrained("google/medsiglip-448")
32processor = SiglipImageProcessor.from_pretrained("google/medsiglip-448")
33
34# Extract features
35img = Image.open("chest_xray.png").convert("RGB")
36inputs = processor(images=img, return_tensors="pt")
37with torch.no_grad():
38 embedding = model.get_image_features(**inputs)
39
40# Predict
41with torch.no_grad():
42 logit = probe(embedding).item()
43 calibrated_prob = calibrator.predict_proba([[logit]])[0, 1]
44
45print(f"TB Probability: {calibrated_prob:.3f}")
1@software{hades_hellix_linear_probe_2026,
2 title={Hades Hellix TB Linear Probe v4},
3 author={Hades Hellix Team},
4 year={2026},
5 note={Calibrated classification head for MedSigLIP-based TB screening}
6}
