KneeSense-cnn/kneesense-cnn

KneeSense CNN: knee osteoarthritis severity (KL grade 0-4)

Predicts the Kellgren-Lawrence (KL) grade, 0 to 4, from a single-knee frontal X-ray, and returns a Grad-CAM heatmap showing which image regions drove the score. It is the vision component of KneeSense, a college major project.

Author: Tejaswi Pise (AI / Deep Learning lead), KneeSense team project, with Lahari (NLP / GenAI) and Jyothi (full-stack and integration).

Not a medical device. This is an educational baseline. Do not use it for diagnosis or treatment decisions. See Limitations.

Files

File	Purpose
best_model.onnx	The model (about 16 MB), CPU inference with ONNX Runtime
README.md / model_card.md	This card

Model

• Backbone: EfficientNet-B0, ImageNet-pretrained, fine-tuned end to end.
• Head: CORAL ordinal head. Instead of a 5-way softmax, the model answers four ordered questions ("is the grade > 0? > 1? > 2? > 3?") with shared weights. The predicted grade is the number of "yes" answers, so being 2 grades wrong costs more than being 1 grade wrong.
• Format: ONNX, opset 17, dynamic batch size.

	Name	Shape	Meaning
input	image	(batch, 3, 224, 224) float32	preprocessed X-ray (see Preprocessing)
output	logits	(batch, 4) float32	ordinal logits, already divided by the calibration temperature
output	cam	(batch, 7, 7) float32	raw class-activation map for the heatmap

The calibration temperature T = 0.5758280388804236 is built into the logits output and stored in the file's metadata (temperature).

Class labels

KL grade	Label
0	Normal
1	Doubtful
2	Mild
3	Moderate
4	Severe

Preprocessing (must match exactly)

1. Read the image as grayscale (uint8).
2. CLAHE contrast enhancement: clip limit 2.0, 8x8 tile grid (OpenCV).
3. Repeat the channel to 3 channels (RGB).
4. Resize to 224 x 224, bilinear interpolation.
5. Scale to [0, 1], then normalise with ImageNet mean (0.485, 0.456, 0.406) and std (0.229, 0.224, 0.225).
6. Arrange as (1, 3, 224, 224), float32.

The model expects a single, cropped, frontal knee X-ray like the training images, not a full bilateral film.

Output

Per image the service returns:

Field	Type	Meaning
kl_grade	int, 0-4	number of ordinal thresholds with probability above 0.5
confidence	float, 0.0-1.0	the model's probability for exactly that grade
heatmap_base64	string	PNG of the Grad-CAM overlay, base64 (no data: prefix)

confidence is the probability of the exact predicted grade. Because the grades are ordered, probability is shared with neighbouring grades, so values around 0.4 to 0.6 are common even for correct predictions. It is not a calibrated clinical probability (see Calibration).

Quick start

import cv2, numpy as np, onnxruntime as ort
from huggingface_hub import hf_hub_download

path = hf_hub_download("KneeSense-cnn/kneesense-cnn", "best_model.onnx")
session = ort.InferenceSession(path, providers=["CPUExecutionProvider"])

