README.md

# 🧠 Text-Conditioned Latent Diffusion for Contrast-Enhanced CT Synthesis

**Model Name**: `mlii0117/sd1.5_MPECT`  
**Model Type**: Fine-tuned `Stable Diffusion v1.5` for medical image-to-image translation  
**Paper**: _Text-Conditioned Latent Diffusion Model for Synthesis of Contrast-Enhanced CT from Non-Contrast CT_  
**Conference**: AAPM 2025 (Oral)  
**Authors**: Mingjie Li, Yizheng Chen, Lei Xing, Michael F. Gensheimer  
**Affiliation**: Department of Radiation Oncology - Medical Physics Divison, Stanford University

---

## 🧬 Model Description

This model is a fine-tuned version of **Stable Diffusion v1.5**, specialized for converting **non-contrast CT images** into **contrast-enhanced CT images**, guided by **textual phase prompts** (e.g., *venous phase*, *arterial phase*). It utilizes the `InstructPix2Pix` framework to enable flexible prompt-conditioned generation, enabling control over contrast timing without requiring explicit paired data.

---

## 💡 Key Features

- 🧾 **Text-guided control** over contrast phase (arterial vs. venous)
- 🖼️ Processes **2D CT slices** in image format (converted from DICOM)
- 🏥 Focused on **clinical realism and anatomical fidelity**
- 🧠 Reconstructs full 3D volume with NIfTI output support
- ✅ Evaluated and presented as **Oral at AAPM 2025**

---

## 🛠️ Usage

### 🔧 Requirements
```bash
pip install diffusers==0.25.0 nibabel pydicom tqdm pillow
```

### 📦 Load the Model
```python
from diffusers import StableDiffusionInstructPix2PixPipeline
import torch

pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    "mlii0117/sd1.5_MPECT", torch_dtype=torch.float16
).to("cuda")
generator = torch.Generator("cuda").manual_seed(0)
```

### 📝 Example Prompts

- **Arterial Phase**
  ```
  Convert this non-contrast CT slice to mimic an arterial-phase contrast-enhanced CT.
  Brighten and enhance the aorta, major arteries, and adjacent organ boundaries to emphasize arterial flow,
  focusing on clarity and contrast in these areas while maintaining other features unchanged.

  ```

- **Venous Phase**
  ```
  Convert this non-contrast CT slice to mimic a venous-phase contrast-enhanced CT.
  Brighten and enhance the veins, especially the portal and hepatic veins,
  and emphasize organ boundaries to mimic venous flow, focusing on brightness and contrast in these areas while maintaining other features unchanged."
  ```

### 🧪 Full Pipeline Example

```python
import os
import numpy as np
import nibabel as nib
from PIL import Image
from glob import glob
from tqdm import tqdm
from pydicom import dcmread
from diffusers import StableDiffusionInstructPix2PixPipeline

pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    "mlii0117/sd1.5_MPECT", torch_dtype=torch.float16
).to("cuda")
generator = torch.Generator("cuda").manual_seed(0)

prompt_art = "Convert this non-contrast CT slice to mimic an arterial-phase contrast-enhanced CT. Brighten and enhance the aorta, major arteries, and adjacent organ boundaries to emphasize arterial flow, focusing on clarity and contrast in these areas while maintaining other features unchanged."
prompt_ven = "Convert this non-contrast CT slice to mimic a venous-phase contrast-enhanced CT. Brighten and enhance the veins, especially the portal and hepatic veins, and emphasize organ boundaries to mimic venous flow, focusing on brightness and contrast in these areas while maintaining other features unchanged."

# read all dicoms
def load_dicom_folder(dicom_folder):
    dicom_folder = os.path.join(dicom_folder, 'DICOM')
    dicom_files = sorted(glob(os.path.join(dicom_folder, "*")))

    slices = []
    for dicom_file in dicom_files:
        ds = dcmread(dicom_file)
        slices.append(ds.pixel_array.astype(np.float32))
    
    dicom_array = np.stack(slices, axis=0)
    dicom_array += ds.RescaleIntercept
    dicom_array = np.clip(dicom_array, -1000, 1000)
    dicom_array = (dicom_array + 1000) / 2000.0
    return dicom_array

# transfer to RGB and send to diffusion
def process_slices(dicom_array):
    outputs = []
    for i in tqdm(range(dicom_array.shape[0])):
        slice_img = (dicom_array[i] * 255).astype(np.uint8)
        pil_img = Image.fromarray(slice_img).convert("RGB")

        edited_image = pipe(
            prompt, #### chose prompt_art or prompt_ven
            image=pil_img,
            num_inference_steps=20,
            image_guidance_scale=1.5,
            guidance_scale=10,
            generator=generator,
        ).images[0]

        gray = edited_image.convert("L")
        gray_np = np.array(gray).astype(np.float32) / 255.0
        gray_np = gray_np * 2000 - 1000  # scale back to [-1000, 1000]
        outputs.append(gray_np)

    volume = np.stack(outputs, axis=0)
    return volume

# save to nii.gz
def save_nifti(volume, output_path):
    affine = np.eye(4)
    nii = nib.Nifti1Image(volume, affine)
    nib.save(nii, output_path)

# main function
def main(input_dicom_path, output_nifti_path):
    dicom_array = load_dicom_folder(input_dicom_path)
    edited_volume = process_slices(dicom_array)
    save_nifti(edited_volume, output_nifti_path)

# DMEO 
# main("/path/to/dicom_folder", "/path/to/output.nii.gz")
```

---

## 🧠 Intended Use

- Medical research and simulation
- Data augmentation for contrast-enhanced imaging
- Exploratory analysis in non-contrast → contrast CT enhancement

> ⚠️ **Disclaimer**: This model is for research purposes only. It is not intended for clinical decision-making or diagnostic use.

---

## 📝 Citation

```
@inproceedings{li2025text,
  title={Text-Conditioned Latent Diffusion Model for Synthesis of Contrast-Enhanced CT from Non-Contrast CT},
  author={Li, Mingjie and Chen, Yizheng and Xing, Lei and Gensheimer, Michael},
  booktitle={AAPM Annual Meeting (Oral)},
  year={2025}
}
```

---

## 🧾 License

This model is released for **non-commercial research purposes only**. Please contact the authors if you wish to use it in clinical or commercial settings.