NeumonIA

Motivation

Pneumonia causes approximately 2.5 million deaths per year globally. Chest X-rays are the most widely used diagnostic tool, but their interpretation requires specialized radiologists — a scarce resource in many regions.

The goal was to explore how AI can assist in pneumonia screening directly in the browser, without relying on a remote server or compromising patient data privacy. The image never leaves the device.

This project also builds on my earlier Pneumonia Prediction in Chest X-Ray Images (a DenseNet-201 model), carrying that goal into the browser.

---

Solution

NeumonIA is a PWA that loads an image classification model in the browser via ONNX Runtime Web and performs inference entirely on the client side. The user drags in a chest X-ray; the model processes it in seconds and shows the result with confidence level and inference time.

Includes sample image gallery, prediction history, dark mode, responsive interface, and Spanish language support.

---

Technical Architecture

Model Training

The model was trained on Google Colab with an NVIDIA RTX PRO 6000 Blackwell GPU (97 GB VRAM). YOLO26n-cls (47 layers, 1.5M parameters) was used on the Chest X-Ray Pneumonia Balanced Dataset with 6,800 training and 1,700 validation images (balanced 50/50).

YOLO was chosen over building a custom CNN from scratch for three reasons:

1. Native ONNX export via Ultralytics — essential for browser inference with onnxruntime-web.
2. Compact footprint — YOLO26n-cls has only 1.5M parameters and produces a 5.9 MB ONNX file, keeping load times reasonable on the web.
3. Proven classification head — YOLO’s -cls variant is purpose-built for image classification, avoiding the complexity of adapting a detection architecture.

This contrasts with the earlier Pneumonia Prediction project, which used DenseNet-201 (18M+ parameters) — a powerful but much heavier architecture. According to Ultralytics benchmarks, YOLO26n achieves just 1.7 ms latency on T4 TensorRT and 38.9 ms on CPU ONNX (up to 43% faster than its predecessor), so it suits real-time client-side inference.

Hyperparameter	Value
Epochs	150 (early stopping, patience=20)
Batch size	512
Image size	224×224
Learning rate	Cosine scheduler + AMP
Validation accuracy	98.5% top-1
Test accuracy	100% (30/30 images)

ONNX Export

The PyTorch checkpoint (3.0 MB) was exported to ONNX with opset 17, FP32 precision, and simplified graph. Final size is 5.9 MB, compatible with onnxruntime-web.

Browser Inference

The inference.ts class (126 lines) implements the pipeline in three phases:

1. Model loading: cascading detection of WebGPU → WebGL → WASM to ensure compatibility with any modern browser.
2. Preprocessing: the image is resized to 224×224 on a canvas, R, G, B channels are extracted separately, and organized in NCHW (planar) format as YOLO expects.
3. Prediction: the tensor is fed to the ONNX session and two probabilities [normal, pneumonia] are returned.

Vite Plugin for WASM Files

The biggest integration challenge was that onnxruntime-web requires 4 WASM files at runtime that Vite doesn’t serve automatically from node_modules. I wrote a custom plugin that:

• In development: serves .wasm files with Content-Type: application/wasm header
• In production: copies WASM files to dist/ during build

Additionally, onnxruntime-web was excluded from optimizeDeps to prevent Vite from breaking WASM imports.

---

Frontend

• React 19.2 with custom hooks: useModel (ONNX model lifecycle), useTheme (light/dark/system), and useHistory (up to 50 in-memory predictions)
• Router-less navigation: useState<View> with conditional rendering to keep the bundle lightweight
• Theme system: CSS custom properties on :root, .dark, and @media (prefers-color-scheme: dark), persisted in localStorage
• 14 components: DropZone, SampleGallery, ImagePreview, ResultCard, ProgressBar, Sidebar, HistoryList, among others

---

Results

Metric	Value
Top-1 accuracy (validation)	98.5%
Top-1 accuracy (test)	100.0%
WebGPU latency	~5–20 ms
WebGL latency	~20–80 ms
WASM latency	~100–500 ms
ONNX model size	5.88 MB
Model load time	~2–8 s

The model maintains perfect accuracy on the test set and performs real-time inference even on consumer hardware.

---

Challenges & Learnings

ONNX + Vite Integration

WASM files are not automatically resolved from node_modules. I fixed it by writing a custom plugin and excluding onnxruntime-web from Vite’s dependency pre-bundling.

Execution Provider Fallback

WebGPU is not yet universal. Cascading detection (WebGPU → WebGL → WASM) was essential to ensure functionality in any browser.

NCHW Preprocessing

YOLO expects channels in planar format (R plane, G plane, B plane), not interleaved. Manual implementation from getImageData() required care to maintain CPU performance.

Medical Context Design

The UI includes a visible disclaimer, confidence bar (not just binary verdict), and ephemeral in-memory history to avoid storing medical data without consent.

---

References

• Ultralytics — YOLO26n-cls. Official documentation.
• ONNX Runtime Web — API and usage guide.
• Kaggle — Chest X-Ray Pneumonia Balanced Dataset.
• ONNX — Opset 17 specification.
• Vite — Plugin API.
• WebGPU — W3C Specification.
• Tailwind CSS v4 — CSS-based configuration.