gray = cv2.imread("knee.png", cv2.IMREAD_GRAYSCALE)
rgb = cv2.cvtColor(cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(gray), cv2.COLOR_GRAY2RGB)
rgb = cv2.resize(rgb, (224, 224), interpolation=cv2.INTER_LINEAR)
x = (rgb.astype(np.float32) / 255 - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
x = x.astype(np.float32).transpose(2, 0, 1)[None]

logits, cam = session.run(["logits", "cam"], {"image": x})
p = 1 / (1 + np.exp(-logits[0].astype(np.float64)))          # P(grade > 0), ..., P(grade > 3)
kl_grade = int((p > 0.5).sum())
probs = np.clip(np.r_[1, p] - np.r_[p, 0], 0, None); probs /= probs.sum()
confidence = float(probs[kl_grade])
print(kl_grade, confidence)

The KneeSense repository's ml-service/model/inference.py also builds the heatmap overlay and the base64 PNG.

Training data

Knee X-rays with KL grades 0-4 from the Kaggle dataset Knee Osteoarthritis Dataset with Severity (shashwatwork/knee-osteoarthritis-dataset-with-severity), used with its provided train / val / test split. It mirrors Pingjun Chen's Knee Osteoarthritis Severity Grading Dataset (Mendeley Data, V1, 2018, DOI 10.17632/56rmx5bjcr.1), which was organised from the Osteoarthritis Initiative (OAI) at the University of Florida. Licence of the Kaggle dataset: Attribution 4.0 International (CC BY 4.0). Please credit the dataset authors when reusing this work.

Split	Images	Grade 0	1	2	3	4
train	5,778	2,286	1,046	1,516	757	173
val	826	328	153	212	106	27
test	1,656	639	296	447	223	51

Training

• Class imbalance: weighted random sampler on the training set (each grade drawn about equally often). Validation and test sets are not resampled.
• Augmentation: rotation up to 10 degrees, zoom 90-110%, brightness/contrast jitter. No horizontal flips, because knee laterality matters.
• Optimiser: AdamW, learning rate 3e-4, weight decay 1e-4, batch size 32, cosine schedule planned for 25 epochs.
• Loss: CORAL (binary cross-entropy on the four ordered thresholds).
• Model selection: best validation quadratic weighted kappa (best: 0.8124). Early stopping after 7 epochs without improvement, which stopped the run at epoch 11.
• Run: a single run, seed 42, on a Kaggle GPU. Exact package versions are in requirements-train.txt.

Evaluation (held-out test set, 1,656 images)

Metric	Value
Accuracy	0.630
Macro F1	0.624
Quadratic weighted kappa	0.830
Mean absolute error	0.41 grades
Within one grade of the true grade	about 96% (1,588 of 1,656)

Grade	Precision	Recall	F1	Support
0 Normal	0.74	0.79	0.77	639
1 Doubtful	0.31	0.38	0.34	296
2 Mild	0.68	0.51	0.58	447
3 Moderate	0.77	0.65	0.71	223
4 Severe	0.58	0.96	0.72	51

Confusion matrix (rows = true grade, columns = predicted grade):

          pred0 pred1 pred2 pred3 pred4
true 0     508   115    16     0     0
true 1     128   113    54     1     0
true 2      47   132   227    41     0
true 3       1     2    38   146    36
true 4       0     0     1     1    49

Calibration

confidence was calibrated with temperature scaling: one number T divides the logits before the sigmoid. T was fitted on the validation set only and evaluated on the test set. T never changes the predicted grade, so all metrics above are unchanged.

Test set	Mean confidence	NLL (lower is better)	ECE (lower is better)
before (T = 1)	0.491	1.063	0.169
after (T = 0.576)	0.613	0.944	0.114

Accuracy on the test set is 0.630, so average confidence now matches average accuracy. A calibration error remains (ECE 0.114): treat confidence as a rough score, not a probability.

Explainability

The cam output equals Grad-CAM for this architecture (the head is linear after global average pooling) and was checked against the pytorch-grad-cam library on the trained model. The heatmap is coarse (7x7 upsampled) and is drawn on the CLAHE-enhanced image the model saw. A visual check of about 10 test images showed heat on or near the joint line and joint margins. That is a plausibility check only, not a clinical validation of what the model relies on.

Export checks

• PyTorch and ONNX give the same grade, confidence, logits and activation map within 1e-3 on 10 test images, including a batch run.
• The serving environment (ml-service/requirements.txt, Python 3.11 slim, CPU) reproduces the test metrics above exactly (accuracy 0.6298, macro F1 0.6241, QWK 0.8301). Prediction takes about 30 ms per image on CPU.

Limitations

• Modest accuracy. About 63% of grades are exactly right. Grade 1 (Doubtful) is the weakest (F1 0.34); this boundary is also hard for human readers.
• Leans toward Severe. Of 223 Moderate (grade 3) test images, 36 were predicted Severe. The weighted sampler shows the model Grade 4 far more often than its real frequency.
• Very few Grade 4 images (27 validation, 51 test), so Grade 4 results and its calibration are uncertain.
• Possible optimistic metrics. The public Kaggle copy has no patient identifiers, so we could not verify that the splits never share a patient between training and test.
• One data source, no external validation. Images come from one research cohort. Other scanners, projections, patient groups, crops or image quality are untested.
• KL grading is subjective. Labels carry reader variability that the model inherits.

Intended use

Educational and research demonstration of an explainable, ordinal CNN for KL grading. Out of scope: clinical diagnosis, screening, treatment decisions, or any use where an error could harm a person.