release some working code

Author: lucidrains

README.md

@@ -1,15 +1,56 @@
 <img src="./med-seg-diff.png" width="450px"></img>
 
-## MedSegDiff - Pytorch (wip)
+## MedSegDiff - Pytorch
 
 Implementation of <a href="https://arxiv.org/abs/2211.00611">MedSegDiff</a> in Pytorch - SOTA medical segmentation out of Baidu using DDPM and enhanced conditioning on the feature level, with filtering of features in fourier space.
 
 I will also add attention and introduce an extended type of cross modulation on the attention matrices, alphafold2 style.
 
+## Install
+
+```bash
+$ pip install med-seg-diff-pytorch
+```
+
+## Usage
+
+```python
+import torch
+from med_seg_diff_pytorch import Unet, MedSegDiff
+
+model = Unet(
+    dim = 64,
+    dim_mults = (1, 2, 4, 8)
+)
+
+diffusion = MedSegDiff(
+    model,
+    image_size = 128,
+    timesteps = 1000
+).cuda()
+
+segmented_imgs = torch.rand(8, 3, 128, 128)  # inputs are normalized from 0 to 1
+input_imgs = torch.rand(8, 3, 128, 128)
+
+loss = diffusion(segmented_imgs, input_imgs)
+loss.backward()
+
+# after a lot of training
+
+pred = diffusion.sample(input_imgs)     # pass in your unsegmented images
+pred.shape                              # predicted segmented images - (8, 3, 128, 128)
+```
+
 ## Appreciation
 
 - <a href="https://stability.ai/">StabilityAI</a> for the generous sponsorship, as well as my other sponsors out there
 
+## Todo
+
+- [ ] add a cross attention variant for generating the attentive map (A)
+- [ ] modulate attention matrices in middle and other self attention layers, wherever full attention is used
+- [ ] some basic training code, with Trainer taking in custom dataset tailored for medical image formats
+
 ## Citations
 
 ```bibtex