# CMDM

**Repository Path**: Louispine/cmdm

## Basic Information

- **Project Name**: CMDM
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-09-14
- **Last Updated**: 2026-09-16

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# CMDM: Controllable Mask Diffusion Model for Medical Annotation Synthesis with Semantic Information Extraction

[[Paper on ScienceDirect](https://www.sciencedirect.com/science/article/pii/S0010482525011588)] 

This is the official implementation of the paper “Controllable Mask Diffusion Model for Medical Annotation Synthesis with Semantic Information Extraction”, published in **Computers in Biology and Medicine**.

<p align="center">
<img src=/assets/cibm_git.png />
</p>

## Table of Contents
- [Requirements](#requirements)
- [Dataset Preparation](#dataset-preparation)
- [Training Your Own CMDM](#training-your-own-cmdm)
- [Sampling with CMDM](#sampling-with-cmdm)
- [Acknowledgement](#acknowledgement)
- [Citations](#citations)

## Requirements
```bash
conda create -n CMDM python=3.8.10
conda activate CMDM
conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch
```
  
## Dataset Preparation
The proposed framework requires medical annotatoin data.

Please organize the dataset with the following structure:
```angular2
├── ${data_root}
│ ├── ${train_maskdataset_dir}
│ │ ├── mask_1.png
│ │ ├── mask_2.png
│ │ ├── mask_3.png
```

## Training Your Own CMDM
To train your own CMDM, Run the following command:

```bash
python train.py --data_path ./TrainDataset \
               --image_size 256 \
               --n_epoch 2000 \
               --n_T 500 \
               --batch_size 2 \
```

### Mask-stage ablation variants

**Correction:** the historical `--variant original_cmdm` below is a four-stage,
shared-backbone CMDM-conditioning comparator. It is NOT the official default
256px CMDM network. It is retained unchanged so old checkpoints/results remain
reproducible. Use the separate upstream entry below for original-model training.

Use one `--variant` value for each independent training run. Named variants
override the related low-level flags and are stored in the checkpoint.

| Variant | Heatmap | Morphology conditioning | Foreground loss | Controls |
| --- | --- | --- | --- | --- |
| `original_cmdm` | no | original size/location multiplication | 1 | area, centroid x/y |
| `ours_wo_heatmap` | no | full-network FiLM | 10 | paper 8-D descriptor |
| `ours_wo_film` | yes | none | 10 | heatmap position only |
| `ours_wo_fg` | yes | full-network FiLM | 1 | paper 8-D descriptor |
| `ours_full` | yes | full-network FiLM | 10 | paper 8-D descriptor |
| `legacy` | CLI flags | historical repository path | CLI flag | shape ID + descriptor |

### Pinned official CMDM training (separate from Ours)

`train_original_cmdm.py` directly imports the upstream model, diffusion, dataset,
preprocessing and defaults from `official_cmdm/`, pinned to
[DSLab-MJU/CMDM commit 15f69578e6de74b26b447da2d0209340d6e84044](https://github.com/DSLab-MJU/CMDM/tree/15f69578e6de74b26b447da2d0209340d6e84044).
`official_cmdm/UPSTREAM.json` records normalized-source SHA256 hashes; these are
checked before every invocation. No downloads are needed on the cloud server.
Do not run the vendored `train.py` directly: the upstream script does not save
checkpoints. Deploy the entire `official_cmdm/` directory with the wrapper.

The default model is 256px, base channels 128, two residual blocks, channel
multipliers `1,1,2,2,4,4`, and attention resolutions `16,8`. Controls are
`[area/(H*W), x/W, y/H]`. The upstream size/location multiplication, ordinary
noise MSE, Adam at `1e-4`, linear epoch learning-rate decay, batch size 2,
500 diffusion steps and condition-drop probability 0.1 are retained. No Ours
FiLM/heatmap/foreground weighting, augmentation or validation-based selection
is inserted. The wrapper only adds seeding, preflight, provenance, epoch-boundary
resume and atomic checkpoint writes. Runtime versions are recorded; identical
results across different CUDA/PyTorch versions or seeds are not promised.

The upstream `range(n_epoch+1)` is intentionally retained: `--n_epoch 2000`
prints 2001 epochs, of which the last has LR=0. `latest.pt` is the terminal model;
`best.pt` is an additional epoch-mean-training-loss choice, not an upstream
validated controllability selection. Report which checkpoint you evaluate.

**Data gate:** upstream produces one label per valid connected component but
indexes images by image number. For this published single-condition training
entry to have aligned pairs, every image must contain exactly one component
with area >= `st=100`, AFTER upstream processing (BGR green channel, bilinear
resize, threshold 127). The wrapper rejects zero/multiple valid components,
without dropping images, selecting a largest component, or changing masks.
Additional sub-threshold foreground remains exactly as upstream. This is a
safety gate, NOT a claim that the whole CMDM paper cannot handle multiple objects.
If preflight fails, agree on and disclose a data adaptation before training;
do not reproduce the upstream flattening/misalignment bug as an experiment.

One additional dependency relative to the current project is `blobfile`,
imported by the unchanged upstream preprocessing module:

```bash
python -m pip install blobfile
python train_original_cmdm.py --data_path dataset/scratch \
  --out_dir checkpoints/scratch_official_cmdm_seed0 --check_only
```

Only after preflight passes, start the official-size model on GPU 0:

```bash
CUDA_VISIBLE_DEVICES=0 python -u train_original_cmdm.py \
  --data_path dataset/scratch \
  --out_dir checkpoints/scratch_official_cmdm_seed0 --seed 0
```

To test one full epoch without changing the 2000-epoch LR budget, add
`--stop_after_epoch 1`, then continue in the same directory with
`--resume_path checkpoints/scratch_official_cmdm_seed0/latest.pt` and remove
`--stop_after_epoch`. For a second independent seed, use GPU 1, `--seed 1` and a
different output directory. Never resume the historical four-stage comparator.
`--smoke` is a separate tiny 64px/32-channel/2-diffusion-step integration test;
it is NOT an original-model result and cannot resume into the full protocol.

Only `latest.pt`, `best.pt`, and `provenance.json` are retained by default.
Saving uses a temporary file followed by replacement; budget space for these
checkpoints PLUS one temporary checkpoint per running job. A failed write keeps
the previous complete target. If no epoch was successfully saved, no recovery
checkpoint exists. Interruptions lose progress since the last completed epoch.

The existing `experiment_cont_sweep.py` can load these checkpoints and use the
saved upstream three-dimensional training bank (no nearest-neighbor relabeling
of training images). Evaluate only `area,centroid_x,centroid_y`. This tests the
generator; it does NOT reproduce the paper's medical datasets, downstream
segmentation experiments or KDE/GPR population sampling protocol.

For the FiLM variants, each descriptor dimension is fitted on the training set
and mapped to `[0, 1]` before it enters the FiLM MLP. The canonical
centroid coordinates are kept unchanged for Gaussian-heatmap construction, and
classifier-free conditioning dropout explicitly replaces the model condition
with an all-zero vector after normalization. The
fitted ranges are stored in `condition_norm_stats` inside every checkpoint.
Checkpoints trained before this change still load in their original unscaled
condition space, but must not be resumed as a newly normalized run.

The paper descriptor is ordered as `centroid_x, centroid_y, area,
eccentricity, solidity, circularity, orientation_sin, orientation_cos`.
`aspect_ratio`, `convexity`, `extent`, and `perimeter` are still extracted for
post-generation diagnostics but are not model conditions. Named variants also
fix a four-stage U-Net (`channel_mult=1,2,3,4`) with attention in the bottleneck,
matching the mask-generator diagram in the paper.

Example:

```bash
python train.py --data_path dataset/scratch \
    --variant ours_full --seed 0 \
    --image_size 256 --n_epoch 2000 --n_T 500 --batch_size 2 \
    --save_dir checkpoints/scratch/ours_full_seed0.pt
```

For repeated experiments, run every named variant independently with seeds
`0`, `1`, and `2`; use a unique checkpoint filename for every pair.

The controllability sweep reads the variant, normalization mode, image size,
and component threshold from the checkpoint:

```bash
python experiment_cont_sweep.py \
    --checkpoint checkpoints/scratch/ours_full_seed0.pt \
    --data_path dataset/scratch \
    --sweep_features area,centroid_x,centroid_y,eccentricity,orientation \
    --n_steps 7 --n_repeats 50 --seed 42 --paired_noise \
    --out_dir paper_figures/scratch/ours_full_seed0
```

The sweep uses an observed training descriptor medoid as its baseline by
default, avoiding the potentially invalid synthetic vector produced by taking
each descriptor's mean independently. `--paired_noise` reuses the same initial
and reverse-process noise for each column across condition values, making the
conditional effect easier to isolate.

For `original_cmdm`, only sweep `area,centroid_x,centroid_y`.

## Corrected labels and controllability training protocol

For revision work, the augmentation/validation recipe below is **optional** and
does not reproduce historical experimental settings. The two-GPU launcher in
the next section leaves those additions off. Both launchers use the current
corrected implementation; neither makes old 12-D results equivalent to new 8-D results.

The model still consumes the paper's eight values. Perimeter is only an
intermediate measurement for circularity (and an optional diagnostic). Feature
extraction version 2 computes all quantities on the same largest pixel-connected
object. Binary masks are resized with nearest-neighbor interpolation.

Choose the target explicitly: `--component_policy main` trains only the largest
object, matching its descriptor; `--component_policy all` retains all foreground
but the descriptor still describes the largest object. The default is `all` for
existing command-line compatibility. Original input files are never rewritten.
Validation reports both total area and largest-object area, and component counts.

Orientation is an unoriented major axis in `[-90, 90)` degrees. Its error uses a
180-degree period. Sampling perturbs the angle and reconstructs its sin/cos pair.
Angle perturbations outside fitted sin/cos support fall back to the source angle
instead of clipping the two components independently. Linear trend diagnostics
use one shared angular branch, avoiding a false slope from per-target wrapping;
axial-angle MAE is the primary direction-accuracy measure.
Sweep plots default to the observed training angle P10-P90 (override with
`--orientation_range=-60,60`). Generated shapes with eccentricity below
`--orientation_min_ecc 0.1` have unstable axes and are reported separately.

Training can split sources before augmentation using `--val_fraction 0.2` and
`--split_seed 0`. Identical preprocessed masks stay in the same split. If files
already contain crops, rotations, or related acquisitions, supply
`--split_groups groups.json`: a JSON mapping **every filename** to its original
source/family ID. Related images cannot be inferred reliably from arbitrary
filenames. Keep a separate final test set outside this training directory.

`--augment_copies 4` retains originals and creates four deterministic affine
copies per training source. Controls are remeasured after nearest-neighbor
rotation, isotropic scaling, and translation. Transforms that crop foreground
are retried; an untransformable sample falls back to its original, and the
fallback count is recorded. This is a fixed bank, not a newly sampled transform
each epoch. FiLM min/max statistics are fitted **only on this training bank**;
validation sources are never used. The bank, split, hashes, feature order, and
normalization contract are stored in the checkpoint. Split and augmentation
seeds are independent of the model seed, allowing identical data across ablations.
New checkpoints supply their stored training bank to automatic sampling and
exploratory sweeps even without `--data_path`; this prevents accidental use of
held-out files when recomputing descriptor ranges.

For a first full-model pilot on a single GPU (change the dataset path):

```bash
python train.py --data_path dataset/scratch --variant ours_full --seed 0 \
  --component_policy main --image_size 256 --batch_size 2 --n_T 500 \
  --val_fraction 0.2 --split_seed 0 \
  --augment_copies 4 --augment_seed 0 \
  --augment_scale_min 0.7 --augment_scale_max 1.3 \
  --augment_rotate_deg 90 --augment_translate_frac 0.2 \
  --max_steps 50000 --lrate 1e-4 --dp 0.1 \
  --val_every_steps 5000 --val_baselines 1 --val_steps 3 --val_repeats 1 \
  --val_guide_w 1 --best_metric control --early_stop_patience 0 \
  --save_dir checkpoints/scratch_v2/ours_full_seed0.pt
```

This enables a **small, noisy pilot validation**: one held-out baseline, three
values for each of seven physical controls, and one paired-noise repeat (22
masks per check for paper-8). Increase `--val_baselines` and `--val_repeats` when
selecting the final experimental protocol; validation cost grows proportionally
and each mask uses the full trained diffusion schedule. P10-P90 one-factor
sweeps can still form out-of-distribution *joint* combinations; the observed
held-out descriptors are evaluated separately and range violations are logged.

With `--max_steps`, the optimizer-update budget and step-based linear LR decay
override the epoch limit. Without it, epoch-based training is retained. Training
loss is now a sample-weighted epoch mean; the progress-bar EMA is diagnostic only.
Validation saves/restores RNG states and model mode, so running validation does
not change later training noise or conditional-dropout draws. Checkpoints store
optimizer step, epoch/batch position, loader state, and RNG states for resuming.
`--save_every_steps N` additionally saves periodic step checkpoints. Resuming
requires unchanged data, normalization, training budget, and validation settings.

Outputs for `ours_full_seed0.pt`:

- `ours_full_seed0.pt`: latest training state.
- `ours_full_seed0_best_control.pt`: lowest validation control score.
- `ours_full_seed0_best.pt`: the selected criterion (`--best_metric control`
  uses the same checkpoint as `best_control`; default `train_loss` uses epoch loss).
- `ours_full_seed0_data_manifest.json`: source split, hashes, augmentation settings.
- `ours_full_seed0_validation/step_*.json`: targets, measured descriptors, errors,
  total/main area, failed masks, and control summaries.
- `ours_full_seed0_validation/step_*_masks/`: generated validation masks, without
  component filtering or position correction.

The selection score is an **evaluation-only heuristic**, not an additional loss:
average capped per-feature absolute error divided by training min/max span,
with axial-angle error divided by 90 degrees and invalid masks penalized by 1;
observed and each sweep group receive equal weight. Constant scalar dimensions
are effectively exact-match constraints. Record this rule in the experiment
protocol; scores with different descriptor sets (Original CMDM versus Ours) are
not directly comparable. Sidecar files also report target MAE, trend, off-target
errors, failures, and the rate within 10% of the normalization scale. The old
P10-P90 range-occupancy metric is **not** target success. Reference lines now
default to `y=x`. `--early_stop_patience`, if used with control selection, counts
validation checks; with train-loss selection it counts epochs.

Compare existing CFG weights without retraining, using exactly the checkpoint's
held-out conditions and noise seeds:

```bash
python experiment_cfg_sweep.py \
  --checkpoint checkpoints/scratch_v2/ours_full_seed0_best_control.pt \
  --guide_weights 0,0.5,1,2,3 --out_dir paper_figures/scratch_v2/cfg_seed0
```

`guide_w=0` is conditional sampling; the effective conventional CFG scale is
`1 + guide_w`. The runner writes `cfg_summary.csv`, detailed JSON, and masks.
Use validation to choose a weight and report performance on untouched test data.

To prepare five one-factor pilot jobs (base LR/dropout, two alternative LRs, two
alternative dropout probabilities), first preview commands:

```bash
python run_control_experiments.py --mode pilot --data_path dataset/scratch \
  --component_policy main --seeds 0 --out_dir checkpoints/scratch_v2/pilot
```

Add `--execute` to run sequentially. After fixing the shared protocol, use
`--mode ablation --seeds 0,1,2` for the 15 variant/seed jobs, supplying the chosen
`--lrate`, `--dp`, and `--val_guide_w`. Every job uses the same split, augmentation
bank, batch size, and optimizer-update budget; output names include variant,
seed, LR, and dropout. Existing checkpoints are not overwritten by the launcher.

Corrected feature extraction changes some labels even for an eight-dimensional
checkpoint. Older checkpoints remain loadable for inference, but resuming them
under the new data protocol is rejected; start a new training run and output
directory. Structural/route tests do not establish achieved controllability.

## One-seed controlled improvement pilot

For exactly one new `ours_full / seed 0` run with unchanged model and loss,
use `run_single_seed_pilot.py`, not the five-job hyperparameter pilot above.
It verifies the reference checkpoint and source-mask hashes, performs CPU
preflight, preserves model/optimizer settings, enables train-only relabeled
augmentation and held-out selection, and disables periodic checkpoint archives.
See [SINGLE_SEED_PILOT.md](SINGLE_SEED_PILOT.md) for the frozen protocol,
cloud commands, output paths, and same-plan old/new CFG comparison caveats.

## Two-GPU revision experiments

Use one independent training job per GPU; no DDP or global-batch change is
required. `train.py` respects external `CUDA_VISIBLE_DEVICES`. For example,
`CUDA_VISIBLE_DEVICES=1 python train.py ...` uses physical/assigned GPU 1 as
logical `cuda:0`. Use the IDs/UUIDs assigned by your cloud environment.

The revision launcher defaults to two short jobs: Full seed 0 and Original CMDM
seed 0, each for 10 epochs. Model size, diffusion steps, and per-job batch match
the normal run. These only test GPU memory, training, and checkpoint I/O.

```bash
python run_revision_experiments.py --data_path dataset/scratch \
  --out_dir checkpoints/scratch_revision_smoke --gpus 0,1 --preset smoke
```

This previews the commands. Add `--execute` to start both jobs. The launcher
checks GPU visibility, writes a log per experiment under `logs/`, and schedules
at most one training process per GPU. Failed jobs stop further dispatch; jobs
already running on other GPUs are allowed to finish. Use a new/empty output
directory to prevent overwriting older experiments. `experiment_plan.json`
records the commands, and `experiment_status.json` records completion states.

After the smoke tests, run only the first two full-length jobs if desired:

```bash
python run_revision_experiments.py --data_path dataset/scratch \
  --out_dir checkpoints/scratch_revision_first_pair --gpus 0,1 \
  --preset ablation --variants ours_full,original_cmdm --seeds 0 --execute
```

The ablation preset defaults to 2000 epochs. Without `--variants`/`--seeds`, it
queues all five variants and seeds 0,1,2 (15 jobs). Do not run that full queue
automatically if the first pair still needs inspection. To run only the remaining
13 jobs after accepting that pair, finish seed 0's three missing variants in one
fresh directory, then all five variants at seeds 1,2 in another:

```bash
python run_revision_experiments.py --data_path dataset/scratch \
  --out_dir checkpoints/scratch_revision_remaining_seed0 --gpus 0,1 \
  --preset ablation --variants ours_wo_heatmap,ours_wo_film,ours_wo_fg \
  --seeds 0 --execute
python run_revision_experiments.py --data_path dataset/scratch \
  --out_dir checkpoints/scratch_revision_seeds12 --gpus 0,1 \
  --preset ablation --seeds 1,2 --execute
```

Every revision-launcher job explicitly disables new augmentation, validation
splitting, and the step-budget override. It retains all mask components by
default, trains by epoch, and selects `best.pt` using the current epoch-mean
training loss. This is a **controlled supplementary run of corrected code**,
not exact reproduction of the historical EMA-selected/12-D experiment. Keep
existing image-generation and downstream results separate until their model
provenance is checked. Short smoke checkpoints cannot resume as 2000-epoch
runs because their LR schedules differ; use the new full-length output directory.

Monitor on Linux with `nvidia-smi` and, for example,
`tail -f checkpoints/scratch_revision_smoke/logs/ours_full_seed0.log`.
GPU utilization and speed must be verified on the cloud host; local tests of the
launcher do not validate the availability of two physical cloud GPUs.

## Sampling with CMDM
To sample with CMDM, set the number of samples you need via ```n_samples``` in ```sampling_si.py```, then run:
```bash
python sampling_mask.py
```
As proposed in the paper, ```sampling_si.py``` first analyzes the correlations among semantic information (e.g., size and location) within the given annotation masks, and ```sampling_mask.py``` then generates the corresponding annotation masks.

## Acknowledgement
This repository is based on [LDM](https://github.com/CompVis/latent-diffusion), [guided-diffusion](https://github.com/openai/guided-diffusion), [CFG](https://github.com/TeaPearce/Conditional_Diffusion_MNIST) and [SDM](https://github.com/WeilunWang/semantic-diffusion-model). We sincerely thank the original authors for their valuable contributions and outstanding work.


## Citations
```
@article{heo2025controllable,
  title={Controllable Mask Diffusion Model for medical annotation synthesis with semantic information extraction},
  author={Heo, Chanyeong and Jung, Jaehee},
  journal={Computers in Biology and Medicine},
  volume={196},
  pages={110807},
  year={2025},
  publisher={Elsevier}
  doi={https://doi.org/10.1016/j.compbiomed.2025.110807}}
```