CLAUDE.md

# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## Project Overview

OmniMorph is a medical image framework for generation, restoration, and registration using a conditional Deformation-Recovery Diffusion Model (DeformDDPM). It supports 2D and 3D multi-modal medical imaging (CT, MRI, PET) with text-conditioned generation via BERT embeddings.

## Common Commands

```bash
# Training (single-mode diffusion)
CUDA_VISIBLE_DEVICES=0 python OM_train.py -C Config/config_om.yaml

# Training (dual-mode: diffusion + registration)
CUDA_VISIBLE_DEVICES=0,1 python OM_train_2modes.py -C Config/config_om.yaml

# Contrastive learning (text-image alignment)
python OM_contrastive.py -C Config/config_om_contrastive.yaml

# XPU testing with dummy data (no real dataset needed)
python OM_contrastive_xpu.py --dummy-samples 20

# Augmentation / inference with a trained model
python OM_aug.py -C Config/config_om.yaml

# Background training (production style)
nohup python -u OM_train_2modes.py -C Config/config_om.yaml > train_log.txt 2>&1 &
```

## Architecture

### Core Pipeline

```
Config YAML → DataLoader(s) → DeformDDPM(Network, STN) → Loss → Checkpoint
```

### Diffusion Module (`Diffusion/`)

- **diffuser.py** — `DeformDDPM`: forward/reverse diffusion over deformation vector fields (DVFs). Generates multi-scale DDFs via control points at ratios [4, 8, 16, 32, 64]. Key methods: `diffuse()`, `recover()`.
- **networks.py** — Network architectures selected by `get_net(net_name)`:
  - `recresacnet` — Atrous convolution UNet (2D CMR)
  - `recmutattnnet` — Multi-head attention network (main 3D, channels [1,16,32,64,128,256])
  - `recmutattnnet_contrastive` — Outputs 1024-dim image embeddings for contrastive training
  - `defrecmutattnnet` — Deformable variant
- **networks.py: `STN`** — Spatial Transformer Network for differentiable image warping via DDFs. Composes deformations: `comp_ddf = dvf + stn(ddf, dvf)`.
- **losses.py** — `Grad` (L1 + negative Jacobian determinant + range penalties), `LNCC` (local normalized cross-correlation), `LMSE` (labeled MSE), `NCC`, `MRSE`.

### Training Modes

| Script | Purpose | DataLoader | Key Loss |
|--------|---------|------------|----------|
| `OM_train.py` | Single diffusion | `OminiDataset_v1` | Grad + MRSE + NCC |
| `OM_train_2modes.py` | Diffusion + registration | `OMDataset_indiv` + `OMDataset_pair` | Above + LNCC + LMSE |
| `OM_train_3modes.py` | Extended dual-mode | Same as 2modes | Different loss weights |
| `OM_contrastive.py` | Text-image alignment | `OMDataset_indiv` | Cosine similarity |
| `OM_reg.py` | Registration only | Paired data | Registration losses |
| `OM_train_uncon.py` | Unconditional generation | Generic | Standard |

All DDP-enabled training scripts use NCCL backend on `localhost:12355`.

### DataLoader (`Dataloader/`)

- **dataLoader.py** — All dataset classes. Data comes from JSON mapping files in `nifty_mappings/` that map NIfTI file paths to metadata (Modality, ROI, Size, Spacing_mm, BERT embeddings).
  - `OMDataset_indiv` → returns `[volume, embd]` (shape: `[1,sz,sz,sz]`, `[1024]`)
  - `OMDataset_pair` → returns `[volume_A, volume_B, embd_A, embd_B]`
  - `DummyOMDataset_indiv` / `DummyOMDataset_pair` → random tensors for XPU testing without data
- **dataloader_utils.py** — `get_sizeRange_dict()` for ROI-based filtering, image thresholding, DICOM reading.
- **bert_helper.py** / **embding_gen.py** — BERT text embedding generation.
- Filtering chain: min dimension → modality → ROI → label presence.

### Config (`Config/`)

YAML files with keys: `data_name`, `net_name`, `ndims` (2 or 3), `img_size`, `batchsize`, `timesteps` (default 80), `v_scale`, `lr`, `epoch`, `noise_scale`, `condition_type` (`'uncon'`, `'adding'`, `'project'`, etc.), augmentation params (`start_noise_step`, `noise_step`, `aug_coe`).

### Augmentation (`OM_aug.py`, `OM_aug_highres.py`)

Loads a trained checkpoint and generates augmented samples. Controlled by `start_noise_step` (higher = less deformation), `aug_coe` (samples per input, typically 32-64). Outputs saved to `Data/Aug_data/{dataset}/img|msk|ddf/`.

## Key Conventions

- Models saved as `{epoch:06d}_{data_name}_{net_name}.pth` containing `model_state_dict`, `optimizer_state_dict`, `epoch`.
- CT images clamped to [-400, 400] HU before normalization.
- SimpleITK axis order is reversed from NumPy (`reverse_axis_order()`).
- Mapping JSON files in `nifty_mappings/` are Git LFS tracked (large files).
- `utils.py` provides `get_transformer()` for random affine augmentations and `get_random_deformed_mask()` for blind masks.

## Dependencies

PyTorch 1.12+ with CUDA, SimpleITK, nibabel, scikit-image, einops, pydicom, transformers (HuggingFace), swanlab (optional, for experiment tracking). See `requirements.txt`.