-10b981){kind=icon}
Live Demo · API Docs · Model Weights · Changelog
ThoraxNet is a production-grade, full-stack AI diagnostic platform that detects 14 thoracic pathologies from chest X-rays. It combines a fine-tuned BioMedCLIP vision-language foundation model with Monte Carlo Dropout uncertainty quantification, ViT-GradCAM explainability, and automated radiology report generation via Groq's LLaMA-3.3-70b.
For research use only. Not FDA cleared. Not a substitute for clinical radiologist interpretation.
Evaluated on the NIH ChestX-ray14 official validation split (224×224 resolution, per-class threshold tuning).
| Pathology | AUC | Threshold |
|---|---|---|
| Cardiomegaly | 0.888 | 0.74 |
| Hernia | 0.872 | 0.62 |
| Edema | 0.851 | 0.75 |
| Effusion | 0.834 | 0.66 |
| Emphysema | 0.823 | 0.61 |
| Pneumothorax | 0.793 | 0.66 |
| Fibrosis | 0.782 | 0.60 |
| Mass | 0.776 | 0.64 |
| Nodule | 0.754 | 0.58 |
| Atelectasis | 0.745 | 0.63 |
| Consolidation | 0.736 | 0.67 |
| Pleural Thickening | 0.728 | 0.61 |
| Infiltration | 0.704 | 0.58 |
| Pneumonia | 0.695 | 0.67 |
| Mean | 0.8215 | — |
Compared to the original NIH paper (Wang et al., 2017) mean AUC of 0.745, ThoraxNet achieves a +7.65% absolute improvement by leveraging BioMedCLIP's medical vision-language pretraining on 15 million biomedical image-text pairs.
- Multi-label classification — detects 14 pathologies simultaneously with per-class calibrated thresholds
- Uncertainty quantification — Monte Carlo Dropout with 20 stochastic forward passes; flags high-uncertainty predictions for radiologist review
- Explainability — ViT-GradCAM heatmaps overlaid on the original X-ray per detected pathology
- Automated radiology reports — structured FINDINGS / IMPRESSION / RECOMMENDATION sections generated by Groq LLaMA-3.3-70b-versatile
- Fairness analysis — per-demographic subgroup AUC evaluation across age and sex
- Production API — async FastAPI with singleton model loading, request validation, GradCAM session store
- PWA frontend — installable Next.js app with Google OAuth, scan history, analytics dashboard
┌─────────────────────────────────────────────────────────┐
│ Next.js Frontend (Vercel) │
│ Login → Home Dashboard → Scan → Stats → Profile │
└────────────────────────┬────────────────────────────────┘
│ HTTPS (multipart/form-data)
┌────────────────────────▼────────────────────────────────┐
│ FastAPI Backend (HuggingFace Spaces) │
│ │
│ POST /api/v1/predict │
│ GET /api/v1/gradcam/{session_id}/{class} │
│ GET /health │
└────────────────────────┬────────────────────────────────┘
│
┌──────────────▼──────────────────┐
│ InferencePipeline │
│ │
│ 1. Preprocess (224×224) │
│ 2. BioMedCLIP ViT-B/16 │
│ 3. MC Dropout (20 passes) │
│ 4. Per-class thresholds │
│ 5. ViT-GradCAM │
│ 6. Groq report generation │
└─────────────────────────────────┘
│
┌──────────────▼──────────────────┐
│ HuggingFace Hub │
│ Sowaiba01/chestai-model │
│ chestai_best.pt │
└─────────────────────────────────┘
| Component | Detail |
|---|---|
| Backbone | BioMedCLIP ViT-B/16 (Microsoft) — pretrained on 15M biomedical image-text pairs |
| Classification head | LayerNorm → Dropout(0.3) → Linear(512→512) → GELU → Dropout(0.3) → Linear(512→14) |
| Uncertainty | Monte Carlo Dropout — 20 stochastic passes at inference; entropy + per-class std |
| Explainability | ViT-GradCAM — gradient-weighted attention map from final transformer block |
| Training data | NIH ChestX-ray14 — 112,120 frontal-view X-rays, 30,805 patients |
| Input | RGB 224×224, normalized to ImageNet mean/std |
| Loss | Weighted binary cross-entropy with class imbalance correction |
| Optimizer | AdamW, lr=1e-4, weight decay=1e-2 |
| Report generation | Groq API — llama-3.3-70b-versatile, T=0.2, max_tokens=512 |
Backend
- Python 3.11, FastAPI, Uvicorn
- PyTorch 2.x, open_clip (BioMedCLIP), transformers
- HuggingFace Hub (model weights), HuggingFace Spaces (deployment)
- Groq Python SDK
Frontend
- Next.js 14, TypeScript, Tailwind CSS
- NextAuth.js (Google OAuth)
- Framer Motion, Lucide React
- Vercel (deployment)
ML / Data
- NIH ChestX-ray14 dataset
- Kaggle (training environment)
- Monte Carlo Dropout, ViT-GradCAM
ThoraxNet/
├── api/ # FastAPI backend
│ ├── main.py # App entrypoint, CORS, lifespan
│ ├── inference.py # InferencePipeline singleton
│ ├── schemas.py # Pydantic v2 request/response models
│ └── routes/
│ └── predict.py # /api/v1/predict endpoint
├── models/
│ ├── backbone.py # BioMedCLIP ViT-B/16 wrapper
│ ├── classifier.py # ChestAIClassifier
│ └── uncertainty.py # Monte Carlo Dropout inference
├── data/
│ ├── dataset.py # NIH ChestX-ray14 dataset + CLASSES
│ └── transforms.py # Train/val image transforms
├── training/
│ ├── trainer.py # Training loop
│ ├── losses.py # Weighted BCE
│ └── metrics.py # Per-class AUC, threshold tuning
├── explainability/
│ └── gradcam.py # ViT-GradCAM implementation
├── fairness/
│ └── evaluator.py # Subgroup fairness metrics
├── report_generation/
│ └── generator.py # Groq LLaMA radiology report
├── frontend/ # Next.js PWA
│ └── src/
│ ├── pages/
│ │ ├── index.tsx # Home / Dashboard
│ │ ├── login.tsx # Google OAuth login
│ │ ├── scan.tsx # X-ray upload + results
│ │ ├── stats.tsx # Model analytics
│ │ ├── profile.tsx # User settings
│ │ └── api/auth/ # NextAuth.js API routes
│ ├── components/
│ │ ├── Layout.tsx # Shared header + bottom nav
│ │ ├── XRayUploader.tsx
│ │ ├── FindingsPanel.tsx
│ │ ├── ReportViewer.tsx
│ │ └── UncertaintyChart.tsx
│ └── lib/
│ └── api.ts # Typed API client + scan history
├── Dockerfile # HuggingFace Spaces Docker config
├── requirements.txt
└── README.md
Base URL: https://Sowaiba01-ThoraxNet.hf.space
Analyze a chest X-ray image.
Request (multipart/form-data)
| Field | Type | Required | Description |
|---|---|---|---|
file |
File |
✓ | PNG or JPEG, max 10 MB |
patient_age |
float |
✗ | Patient age in years |
patient_gender |
string |
✗ | "M" or "F" |
generate_report |
bool |
✗ | Default true. Set false to skip the LLM narrative and return in ~85 ms. |
generate_gradcam |
bool |
✗ | Default true. Set false to skip heatmap generation. |
Response
{
"findings": [
{
"name": "Effusion",
"probability": 0.724,
"uncertainty": 0.043,
"present": true,
"high_uncertainty": false
}
],
"report": "FINDINGS:\n...\nIMPRESSION:\n...\nRECOMMENDATION:\n...",
"entropy": 0.312,
"inference_time_ms": 3145.8,
"stage_timings_ms": {
"preprocess": 15.7,
"mc_dropout": 2770.5,
"gradcam": 5.4,
"report": 42.4
},
"model_version": "1.1.0",
"gradcam_available": true,
"gradcam_classes": ["Effusion"],
"gradcam_session_id": "9f2c1e40-5b3a-4d81-b7e6-2a4c8d1f0e33"
}report is null when generate_report=false.
stage_timings_ms gives the per-stage latency breakdown in milliseconds — use
it to attribute latency regressions to a specific stage in production.
Returns the GradCAM heatmap overlay as a PNG for a specific pathology. Use the
gradcam_session_id returned by /predict; sessions are held in memory and
evicted after 100 newer scans.
Returns model load status and device info.
git clone https://github.com/Sowaiba-01/ThoraxNet.git
cd chestai
pip install -r requirements.txt
# Set environment variables
export MODEL_HUB_REPO=Sowaiba01/chestai-model
export GROQ_API_KEY=your_groq_api_key
uvicorn api.main:app --host 0.0.0.0 --port 7860 --reloadAPI docs available at http://localhost:7860/docs
cd frontend
npm install
# Create .env.local from template
cp .env.local.example .env.local
# Fill in your values
npm run devThe API is deployed as a Docker container on HuggingFace Spaces. Model weights are downloaded at startup from Sowaiba01/chestai-model via huggingface_hub.
Required Space secrets:
GROQ_API_KEYMODEL_HUB_REPO(optional, defaults toSowaiba01/chestai-model)
cd frontend
vercel --prodRequired environment variables in Vercel dashboard:
NEXTAUTH_SECRETNEXTAUTH_URLGOOGLE_CLIENT_IDGOOGLE_CLIENT_SECRETNEXT_PUBLIC_API_URL
ThoraxNet demonstrates that medical vision-language foundation models (BioMedCLIP) significantly outperform task-specific CNN architectures on chest X-ray multi-label classification when fine-tuned with appropriate regularization. Key observations:
- Cardiomegaly achieves the highest AUC (0.888) due to its distinct visual morphology
- Pneumonia is the hardest class (AUC 0.695), consistent with literature — its appearance overlaps heavily with Consolidation and Infiltration
- Monte Carlo Dropout provides well-calibrated uncertainty estimates; predictions with std > 0.15 correlate strongly with ambiguous or borderline cases
- Per-class threshold tuning on the validation set improves macro F1 by ~4% over using a uniform 0.5 threshold
Single-scan latency, measured client-side against the deployed Space with
scripts/benchmark.py. Reproduce with:
python scripts/benchmark.py --image tests/fixtures/sample_cxr.png \
--requests 200 --concurrency 1,4,16Environment. HuggingFace Spaces free tier, CPU only (device: cpu,
2 vCPU). Latency is wall-clock from the client and therefore includes network
round-trip to the Space. All figures below are from real runs; nothing here is
estimated.
Identical hardware, identical image, identical protocol: 30 requests, 3 warmup requests discarded, 0 failures. The only variable is the code.
| Metric | v1.0.0 | v1.1.0 | Improvement |
|---|---|---|---|
| p50 latency | 6,387 ms | 3,327 ms | 1.9× faster |
| p95 latency | 7,525 ms | 3,862 ms | 1.9× faster |
| p99 latency | 8,026 ms | 4,768 ms | 1.7× faster |
| Throughput | 0.15 req/s | 0.30 req/s | 2.0× |
Every percentile improved by ~2×, with no change to model weights or accuracy — this is purely a change in how the model is executed.
Three changes to the request path, no change to the model:
| Change | Effect |
|---|---|
| Batched MC Dropout | 20 sequential batch-1 forward passes → one batched forward pass. The dominant win. |
| Async report | Groq call moved off the critical path onto a worker thread. |
| GradCAM cache | Heatmaps were recomputed every call; now LRU-cached per (image, class). |
v1.1.0 exposes a per-stage breakdown on every response (stage_timings_ms),
measured in production — MC Dropout is now 2,868 ms of ~3,000 ms server-side,
preprocess 16 ms, GradCAM 6 ms (cached), report 92 ms (async). (v1.0.0 had no
per-stage instrumentation, so only its end-to-end total is comparable.)
The single biggest win: the 20 Monte Carlo Dropout passes were running sequentially at batch size 1. Tiling the input into one batched forward pass — while preserving the independent per-sample dropout masks the estimator requires — roughly halved end-to-end latency.
The gradcam figure is a cache-hit number: the benchmark reuses one image, so every request after the first hits the cache. Cold heatmaps cost more.
GradCAM retrieval had never worked. api/routes/predict.py read
pipeline.gradcam._last_overlays to populate its session store. That
attribute did not exist — InferencePipeline.predict() built the overlays
into a local variable and let them go out of scope. Every
GET /api/v1/gradcam/{session_id}/{class} returned 404. Nothing logged an
error, and the frontend rendered an empty panel rather than a failure state,
so it shipped and stayed broken. Found while reading the request path
end-to-end for the latency work, not from a bug report — which is the
uncomfortable part. Fixed by having ViTGradCAM.generate_overlays() record
overlays on the instance, and pinned with a regression test.
MC Dropout was 65% of request time and nobody had measured it. The
implementation looked reasonable — a list comprehension over 20 forward
passes. But at batch size 1 a T4 is almost entirely idle during each pass;
the cost was kernel launch overhead and Python dispatch, repeated 20 times.
Tiling the input along the batch dimension gives the same 20 independent
dropout samples in one pass. The correctness argument matters here: dropout
masks are sampled per batch element, so the tiled copies are independent
samples, not correlated ones. test_batched_matches_sequential_in_distribution
verifies this empirically at n=400 rather than taking it on faith.
std returned NaN at n_samples=1. torch.std applies Bessel's
correction by default, so a single sample divides by zero. Only surfaced when
writing the chunking tests — no production request uses T=1, but the failure
mode would have been a NaN silently propagating into the uncertainty field
and rendering as a blank badge in the UI.
Session-store eviction leaked memory. The eviction check was if len(store) > MAX rather than while, so a burst of concurrent requests
could only ever evict one entry per request while adding one — the store
grew monotonically under load.
A test was passing on a vacuous truth. The stub backbone used to avoid
downloading BioMedCLIP in CI returned a hard-coded torch.zeros(2, 512),
ignoring its input entirely. All-zero features stay zero through LayerNorm
and zero-initialised Linear layers, so every MC Dropout sample was identical
and assert std.mean() > 0 was asserting nothing about dropout. The stub also
ignored the input batch size, so the shape assertions in the same file could
never have held. Both tests were red before this work started — the batching
rewrite surfaced them rather than caused them, because reshaping into
(T, B, C) forces the batch dimension to actually be correct. A test that
cannot fail is worse than no test: it reports safety it isn't providing.
An optional dependency was a hard import. report_generation/generator.py
imported groq at module scope, which made the whole inference pipeline
unimportable without it — even though the code already carried a template
fallback for precisely the case where Groq isn't available. The fallback
existed; the import made it unreachable.
torch.jit.script on the full model. The BioMedCLIP backbone contains
constructs TorchScript's compiler rejects (dynamic attribute access in
open_clip's attention implementation). Switched to torch.jit.trace, which
is valid here only because inference has no data-dependent control flow — a
worse tool in general, the right one for this specific graph.
Static (calibrated) INT8 quantization. Quantizing activations as well as weights meant quantizing the attention softmax path, and ViT accuracy is notably sensitive there. Early runs showed AUC regressions above 0.01 on the low-prevalence classes, which is not a trade I'm willing to make for a medical model. Fell back to dynamic quantization, which touches only Linear weights. The measured delta is published in the quantization table rather than asserted to be negligible.
Caching MC Dropout results by image hash. Tempting, and wrong: the whole
point of MC Dropout is that repeated evaluation draws fresh samples. Caching
the result would return a stale point estimate and quietly destroy the
uncertainty semantics. Only GradCAM — which is deterministic given
(image, class) — is safe to cache.
Raising max_chunk above 32. No measurable throughput gain past 32
samples per pass on a T4, and it pushed peak memory high enough to OOM when
two requests overlapped. The lock serialises GPU work, but the allocator
still has to hold both.
- INT8 weights are exported but not yet serving traffic; the accuracy table has to be published before that promotion.
- The GradCAM session store is process-local in-memory. It does not survive a restart and will not work behind more than one replica — Redis is the correct fix and is not done.
- The benchmark measures client-side latency against a single Space instance. It is not a load test of a horizontally scaled deployment.
If you use ChestAI in your research, please cite:
@software{thoraxnet2026,
author = {Arshad, Sowaiba},
title = {ThoraxNet: AI Chest X-Ray Diagnostic Platform with Uncertainty Quantification},
year = {2026},
url = {https://github.com/Sowaiba-01/ThoraxNet},
}Acknowledgements
- NIH ChestX-ray14 — Wang et al., 2017
- BioMedCLIP — Microsoft Research
- Groq — LLaMA inference API
MIT License — see LICENSE for details.
Built by Sowaiba Arshad