Add README.md

README.md

@@ -0,0 +1,96 @@
+# Conditional Diffusion Model for Medical Image Generation
+
+This repository contains a conditional diffusion model trained to generate 3D medical CT scan images based on segmentation masks. The model uses a U-Net architecture with score-based diffusion for high-quality medical image synthesis.
+
+## Model Architecture
+
+- **Base Model**: U-Net with 5-level encoder-decoder
+- **Input**: 4-channel 256x256 CT scan images
+- **Conditioning**: Segmentation masks (4-channel 256x256)
+- **Output**: 4-channel 256x256 generated images
+- **Sampling**: Euler-Maruyama sampler with 250 steps
+- **Training**: Score matching loss with conditional generation
+
+## Model Details
+
+- **Training Data**: 3,346 medical CT scan examples
+- **Lambda Parameter**: 25.0 (diffusion coefficient)
+- **Embedding Dimension**: 256
+- **Channels**: [32, 64, 128, 256, 512]
+- **Activation**: SiLU (Swish)
+
+## Usage
+
+### Using the Hugging Face API
+
+```python
+from transformers import AutoModelForImageGeneration
+import torch
+
+# Load the model
+model = AutoModelForImageGeneration.from_pretrained("your-username/your-model-name")
+
+# Generate images
+conditioning_mask = torch.randn(1, 4, 256, 256)  # Your segmentation mask
+generated_image = model.generate(conditioning_mask)
+```
+
+### Local Usage
+
+```python
+import torch
+from model import UNet, marginal_prob_std, diffusion_coeff, Euler_Maruyama_sampler
+
+# Load model
+Lambda = 25.0
+marginal_prob_std_fn = lambda t: marginal_prob_std(t, Lambda=Lambda, device='cuda')
+score_model = UNet(marginal_prob_std=marginal_prob_std_fn)
+score_model.load_state_dict(torch.load("ckpt_3D_v2.pth"))
+score_model.eval()
+
+# Generate sample
+conditioning_mask = torch.randn(1, 4, 256, 256)
+samples = Euler_Maruyama_sampler(
+    score_model,
+    marginal_prob_std_fn,
+    lambda t: diffusion_coeff(t, Lambda=Lambda, device='cuda'),
+    batch_size=1,
+    x_shape=(4, 256, 256),
+    num_steps=250,
+    device='cuda',
+    y=conditioning_mask
+)
+```
+
+## Training
+
+The model was trained for 5000 epochs with:
+- Learning rate: 2e-4 (with decay)
+- Batch size: 1
+- Optimizer: Adam
+- Loss: Score matching loss
+
+## Dataset
+
+The model was trained on medical CT scan data with corresponding segmentation masks. The dataset contains 3,346 training examples with 80-20 train/validation split.
+
+## Citation
+
+If you use this model in your research, please cite:
+
+```bibtex
+@misc{conditional_diffusion_medical,
+  title={Conditional Diffusion Model for Medical Image Generation},
+  author={Your Name},
+  year={2024},
+  url={https://huggingface.co/your-username/your-model-name}
+}
+```
+
+## License
+
+[Add your license here]
+
+## Contact
+
+For questions or issues, please open an issue on this repository.