Fetal brain sub-plane CNN

This repository contains a custom PyTorch convolutional neural network trained to classify a fetal-brain ultrasound image into one of four classes:

Trans-thalamic
Trans-cerebellum
Trans-ventricular
Other

The project is university coursework by Drileba Shem Remo (student number 2400704632), Bachelor of Science in Electrical and Electronics Engineering.

Model architecture

The network accepts one-channel 224 x 224 images. It contains four convolutional blocks with 32, 64, 128 and 256 channels. Each block contains two 3 x 3 convolutions, batch normalisation, ReLU activations and max pooling. Adaptive average pooling reduces each feature map to one value, followed by dropout 0.5 and a four-output linear classifier. The model has 1,173,668 trainable parameters.

Training and model selection

The training workflow used Adam optimisation, weight decay 0.0001, geometric and intensity augmentation without horizontal flipping, cosine learning-rate decay and early stopping. Macro-F1 was used for model selection because the four classes are highly imbalanced.

Three final random-seed runs produced validation macro-F1 values of 0.6272, 0.6426 and 0.6288, giving a mean of 0.6329 with a standard deviation of 0.0069. The checkpoint in this repository is the seed-42 model used for the single final test evaluation. Its best validation macro-F1 was 0.6272 at epoch 38.

Final test results

Metric	Result
Accuracy	0.7126
Macro-F1	0.5421
Balanced accuracy	0.5410
Weighted F1	0.6953

Class	Recall	F1
Trans-thalamic	0.8078	0.7658
Trans-cerebellum	0.7729	0.7391
Trans-ventricular	0.5530	0.6073
Other	0.0303	0.0563

The model recognises the three named anatomical planes much better than the rare Other class. Only 2 of 66 Other test images were correctly classified. Accuracy must therefore not be quoted without macro-F1 and class-wise results.

Preprocessing

Inference must reproduce the training preprocessing:

Convert the image to grayscale.
Resize to 224 x 224 pixels.
Convert pixels to a tensor in the range 0 to 1.
Normalise using mean 0.17076 and standard deviation 0.17093.

The checkpoint contains the state dictionary, training configuration and ordered class list.

Experiment record

W&B was used to record hyperparameters, learning curves, validation results, confusion matrices, misclassified examples and versioned model artifacts.

Intended use

This model is intended for education, reproducibility demonstrations and exploratory research on the FETAL_PLANES_DB dataset.

Limitations and safety

This is not a medical device and must not be used for diagnosis, clinical decisions or patient care. The model was evaluated on one published dataset and has not been validated across hospitals, populations or clinical workflows. Performance on Other is extremely poor, and the confidence score is not a guarantee that a prediction is correct. Machine- and operator-related acquisition differences may also affect predictions.

Files

FINAL-test-evaluation.pt: trained checkpoint
model.py: exact network architecture
config.json: preprocessing and model configuration
dataset_stats.json: dataset statistics used by the notebook
final_test_metrics.json: machine-readable final results

Dataset attribution

The model was trained using FETAL_PLANES_DB, released by Burgos-Artizzu et al. under CC BY 4.0:

Zenodo record: https://doi.org/10.5281/zenodo.3904280
Paper: Burgos-Artizzu et al., Scientific Reports 10, 10200 (2020), https://doi.org/10.1038/s41598-020-67076-5

Downloads last month: 13