NoaCahan/pe_classification_model

Pulmonary Embolism Classifier 🧠

This is a classification model for pulmonary embolism classification from CTPA scans. This classification model is used to aid a classification task to a diffusion model that generated 3D CTPA scans from a single 2D chest X-ray.

The input to the classification model are CTPA scan latents of size 32 x 32 x 64. This model was trained on CTPAs that were downscaled from size 256 x 256 x 64 in a 2D slice-by-slice manner using a VAE.

This model is used in this git repository: https://github.com/NoaCahan/X-ray2CTPA

🧩 Model Details

• Architecture: 3D 121-DenseNet
• Framework: PyTorch
• Input: 3D volumes (64×32×32)
• Output: PE Binary Classification probability

🧠 How to Use

from huggingface_hub import hf_hub_download
import torch

model_path = hf_hub_download("username/medical-diffusion-classifier", "classification_model_256.pth.tar")
model = torch.load(model_path, weights_only=False)

🧾 Citation

If you use this model, please cite our work:

Cahan, N. et al. X-ray2CTPA: leveraging diffusion models to enhance pulmonary embolism classification. npj Digit. Med. 8, 439 (2025). https://doi.org/10.1038/s41746-025-01857-y