For hepatocellular carcinoma (HCC) patients, early-stage diagnosis continues to be the most important predictor of survival. However, lesion detection typically only occurs at an advanced stage and is largely affected by radiologists' experience. Therefore, this repository implements HCCNet - a spatio-temporal neural network that utilizes a 3D ConvNeXt 1 backbone combined with a transformer encoder 2 - to predict future cancer development based on past MRI examinations.
The model embeds the raw MR images in higher-dimensional feature space, adds time-based positional encodings, and subsequently processes the sequence of embedded image representations to predict the likelihood of future HCC development.
Model training follows a step-wise approach. That is, we first pretrain both CNN backbone and Transformer encoder using self-supervised learning and then fine-tune the full model on our downstream task.
We pretrain the CNN backbone adapting the DINO pretraining framework proposed in 3 to 3D medical images. Furthermore, for the Transformer encoder, we employ a custom pre-training approach, where we randomly shuffle 50% of embedded image sequences and train the model to differentiate shuffled from non-shuffled sequences. Note that during Transformer pre-training, we keep the parameters of the CNN backbone frozen.
After pre-training, we initialize the model with the weights obtained after pre-training, add a linear pooling layer to the architecture, and fine-tune the full model on our downstream task. We train the model over 10 runs with different random seeds and ensemble the runs' predictions by averaging over the predicted probabilities.
| Random init | Pretrained init | ||||||
|---|---|---|---|---|---|---|---|
| Modality | Architecture | Params | AUPRC | AUROC | AUPRC | AUROC | Downloads |
| DW-MRI | HCCNet-F | 12.4M | 0.26 | 0.72 | 0.54 | 0.89 | model weights |
| HCCNet-P | 22.0M | 0.29 | 0.72 | 0.74 | 0.92 | model weights | |
| HCCNet-N | 45.9M | 0.31 | 0.72 | 0.79 | 0.95 | model weights | |
| HCCNet-T | 72.4M | 0.30 | 0.73 | 0.73 | 0.93 | model weights | |
| T1 DCE-MRI | HCCNet-F | 12.4M | 0.30 | 0.74 | 0.53 | 0.76 | model weights |
| HCCNet-P | 22.0M | 0.35 | 0.78 | 0.59 | 0.81 | model weights | |
| HCCNet-N | 45.9M | 0.39 | 0.78 | 0.48 | 0.76 | model weights | |
| HCCNet-T | 72.4M | 0.56 | 0.77 | 0.65 | 0.88 | model weights | |
Diffusion weighted MRI's (i.e., DW-MRI) compose a four channel image with diffusion coefficients b = 0, 150, 400, and 800. Contrarily, dynamic contrast enhanced T1 weighted MRIs comprise of one pre-contrast and three post-contrast (i.e., late aterial, portal venous, and delayed phase) scans.
- Add results for diffusion MRI
- Add results for T1 contrast enhanced MRI
- Add results for T1 in- and out-of-phase MRI
- Add results for T2 MRI
The repository is licensed under the apache-2.0 license.
Footnotes
-
Liu, Z., Mao, H., Wu, C.Y., Feichtenhofer, C., Darrell, T., Xie, S., 2022. A ConvNet for the 2020s. URL: http://arxiv.org/abs/2201.03545. arXiv:2201.03545. ↩
-
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., Polosukhin, I., 2023. Attention Is All You Need. URL: http://arxiv.org/abs/1706.03762. arXiv:1706.03762. ↩
-
Caron, M., Touvron, H., Misra, I., J ́egou, H., Mairal, J., Bojanowski, P., Joulin, A., 2021. Emerging Properties in Self-Supervised Vision Transformers. URL: http://arxiv.org/abs/2104.14294. arXiv:2104.14294. ↩