NephroUro-AI

A unified, two-tier system for 14 nephrology and urology AI tasks. It covers image, whole-slide, signal, and tabular models built on pretrained CNN backbones, trained under one reproducible 70/15/15 protocol, and served from a model inference host that a separate website consumes over an HTTP/JSON API.

⚠️ Research prototype. Not a medical device. It has not been clinically validated and must not be used for clinical decisions. Several tasks ship with synthetic data and their models carry no clinical validity. See the methodology paper in docs/ for the full picture.

---

What's in the box

Deliverable	Location
Two-tier architecture (model server + website)	model_server/, website/
14-task registry (single source of truth)	config/tasks.yaml
Training pipeline (70/15/15, transfer learning)	model_server/train.py
Google Colab training notebook	notebooks/NephroUro_AI_Training.ipynb
Dataset download helper	scripts/download_datasets.py
Architecture diagram	docs/architecture.svg / .png
Methodology research paper	docs/NephroUro-AI_Methodology_Paper.docx / .pdf

The two tiers

Browser  ──HTTP──▶  Website (Flask, no ML)  ──HTTP/JSON──▶  Model Server (Flask + PyTorch)

• Model inference server (model_server/): hosts every model, the only component that needs PyTorch. Designed to run on a GPU host.
• Website (website/): holds no models and imports no ML libraries; it forwards every request to the model server. The two tiers can run on different machines.

The 14 tasks

Nephrology: kidney biopsy pathology, CKD progression, AKI detection, renal US/CT interpretation, dialysis optimisation.

Urology: kidney-stone CT, prostate mpMRI PI-RADS, prostate Gleason/ISUP, PSA decision support, bladder cystoscopy, urine cytology, urodynamics, hematuria triage, renal mass characterisation.

Each task is one entry in config/tasks.yaml. Adding a 15th is a config edit.

---

Setup

Requires Python 3.10–3.12 (PyTorch wheels may lag the newest Python).

# Option A: one environment for everything
pip install -r requirements.txt

# Option B: split deployment (recommended for two hosts)
pip install -r requirements-server.txt   # on the GPU / inference host
pip install -r requirements-web.txt      # on the website host

Running locally

Open two terminals from the project root:

# Terminal 1: model inference server  (http://localhost:8001)
python -m model_server.app

# Terminal 2: website                 (http://localhost:8000)
python -m website.app

Then open http://localhost:8000. The six small models (the five tabular tasks plus urodynamics) ship pre-trained on synthetic data, so the site is demoable immediately. The eight image-CNN checkpoints are large (20 to 100 MB each) and are not committed to the repo; run python -m model_server.train --task all to reproduce them locally in a few minutes, or use the Colab notebook. All fourteen tasks train cleanly.

Running the two tiers on separate servers

The website reaches the model server purely through one environment variable:

# on the website host: point it at the remote inference server
export NEPHROURO_MODEL_SERVER=http://<model-server-host>:8001
python -m website.app

Other environment variables (all optional):

Variable	Purpose	Default
NEPHROURO_MODEL_SERVER	URL the website calls	http://127.0.0.1:8001
NEPHROURO_MS_HOST / NEPHROURO_MS_PORT	model server bind	0.0.0.0 / 8001
NEPHROURO_WEB_HOST / NEPHROURO_WEB_PORT	website bind	0.0.0.0 / 8000
NEPHROURO_API_KEY	shared secret for the model API (X-API-Key)	disabled
NEPHROURO_DEVICE	auto / cpu / cuda	auto
NEPHROURO_CKPT_DIR	where checkpoints are read/written	./checkpoints
NEPHROURO_DATA_DIR	where datasets live	./data

---

Training the models

Locally

python -m model_server.train --task all --epochs 8          # train everything
python -m model_server.train --task kidney_stone_ct --epochs 20 --device cuda

Each run prints the 70/15/15 split sizes and the held-out test metrics, and writes checkpoints/<task_id>.pt. The model server picks up checkpoints automatically. No restart logic is needed beyond a process restart.

On Google Colab (GPU)

Open notebooks/NephroUro_AI_Training.ipynb in Colab, set the runtime to GPU, and follow the cells: upload the repo as a zip → install → train → download nephrouro_checkpoints.zip. Unzip the .pt files into checkpoints/.

Datasets

python scripts/download_datasets.py --list                  # show sources
python scripts/download_datasets.py --task ckd_progression  # fetch one

Public datasets are fetched directly or via the Kaggle CLI; credentialed sets (e.g. MIMIC-IV) print instructions. Any task whose data/<task_id>/ folder is empty falls back to a synthetic generator, so training always runs.

---

Model API (model server)

Method & route	Purpose
GET /health	liveness probe
GET /tasks	metadata for all 14 tasks
GET /tasks/<id>	one task's metadata
GET /status	per-task: is a trained checkpoint present?
POST /predict/<id>	run inference

POST /predict/<id> accepts:

• image / WSI tasks: multipart form, file field image
• tabular tasks: JSON {"features": { ... }}
• signal tasks: JSON {"signal": [[...], ...]}

Example:

curl -X POST localhost:8001/predict/ckd_progression \
  -H 'Content-Type: application/json' \
  -d '{"features":{"age":71,"serum_creatinine":3.4,"uacr":420}}'

Every response includes a trained flag. When false, the model has no checkpoint and its output is meaningless.

Image and WSI-tile responses also include a heatmap field: a Grad-CAM saliency overlay (a PNG data URI) showing which regions of the image drove the predicted class. The website displays it under the prediction.

---

Repository layout

nephrouro-ai/
├── config/tasks.yaml          # the 14-task registry, the single source of truth
├── common/                    # config + registry loader (shared by both tiers)
├── model_server/              # inference server + training pipeline (PyTorch)
│   ├── app.py                 #   JSON API
│   ├── backbones.py           #   pretrained CNN loader
│   ├── models.py              #   multi-head model builder
│   ├── datasets.py            #   preprocessing + superset logic + synth data
│   ├── inference.py           #   checkpoint loading + prediction
│   └── train.py               #   70/15/15 training pipeline
├── website/                   # Flask front-end (no ML libraries)
├── scripts/download_datasets.py
├── notebooks/                 # Colab training notebook
├── docs/                      # architecture diagram + research paper
└── checkpoints/               # trained models (<task_id>.pt)

Key design choices

• One registry, not 14 codebases. One training loop, one serving path.
• Superset of parameters. Tabular models train on the union of every dataset's columns; missing values are imputed, so adding a source never drops a feature.
• Transfer learning. ImageNet-pretrained ResNet-50 / DenseNet-121 / EfficientNet-B0 backbones with task-specific multi-head outputs.
• Uniform 70/15/15 split with a fixed seed, reproducible and comparable across all tasks.
• Honest untrained state. Untrained models are flagged, never served silently as if meaningful.
• Built-in explainability. Every image prediction returns a Grad-CAM heatmap so the result can be sanity-checked, not taken on trust.

See docs/NephroUro-AI_Methodology_Paper.pdf for the full rationale.

Limitations

This is a methodology/architecture prototype. It is not a medical device, has no regulatory clearance, and has not been clinically validated. Six tasks use synthetic data; whole-slide pathology is handled at the tile level; label spaces are simplified. Do not use any output for patient care.