pymite6941's picture
Add model card
1496b47 verified
|
Raw
History Blame Contribute Delete
10.2 kB
---
license: apache-2.0
library_name: transformers
pipeline_tag: image-text-to-text
tags:
- medical
- chest-xray
- radiology
- onnx
- clip
- blip
- multimodal
- cpu
base_model:
- openai/clip-vit-base-patch32
- emilyalsentzer/Bio_ClinicalBERT
- Salesforce/blip-image-captioning-base
language:
- en
---
# MedicalAI β€” Light Weight
Chest X-ray analysis that runs on consumer hardware β€” a laptop CPU, no GPU, no cloud.
> **⚠️ Not a medical device.** This is a research and educational project. It is **not** FDA/CE cleared, has not been clinically validated, and must not be used to diagnose, treat, or make any decision about a real patient. Outputs are frequently wrong. See [Limitations](#limitations) β€” they are substantial and you should read them before using anything here.
## What's in this repo
> **The X-ray and the symptoms go into one model, not two.** The fusion model is a single network that consumes the radiograph *and* the symptom text together and emits one set of logits β€” `fusion_full.onnx` is one graph with three inputs (`pixel_values`, `input_ids`, `attention_mask`). There is no separate image classifier and text classifier whose outputs get merged afterwards; the two modalities are fused inside the model, before the classifier head. The BLIP captioner below is a **separate, optional** model that only writes a text description of the image β€” it takes no symptom input and plays no part in the diagnosis path.
| Component | Path | Size | What it does |
|---|---|---|---|
| **Fusion model** (ONNX, end-to-end) β€” *the main model* | `checkpoints/onnx_full/fusion_full.onnx` | 787 MB | X-ray **and** symptom text β†’ diagnosis logits, in one graph. Runs with `onnxruntime` alone β€” no PyTorch. |
| Fusion classifier head (PyTorch) | `checkpoints/fusion_model.pth` | 5.4 MB | Trained classifier head only; needs CLIP + Bio_ClinicalBERT at runtime. |
| Fusion classifier head (ONNX) | `checkpoints/onnx/fusion_classifier.onnx` | 4.5 MB | Head-only ONNX; encoders still run in PyTorch. |
| **BLIP X-ray captioner** | `blip-xray-finetuned/` | 896 MB | `Salesforce/blip-image-captioning-base` fine-tuned on IU-Xray reports β†’ radiology-style caption. |
| Default/demo classifier | `models/default/fusion_classifier.onnx` | 1.3 MB | 15 NIH classes, **random weights**. Ships so the app runs before training. Not predictive. |
| Application code | `*.py`, `launch.*`, `config.json` | β€” | CLI, Gradio web UI, batch predictor, training and ONNX export scripts. |
The training dataset is **not** included β€” see [Data](#data).
## Architecture
**Fusion (Symptom Check)** β€” one multimodal classifier over both inputs. Both encoders feed a single shared head, so the prediction is a joint function of the image and the symptoms; neither modality is scored on its own:
```
image ──► CLIP ViT-B/32 vision tower ──► visual_projection ──► L2-normalize ──┐
β”œβ”€β–Ί concat ──► MLP classifier ──► logits
symptom text ──► Bio_ClinicalBERT ──► mean-pool last_hidden_state β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
```
Encoders are **frozen**; only the MLP head is trained. `fusion_full.onnx` bakes the whole graph β€” encoders included β€” into one file, which is why it is 787 MB.
ONNX signature (opset 14, dynamic batch and sequence length):
| | Name | Shape | dtype |
|---|---|---|---|
| in | `pixel_values` | `[batch, 3, 224, 224]` | float32 |
| in | `input_ids` | `[batch, seq_len]` | int64 |
| in | `attention_mask` | `[batch, seq_len]` | int64 |
| out | `logits` | `[batch, 3018]` | float32 |
Preprocess with `CLIPProcessor` (`openai/clip-vit-base-patch32`) for the image and `AutoTokenizer` (`emilyalsentzer/Bio_ClinicalBERT`) for the text. Class names are in `checkpoints/onnx_full/labels.json`, index-aligned to the logits.
**Vision (captioning)** β€” a separate BLIP model, image-only, loadable with `BlipForConditionalGeneration.from_pretrained`. It does not see the symptoms and does not feed the fusion model; it exists to write a human-readable description alongside the diagnosis.
## Usage
### ONNX, no PyTorch
```python
import json
import numpy as np
import onnxruntime as ort
from PIL import Image
from transformers import CLIPProcessor, AutoTokenizer
from huggingface_hub import hf_hub_download
repo = "GAD-Research-Lab/MedicalAI-Light-Weight"
onnx_path = hf_hub_download(repo, "checkpoints/onnx_full/fusion_full.onnx")
labels = json.load(open(hf_hub_download(repo, "checkpoints/onnx_full/labels.json")))
clip = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
tok = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
image = Image.open("xray.jpg").convert("RGB")
pixel_values = clip(images=image, return_tensors="np")["pixel_values"]
text = tok("cough and fever", return_tensors="np", padding="max_length",
truncation=True, max_length=64)
sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
logits = sess.run(["logits"], {
"pixel_values": pixel_values.astype(np.float32),
"input_ids": text["input_ids"].astype(np.int64),
"attention_mask": text["attention_mask"].astype(np.int64),
})[0]
probs = np.exp(logits - logits.max()) / np.exp(logits - logits.max()).sum()
top = probs[0].argmax()
print(labels[top], float(probs[0][top]))
```
### BLIP captioning
```python
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
repo = "GAD-Research-Lab/MedicalAI-Light-Weight"
processor = BlipProcessor.from_pretrained(repo, subfolder="blip-xray-finetuned")
model = BlipForConditionalGeneration.from_pretrained(repo, subfolder="blip-xray-finetuned")
inputs = processor(Image.open("xray.jpg").convert("RGB"), return_tensors="pt")
print(processor.decode(model.generate(**inputs, max_new_tokens=64)[0],
skip_special_tokens=True))
```
### Full application
```bash
git clone https://huggingface.co/GAD-Research-Lab/MedicalAI-Light-Weight
cd MedicalAI-Light-Weight
pip install -r requirements.txt gradio
python web_ui.py # http://127.0.0.1:7860
```
Or `launch.ps1` (Windows) / `launch.sh` (Linux/macOS) to set up a venv and start the UI in one step. `python run.py` for the interactive CLI, `python batch_predict.py <dir> -o out.csv` for batch.
## Training
- **Fusion head** β€” cross-entropy over the label set, 85/15 random train/val split, AdamW, frozen encoders. `python training.py --mode train --epochs 10 --batch_size 8`.
- **BLIP** β€” fine-tuned on IU-Xray image/report pairs. `python xray_training.py --mode train --epochs 3 --batch_size 4 --max_samples 500`.
Sources: IU-Xray (~6,687 rows), NIH Chest X-ray (~3,000 rows), plus 160 synthetic rare-finding rows from `expand_dataset.py` β€” ~9,847 total.
## Limitations
Read these. They are not boilerplate.
- **The label space is degenerate.** `labels.json` has **3,018 classes**, and most are not diagnoses β€” they are raw, deduplicated report strings scraped from IU-Xray, e.g. `"findings: . impression: 1. all lines and tubes in stable , xxxx position..."`. Only a handful (`atelectasis`, `cardiomegaly`, `consolidation`, `edema`, `effusion`, `emphysema`, `fibrosis`, …) are clean condition names. With ~9.8k training rows across 3,018 classes there are roughly **3 examples per class**, and the reported "confidence" is a softmax over that space β€” it is not calibrated and should not be read as a probability of disease. Treat the classifier as a demonstration of the architecture, not as a working diagnostic.
- **No held-out evaluation is published.** Training reports validation *loss* only. There is no accuracy, AUROC, sensitivity/specificity, or per-class breakdown in this repo, and no evaluation on an external site or scanner. Nothing here supports a claim about real-world performance.
- **Encoders are frozen and general-purpose.** CLIP ViT-B/32 was pretrained on web images, not radiographs. Only a small MLP adapts it to this domain.
- **Narrow data.** Two US datasets, adult chest radiographs, frontal views, English free-text reports. Expect degradation on pediatric films, lateral views, other modalities, other populations, and other equipment. Known demographic and label-noise problems in IU-Xray and NIH ChestX-ray14 are inherited wholesale β€” NIH labels were themselves NLP-mined from reports and are noisy.
- **BLIP captions are fluent, not faithful.** The captioner will produce confident, plausible, well-formed radiology prose for an image it has no ability to read correctly, including for non-X-ray inputs. Fluency here carries no signal about correctness.
- **`models/default/` is random weights** by construction, so the app can start before training. Its predictions are noise.
- **Automation bias is the main risk.** The most likely harm from this repo is a person believing a confident-looking output. Do not put it in front of patients or in any workflow where a wrong answer reaches one.
## Data
The training data is **not redistributed here** β€” download it yourself with `python training.py --mode prepare-data`.
- **IU-Xray** (Indiana University / Open-i, NLM) β€” de-identified, public.
- **NIH ChestX-ray14** (NIH Clinical Center) β€” de-identified, public.
Both are de-identified at source; no PHI is contained in this repo. The BLIP model was fine-tuned on IU-Xray report text and can emit dataset artifacts such as `xxxx` anonymization tokens. Check each dataset's own terms before redistributing derivatives.
## License
Apache-2.0 for the code and the trained weights in this repo. Upstream components carry their own licenses β€” CLIP (MIT), BLIP (BSD-3-Clause), Bio_ClinicalBERT (MIT) β€” and the source datasets carry their own terms. The license permits use; it does not make the model safe or fit for clinical use.
## Citation
```bibtex
@software{medicalai_light_weight,
title = {MedicalAI β€” Light Weight: CPU-friendly chest X-ray analysis},
author = {GAD Research Lab},
year = {2026},
url = {https://huggingface.co/GAD-Research-Lab/MedicalAI-Light-Weight}
}
```