Add dummy datasets for XPU testing, XPU contrastive training script, and CLAUDE.md

- DummyOMDataset_indiv/pair in dataLoader.py for testing without real data
- OM_contrastive_xpu.py with XPU/CUDA/CPU auto-detection
- CLAUDE.md for codebase guidance

CLAUDE.md

@@ -0,0 +1,91 @@
+# 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`.

Dataloader/dataLoader.py

@@ -74,6 +74,47 @@ def sample_random_uniform_multi_order(high=1., low=0., order_num=2, type='high')
             sample_value = np.random.uniform(low, high=sample_value)
     return sample_value
 
+class DummyOMDataset_indiv(Dataset):
+    """Dummy dataset that generates random 3D volumes and embeddings for XPU testing."""
+    def __init__(self, out_sz=128, num_samples=100, embd_dim=1024, transform=None):
+        self.out_sz = out_sz
+        self.num_samples = num_samples
+        self.embd_dim = embd_dim
+        self.transform = transform
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, idx):
+        volume = np.random.rand(1, self.out_sz, self.out_sz, self.out_sz).astype(np.float64)
+        embd = np.random.randn(self.embd_dim).astype(np.float32)
+        if self.transform is not None:
+            volume = self.transform(volume)
+        return volume, embd
+
+
+class DummyOMDataset_pair(Dataset):
+    """Dummy dataset that generates random paired 3D volumes and embeddings for XPU testing."""
+    def __init__(self, out_sz=128, num_samples=100, embd_dim=1024, transform=None):
+        self.out_sz = out_sz
+        self.num_samples = num_samples
+        self.embd_dim = embd_dim
+        self.transform = transform
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, idx):
+        volume_A = np.random.rand(1, self.out_sz, self.out_sz, self.out_sz).astype(np.float64)
+        volume_B = np.random.rand(1, self.out_sz, self.out_sz, self.out_sz).astype(np.float64)
+        embd_A = np.random.randn(self.embd_dim).astype(np.float32)
+        embd_B = np.random.randn(self.embd_dim).astype(np.float32)
+        if self.transform is not None:
+            volume_A = self.transform(volume_A)
+            volume_B = self.transform(volume_B)
+        return [volume_A, volume_B, embd_A, embd_B]
+
+
 class OminiDataset(object):
     """Base class for OmniMorph datasets."""
     def init(self, out_sz, transform, clamp_range, min_crop_ratio, ROIs, modality,reverse_axis_order ,min_dim,mapping_files):

OM_contrastive_xpu.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn.functional as F
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+from Diffusion.networks import get_net
+from Dataloader.dataLoader import DummyOMDataset_indiv
+import argparse
+import yaml
+import os
+import time
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--config", "-C", type=str, default="Config/config_om_contrastive.yaml")
+parser.add_argument("--dummy-samples", type=int, default=100, help="Number of dummy samples")
+args = parser.parse_args()
+
+with open(args.config, 'r') as file:
+    hyp = yaml.safe_load(file)
+
+# Setup device: prefer XPU, fallback to CUDA, then CPU
+if hasattr(torch, 'xpu') and torch.xpu.is_available():
+    device = torch.device('xpu')
+    print(f"Using XPU device: {torch.xpu.get_device_name(0)}")
+elif torch.cuda.is_available():
+    device = torch.device(hyp['device'])
+    print(f"Using CUDA device")
+else:
+    device = torch.device('cpu')
+    print(f"Using CPU device")
+
+data_name = hyp['data_name']
+net_name = hyp['net_name']
+ndims = hyp['ndims']
+img_size = hyp['img_size']
+model_save_path = os.path.join('Models', f'{data_name}_{net_name}/')
+os.makedirs(model_save_path, exist_ok=True)
+
+# Model
+Net = get_net(net_name)
+model = Net(n_steps=hyp['timesteps'], ndims=ndims, num_input_chn=hyp['num_input_chn'], res=img_size).to(device)
+optimizer = Adam(model.parameters(), lr=hyp['lr'])
+
+# Data - dummy dataset for XPU testing
+dataset = DummyOMDataset_indiv(out_sz=img_size, num_samples=args.dummy_samples)
+train_loader = DataLoader(dataset, batch_size=hyp['batchsize'], shuffle=True, drop_last=True)
+
+# Training
+print(f'Start training on {device} with {len(dataset)} dummy samples...')
+for epoch in range(hyp['epoch']):
+    epoch_loss = 0.0
+
+    for i, (volume, embd) in enumerate(train_loader):
+        t0 = time.time()
+        volume = volume.float().to(device)
+        embd = embd.to(device)  # [B, 1024] GT text embedding
+        t = torch.randint(0, hyp['timesteps'], (volume.shape[0],)).to(device)
+
+        _, img_embd = model(x=volume, y=volume, t=t)  # img_embd: [B, 1024]
+
+        # Cosine similarity loss: align img_embd with GT text embedding
+        loss = 1 - F.cosine_similarity(img_embd, embd, dim=-1).mean()
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        epoch_loss += loss.item()
+        t1 = time.time()
+        dt = t1 - t0
+        print(f"  Batch {i:04d} | Loss: {loss.item():.6f} | Time: {dt:.2f}s")
+    avg_loss = epoch_loss / max(len(train_loader), 1)
+    print(f"Epoch {epoch:04d} | Avg Loss: {avg_loss:.6f}")