Initial upload: OmniMorph codebase

.gitattributes

@@ -33,3 +33,16 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Dataloader/PSMA-CT_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/AbdomenAtlas_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/AbdomenCT1k_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/Brats2019_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/Brats2020_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/Brats2021_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/CIA_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/MSD_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/OASIS_1_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/OASIS_2_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/PSMA-CT-Longitud_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/PSMA-FDG-PET-CT-LESION_mappings.json filter=lfs diff=lfs merge=lfs -text
+Dataloader/nifty_mappings/TotalSegmentorCT_MRI_mappings.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,29 @@
+# Model checkpoints
+Models/
+
+# Data files
+Data/
+
+# Python cache
+__pycache__/
+
+# Virtual environment
+ominenv/
+
+# External libraries
+External/
+
+# Logs
+Log/
+swanlog/
+train_log.txt
+aug_log.txt
+
+# Reference implementation
+def_diff_rec/
+
+# IDE
+.vscode/
+
+# Misc
+CLAUDE.md

Config/config_cmr.yaml

@@ -0,0 +1,29 @@
+data_name: cmr
+net_name: recresacnet
+ndims: 2
+img_size: 256
+batchsize: 1
+ddf_pad_mode: border
+device: cuda
+img_pad_mode: zeros
+num_input_chn: 1
+padding_mode: zeros
+resample_mode: bicubic
+timesteps: 80
+v_scale: 4.0e-05
+# =========================
+# TRAINING SETTING
+epoch: 10000
+epoch_per_save: 1
+lr: 0.0001
+noise_scale: 0.1
+# =========================
+# AUGMENTATION SETTING
+patients_list: []
+model_id_str: '000000'
+start_noise_step: 48
+noise_step: 2
+aug_coe: 32                  # how many times each sample will be augmented
+aug_img_savepath: Data/Aug_data/cmr/img/
+aug_msk_savepath: Data/Aug_data/cmr/msk/
+aug_ddf_savepath: Data/Aug_data/cmr/ddf/

Config/config_lct.yaml

@@ -0,0 +1,31 @@
+data_name: lct
+net_name: recmutattnnet
+# net_name: recresacnet
+ndims: 3
+img_size: 128 #was 128
+batchsize: 2
+ddf_pad_mode: border
+device: cuda
+img_pad_mode: zeros
+num_input_chn: 1
+padding_mode: border
+resample_mode: bilinear
+timesteps: 80
+v_scale: 4.0e-05
+# =========================
+# TRAINING SETTING
+epoch: 10000
+epoch_per_save: 1
+lr: 0.00001
+noise_scale: 0.1
+# =========================
+# AUGMENTATION SETTING
+patients_list: []
+model_id_str: '001157'
+start_noise_step: 64        
+noise_step: 1
+aug_coe: 32                  # how many times each sample will be augmented
+condition_type: 'project'       # 'None', 'none', 'adding','independ', 'downsample', 'slice', 'project', 'uncon'
+aug_img_savepath: Data/Aug_data/lct/img/
+aug_msk_savepath: Data/Aug_data/lct/msk/
+aug_ddf_savepath: Data/Aug_data/lct/ddf/

Config/config_om.yaml

@@ -0,0 +1,53 @@
+data_name: all
+# net_name: recresacnet
+net_name: recmutattnnet
+# net_name: recmutattnnet1
+# net_name: defrecmutattnnet
+ndims: 3
+img_size: 128 
+batchsize: 2
+ddf_pad_mode: border
+device: cuda
+img_pad_mode: zeros
+num_input_chn: 1
+padding_mode: border
+resample_mode: bilinear
+timesteps: 80
+v_scale: 5.0e-05
+# =========================
+# TRAINING SETTING
+epoch: 10000
+epoch_per_save: 1
+lr: 0.00001
+noise_scale: 0.1
+# =========================
+# AUGMENTATION SETTING
+patients_list: []
+# model_id_str: '000000'
+# model_id_str: '000180' # before registration training
+# model_id_str: '000353'  # good augmentation results on msd
+model_id_str: '000354'  # 
+# model_id_str: '000157'
+# model_id_str: '000171'
+start_noise_step: 48      # starting from which noise step to add noise
+noise_step: 1
+aug_coe: 64                  # how many times each sample will be augmented
+# start_noise_step: 56      # starting from which noise step to add noise
+# noise_step: 4
+# aug_coe: 4                  # how many times each sample will be augmented
+condition_type: 'uncon'       # 'None', 'none', 'adding','independ', 'downsample', 'slice', 'project', 'uncon'
+# aug_img_savepath: Data/Aug_data/totseg/img/
+# aug_msk_savepath: Data/Aug_data/totseg/msk/
+# aug_ddf_savepath: Data/Aug_data/totseg/ddf/
+# aug_img_savepath: Data/Aug_data/om/img/
+# aug_msk_savepath: Data/Aug_data/om/msk/
+# aug_ddf_savepath: Data/Aug_data/om/ddf/
+reg_img_savepath: Data/Reg_data/om/img/
+reg_msk_savepath: Data/Reg_data/om/msk/
+reg_ddf_savepath: Data/Reg_data/om/ddf/
+# aug_img_savepath: Data/Aug_data/msd/img/
+# aug_msk_savepath: Data/Aug_data/msd/msk/
+# aug_ddf_savepath: Data/Aug_data/msd/ddf/
+aug_img_savepath: Data/Aug_data/mnms/img/
+aug_msk_savepath: Data/Aug_data/mnms/msk/
+aug_ddf_savepath: Data/Aug_data/mnms/ddf/

Config/config_om_contrastive.yaml

@@ -0,0 +1,51 @@
+data_name: all
+# net_name: recresacnet
+# net_name: recmutattnnet
+net_name: recmutattnnet_contrastive
+# net_name: recmutattnnet1
+# net_name: defrecmutattnnet
+ndims: 3
+img_size: 128 
+batchsize: 1 #1 for testing
+ddf_pad_mode: border
+device: cuda
+img_pad_mode: zeros
+num_input_chn: 1
+padding_mode: border
+resample_mode: bilinear
+timesteps: 80
+v_scale: 5.0e-05
+# =========================
+# TRAINING SETTING
+epoch: 10000
+epoch_per_save: 1
+lr: 0.00001
+noise_scale: 0.1
+# =========================
+# AUGMENTATION SETTING
+patients_list: []
+# model_id_str: '000000'
+# model_id_str: '000180' # before registration training
+# model_id_str: '000353'  # good augmentation results on msd
+model_id_str: '000354'  # 
+# model_id_str: '000157'
+# model_id_str: '000171'
+start_noise_step: 48      # starting from which noise step to add noise
+noise_step: 1
+aug_coe: 64                  # how many times each sample will be augmented
+# start_noise_step: 56      # starting from which noise step to add noise
+# noise_step: 4
+# aug_coe: 4                  # how many times each sample will be augmented
+condition_type: 'uncon'       # 'None', 'none', 'adding','independ', 'downsample', 'slice', 'project', 'uncon'
+# aug_img_savepath: Data/Aug_data/totseg/img/
+# aug_msk_savepath: Data/Aug_data/totseg/msk/
+# aug_ddf_savepath: Data/Aug_data/totseg/ddf/
+# aug_img_savepath: Data/Aug_data/om/img/
+# aug_msk_savepath: Data/Aug_data/om/msk/
+# aug_ddf_savepath: Data/Aug_data/om/ddf/
+reg_img_savepath: Data/Reg_data/om/img/
+reg_msk_savepath: Data/Reg_data/om/msk/
+reg_ddf_savepath: Data/Reg_data/om/ddf/
+aug_img_savepath: Data/Aug_data/msd/img/
+aug_msk_savepath: Data/Aug_data/msd/msk/
+aug_ddf_savepath: Data/Aug_data/msd/ddf/

Dataloader/PSMA-CT_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4fbbdc9b4b48688a37c4f828eea2823820a1ee27f954d5987d8cbf3b67d6d9bf
+size 179285490

Dataloader/bert_helper.py

@@ -0,0 +1,258 @@
+import os
+import numpy as np
+import math
+from torch.nn import Tanh, BatchNorm1d
+from typing import Optional
+import torch.nn as nn
+import torch
+from transformers import BertModel, BertForSequenceClassification
+from transformers import BertTokenizer
+from transformers import AutoTokenizer, AutoModel
+
+from torch.utils.data import Dataset as Dataset_n
+from torch.utils.data import DataLoader as DataLoader_n
+from torch.utils.data import WeightedRandomSampler
+
+def _freeze_bert(
+        bert_model: BertModel, freeze_bert=True, freeze_layer_count=-1
+):
+    """Freeze parameters in BertModel (in place)
+    Args:
+        bert_model: HuggingFace bert model
+        freeze_bert: Bool whether to freeze the bert model
+        freeze_layer_count: If freeze_bert, up to what layer to freeze.
+    Returns:
+        bert_model
+    """
+    if freeze_bert:
+        # freeze the entire bert model
+        for param in bert_model.parameters():
+            param.requires_grad = False
+    else:
+        # freeze the embeddings
+        for param in bert_model.embeddings.parameters():
+            param.requires_grad = False
+        if freeze_layer_count != -1:
+            if freeze_layer_count > 0 :
+                # freeze layers in bert_model.encoder
+                for layer in bert_model.encoder.layer[:freeze_layer_count]:
+                    for param in layer.parameters():
+                        param.requires_grad = False
+
+            if freeze_layer_count < 0 :
+                # freeze layers in bert_model.encoder
+                for layer in bert_model.encoder.layer[freeze_layer_count:]:
+                    for param in layer.parameters():
+                        param.requires_grad = False
+    return None
+
+def get_frozen_embeder(key_word="bert-large-uncased"):
+    tokenizer = AutoTokenizer.from_pretrained(key_word, do_lower_case=False)
+    model = AutoModel.from_pretrained(key_word)
+
+    _freeze_bert(model, freeze_bert=True, freeze_layer_count=None)
+    return model, tokenizer
+
+
+def str2emb(string, max_words_num=100, embeder=None, tokenizer=None, reduce_method='mean'):
+    string = string.lower()
+    str_token = tokenizer(string, return_tensors='pt', max_length=max_words_num,
+                          padding='max_length', truncation=True)    
+    embeder_output = embeder(**str_token)
+    if reduce_method == 'mean':
+        embeder_output = torch.mean(embeder_output.last_hidden_state, dim=1)
+    elif reduce_method == 'max':
+        embeder_output = torch.max(embeder_output.last_hidden_state, dim=1)[0]
+    else:
+        embeder_output = embeder_output.last_hidden_state
+    return embeder_output
+
+def get_synonyms_dict(dict_type=None):
+    '''
+    Get the dictionary of synonyms for the specified dictionary type
+    '''
+    if dict_type == 'ROI':
+        dict_synonyms = {
+            'whole-body': ['whole-body', 'whole body', 'wholebody', 'whole body', 'whole-body', 'whole body', 'wholebody','polytrauma','head-neck-thorax-abdomen-pelvis-leg','head-neck-thorax-abdomen-pelvis'],
+            'neck-thorax-abdomen-pelvis-leg': ['neck-thorax-abdomen-pelvis-leg','neck-thx-abd-pelvis-leg', 'angiography neck-thx-abd-pelvis-leg', 'neck thorax abdomen pelvis leg', 'neck and thorax and abdomen and pelvis and leg', 'neck, thorax, abdomen, pelvis & leg', 'neck/thorax/abdomen/pelvis/leg', 'neck, thorax, abdomen, pelvis and leg', 'neck thorax abdomen pelvis leg'],
+            'neck-thorax-abdomen-pelvis': ['neck-thorax-abdomen-pelvis', 'neck-thx-abd-pelvis', 'neck thorax abdomen pelvis', 'neck and thorax and abdomen and pelvis', 'neck, thorax, abdomen & pelvis', 'neck/thorax/abdomen/pelvis', 'neck, thorax, abdomen and pelvis', 'neck thorax abdomen & pelvis'],
+            'thorax-abdomen-pelvis-leg': ['thorax-abdomen-pelvis-leg','thx-abd-pelvis-leg', 'angiography thx-abd-pelvis-leg', 'thorax abdomen pelvis leg', 'thorax and abdomen and pelvis and leg', 'thorax, abdomen, pelvis & leg', 'thorax/abdomen/pelvis/leg', 'thorax, abdomen, pelvis and leg', 'thorax abdomen pelvis leg'],
+            'neck-thorax-abdomen': ['neck-thorax-abdomen', 'neck-thorax-abdomen', 'neck thorax abdomen', 'neck and thorax and abdomen', 'neck, thorax, abdomen', 'neck/thorax/abdomen', 'neck, thorax, abdomen', 'neck thorax abdomen'],
+            'head-neck-thorax-abdomen': ['head-neck-thorax-abdomen', 'head-neck-thorax-abdomen', 'head neck thorax abdomen', 'head and neck and thorax and abdomen', 'head, neck, thorax, abdomen', 'head/thorax/abdomen', 'head, thorax, abdomen', 'head thorax abdomen'],
+            'head-neck-thorax': ['head-neck-thorax', 'head neck thorax', 'head and neck and thorax', 'head, neck, thorax', 'head/thorax', 'head, thorax', 'head thorax'],
+            'thorax-abdomen-pelvis': ['thorax-abdomen-pelvis', 'thx-abd-pelvis', 'polytrauma', 'thorax abdomen pelvis', 'thorax and abdomen and pelvis', 'thorax, abdomen & pelvis', 'thorax/abdomen/pelvis', 'thorax, abdomen and pelvis', 'thorax abdomen & pelvis'],
+            'abdomen-pelvis-leg': ['abdomen-pelvis-leg', 'angiography abdomen-pelvis-leg', 'abd-pelvis-leg', 'abdomen pelvis leg', 'abdomen and pelvis and leg', 'abdomen, pelvis & leg', 'abdomen/pelvis/leg', 'abdomen, pelvis, leg', 'abdomen pelvis leg'],
+            'neck-thorax': ['neck-thorax', 'neck thorax', 'neck and thorax', 'neck, thorax', 'thorax-neck', 'thorax neck', 'thorax and neck', 'thorax, neck','thorax/neck'],
+            'thorax-abdomen': ['thorax-abdomen', 'thorax abdomen', 'thorax and abdomen', 'thorax, abdomen', 'aortic valve'],
+            'abdomen-pelvis': ['abdomen-pelvis', 'abdomen pelvis', 'abdomen and pelvis', 'abdomen & pelvis', 'abdomen/pelvis', 'abdomen-pelvis', 'abdomen pelvis', 'abdomen and pelvis', 'abdomen & pelvis', 'abdomen/pelvis'],
+            'pelvis-leg': ['pelvis-leg', 'pelvis leg', 'pelvis and leg', 'pelvis, leg', 'pelvis/leg', 'pelvis-leg', 'pelvis leg', 'pelvis and leg', 'pelvis, leg', 'pelvis/leg'],
+            'head-neck': ['head-neck', 'head neck', 'head and neck', 'head, neck', 'head/neck', 'head-neck', 'head neck', 'head and neck', 'head, neck', 'head/neck'],
+            'abdomen': ['abdomen', 'abdominal', 'belly', 'stomach', 'tummy', 'gut', 'guts', 'viscera', 'bowels', 'intestines', 'gastrointestinal', 'digestive', 'peritoneum','gastric', 'liver', 'spleen', 'pancreas','kidney','lumbar','renal','hepatic','splenic','pancreatic','intervention'],
+            'thorax': ['chest', 'thorax', 'breast', 'lung', 'heart','heart-thorakale aorta', 'heart-thorakale', 'mediastinum', 'pleura', 'bronchus', 'bronchi', 'trachea', 'esophagus', 'diaphragm', 'rib', 'sternum', 'clavicle', 'scapula', 'axilla', 'armpit','breast biopsy','thoracic','mammary','caeiothoracic','mediastinal','pleural','bronchial','bronchial tree','tracheal','esophageal','diaphragmatic','costal','sternal','clavicular','scapular','axillary','axillar','cardiac','pericardial','pericardiac','pericardium'],
+            'head': ['head', 'headbasis', 'brain', 'skull', 'face','nose','ear','eye','mouth','jaw','cheek','chin','forehead','temporal','parietal','occipital','frontal','mandible','maxilla','mandibular','maxillary','nasal','orbital','orbita','ocular','auricular','otic','oral','buccal','labial','lingual','palatal'],
+            'neck': ['neck', 'throat', 'cervical', 'thyroid', 'trachea', 'larynx', 'pharynx', 'esophagus','pharyngeal','laryngeal','cervical','thyroid','trachea','esophagus','carotid','jugular'],
+            'hand': ['hand', 'finger', 'thumb', 'palm', 'wrist', 'knuckle', 'fingernail', 'phalanx', 'metacarpal', 'carpal', 'radius'],
+            'arm': ['arm', 'forearm', 'upper arm', 'bicep', 'tricep', 'brachium', 'brachial', 'humerus', 'radius', 'ulna', 'elbow', 'shoulder', 'armpit''clavicle', 'scapula', 'acromion', 'acromioclavicular'],
+            'leg': ['leg', 'felsenleg','thigh', 'calf', 'shin', 'knee', 'foot', 'ankle', 'toe', 'heel', 'sole', 'arch', 'instep', 'metatarsal', 'phalanx', 'tibia', 'fibula', 'femur', 'patella', 'kneecap','achilles tendon','achilles'],
+            'pelvis': ['pelvis', 'hip', 'groin', 'buttock', 'gluteus', 'gluteal', 'ischium', 'pubis', 'sacrum', 'coccyx', 'acetabulum', 'iliac', 'iliac crest', 'iliac spine', 'iliac wing', 'sacroiliac', 'sacroiliac joint', 'sacroiliac ligament', 'sacroiliac spine', 'ureter', 'bladder', 'urethra', 'prostate', 'testicle', 'ovary', 'uterus',],
+            'skeleton': ['skeleton','bone','spine', 'back', 'vertebra', 'sacrum', 'coccyx'],
+        }
+    elif dict_type == 'Label_tissue':
+        dict_synonyms = {
+            'liver': ['liver','hepatic'],
+            'spleen': ['spleen','splenic'],
+            'kidney': ['kidney','renal'],
+            'pancreas': ['pancreas','pancreatic'],
+            'stomach': ['stomach','gastric'],
+            'intestine': ['large intestine', 'small intestine','large bowel','small bowel'],
+            'gallbladder': ['gallbladder'],
+            'adrenal_gland': ['adrenal_gland','adrenal gland'],
+            'bladder': ['bladder'],
+            'prostate': ['prostate'],
+            'uterus': ['uterus'],
+            'ovary': ['ovary'],
+            'testicle': ['testicle'],
+            'lymph_node': ['lymph_node','lymph node'],
+            'bone': ['bone'],
+            'lung': ['lung'],
+            'heart': ['heart'],
+            'esophagus': ['esophagus'],
+            'muscle': ['muscle'],
+            'fat': ['fat'],
+            'skin': ['skin'],
+            'vessel': ['vessel'],
+            'tumor': ['tumor'],
+            'other': ['other']
+        }
+    elif dict_type == 'Task':
+        dict_synonyms = {
+            'segmentation': ['segmentation', 'seg', 'mask'],
+            'classification': ['classification', 'class', 'diagnosis','identify','identification'],
+            'localization': ['localization', 'locate', 'location', 'position'],
+            'registration': ['registration', 'register', 'align', 'alignment'],
+            'detection': ['detection', 'detect', 'find', 'locate'],
+            'quantification': ['quantification', 'quantify', 'measure', 'measurement'],
+        }
+    elif dict_type == 'Modality':
+        dict_synonyms = {
+            'CT': ['CT', 'computed tomography'],
+            'MRI': ['MRI', 'MR', 'magnetic resonance imaging'],
+            'PET': ['PET', 'positron emission tomography'],
+            'US': ['US', 'ultrasound'],
+            'X-ray': ['X-ray', 'radiography'],
+            'SPECT': ['SPECT', 'single-photon emission computed tomlogy'],
+        }
+    else:
+        dict_synonyms = {
+            '\'gender\'': ['\'gender\'', '\'sex\'', '\'M/F\'', '\'m/f\''],
+            '\'modality\'': ['\'modality\'', '\'modal\''],
+            '\'male\'': ['\'male\'', '\'m\''],
+            '\'female\'': ['\'female\'', '\'f\'','\'woman\''],
+            '\'high-grade glioma\'': ['\'high-grade glioma\'', '\'high grade glioma\'', '\'HGG\''],
+            '\'low-grade glioma\'': ['\'low-grade glioma\'', '\'low grade glioma\'', '\'LGG\''],
+            '\'atlas scaling factor\'': ['\'atlas scaling factor\'', '\'asf\''],
+            '\'age\'': ['\'age\'', '\'years\'', '\'year\'', '\'y/o\'', '\'y.o.\''],
+            '\'education\'': ['\'educ\'', '\'educat\'', '\'education\''],
+            '\'roi\'': ['\'roi\'', '\'region of interest\'', '\'region\''],
+            '\'mini-mental state examination\'': ['\'mini-mental state examination\'', '\'mmse\''],
+            '\'clinical dementia rating\'': ['\'clinical dementia rating\'', '\'cdr\''],
+            '\'socio-economic status\'': ['\'socio-economic status\'', '\'ses\''],
+            '\'unknown\'': ['\'unknown\'', '\'unkn\'', '\'not available\'', '\'nan\'', '\'n/a\'', '\'none\'', '\'n.a.\'', '\'not applicable\'','\'not specified\'', '\'unspecified\'', '\'not given\'', '\'null\''],
+            '': [' segmentation', '\'seg\'', '\'registration\''],
+        }
+    return dict_synonyms
+
+def replace_text(text, dict_synonyms):
+    '''
+    Replace the text in the text with the standard term
+    '''
+    if isinstance(text, str):
+        for key, value in dict_synonyms.items():
+            for v in value:
+                if v.lower() in text.lower():
+                    text = text.replace(v, key)
+        return text
+    elif isinstance(text, list):
+        text = [replace_text(t, dict_synonyms) for t in text]
+    elif isinstance(text, dict):
+        for key in text.keys():
+            # replace values in dict
+            text[key] = replace_text(text[key], dict_synonyms)
+            # replace keys in dict
+            for k in dict_synonyms.keys():
+                if k.lower() in key.lower():
+                    text[dict_synonyms[k]] = text.pop(key)
+    return text
+
+
+def replace_synonyms(text, dict_synonyms):
+    '''
+    Replace the synonyms in the text with the standard term
+    '''
+    if isinstance(text,str):
+        for key, value in dict_synonyms.items():
+            for v in value:
+                if v.lower() in text.lower():
+                    return key
+        Warning(f"Value {text} is not in the correct format")
+    elif isinstance(text,list):
+        text = [replace_synonyms(t, dict_synonyms) for t in text]
+    elif isinstance(text,dict):
+        for key in text.keys():
+            # replace values in dict
+            text[key] = replace_synonyms(text[key], dict_synonyms)
+            # replace keys in dict
+            for k in dict_synonyms.keys():
+                text[dict_synonyms[k]] = text.pop(key)
+    return text    
+
+if __name__ == "__main__":
+    # model_name = "bert-base-uncased"
+    # model_name = "bert-large-uncased"
+    model_name = "/home/jachin/data/Github/OmniMorph/External/Models/bert_large_uncased"
+    # model_name = "Rostlab/prot_bert"
+    # model_name = "fspanda/Medical-Bio-BERT2"
+    # model_name = "GerMedBERT/medbert-512"
+
+    reduce_method = 'mean'
+    max_words_num = 32  # max number of words in the caption > 2
+    
+    embeder, tokenizer = get_frozen_embeder(model_name)
+
+    # string1 = ["mri", "female"]
+    string1 = "modality: ct, gender: female, age: 51, roi: abdomen"
+    # string1 = "modality: Magnetic Resonance, gender: female"
+    embeder_output1 = str2emb(string1, max_words_num, embeder, tokenizer, reduce_method=reduce_method)
+
+    # string2 = "Hello world!"
+    # string2 = ["ct", "male"]
+    # string2 = "modality: mri, gender: female, roi: head"
+    string2 = "modality: ct, gender: female, age: 50, roi: head"
+    # string2 = "modality: ct, gender: male, roi: head"
+    embeder_output2 = str2emb(string2, max_words_num, embeder, tokenizer, reduce_method=reduce_method)
+
+    input_size = embeder.config.vocab_size
+    in_size = embeder.config.hidden_size
+    
+    print(embeder, input_size, in_size)
+    print(tokenizer)
+    
+    # embeder_output1 shape: [batch_size, max_words_num, hidden_size]
+    print(embeder_output1)
+    print(embeder_output1.shape)  # torch.Size([1, 8, 768])
+    
+    # embeder_output2 shape: [batch_size, max_words_num, hidden_size]
+    print(embeder_output2)
+    print(embeder_output2.shape)  # torch.Size([1, 8, 768])
+    
+    # check the difference between the two sentences in embedding space
+    # embeder_output1[0, :, :] shape: [max_words_num, hidden_size]
+    # embeder_output2[0, :, :] shape: [max_words_num, hidden_size]
+    # error = torch.max(torch.abs(embeder_output1[0, :, :] - embeder_output2[0, :, :]), dim=-1)
+    error = torch.abs(embeder_output1 - embeder_output2)
+    print(error)
+    print("Embedding distance between the two sentences: ")
+    print(f"String1: {string1}")
+    print(f"String2: {string2}")
+    print(torch.mean(error))
+    exit()

Dataloader/dataLoader.py

@@ -0,0 +1,1473 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+import json
+import SimpleITK as sitk
+import numpy as np
+from skimage.transform import rescale, resize, downscale_local_mean
+# from torchvision.transforms import v2
+import sys
+sys.path.append('./')
+from Dataloader.dataloader_utils import *
+import random
+
+# add your mapping files here
+# mapping_files = {
+#     'TotalSegmentor': '/home/data/Github/data/data_gen_def/DATASETS_processed/TotalSegmentorCT_MRI/nifti_mappings.json',
+#     'MSD': '/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/MSD_processed/nifti_mappings_updated.json',
+#     # 'CancerImageArchive': '/home/data/Github/data/data_gen_def/DATASETS_processed/CancerImageArchive_1/nifti_mappings.json',
+# }
+
+
+mapping_files = {
+    'MSD': '/home/data/Github/OmniMorph/Dataloader/nifty_mappings/MSD_mappings.json',
+    'TotalSegmentor': '/home/data/Github/OmniMorph/Dataloader/nifty_mappings/TotalSegmentorCT_MRI_mappings.json',
+    'Kaggle_osic': '/home/data/Github/OmniMorph/Dataloader/nifty_mappings/Kaggle_osic_mappings.json',
+    'CancerImageArchive': '/home/data/Github/OmniMorph/Dataloader/nifty_mappings/CIA_mappings.json',
+    'MnMs': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/MnMs_mappings.json',
+    # 'Brats2019': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/Brats2019_mappings.json',
+    'Brats2020': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/Brats2020_mappings.json',
+    'Brats2021': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/Brats2021_mappings.json',
+    'OASIS_1': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/OASIS_1_mappings.json',
+    'OASIS_2': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/OASIS_2_mappings.json',
+    'PSMA-FDG-PET-CT-LESION':'/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/PSMA-FDG-PET-CT-LESION_mappings.json',
+    'PSMA-CT':'/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/PSMA-CT-Longitud_mappings.json',
+    'AbdomenAtlas':'/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/AbdomenAtlas_mappings.json',
+    'AbdomenCT1k':'/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/AbdomenCT1k_mappings.json',
+}
+
+CLAMP_RANGE = [-400, 400]  # default clamp range for the images
+
+indivi_ROI_list = ['abdomen','arm','brain','hand','head','leg','neck','pelvis','skeleton','thorax']
+
+def reverse_axis_order(arr):
+    """SimpleITK to NumPy axis order conversion."""
+    # For 3D or 4D arrays, this is just a fast view, not a copy.
+    return np.ascontiguousarray(arr.transpose(tuple(range(arr.ndim)[::-1])))
+
+def sample_random_uniform_multi_order(high=1., low=0., order_num=2, type='high'):
+    """Sample a random value from a uniform distribution with multiple orders.
+
+    Args:
+        high (float): Upper bound of the uniform distribution.
+        low (float): Lower bound of the uniform distribution.
+        order_num (int): Number of times to sample.
+        type (str): 'high' or 'low', determines the sampling direction.
+
+    Returns:
+        sample_value (float): The sampled value after multiple orders.
+
+    Notes:
+        - If type is 'high', samples are drawn iteratively from [low, high], each time using the previous sample as the new lower bound.
+        - If type is 'low', samples are drawn iteratively from [low, high], each time using the previous sample as the new upper bound.
+        - If order_num is 0, returns the low value.
+        - If order_num is 1, returns a single random value from the uniform distribution.
+        - If order_num is 2, returns a value from a linear distribution.
+        - If order_num is 3, returns a value from a quadratic distribution.
+    """
+    if type == 'high':
+        sample_value = low
+        for _ in range(order_num):
+            sample_value = np.random.uniform(low=sample_value, high=high)
+    elif type == 'low':
+        sample_value = high
+        for _ in range(order_num):
+            sample_value = np.random.uniform(low, high=sample_value)
+    return sample_value
+
+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):
+      
+        # self.mappings = mapping_files
+        self.ALLdata = self.combine_data(mappings = mapping_files)
+        self.out_sz = out_sz
+        self.reverse_axis_order = reverse_axis_order
+        self.min_dim = min_dim
+        self.clamp_range = clamp_range
+        self.min_crop_ratio = min_crop_ratio
+        self.transform = transform
+        self.ndims = 3
+        
+    def get_ALLdata(self):
+        return self.ALLdata
+        
+    def get_all_ROI(self):
+        # Get all the ROI options. and remove the reduntant ones
+        ROIs = []
+        # ALLdata_filtered = data
+        for k in self.ALLdata_filtered.keys():
+            ROIs.append(self.ALLdata[k]['ROI'])
+        ROIs = set(ROIs)
+        return ROIs
+    
+    def get_filter_ROIs(self,keep_single_roi=False):
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        # if keep_single_roi == True:
+        #     for k in self.ALLdata_filtered.keys():
+        #         if '-' in self.ALLdata_filtered[k]['ROI']:
+        #             del ALLdata_filtered[k]
+        # d = {k: v for k, v in ALLdata_filtered.items() if v['ROI'] in self.ROIs}
+        for k in ALLdata_filtered.keys():
+            if self.ALLdata_filtered[k]['ROI'] not in self.ROIs:
+                del ALLdata_filtered[k]
+        return ALLdata_filtered
+
+    def combine_data(self, mappings = mapping_files):
+        ALLdata = {}
+        for j in mappings.keys():
+            with open(mappings[j], 'r') as f:
+                mappings_tmp = json.load(f)
+                ALLdata.update(mappings_tmp)
+        return ALLdata
+        
+    def get_3D_volume(self, volume, select_channel = None):
+        # Get a 3D volume from the 4D volume, sometime the input image may have 4 dimensions
+        if self.reverse_axis_order:
+            volume = reverse_axis_order(volume)
+        if volume.ndim == 4:
+            if select_channel is None:
+                select_channel = np.random.randint(0, volume.shape[3] - 1)
+            volume = volume[:, :, :, select_channel]
+        return volume
+        
+    def get_filter_mindim(self):
+        # Filter out images with dimensions less than min_dim
+        # Top priority is to filter out images with dimensions less than min_dim
+        ALLdata = self.ALLdata.copy()
+        for k in self.ALLdata.keys():
+            if min(self.ALLdata[k]['Size'][:self.ndims]) < self.out_sz/2:
+                del ALLdata[k]
+        return ALLdata
+
+    def normalize(self, volume, eps=1e-7):
+        # Normalize the image (0-1)
+        volume = volume.astype(np.float64)
+        volume = (volume - np.min(volume)) / (np.ptp(volume) + eps)
+        return volume
+    
+    def random_crop_3d(self, volume, crop_size=None):
+        # Fast random crop with optional padding using NumPy
+        d, h, w = volume.shape
+        if crop_size is None:
+            crop_size = self.out_sz
+        crop_d, crop_h, crop_w = crop_size, crop_size, crop_size
+
+        # Only pad if needed (avoid np.pad if not necessary)
+        pad_d = max(0, crop_d - d)
+        pad_h = max(0, crop_h - h)
+        pad_w = max(0, crop_w - w)
+        if pad_d or pad_h or pad_w:
+            pad_width = (
+                (np.random.randint(0, pad_d + 1), pad_d - np.random.randint(0, pad_d + 1)),
+                (np.random.randint(0, pad_h + 1), pad_h - np.random.randint(0, pad_h + 1)),
+                (np.random.randint(0, pad_w + 1), pad_w - np.random.randint(0, pad_w + 1)),
+            )
+            volume = np.pad(volume, pad_width, mode='constant', constant_values=0)
+            d, h, w = volume.shape
+
+        # Crop indices
+        start_d = np.random.randint(0, d - crop_d + 1) if d > crop_d else 0
+        start_h = np.random.randint(0, h - crop_h + 1) if h > crop_h else 0
+        start_w = np.random.randint(0, w - crop_w + 1) if w > crop_w else 0
+
+        # Use NumPy slicing (very fast)
+        return volume[start_d:start_d + crop_d, start_h:start_h + crop_h, start_w:start_w + crop_w]
+
+class OminiDataset_v1(Dataset):
+    def __init__(self, out_sz = 128, transform=None, clamp_range = CLAMP_RANGE, min_crop_ratio = 0.2, reverse_axis_order = False):
+        self.mappings = mapping_files
+        self.ALLdata = self.combine_data()
+        self.out_sz = out_sz
+        self.reverse_axis_order = reverse_axis_order
+        self.min_crop_ratio = min_crop_ratio
+        self.crop_ratio_sample_order = 2
+        self.transform = transform
+        self.clamp_range = clamp_range
+        self.ndims = 3
+        # Start you filtering here
+        self.ALLdata_filtered = self.get_filter_mindim()
+        
+
+        # self.min_dim = self.find_min_dim()
+        
+    def find_min_dim(self):
+        # Find the minimum dimension of the images
+        min_dim = 100000
+        for k in self.ALLdata.keys():
+            value = self.ALLdata[k]
+            if min(value['Size']) < min_dim:
+                min_dim = min(value['Size'])
+        return min_dim
+    
+    def random_crop_3d(self, volume, crop_size=None):
+        # Fast random crop with optional padding using NumPy
+        d, h, w = volume.shape
+        if crop_size is None:
+            crop_size = self.out_sz
+        crop_d, crop_h, crop_w = crop_size, crop_size, crop_size
+
+        # Only pad if needed (avoid np.pad if not necessary)
+        pad_d = max(0, crop_d - d)
+        pad_h = max(0, crop_h - h)
+        pad_w = max(0, crop_w - w)
+        if pad_d or pad_h or pad_w:
+            pad_width = (
+                (np.random.randint(0, pad_d + 1), pad_d - np.random.randint(0, pad_d + 1)),
+                (np.random.randint(0, pad_h + 1), pad_h - np.random.randint(0, pad_h + 1)),
+                (np.random.randint(0, pad_w + 1), pad_w - np.random.randint(0, pad_w + 1)),
+            )
+            volume = np.pad(volume, pad_width, mode='constant', constant_values=0)
+            d, h, w = volume.shape
+
+        # Crop indices
+        start_d = np.random.randint(0, d - crop_d + 1) if d > crop_d else 0
+        start_h = np.random.randint(0, h - crop_h + 1) if h > crop_h else 0
+        start_w = np.random.randint(0, w - crop_w + 1) if w > crop_w else 0
+
+        # Use NumPy slicing (very fast)
+        return volume[start_d:start_d + crop_d, start_h:start_h + crop_h, start_w:start_w + crop_w]
+    
+    def get_ALLdata(self):
+        # Return all data
+        return self.ALLdata
+    
+    def get_3D_volume(self, volume, select_channel = None):
+        if self.reverse_axis_order:
+            volume = reverse_axis_order(volume)
+        if volume.ndim == 4:
+            if select_channel is None:
+                select_channel = np.random.randint(0, volume.shape[3] - 1)
+            volume = volume[:, :, :, select_channel]
+            # print(f"Volume shape: {volume.shape}, selected channel: {select_channel}")
+        return volume
+        
+    def get_filter_ROI(self, key_word):
+        # Filter out images with a key word
+        ALLdata = self.ALLdata.copy()
+        for k in self.ALLdata.keys():
+            if key_word not in k["ROI"]:
+                del ALLdata[k]
+        return ALLdata
+
+    def get_filter_mindim(self):
+        # Filter out images with dimensions less than min_dim
+        # Top priority is to filter out images with dimensions less than min_dim
+        ALLdata = self.ALLdata.copy()
+        for k in self.ALLdata.keys():
+            if min(self.ALLdata[k]['Size'][:self.ndims]) < self.out_sz/2:
+                del ALLdata[k]
+        return ALLdata
+    
+    def combine_data(self):
+        ALLdata = {}
+        for j in self.mappings.keys():
+            with open(self.mappings[j], 'r') as f:
+                mappings = json.load(f)
+                ALLdata.update(mappings)
+        return ALLdata
+
+    def __len__(self):
+        return len(self.ALLdata_filtered.keys())
+
+    def normalize(self, volume, eps=1e-7):
+        # Normalize the image (0-1)
+        volume = volume.astype(np.float64)
+        volume = (volume - np.min(volume)) / (np.ptp(volume) + eps)
+        return volume
+
+    def __getitem__(self, idx):  
+        key = list(self.ALLdata_filtered.keys())[idx]
+        if 0:
+            print(key)
+        volume = sitk.ReadImage(key)
+        volume = sitk.GetArrayFromImage(volume)
+        # if volume.ndim == 4:
+        volume = self.get_3D_volume(volume)
+        
+        if self.clamp_range is not None:
+            modality = self.ALLdata_filtered[key].get("Modality", None)
+            if modality == "CT":
+                volume = np.clip(volume, self.clamp_range[0], self.clamp_range[1])
+        volume = self.normalize(volume)
+        
+        if self.min_crop_ratio is not None:
+            # print(f'before volume_shape: {volume.shape}')
+            # crop_ratio = np.random.uniform(self.min_crop_ratio, 1)
+            crop_ratio = sample_random_uniform_multi_order(high=1., low=self.min_crop_ratio, order_num=self.crop_ratio_sample_order, type='high')
+            # crop_size = int(min(volume.shape) * crop_ratio)
+            crop_size = int(max(volume.shape) * crop_ratio)
+            volume = self.random_crop_3d(volume, crop_size)
+            volume = resize(volume, [self.out_sz]*self.ndims, anti_aliasing = True, preserve_range = True)
+
+            
+        else:
+            volume = self.random_crop_3d(volume, self.out_sz)
+        volume = volume[None, :, :, :]
+
+        if self.transform is not None:
+            return  self.transform(volume)
+
+        return volume
+        
+class OMDataset_indiv(Dataset):
+    def __init__(self, out_sz = 128, transform=None, clamp_range = CLAMP_RANGE, min_crop_ratio = 0.3, reverse_axis_order = False):
+        # self.mappings = mapping_files
+        self.ALLdata = self.combine_data(mappings=mapping_files)
+        self.out_sz = out_sz
+        self.max_sz = out_sz*8
+        self.reverse_axis_order = reverse_axis_order
+        self.min_crop_ratio = min_crop_ratio
+        self.crop_ratio_sample_order = 2
+        self.transform = transform
+        self.clamp_range = clamp_range
+        self.ndims = 3
+
+        # Start you filtering here
+        # print(f"Filtering data with out_sz: {self.out_sz}, min_crop_ratio: {min_crop_ratio}")
+        print(f"Diffusion mode: Total data size before filtering: {len(self.ALLdata)}")
+        self.ALLdata_filtered = self.get_filter_mindim()
+        print(f"Diffusion mode: Filtered data size: {len(self.ALLdata_filtered)}")
+        
+        
+        # self.min_dim = self.find_min_dim()
+        
+    def find_min_dim(self):
+        # Find the minimum dimension of the images
+        min_dim = 100000
+        for k in self.ALLdata.keys():
+            value = self.ALLdata[k]
+            if min(value['Size']) < min_dim:
+                min_dim = min(value['Size'])
+        return min_dim
+    
+    def random_crop_3d(self, volume, crop_size=None):
+        # Fast random crop with optional padding using NumPy
+        d, h, w = volume.shape
+        if crop_size is None:
+            crop_size = self.out_sz
+        crop_d, crop_h, crop_w = crop_size, crop_size, crop_size
+
+        # Only pad if needed (avoid np.pad if not necessary)
+        pad_d = max(0, crop_d - d)
+        pad_h = max(0, crop_h - h)
+        pad_w = max(0, crop_w - w)
+        if pad_d or pad_h or pad_w:
+            pad_width = (
+                (np.random.randint(0, pad_d + 1), pad_d - np.random.randint(0, pad_d + 1)),
+                (np.random.randint(0, pad_h + 1), pad_h - np.random.randint(0, pad_h + 1)),
+                (np.random.randint(0, pad_w + 1), pad_w - np.random.randint(0, pad_w + 1)),
+            )
+            volume = np.pad(volume, pad_width, mode='constant', constant_values=0)
+            d, h, w = volume.shape
+
+        # Crop indices
+        start_d = np.random.randint(0, d - crop_d + 1) if d > crop_d else 0
+        start_h = np.random.randint(0, h - crop_h + 1) if h > crop_h else 0
+        start_w = np.random.randint(0, w - crop_w + 1) if w > crop_w else 0
+
+        # Use NumPy slicing (very fast)
+        return volume[start_d:start_d + crop_d, start_h:start_h + crop_h, start_w:start_w + crop_w]
+    
+    def get_ALLdata(self):
+        # Return all data
+        return self.ALLdata
+    
+    def get_3D_volume(self, volume, select_channel = None):
+        if self.reverse_axis_order:
+            volume = reverse_axis_order(volume)
+        if volume.ndim == 4:
+            if select_channel is None:
+                select_channel = np.random.randint(0, volume.shape[3] - 1)
+            volume = volume[:, :, :, select_channel]
+            # print(f"Volume shape: {volume.shape}, selected channel: {select_channel}")
+        return volume
+        
+    def get_filter_ROI(self, key_word):
+        # Filter out images with a key word
+        ALLdata = self.ALLdata.copy()
+        for k in self.ALLdata.keys():
+            if key_word not in k["ROI"]:
+                del ALLdata[k]
+        return ALLdata
+
+    def get_filter_mindim(self):
+        # Filter out images with dimensions less than min_dim
+        # Top priority is to filter out images with dimensions less than min_dim
+        ALLdata = self.ALLdata.copy()
+        for k in self.ALLdata.keys():
+            if min(self.ALLdata[k]['Size'][:self.ndims]) < self.out_sz/2:
+                del ALLdata[k]
+        return ALLdata
+    
+    def combine_data(self, mappings = mapping_files):
+        ALLdata = {}
+        for j in mappings.keys():
+            with open(mappings[j], 'r') as f:
+                mappings_tmp = json.load(f)
+                ALLdata.update(mappings_tmp)
+        return ALLdata
+
+    def __len__(self):
+        return len(self.ALLdata_filtered.keys())
+
+    def normalize(self, volume, eps=1e-7):
+        # Normalize the image (0-1)
+        volume = volume.astype(np.float64)
+        volume = (volume - np.min(volume)) / (np.ptp(volume) + eps)
+        return volume
+
+    def __getitem__(self, idx):  
+        key = list(self.ALLdata_filtered.keys())[idx]
+        embd = self.ALLdata_filtered[key]['embd']
+        embd = np.array(embd, dtype=np.float32)
+
+        if 0:
+            print(key)
+        volume = sitk.ReadImage(key)
+        volume = sitk.GetArrayFromImage(volume)
+        # if volume.ndim == 4:
+        volume = self.get_3D_volume(volume)
+        
+        if self.clamp_range is not None:
+            modality = self.ALLdata_filtered[key].get("Modality", None)
+            if modality == "CT":
+                volume = np.clip(volume, self.clamp_range[0], self.clamp_range[1])
+        volume = self.normalize(volume)
+        
+        if self.min_crop_ratio is not None:
+            # print(f'before volume_shape: {volume.shape}')
+            # crop_ratio = np.random.uniform(self.min_crop_ratio, 1)
+            crop_ratio = sample_random_uniform_multi_order(high=1., low=self.min_crop_ratio, order_num=self.crop_ratio_sample_order, type='high')
+            # crop_size = int(min(volume.shape) * crop_ratio)
+            crop_size = int(max(volume.shape) * crop_ratio)
+            crop_size = min(crop_size, self.max_sz)
+            volume = self.random_crop_3d(volume, crop_size)
+            volume = resize(volume, [self.out_sz]*self.ndims, anti_aliasing = True, preserve_range = True)
+
+            
+        else:
+            volume = self.random_crop_3d(volume, self.out_sz)
+        volume = volume[None, :, :, :]
+
+        if self.transform is not None:
+            return  self.transform(volume)
+
+        return [volume, embd]
+        
+class OminiDataset_paired(Dataset):
+    def __init__(self, out_sz = 128, transform=None, clamp_range = CLAMP_RANGE, min_crop_ratio = 0.9, ROIs = None, modality = None, reverse_axis_order = False):
+        # self.mappings = mapping_files
+        self.ALLdata = self.combine_data(mappings=mapping_files)
+        self.out_sz = out_sz
+        self.sz_range = get_sizeRange_dict()
+        self.min_dim_ratio = 0.5
+        self.reverse_axis_order = reverse_axis_order
+        self.min_crop_ratio = min_crop_ratio
+        self.transform = transform
+        self.clamp_range = clamp_range
+        self.ndims = 3
+        # Start you filtering here
+        # print(f"Number of images before filtering: {len(self.ALLdata.keys())}")
+        self.ALLdata_filtered = self.get_filter_mindim()
+        # print(f"Number of images after filtering: {len(self.ALLdata_filtered.keys())}")
+        self.ALLdata_filtered = self.get_filter_modality(modality)
+        # print(f"Number of images after modality filtering: {len(self.ALLdata_filtered.keys())}")
+        if ROIs is None:# if no ROIs are provided, get all the ROIs from filtered data
+            self.ROIs = self.get_all_ROI()
+        else:
+            self.ROIs = ROIs
+        self.ALLdata_filtered = self.get_filter_ROIs()
+        # print(f"Number of images after ROI filtering: {len(self.ALLdata_filtered.keys())}")
+        # filtering ends here
+        
+
+
+    def combine_data(self, mappings = mapping_files):
+        ALLdata = {}
+        for j in mappings.keys():
+            with open(mappings[j], 'r') as f:
+                mappings_tmp = json.load(f)
+                ALLdata.update(mappings_tmp)
+        return ALLdata
+    
+    def normalize(self, volume, eps=1e-7):
+        # Normalize the image (0-1)
+        volume = volume.astype(np.float64)
+        volume = (volume - np.min(volume)) / (np.ptp(volume) + eps)
+        return volume
+    
+    def random_crop_3d(self, volume, crop_size=None):
+        # Fast random crop with optional padding using NumPy
+        d, h, w = volume.shape
+        if crop_size is None:
+            crop_size = self.out_sz
+        crop_d, crop_h, crop_w = crop_size, crop_size, crop_size
+
+        # Only pad if needed (avoid np.pad if not necessary)
+        pad_d = max(0, crop_d - d)
+        pad_h = max(0, crop_h - h)
+        pad_w = max(0, crop_w - w)
+        if pad_d or pad_h or pad_w:
+            pad_width = (
+                (np.random.randint(0, pad_d + 1), pad_d - np.random.randint(0, pad_d + 1)),
+                (np.random.randint(0, pad_h + 1), pad_h - np.random.randint(0, pad_h + 1)),
+                (np.random.randint(0, pad_w + 1), pad_w - np.random.randint(0, pad_w + 1)),
+            )
+            volume = np.pad(volume, pad_width, mode='constant', constant_values=0)
+            d, h, w = volume.shape
+
+        # Crop indices
+        start_d = np.random.randint(0, d - crop_d + 1) if d > crop_d else 0
+        start_h = np.random.randint(0, h - crop_h + 1) if h > crop_h else 0
+        start_w = np.random.randint(0, w - crop_w + 1) if w > crop_w else 0
+
+        # Use NumPy slicing (very fast)
+        return volume[start_d:start_d + crop_d, start_h:start_h + crop_h, start_w:start_w + crop_w]
+    
+    # def random_crop_3d(self, volume, crop_size=None):
+    #     # Randomly crop the image
+    #     d, h, w = volume.shape
+    #     if crop_size is None:
+    #         crop_size = self.out_sz
+    #     crop_d, crop_h, crop_w = crop_size, crop_size, crop_size
+
+    #     if crop_d > d or crop_h > h or crop_w > w:
+    #         raise ValueError("Crop size must be smaller than the original array size")
+        
+    #     start_d = np.random.randint(0, d - crop_d + 1)
+    #     start_h = np.random.randint(0, h - crop_h + 1)
+    #     start_w = np.random.randint(0, w - crop_w + 1)
+
+    #     cropped_array = volume[start_d:start_d + crop_d, start_h:start_h + crop_h, start_w:start_w + crop_w]
+
+    #     return cropped_array
+    
+    def get_all_ROI(self):
+        # Get all the ROI options. and remove the reduntant ones
+        ROIs = []
+        for k in self.ALLdata_filtered.keys():
+            ROIs.append(self.ALLdata[k]['ROI'])
+        ROIs = set(ROIs)
+        return ROIs
+        
+    def find_min_dim(self):
+        # Find the minimum dimension of the images
+        min_dim = 100000
+        for k in self.ALLdata.keys():
+            value = self.ALLdata[k]
+            if min(value['Size']) < min_dim:
+                min_dim = min(value['Size'])
+        return min_dim
+            
+    def get_ALLdata(self):
+        # Return all data
+        return self.ALLdata
+    
+    def get_filter_modality(self, key_words=None):
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        if key_words is not None:
+            for k in self.ALLdata_filtered.keys():
+                if ALLdata_filtered[k]["Modality"] not in key_words:
+                    del ALLdata_filtered[k]
+        return ALLdata_filtered
+
+    def get_filter_ROI(self, key_word):
+        # Filter out images with a key word
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        for k in self.ALLdata_filtered.keys():
+            if key_word not in k["ROI"]:
+                del ALLdata_filtered[k]
+        return ALLdata_filtered
+    
+    def get_key_by_ROI(self, key_word):
+        # Get all the keys with a key word
+        keys = []
+        for k in self.ALLdata_filtered.keys():
+            if key_word == self.ALLdata_filtered[k]["ROI"]:
+                keys.append(k)
+        return keys
+    
+    def get_filter_ROIs(self):
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        for k in self.ALLdata_filtered.keys():
+            if self.ALLdata_filtered[k]['ROI'] not in self.ROIs:
+                del ALLdata_filtered[k]
+        return ALLdata_filtered
+
+    def get_3D_volume(self, volume, select_channel = None):
+        if self.reverse_axis_order:
+            volume = reverse_axis_order(volume)
+        if volume.ndim == 4:
+            if select_channel is None:
+                select_channel = np.random.randint(0, volume.shape[3] - 1)
+            volume = volume[:, :, :, select_channel]
+        return volume
+        
+    def get_filter_mindim(self):
+        # Filter out images with dimensions less than min_dim
+        # Top priority is to filter out images with dimensions less than min_dim
+        ALLdata = self.ALLdata.copy()
+        for k in self.ALLdata.keys():
+            img_sz = self.ALLdata[k]['Size'][:self.ndims]
+            del_flag = False
+            del_flag = del_flag or min(img_sz) < self.out_sz
+            # print(f"Size: {self.ALLdata[k]['Size']}, Spacing_mm: {self.ALLdata[k]['Spacing_mm']}, ROI: {self.ALLdata[k]['ROI']}")
+            # print(f"sz_range: {self.sz_range[self.ALLdata[k]['ROI']]}, min_dim_ratio: {self.min_dim_ratio}")
+            del_flag = del_flag or (min(img_sz)*self.ALLdata[k]['Spacing_mm']) < self.sz_range[self.ALLdata[k]['ROI']][0]
+            del_flag = del_flag or (min(img_sz)/max(img_sz) < self.min_dim_ratio)
+            # del_flag = min(self.ALLdata[k]['Size']) < self.out_sz or (min(self.ALLdata[k]['Size'])*self.ALLdata[k]['Spacing_mm']) < self.sz_range[self.ALLdata[k]['ROI']] or (min(self.ALLdata[k]['Size'])/max(self.ALLdata[k]['Size']) < self.min_dim_ratio)
+            if del_flag:
+                del ALLdata[k]
+        return ALLdata
+
+        
+        
+    def __getitem__(self,idx):
+        key = list(self.ALLdata_filtered.keys())[idx]
+        volume_A = sitk.ReadImage(key)
+        volume_A = sitk.GetArrayFromImage(volume_A)
+        
+        paired_keys = self.get_key_by_ROI(self.ALLdata_filtered[key]['ROI'])
+        paired_key = random.choice(paired_keys)
+        
+        volume_B = sitk.ReadImage(paired_key)
+        volume_B = sitk.GetArrayFromImage(volume_B)
+        
+        # if volume_A.ndim == 4 or volume_B.ndim == 4:
+        volume_A = self.get_3D_volume(volume_A)
+        volume_B = self.get_3D_volume(volume_B)
+
+        if self.clamp_range is not None:
+            modality = self.ALLdata_filtered[key].get("Modality", None)
+            if modality == "CT":
+                volume_A = np.clip(volume_A, self.clamp_range[0], self.clamp_range[1])
+                volume_B = np.clip(volume_B, self.clamp_range[0], self.clamp_range[1])
+        volume_A = self.normalize(volume_A)
+        volume_B = self.normalize(volume_B)
+        
+        if self.min_crop_ratio is not None:
+            
+            # print(f'before volume_shape: {volume.shape}')
+            crop_ratio = np.random.uniform(self.min_crop_ratio, 1)
+            crop_size_A = int(min(volume_A.shape) * crop_ratio)
+            crop_size_B = int(min(volume_B.shape) * crop_ratio)
+            # crop_size_A = int(max(volume_A.shape) * crop_ratio)
+            # crop_size_B = int(max(volume_B.shape) * crop_ratio)
+            volume_A = self.random_crop_3d(volume_A, crop_size_A)
+            volume_B = self.random_crop_3d(volume_B, crop_size_B)
+            volume_A = resize(volume_A, [self.out_sz]*self.ndims, anti_aliasing = True, preserve_range = True)
+            volume_B = resize(volume_B, [self.out_sz]*self.ndims, anti_aliasing = True, preserve_range = True)
+
+        else:
+            volume_A = self.random_crop_3d(volume_A, self.out_sz)
+            volume_B = self.random_crop_3d(volume_B, self.out_sz)
+        volume_A = volume_A[None, :, :, :]
+        volume_B = volume_B[None, :, :, :]
+
+        if self.transform is not None:
+            return  self.transform(volume_A), self.transform(volume_B)
+        
+        # print(self.ALLdata_filtered[key]['ROI'],self.ALLdata_filtered[key]['Modality'],self.ALLdata_filtered[key]['Dataset_name'],'---',self.ALLdata_filtered[paired_key]['ROI'], self.ALLdata_filtered[paired_key]['Modality'], self.ALLdata_filtered[paired_key]['Dataset_name'])
+        return volume_A, volume_B
+
+    def __len__(self):
+        return len(self.ALLdata_filtered.keys())
+
+class OMDataset_pair(Dataset):
+    def __init__(self, out_sz = 128, transform=None, clamp_range = CLAMP_RANGE, min_crop_ratio = 0.75, ROIs = indivi_ROI_list, modality = None, reverse_axis_order = False):
+        # self.mappings = mapping_files
+        self.ALLdata = self.combine_data(mappings=mapping_files)
+        self.out_sz = out_sz
+        self.max_sz = out_sz*8
+        self.sz_range = get_sizeRange_dict()
+        self.min_dim_ratio = 0.7
+        self.reverse_axis_order = reverse_axis_order
+        self.min_crop_ratio = min_crop_ratio
+        self.transform = transform
+        self.clamp_range = clamp_range
+        self.ndims = 3
+        # Start you filtering here
+        # print(f"Number of images before filtering: {len(self.ALLdata.keys())}")
+        print(f"Registration mode: Total data size before filtering: {len(self.ALLdata)}")
+
+        self.ALLdata_filtered = self.get_filter_mindim()
+        # print(f"Number of images after filtering: {len(self.ALLdata_filtered.keys())}")
+        self.ALLdata_filtered = self.get_filter_modality(modality)
+        # print(f"Number of images after modality filtering: {len(self.ALLdata_filtered.keys())}")
+        if ROIs is None:# if no ROIs are provided, get all the ROIs from filtered data
+            self.ROIs = self.get_all_ROI()
+        else:
+            self.ROIs = ROIs
+        self.ALLdata_filtered = self.get_filter_ROIs()
+        print(f"Registration mode: Number of images after filtering: {len(self.ALLdata_filtered.keys())}")
+        # filtering ends here
+        
+
+
+    def combine_data(self, mappings = mapping_files):
+        ALLdata = {}
+        for j in mappings.keys():
+            with open(mappings[j], 'r') as f:
+                mappings_tmp = json.load(f)
+                ALLdata.update(mappings_tmp)
+        return ALLdata
+    
+    def normalize(self, volume, eps=1e-7):
+        # Normalize the image (0-1)
+        volume = volume.astype(np.float64)
+        volume = (volume - np.min(volume)) / (np.ptp(volume) + eps)
+        return volume
+    
+    def random_crop_3d(self, volume, crop_size=None):
+        # Fast random crop with optional padding using NumPy
+        d, h, w = volume.shape
+        if crop_size is None:
+            crop_size = self.out_sz
+        crop_d, crop_h, crop_w = crop_size, crop_size, crop_size
+
+        # Only pad if needed (avoid np.pad if not necessary)
+        pad_d = max(0, crop_d - d)
+        pad_h = max(0, crop_h - h)
+        pad_w = max(0, crop_w - w)
+        if pad_d or pad_h or pad_w:
+            pad_width = (
+                (np.random.randint(0, pad_d + 1), pad_d - np.random.randint(0, pad_d + 1)),
+                (np.random.randint(0, pad_h + 1), pad_h - np.random.randint(0, pad_h + 1)),
+                (np.random.randint(0, pad_w + 1), pad_w - np.random.randint(0, pad_w + 1)),
+            )
+            volume = np.pad(volume, pad_width, mode='constant', constant_values=0)
+            d, h, w = volume.shape
+
+        # Crop indices
+        start_d = np.random.randint(0, d - crop_d + 1) if d > crop_d else 0
+        start_h = np.random.randint(0, h - crop_h + 1) if h > crop_h else 0
+        start_w = np.random.randint(0, w - crop_w + 1) if w > crop_w else 0
+
+        # Use NumPy slicing (very fast)
+        return volume[start_d:start_d + crop_d, start_h:start_h + crop_h, start_w:start_w + crop_w]
+    
+    # def random_crop_3d(self, volume, crop_size=None):
+    #     # Randomly crop the image
+    #     d, h, w = volume.shape
+    #     if crop_size is None:
+    #         crop_size = self.out_sz
+    #     crop_d, crop_h, crop_w = crop_size, crop_size, crop_size
+
+    #     if crop_d > d or crop_h > h or crop_w > w:
+    #         raise ValueError("Crop size must be smaller than the original array size")
+        
+    #     start_d = np.random.randint(0, d - crop_d + 1)
+    #     start_h = np.random.randint(0, h - crop_h + 1)
+    #     start_w = np.random.randint(0, w - crop_w + 1)
+
+    #     cropped_array = volume[start_d:start_d + crop_d, start_h:start_h + crop_h, start_w:start_w + crop_w]
+
+    #     return cropped_array
+    
+    def get_all_ROI(self):
+        # Get all the ROI options. and remove the reduntant ones
+        ROIs = []
+        for k in self.ALLdata_filtered.keys():
+            ROIs.append(self.ALLdata[k]['ROI'])
+        ROIs = set(ROIs)
+        return ROIs
+        
+    def find_min_dim(self):
+        # Find the minimum dimension of the images
+        min_dim = 100000
+        for k in self.ALLdata.keys():
+            value = self.ALLdata[k]
+            if min(value['Size']) < min_dim:
+                min_dim = min(value['Size'])
+        return min_dim
+            
+    def get_ALLdata(self):
+        # Return all data
+        return self.ALLdata
+    
+    def get_filter_modality(self, key_words=None):
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        if key_words is not None:
+            for k in self.ALLdata_filtered.keys():
+                if ALLdata_filtered[k]["Modality"] not in key_words:
+                    del ALLdata_filtered[k]
+        return ALLdata_filtered
+
+    def get_filter_ROI(self, key_word):
+        # Filter out images with a key word
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        for k in self.ALLdata_filtered.keys():
+            if key_word not in k["ROI"]:
+                del ALLdata_filtered[k]
+        return ALLdata_filtered
+    
+    def get_key_by_ROI(self, key_word):
+        # Get all the keys with a key word
+        keys = []
+        for k in self.ALLdata_filtered.keys():
+            if key_word == self.ALLdata_filtered[k]["ROI"]:
+                keys.append(k)
+        return keys
+
+    def filter_keys_by_xx(self, key_word, keys=None, term="ROI"):
+        # Filter out images with a key word
+        filtered_keys = []
+        if keys is None:
+            keys = self.ALLdata_filtered.keys()
+        for k in keys:
+            value = self.ALLdata_filtered[k].get(term, None)
+            if value is not None and key_word == value:
+                filtered_keys.append(k)
+        return filtered_keys
+    
+    def get_filter_ROIs(self):
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        for k in self.ALLdata_filtered.keys():
+            if self.ALLdata_filtered[k]['ROI'] not in self.ROIs:
+                del ALLdata_filtered[k]
+        return ALLdata_filtered
+
+    def get_3D_volume(self, volume, select_channel = None):
+        if self.reverse_axis_order:
+            volume = reverse_axis_order(volume)
+        if volume.ndim == 4:
+            if select_channel is None:
+                select_channel = np.random.randint(0, volume.shape[3] - 1)
+            volume = volume[:, :, :, select_channel]
+        return volume
+        
+    def get_filter_mindim(self):
+        # Filter out images with dimensions less than min_dim
+        # Top priority is to filter out images with dimensions less than min_dim
+        ALLdata = self.ALLdata.copy()
+        for k in self.ALLdata.keys():
+            img_sz = self.ALLdata[k]['Size'][:self.ndims]
+            del_flag = False
+            del_flag = del_flag or min(img_sz) < self.out_sz
+            # print(f"Size: {self.ALLdata[k]['Size']}, Spacing_mm: {self.ALLdata[k]['Spacing_mm']}, ROI: {self.ALLdata[k]['ROI']}")
+            # print(f"sz_range: {self.sz_range[self.ALLdata[k]['ROI']]}, min_dim_ratio: {self.min_dim_ratio}")
+            del_flag = del_flag or (min(img_sz)*self.ALLdata[k]['Spacing_mm']) < self.sz_range[self.ALLdata[k]['ROI']][0]
+            del_flag = del_flag or (min(img_sz)/max(img_sz) < self.min_dim_ratio)
+            # del_flag = min(self.ALLdata[k]['Size']) < self.out_sz or (min(self.ALLdata[k]['Size'])*self.ALLdata[k]['Spacing_mm']) < self.sz_range[self.ALLdata[k]['ROI']] or (min(self.ALLdata[k]['Size'])/max(self.ALLdata[k]['Size']) < self.min_dim_ratio)
+            if del_flag:
+                del ALLdata[k]
+        return ALLdata
+
+        
+        
+    def __getitem__(self,idx):
+        key = list(self.ALLdata_filtered.keys())[idx]
+        volume_A = sitk.ReadImage(key)
+        volume_A = sitk.GetArrayFromImage(volume_A)
+
+        embd_A = self.ALLdata_filtered[key]['embd']
+        embd_A = np.array(embd_A, dtype=np.float32)
+        
+        all_keys = list(self.ALLdata_filtered.keys())
+        paired_keys = self.filter_keys_by_xx(self.ALLdata_filtered[key]['ROI'], all_keys, term="ROI")
+        paired_keys = self.filter_keys_by_xx(self.ALLdata_filtered[key]['Modality'], paired_keys, term="Modality")
+        # paired_keys = self.get_key_by_ROI(self.ALLdata_filtered[key]['ROI'])
+        
+        paired_key = random.choice(paired_keys)
+        
+        print(f"Key: {key}, Paired Key: {paired_key}")
+        print(f"ROI: {self.ALLdata_filtered[key]['ROI']}, {self.ALLdata_filtered[paired_key]['ROI']}; Modality: {self.ALLdata_filtered[key]['Modality']}, {self.ALLdata_filtered[paired_key]['Modality']}")
+
+
+        volume_B = sitk.ReadImage(paired_key)
+        volume_B = sitk.GetArrayFromImage(volume_B)
+
+        embd_B = self.ALLdata_filtered[paired_key]['embd']
+        embd_B = np.array(embd_B, dtype=np.float32)
+
+        # if volume_A.ndim == 4 or volume_B.ndim == 4:
+        volume_A = self.get_3D_volume(volume_A)
+        volume_B = self.get_3D_volume(volume_B)
+
+        if self.clamp_range is not None:
+            modality = self.ALLdata_filtered[key].get("Modality", None)
+            if modality == "CT":
+                volume_A = np.clip(volume_A, self.clamp_range[0], self.clamp_range[1])
+                volume_B = np.clip(volume_B, self.clamp_range[0], self.clamp_range[1])
+        volume_A = self.normalize(volume_A)
+        volume_B = self.normalize(volume_B)
+        
+        if self.min_crop_ratio is not None:
+            
+            # print(f'before volume_shape: {volume.shape}')
+            crop_ratio = np.random.uniform(self.min_crop_ratio, 1)
+            # crop_size_A = int(min(volume_A.shape) * crop_ratio)
+            # crop_size_B = int(min(volume_B.shape) * crop_ratio)
+            crop_size_A = int(max(volume_A.shape) * crop_ratio)
+            crop_size_B = int(max(volume_B.shape) * crop_ratio)
+            crop_size_A = min(crop_size_A, self.max_sz)
+            crop_size_B = min(crop_size_B, self.max_sz)
+            volume_A = self.random_crop_3d(volume_A, crop_size_A)
+            volume_B = self.random_crop_3d(volume_B, crop_size_B)
+            volume_A = resize(volume_A, [self.out_sz]*self.ndims, anti_aliasing = True, preserve_range = True)
+            volume_B = resize(volume_B, [self.out_sz]*self.ndims, anti_aliasing = True, preserve_range = True)
+
+        else:
+            volume_A = self.random_crop_3d(volume_A, self.out_sz)
+            volume_B = self.random_crop_3d(volume_B, self.out_sz)
+        volume_A = volume_A[None, :, :, :]
+        volume_B = volume_B[None, :, :, :]
+
+        
+        if self.transform is not None:
+            return  self.transform(volume_A), self.transform(volume_B)
+        
+        # print(self.ALLdata_filtered[key]['ROI'],self.ALLdata_filtered[key]['Modality'],self.ALLdata_filtered[key]['Dataset_name'],'---',self.ALLdata_filtered[paired_key]['ROI'], self.ALLdata_filtered[paired_key]['Modality'], self.ALLdata_filtered[paired_key]['Dataset_name'])
+        return [volume_A, volume_B, embd_A, embd_B]
+
+    def __len__(self):
+        return len(self.ALLdata_filtered.keys())
+
+class OminiDataset_paired_inf(object):
+    def __init__(self, out_sz = 128, transform=None, clamp_range = CLAMP_RANGE, min_crop_ratio = 0.3, ROIs = None):
+        # self.mappings = mapping_files
+        self.ALLdata = self.combine_data(mappings=mapping_files)
+        self.out_sz = out_sz
+        self.min_crop_ratio = min_crop_ratio
+        self.transform = transform
+        self.clamp_range = clamp_range
+        self.ndims = 3
+        # Start you filtering here:
+        # filter out images with dimensions less than min_dim
+        self.ALLdata_filtered = self.get_filter_mindim()
+        # filter out images with ROIs that are not in the provided ROIs
+        if ROIs is None:
+            self.ROIs = self.get_all_ROI()
+        else:
+            self.ROIs = ROIs
+        self.ALLdata_filtered = self.get_filter_ROIs()
+        # filtering ends here
+        
+        self.roi_scan_mapping = self.build_ROI_scan_mapping()
+        self.keys_dist, self.total = self.get_keys_dist()
+        
+        
+
+        
+    def get_all_ROI(self):
+        # Get all the ROI options. and remove the reduntant ones
+        ROIs = []
+        for k in self.ALLdata_filtered.keys():
+            ROIs.append(self.ALLdata[k]['ROI'])
+        ROIs = set(ROIs)
+        return ROIs
+        
+    def get_ALLdata(self):
+        # Return all data
+        return self.ALLdata
+        
+    def combine_data(self, mappings = mapping_files):
+        ALLdata = {}
+        for j in mappings.keys():
+            with open(mappings[j], 'r') as f:
+                mappings_tmp = json.load(f)
+                ALLdata.update(mappings_tmp)
+        return ALLdata
+    
+    def __len__(self):
+        return len(self.ALLdata_filtered.keys())
+    
+    def random_crop_3d(self, volume, crop_size=None):
+        # Randomly crop the image
+        d, h, w = volume.shape
+        if crop_size is None:
+            crop_size = self.out_sz
+        crop_d, crop_h, crop_w = crop_size, crop_size, crop_size
+
+        if crop_d > d or crop_h > h or crop_w > w:
+            raise ValueError("Crop size must be smaller than the original array size")
+        
+        start_d = np.random.randint(0, d - crop_d + 1)
+        start_h = np.random.randint(0, h - crop_h + 1)
+        start_w = np.random.randint(0, w - crop_w + 1)
+
+        cropped_array = volume[start_d:start_d + crop_d, start_h:start_h + crop_h, start_w:start_w + crop_w]
+
+        return cropped_array
+    
+    def normalize(self, volume, eps=1e-7):
+        # Normalize the image (0-1)
+        volume = volume.astype(np.float64)
+        volume = (volume - np.min(volume)) / (np.ptp(volume) + eps)
+        return volume
+    
+    def get_3D_volume(self, volume, select_channel = None):
+        volume = reverse_axis_order(volume)
+        if volume.ndim == 4:
+            if select_channel is None:
+                select_channel = np.random.randint(0, volume.shape[3] - 1)
+            volume = volume[:, :, :, select_channel]
+        return volume
+    
+    def get_filter_mindim(self):
+        # Filter out images with dimensions less than min_dim
+        # Top priority is to filter out images with dimensions less than min_dim
+        ALLdata = self.ALLdata.copy()
+        for k in self.ALLdata.keys():
+            if min(self.ALLdata[k]['Size'][:self.ndims]) < self.out_sz/2:
+                del ALLdata[k]
+        return ALLdata
+    
+    def get_filter_ROI(self, key_word):
+        # Filter out images with a key word
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        for k in self.ALLdata_filtered.keys():
+            if key_word not in k["ROI"]:
+                del ALLdata_filtered[k]
+        return ALLdata_filtered
+    
+        
+    def get_filter_ROIs(self):
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        for k in self.ALLdata_filtered.keys():
+            if self.ALLdata_filtered[k]['ROI'] not in self.ROIs:
+                del ALLdata_filtered[k]
+        return ALLdata_filtered
+
+    def get_keys_dist(self):
+        ROIs = self.get_all_ROI()
+        keys_dist = {}
+        total = 0
+        for item in self.ALLdata_filtered.keys():
+            if self.ALLdata_filtered[item]['ROI'] not in keys_dist:
+                keys_dist[self.ALLdata_filtered[item]['ROI']] = 0
+            keys_dist[self.ALLdata_filtered[item]['ROI']] += 1
+            
+        return keys_dist, total
+    
+    def build_ROI_scan_mapping(self):
+        # Build a mapping of ROIs to scans
+        ROI_scan_mapping = {}
+        for item in self.ALLdata_filtered.keys():
+            if self.ALLdata_filtered[item]['ROI'] not in ROI_scan_mapping:
+                ROI_scan_mapping[self.ALLdata_filtered[item]['ROI']] = []
+            ROI_scan_mapping[self.ALLdata_filtered[item]['ROI']].append(item)
+        return ROI_scan_mapping
+
+
+    def get_random_2_items(self, mode = 'uniform'):
+        # Get a random pair of items from the dataset with the same ROI    
+        if mode == 'uniform':
+            idx = random.randint(0, len(self.keys_dist.keys()) - 1)
+            key = list(self.keys_dist.keys())[idx]
+            path_1 = random.choice(self.roi_scan_mapping[key])
+            path_2 = random.choice(self.roi_scan_mapping[key])
+            
+            volume_A = sitk.ReadImage(path_1)
+            volume_A = sitk.GetArrayFromImage(volume_A)
+        
+            volume_B = sitk.ReadImage(path_2)
+            volume_B = sitk.GetArrayFromImage(volume_B)
+            
+            if self.clamp_range is not None:
+                modality = self.ALLdata_filtered[key].get("Modality", None)
+                if modality == "CT":
+                    volume_A = np.clip(volume_A, self.clamp_range[0], self.clamp_range[1])
+                    volume_B = np.clip(volume_B, self.clamp_range[0], self.clamp_range[1])
+            volume_A = self.normalize(volume_A)
+            volume_B = self.normalize(volume_B)
+        
+            if self.min_crop_ratio is not None:
+                crop_ratio = np.random.uniform(self.min_crop_ratio, 1)
+                crop_size_A = int(min(volume_A.shape) * crop_ratio)
+                crop_size_B = int(min(volume_B.shape) * crop_ratio)
+                volume_A = self.random_crop_3d(volume_A, crop_size_A)
+                volume_B = self.random_crop_3d(volume_B, crop_size_B)
+                volume_A = resize(volume_A, [self.out_sz]*self.ndims, anti_aliasing = True, preserve_range = True)
+                volume_B = resize(volume_B, [self.out_sz]*self.ndims, anti_aliasing = True, preserve_range = True)
+            else:
+                volume_A = self.radndom_crop_3d(volume_A, self.out_sz)
+                volume_B = self.radndom_crop_3d(volume_B, self.out_sz)
+            volume_A = volume_A[None, :, :, :]
+            volume_B = volume_B[None, :, :, :]
+            if self.transform is not None:
+                return  self.transform(volume_A), self.transform(volume_B)
+            return volume_A, volume_B
+        
+        elif mode == 'original':
+            pass
+    
+    def build_batch(self, batch_size = 2):
+        batch_1 = []
+        batch_2 = []
+        for i in range(batch_size):
+            V_a, V_b = self.get_random_2_items()
+            batch_1.append(V_a)
+            batch_2.append(V_b)
+        return np.array(batch_1), np.array(batch_2)
+            
+class OminiDataset_inference_w_all(object):
+    def __init__(self, out_sz = 128, transform=None, clamp_range = CLAMP_RANGE, min_crop_ratio = 0.75, ROIs = None, label_key = ['brain'], task_key = 'segmentation', database = None, select_channels_dict = {}):
+        self.mappings = mapping_files
+        # database=['MSD', 'TotalSegmentor']
+        if database is not None:
+            self.mappings = {db: self.mappings[db] for db in database if db in self.mappings}
+        # select_channels_dict={
+        #     "ImgDict":["ed","es"]
+        # }
+        self.select_channels_dict = select_channels_dict
+        self.ALLdata = self.combine_data(mappings=self.mappings)
+        self.out_sz = out_sz
+        self.label_key = label_key
+        self.min_crop_ratio = min_crop_ratio
+        self.transform = transform
+        self.clamp_range = clamp_range
+        self.ndims = 3
+        self.is_reverse_axis_order = True # for inference, always reverse axis order (nifty is reverse order than numpy)
+
+        # Start you filtering here:
+        # self.ALLdata_filtered = self.ALLdata.copy()
+        # filter out images with dimensions less than min_dim
+        self.ALLdata_filtered = self.get_filter_mindim()
+        # filter out images with ROIs that are not in the provided ROIs
+        if ROIs is None:
+            self.ROIs = self.get_all_ROI()
+        else:
+            self.ROIs = ROIs
+        self.ALLdata_filtered = self.get_filter_ROIs()
+        self.ALLdata_filtered = self.get_filter_labels(task_key=task_key,label_keys=label_key)
+        # filtering ends here
+        
+        self.roi_scan_mapping = self.build_ROI_scan_mapping()
+        self.keys_dist, self.total = self.get_keys_dist()
+        
+        
+
+    def get_all_ROI(self):
+        # Get all the ROI options. and remove the reduntant ones
+        ROIs = []
+        for k in self.ALLdata_filtered.keys():
+            ROIs.append(self.ALLdata[k]['ROI'])
+        ROIs = set(ROIs)
+        return ROIs
+
+    def get_keys_dist(self):
+        ROIs = self.get_all_ROI()
+        keys_dist = {}
+        total = 0
+        for item in self.ALLdata_filtered.keys():
+            if self.ALLdata_filtered[item]['ROI'] not in keys_dist:
+                keys_dist[self.ALLdata_filtered[item]['ROI']] = 0
+            keys_dist[self.ALLdata_filtered[item]['ROI']] += 1
+            
+        return keys_dist, total
+
+    def build_ROI_scan_mapping(self):
+        # Build a mapping of ROIs to scans
+        ROI_scan_mapping = {}
+        for item in self.ALLdata_filtered.keys():
+            if self.ALLdata_filtered[item]['ROI'] not in ROI_scan_mapping:
+                ROI_scan_mapping[self.ALLdata_filtered[item]['ROI']] = []
+            ROI_scan_mapping[self.ALLdata_filtered[item]['ROI']].append(item)
+        return ROI_scan_mapping
+
+    def get_3D_volume(self, volume, select_channel = None):
+        volume = reverse_axis_order(volume) if self.is_reverse_axis_order else volume
+        if volume.ndim == 4:
+            if select_channel is None:
+                select_channel = np.random.randint(0, volume.shape[3] - 1)
+            volume = volume[:, :, :, select_channel]
+            # print(f"Volume shape: {volume.shape}, selected channel: {select_channel}")
+        return volume
+
+    def get_filter_mindim(self):
+        # Filter out images with dimensions less than min_dim
+        # Top priority is to filter out images with dimensions less than min_dim
+        ALLdata = self.ALLdata.copy()
+        for k in self.ALLdata.keys():
+            if min(self.ALLdata[k]['Size'][:self.ndims]) < self.out_sz/2:
+                del ALLdata[k]
+        return ALLdata
+
+    def find_min_dim(self):
+        # Find the minimum dimension of the images
+        min_dim = 100000
+        for k in self.ALLdata.keys():
+            value = self.ALLdata[k]
+            if min(value['Size']) < min_dim:
+                min_dim = min(value['Size'])
+        return min_dim
+
+    # def combine_data(self):
+    #     ALLdata = {}
+    #     for j in self.mappings.keys():
+    #         with open(self.mappings[j], 'r') as f:
+    #             mappings = json.load(f)
+    #             ALLdata.update(mappings)
+    #     return ALLdata
+    
+    def combine_data(self, mappings = mapping_files):
+        ALLdata = {}
+        for j in mappings.keys():
+            with open(mappings[j], 'r') as f:
+                mappings_tmp = json.load(f)
+                ALLdata.update(mappings_tmp)
+        return ALLdata
+
+    def normalize(self, volume, eps=1e-7):
+        # Normalize the image (0-1)
+        volume = volume.astype(np.float64)
+        volume = (volume - np.min(volume)) / (np.ptp(volume) + eps)
+        return volume
+
+    def get_key_by_ROI(self, key_word):
+        # Get all the keys with a key word
+        keys = []
+        for k in self.ALLdata_filtered.keys():
+            if key_word == self.ALLdata_filtered[k]["ROI"]:
+                keys.append(k)
+        return keys
+    
+    def get_filter_task(self, task_key = 'segmentation'):
+        # Filter out images with task type that are not in the provided labels_path
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        for k in self.ALLdata_filtered.keys():
+            if 'Label_path' not in self.ALLdata_filtered[k] or task_key not in self.ALLdata_filtered[k]['Label_path']:
+                del ALLdata_filtered[k]
+                Warning(f"Label path not found for {k} with task key {task_key}. This image will be removed from the dataset.")
+        return ALLdata_filtered
+
+    def get_filter_labels(self, task_key='segmentation', label_keys=['heart']):
+        # Filter out images where 'Label_path' does not contain any of the label_keys for the given task_key
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        keys_to_remove = []
+        for k in list(ALLdata_filtered.keys()):
+            label_path = ALLdata_filtered[k].get('Label_path', {})
+            task_labels = label_path.get(task_key, {})
+            # Check if any label_keys are present in task_labels
+            # print(f"Checking {k} for task key {task_labels.keys()} with label keys {label_keys}")
+            has_any_label = any((tk in label_keys) for tk in task_labels.keys())
+            # print(f"Has any label: {has_any_label}")
+            if not has_any_label:
+                keys_to_remove.append(k)
+                # print(f"Label path not found for {k} with task key {task_key} and label keys {label_keys}. This image will be removed from the dataset.")
+        for k in keys_to_remove:
+            del ALLdata_filtered[k]
+        return ALLdata_filtered
+
+    def get_random_pad_crop_params(self, volume_shape, crop_size=None, random=True):
+        # Get random padding and cropping parameters for a given shape
+        d, h, w = volume_shape[:3]
+        if crop_size is None:
+            crop_size = self.out_sz
+        crop_d, crop_h, crop_w = crop_size, crop_size, crop_size
+
+        # Calculate padding
+        pad_width = []
+        for size, crop in zip((d, h, w), (crop_d, crop_h, crop_w)):
+            if crop > size:
+                total_pad = crop - size
+                pad_before = np.random.randint(0, total_pad + 1)
+                pad_after = total_pad - pad_before
+                pad_width.append((pad_before, pad_after))
+            else:
+                pad_width.append((0, 0))
+
+        # Update shape after padding
+        d_p, h_p, w_p = d + pad_width[0][0] + pad_width[0][1], h + pad_width[1][0] + pad_width[1][1], w + pad_width[2][0] + pad_width[2][1]
+
+        if random:
+            # Calculate cropping start indices (random crop)
+            start_d = np.random.randint(0, d_p - crop_d + 1) if d_p > crop_d else 0
+            start_h = np.random.randint(0, h_p - crop_h + 1) if h_p > crop_h else 0
+            start_w = np.random.randint(0, w_p - crop_w + 1) if w_p > crop_w else 0
+        else:    
+            # Calculate cropping start indices (center crop)
+            start_d = max((d_p - crop_d) // 2, 0)
+            start_h = max((h_p - crop_h) // 2, 0)
+            start_w = max((w_p - crop_w) // 2, 0)
+
+        crop_slices = (start_d, start_h, start_w, crop_d, crop_h, crop_w)
+        return pad_width, crop_slices
+
+    def apply_pad_crop(self, volume, pad_width, crop_slices):
+        # Apply padding and cropping to the volume
+        if any(pad != (0, 0) for pad in pad_width):
+            volume = np.pad(volume, pad_width, mode='constant', constant_values=0)
+        start_d, start_h, start_w, crop_d, crop_h, crop_w = crop_slices
+        cropped_array = volume[start_d:start_d + crop_d, start_h:start_h + crop_h, start_w:start_w + crop_w]
+        return cropped_array
+    
+    def get_filter_ROIs(self):
+        ALLdata_filtered = self.ALLdata_filtered.copy()
+        for k in self.ALLdata_filtered.keys():
+            if self.ALLdata_filtered[k]['ROI'] not in self.ROIs:
+                del ALLdata_filtered[k]
+        return ALLdata_filtered
+
+    def get_channel_ids(self, key):
+        """
+        Get the indices where ImgDict values match the selected channels (e.g., 'ed', 'es').
+        
+        Returns:
+            list: List of integer indices matching the selected channels
+        """
+        img_dict = self.ALLdata_filtered[key].get("ImgDict", {})
+        selected_values = self.select_channels_dict.get("ImgDict", [])
+        # Build reverse mapping: value -> index
+        value_to_idx = {value: int(idx) for idx, value in img_dict.items()}
+    
+        # Get indices in the order of selected_values
+        indices = [
+            value_to_idx[val] for val in selected_values 
+            if val in value_to_idx
+        ]
+        return indices
+        # return sorted(indices)
+
+    def __len__(self):
+        return len(self.ALLdata_filtered.keys())
+
+    def __getitem__(self, idx):
+        key = list(self.ALLdata_filtered.keys())[idx]
+        return_dict = dict()
+        
+        print(f"Processing key: {key}")
+
+        volume = sitk.ReadImage(key)
+        volume = sitk.GetArrayFromImage(volume)
+        
+        if volume.ndim == 4:
+            channel_ids = self.get_channel_ids(key)
+            if len(channel_ids) == 0:
+                # warning message that this key has no matching channels
+                Warning(f"No matching channels found for key: {key} with ImgDict: {self.ALLdata_filtered[key].get('ImgDict', {})} and selected channels: {self.select_channels_dict.get('ImgDict', [])}. Using random channel.")
+                channel_id = None
+            else:
+                channel_id=channel_ids[0]
+
+        volume = self.get_3D_volume(volume, select_channel = channel_id)
+
+        if self.clamp_range is not None:
+            modality = self.ALLdata_filtered[key].get("Modality", None)
+            if modality == "CT":
+                volume = np.clip(volume, self.clamp_range[0], self.clamp_range[1])
+        volume = self.normalize(volume)
+        
+        crop_ratio = np.random.uniform(self.min_crop_ratio, 1)
+
+        crop_size = int(max(volume.shape) * crop_ratio)
+        pad_width, crop_slices = self.get_random_pad_crop_params(volume.shape, crop_size)
+        # print(f"Pad width: {pad_width}, Crop slices: {crop_slices}, Original shape: {volume.shape}")
+        volume = self.apply_pad_crop(volume, pad_width, crop_slices)
+
+        label_dict = dict()
+        if 'Label_path' in self.ALLdata_filtered[key]:
+            for lk in self.label_key:
+                if lk in self.ALLdata_filtered[key]['Label_path']['segmentation'].keys():
+                    label = sitk.ReadImage(self.ALLdata_filtered[key]['Label_path']['segmentation'][lk])
+                    label = sitk.GetArrayFromImage(label)
+                    # print(f"Label shape: {label.shape}, key: {key}, label key: {lk}")
+                    label = reverse_axis_order(label) if self.is_reverse_axis_order else label
+                    
+                    # print(f"Label shape: {label.shape}, key: {key}, label key: {lk}")
+                    if label.ndim > self.ndims:
+                        if len(channel_ids) != 0:
+                            label = label[...,channel_ids]  # assuming channel last
+                        pad_width_lab = pad_width + [(0,0)]*(label.ndim - self.ndims)
+                        # print(f"Label with channels, pad_width_lab: {pad_width_lab}")
+                    else:
+                        pad_width_lab = pad_width
+                    label = self.apply_pad_crop(label, pad_width_lab, crop_slices)
+                    # print(f"After pad and crop, label shape: {label.shape}, key: {key}, label key: {lk}")
+                    label_dict[lk] = resize(label,[self.out_sz]*self.ndims, anti_aliasing = False, preserve_range = True, order=0)
+                    if label.ndim > self.ndims:
+                        if self.ndims==3:
+                            label_dict[lk] = np.transpose(label_dict[lk], (3,0,1,2))  # assuming channel last
+                        elif self.ndims==4:
+                            label_dict[lk] = np.transpose(label_dict[lk], (4,0,1,2,3))  # assuming channel last
+                    # print(f"After resize, label shape: {label_dict[lk].shape}, key: {key}, label key: {lk}")
+                else:
+                    label_dict[lk] = np.full([self.out_sz]*self.ndims, -1)
+                    Warning(f"Label path not found for {key} with label key {lk}.")
+                label_dict[lk] = label_dict[lk][None, :, :, :] if label_dict[lk].ndim == 3 else label_dict[lk]
+        else:
+           for lk in self.label_key:
+                label_dict[lk] = np.full([self.out_sz]*self.ndims, -1)
+                Warning(f"Label path not found for {key} with label key {lk}.")
+                label_dict[lk] = label_dict[lk][None, :, :, :]
+        
+        volume =resize(volume, [self.out_sz]*self.ndims, anti_aliasing = True, preserve_range = True)
+        # return_dict['labels'] = label_dict
+        return_dict['labels'] = np.concatenate([v for v in label_dict.values()], axis=1)
+        
+        return_dict['img'] = volume[None, :, :, :]
+        return_dict['label_channels'] = list(self.select_channels_dict.get("ImgDict", []))
+        return return_dict
+
+
+class OminiDataset_bertembd(OminiDataset):
+    def __init__(self, 
+                 out_sz = 128, 
+                 transform=None, 
+                 clamp_range = CLAMP_RANGE, 
+                 min_crop_ratio = 0.85, 
+                 ROIs = None, 
+                 modality = None, 
+                 reverse_axis_order = False,
+                 min_dim = 3,
+                 mapping_files = mapping_files):
+        super().init(out_sz = out_sz, 
+                     transform = transform,
+                     clamp_range = clamp_range, 
+                     min_crop_ratio = min_crop_ratio, 
+                     ROIs = ROIs, 
+                     modality = modality, 
+                     reverse_axis_order = reverse_axis_order,
+                     min_dim = min_dim,
+                     mapping_files=mapping_files)
+        # start you filtering here
+        self.ALLdata_filtered = self.get_filter_mindim()
+        if ROIs is None:
+            #  if no ROIs are provided, get all the ROIs from filtered data
+            self.ROIs = self.get_all_ROI()
+        else:
+            self.ROIs = ROIs
+        self.ALLdata_filtered = self.get_filter_ROIs()
+        # self.ALLdata_filtered = self.filter_embd()
+        # self.ALLdata_filtered = self.get_filter_labels(task_key=task_key,label_keys=label_key)
+        # end your filtering here
+    def __getitem__(self, idx):
+        key = list(self.ALLdata_filtered.keys())[idx]
+        embd = self.ALLdata_filtered[key]['embd']
+        if 0:
+            print(key) 
+
+        volume = sitk.ReadImage(key)
+        volume = sitk.GetArrayFromImage(volume)
+        volume = self.get_3D_volume(volume)
+        
+        if self.clamp_range is not None:
+            modality = self.ALLdata_filtered[key].get("Modality", None)
+            if modality == "CT":
+                volume = np.clip(volume, self.clamp_range[0], self.clamp_range[1])
+        volume = self.normalize(volume)
+        
+        if self.min_crop_ratio is not None:
+            crop_ratio = np.random.uniform(self.min_crop_ratio, 1)
+            crop_size = int(max(volume.shape) * crop_ratio)
+            volume = self.random_crop_3d(volume, crop_size)
+            volume = resize(volume, [self.out_sz]*self.ndims, anti_aliasing = True, preserve_range = True)
+        else:
+            volume = self.random_crop_3d(volume, self.out_sz)
+        volume = volume[None, :, :, :]
+
+        if self.transform is not None:
+            return  self.transform(volume)
+
+        return volume,np.array(embd)
+    
+    def __len__(self):
+        return len(self.ALLdata_filtered.keys())
+    
+    def filter_embd(self):
+        for k in self.ALLdata_filtered.keys():
+            if 'BERT_embedding_keys' not in self.ALLdata_filtered[k]['Metadata']:
+                del self.ALLdata_filtered[k]
+        return self.ALLdata_filtered
+                

Dataloader/dataloader0.py

@@ -0,0 +1,421 @@
+import os
+import torch
+import torchvision
+from torch.utils.data import Dataset, DataLoader
+from torchvision import datasets, transforms
+import nibabel as nib
+from skimage.transform import rescale, resize, downscale_local_mean
+from scipy.ndimage import zoom
+import numpy as np
+# import SimpleITK as sitk
+
+# print(os.getcwd())
+import sys
+sys.path.append('./')
+from Dataloader.dataloader_utils import *
+
+EPS = 1e-7
+
+def get_dataloader(data_name='cmr',mode='train'):
+  if data_name=='cmr':
+    if mode=='train':
+      dataloader=CMR_loader
+    elif mode =='aug':
+      dataloader=CMR_tgt_loader
+    else:
+      print('mode not exist')
+  elif data_name=='lct':
+    if mode=='train':
+      dataloader=LCT_loader
+    elif mode =='aug':
+      dataloader=LCT_tgt_loader
+    else:
+      print('mode not exist')
+  else:
+    print('dataloader not exist')
+  return dataloader
+
+class LCT_loader(Dataset):
+  def __init__(self, data_root_path = f'Data/Src_data/CTLung_processed/', target_res = (256, 256),transforms = None, noise_scale=0.0, patient_index = None):
+  # def __init__(self, data_root_path = '/home/data/jzheng/CTLung_processed/', target_res = (256, 256),transforms = None, noise_scale=0.0, patient_index = None):
+    self.files = [data_root_path + f for f in os.listdir(data_root_path) if f.endswith('.npy')]
+    self.transforms = transforms
+    self.noise_scale=noise_scale
+    self.d_p = data_root_path
+  
+  def __getitem__(self, item):
+    array = np.load(self.files[item])
+    if 'process' not in self.d_p:
+      array = (array - array.min()) / (array.max() - array.min() + EPS) # Normalize to 0 to 1
+    array = array[None,:,:,:] # add a channel to array make it (‘C’,H,W,Z)
+    if self.transforms != None:
+      array = self.transforms(array)
+    # print(array.shape)
+    return array, array, item # -> (B, C, H, W, Z)
+    # return array, array # -> (B, C, H, W, Z)
+  
+  def __len__(self):
+    return len(self.files)
+
+class LCT_tgt_loader(Dataset):
+  def __init__(self, data_root_path = "Data/Tgt_data/lct/",noise_scale=0.0, patient_index = None):
+    self.files_gt = [data_root_path + "Gt/" + f for f in os.listdir(data_root_path + "Gt/")]
+    self.files_tr = [data_root_path + 'Tr/' + f for f in os.listdir(data_root_path + "Tr/")]
+
+    self.files_tr.sort()
+    self.files_gt.sort()
+
+    self.transforms = transforms
+    self.noise_scale=noise_scale
+
+  def __getitem__(self, item):
+    img_nib = nib.load(self.files_tr[item])
+    mask_nib = nib.load(self.files_gt[item])
+
+    image = img_nib.get_fdata()
+    mask = mask_nib.get_fdata()
+
+    image = image[None,:,:,:]
+    mask = mask[None,:,:,:]
+
+    print(self.files_tr[item],self.files_gt[item])
+    
+    return image, mask, item 
+
+    
+
+  def __len__(self):
+    assert len(self.files_gt) == len(self.files_tr)
+    return len(self.files_gt)
+
+class LCT_seg(Dataset):
+  def __init__(self, data_root_path = "/home/data/jzheng/CTLung_processed/testset/modality_0001/",noise_scale=0.0, patient_index = None):
+    self.files_gt = [data_root_path + "Gt/" + f for f in os.listdir(data_root_path + "Gt/")]
+    self.files_tr = [data_root_path + 'Tr/' + f for f in os.listdir(data_root_path + "Tr/")]
+
+    self.files_tr.sort()
+    self.files_gt.sort()
+
+    self.transforms = transforms
+    self.noise_scale=noise_scale
+
+  def __getitem__(self, item):
+    img_nib = nib.load(self.files_tr[item])
+    mask_nib = nib.load(self.files_gt[item])
+
+    image = img_nib.get_fdata()
+    mask = mask_nib.get_fdata()
+
+    image = image[None,:,:,:]
+    mask = mask[None,:,:,:]
+
+    print(self.files_tr[item],self.files_gt[item])
+    
+    return image, mask, item 
+
+    
+
+  def __len__(self):
+    assert len(self.files_gt) == len(self.files_tr)
+    return len(self.files_gt)
+
+class CMR_loader_preprocess(Dataset):
+  # This is for pre_processing for CMR. not use for training model   
+  def __init__(self, data_path = 'Data/CTLung_processed/', target_res = (256, 256), transforms = None, noise_scale=0.0):
+  # def __init__(self, data_path = '/home/data/jzheng/CMR_processed/', target_res = (256, 256), transforms = None, noise_scale=0.0):
+    self.d_p = data_path
+    self.target_res = target_res
+    self.files = [self.d_p + x for x in os.listdir(self.d_p)]
+    self.transforms = transforms
+    self.noise_scale=noise_scale
+    
+  def __getitem__(self, item):
+     array = nib.load(self.files[item]).get_fdata()
+     array = resize(array, self.target_res, anti_aliasing = True, preserve_range = True)
+     array = array[None, :, :]
+     array = remove_background(array)      # jzheng 20240228
+     array = (array - array.min()) / (array.max() - array.min() + EPS)
+
+     if self.noise_scale > 0:
+       array = thresh_img(array,[0,self.noise_scale])
+       array = array * (np.random.normal(1, self.noise_scale*2))
+
+     if self.transforms != None:
+       array = self.transforms(array)
+     return array, self.files[item]
+                         
+  def  __len__(self):
+    return len(self.files)
+  
+class CMR_loader(Dataset):
+#   niff format size is (H,W) for CMR
+#   CMR_processed_rmbg_resize means the niif image has been gone throught rmbg and resize offline to make trainig fast
+  def __init__(self, data_path = f'Data/Src_data/CMR_processed_rmbg_resize/', target_res = (256, 256), transforms = None, noise_scale=0.0):
+  # def __init__(self, data_path = '/home/data/jzheng/CMR_processed_rmbg_resize/', target_res = (256, 256), transforms = None, noise_scale=0.0):
+    self.d_p = data_path
+    self.ndims = 2
+    self.target_res = target_res
+    self.files = [self.d_p + x for x in os.listdir(self.d_p)]
+    self.transforms = transforms
+    # self.get_transform()
+    self.noise_scale=noise_scale
+    self.preprocessed='resize' in data_path
+    
+  def __getitem__(self, item):
+     array = nib.load(self.files[item]).get_fdata()
+     if not self.preprocessed:
+        array = resize(array, self.target_res, anti_aliasing = True, preserve_range = True)
+     array = array[None, :, :]
+     if not self.preprocessed:
+         array = remove_background(array)      # jzheng 20240228
+         array = (array - array.min()) / (array.max() - array.min() + EPS)
+
+     # if self.noise_scale > 0:
+     #   array = thresh_img(array,[0,self.noise_scale])
+     #   array = array * (np.random.normal(1, self.noise_scale*2)) + np.random.normal(0, self.noise_scale*2)
+
+     if self.transforms != None:
+       array = self.transforms(array)
+     return array, array, item
+
+  def  __len__(self):
+    return len(self.files)
+
+  def get_transform(self,degrees=np.pi,translate=0.125):
+    # self.transforms = torchvision.transforms.RandomAffine(degrees=degrees,translate=[translate]*self.ndims,interpolation=torchvision.transforms.InterpolationMode.BILINEAR)
+    self.transforms = torchvision.transforms.Compose([
+      # torchvision.transforms.Resize((hyp_parameters['img_size'], hyp_parameters['img_size'])),
+      torchvision.transforms.ToTensor(),
+      torchvision.transforms.RandomAffine(degrees=degrees,translate=[translate]*self.ndims,interpolation=torchvision.transforms.InterpolationMode.BILINEAR),
+      # torchvision.transforms.ToTensor(),
+      # torchvision.transforms.Normalize(0.5, 0.5)
+      # Lambda(lambda x: (x - 0.5) * 2)
+    ])
+    return
+
+class CMR_tgt_loader(Dataset):
+  def __init__(self, 
+               data_path = 'Data/Tgt_data/cmr/',
+              #  gt_path = '/home/data/jzheng/acdc/train_gt/',
+               target_res = (256,256), 
+               is_3d = False,
+               patient_index = [],
+               ):
+    
+    #  parameter initialize
+    self.d_p = os.path.join(data_path,'Tr','') 
+    self.gt_p = os.path.join(data_path,'Gt','')
+    self.img_files = os.listdir(self.d_p)
+    self.gt_files = os.listdir(self.gt_p)
+    self.p_indice = patient_index
+    self.target_res_2d = target_res
+    self.img_files.sort()
+    self.gt_files.sort()
+    self.img_samples = []
+    self.gt_samples = []
+    self.p_id = []
+
+    if len(self.p_indice) == 0:
+      self.p_indice = [x for x in range(1,101)]
+    # build patient-to-file correspondence
+    p2f = {}
+    assert len(self.gt_files) == len(self.img_files)
+    print(self.p_indice)
+    for i in self.p_indice:
+      for gt_f, img_f in zip(self.gt_files, self.img_files):
+        pf_id = gt_f.split('_')[0]
+        pf_id = pf_id[-3:]
+        if i == int(pf_id):
+          img_volume = nib.load(self.d_p + img_f).get_fdata()
+          gt_volume = nib.load(self.gt_p + gt_f).get_fdata()
+          assert img_volume.shape == gt_volume.shape
+          depth = img_volume.shape[2]
+          for si in range(depth):
+            img = resize(img_volume[:, :, si], self.target_res_2d, anti_aliasing=True, preserve_range=True)
+            img = (img - img.min()) / (img.max() - img.min() + EPS)
+
+            gt = gt_volume[:, :, si]
+
+            gt_1_index = gt == 1
+            gt_2_index = gt == 2
+            gt_3_index = gt == 3
+            gt_4_index = gt == 4
+
+            gt_1 = gt * gt_1_index
+            gt_2 = gt * gt_2_index
+            gt_3 = gt * gt_3_index
+            gt_4 = gt * gt_4_index
+
+
+            gt_1 = resize(gt_1, self.target_res_2d, anti_aliasing=True, preserve_range=True)
+            gt_2 = resize(gt_2, self.target_res_2d, anti_aliasing=True, preserve_range=True)
+            gt_3 = resize(gt_3, self.target_res_2d, anti_aliasing=True, preserve_range=True)
+            gt_4 = resize(gt_4, self.target_res_2d, anti_aliasing=True, preserve_range=True)
+
+
+            self.img_samples.append(img[np.newaxis, :, :])
+            self.gt_samples.append(np.array([gt_1, gt_2, gt_3, gt_4]))
+            self.p_id.append(i)
+
+
+  def __getitem__(self, item):
+
+    return self.img_samples[item], self.gt_samples[item], self.p_id[item]
+
+
+  def __len__(self):
+
+    assert len(self.img_samples) == len(self.gt_samples)
+    return len(self.img_samples)
+
+class acdc_seg(Dataset):
+  def __init__(self, 
+               data_path = '/home/data/jzheng/acdc/train_images/',
+               gt_path = '/home/data/jzheng/acdc/train_gt/',
+               target_res = (256,256), 
+               is_3d = False,
+               patient_index = [],
+               ):
+    
+    #  parameter initialize
+    self.d_p = data_path
+    self.gt_p = gt_path
+    self.img_files = os.listdir(self.d_p)
+    self.gt_files = os.listdir(self.gt_p)
+    self.p_indice = patient_index
+    self.target_res_2d = target_res
+    self.img_files.sort()
+    self.gt_files.sort()
+    self.img_samples = []
+    self.gt_samples = []
+    self.p_id = []
+
+    if len(self.p_indice) == 0:
+      self.p_indice = [x for x in range(1,101)]
+    # build patient-to-file correspondence
+    p2f = {}
+    assert len(self.gt_files) == len(self.img_files)
+    print(self.p_indice)
+    for i in self.p_indice:
+      for gt_f, img_f in zip(self.gt_files, self.img_files):
+        pf_id = gt_f.split('_')[0]
+        pf_id = pf_id[-3:]
+        if i == int(pf_id):
+          img_volume = nib.load(self.d_p + img_f).get_fdata()
+          gt_volume = nib.load(self.gt_p + gt_f).get_fdata()
+          assert img_volume.shape == gt_volume.shape
+          depth = img_volume.shape[2]
+          for si in range(depth):
+            img = resize(img_volume[:, :, si], self.target_res_2d, anti_aliasing=True, preserve_range=True)
+            img = (img - img.min()) / (img.max() - img.min() + EPS)
+
+            gt = gt_volume[:, :, si]
+
+            gt_1_index = gt == 1
+            gt_2_index = gt == 2
+            gt_3_index = gt == 3
+            gt_4_index = gt == 4
+
+            gt_1 = gt * gt_1_index
+            gt_2 = gt * gt_2_index
+            gt_3 = gt * gt_3_index
+            gt_4 = gt * gt_4_index
+
+
+            gt_1 = resize(gt_1, self.target_res_2d, anti_aliasing=True, preserve_range=True)
+            gt_2 = resize(gt_2, self.target_res_2d, anti_aliasing=True, preserve_range=True)
+            gt_3 = resize(gt_3, self.target_res_2d, anti_aliasing=True, preserve_range=True)
+            gt_4 = resize(gt_4, self.target_res_2d, anti_aliasing=True, preserve_range=True)
+
+
+            self.img_samples.append(img[np.newaxis, :, :])
+            self.gt_samples.append(np.array([gt_1, gt_2, gt_3, gt_4]))
+            self.p_id.append(i)
+
+
+  def __getitem__(self, item):
+
+    return self.img_samples[item], self.gt_samples[item], self.p_id[item]
+
+
+  def __len__(self):
+
+    assert len(self.img_samples) == len(self.gt_samples)
+    return len(self.img_samples)
+
+class acdc_gan(Dataset):
+  def __init__(self, 
+               train_path = '/home/data/jzheng/acdc/images/',
+               target_res = (32, 256, 256),
+               is_3d = False,
+               transforms = None
+               ):
+    self.t_p = train_path
+    self.files = os.listdir(self.t_p)
+    self.sample_list_2d = []
+    self.is_3d = is_3d
+    self.target_res = target_res
+    self.res_2d = (target_res[1], target_res[2])
+    self.transforms = transforms
+
+    if self.is_3d == False:
+      for f in self.files:
+        img = nib.load(self.t_p + f).get_fdata()
+        depth = img.shape[2]
+        f_i = int(round(depth*0.1))
+        b_i = int(round(depth*0.9))
+        interval_slice = img[:, :, f_i:b_i]
+        for ii in range(interval_slice.shape[2]):
+          single_slice = interval_slice[:,:,ii]
+          single_slice = resize(single_slice, self.res_2d, anti_aliasing=True, preserve_range=True)
+          single_slice = (single_slice -  single_slice.min()) / ( single_slice.max() -  single_slice.min() + EPS)
+          self.sample_list_2d.append(single_slice[None,:,:])
+
+
+  def __len__(self):
+    if self.is_3d == False:
+      return len(self.sample_list_2d)
+    else:
+      return len(self.files )
+
+  def __getitem__(self, index):
+    if self.is_3d == False:
+      return self.sample_list_2d[index], self.sample_list_2d[index]
+    for f in self.files:
+      img = nib.load(self.t_p + f).get_fdata()
+      target_d_ratio = self.target_res[0] / img.shape[2]
+      target_w_ratio = self.target_res[1] / img.shape[0]
+      target_h_ratio = self.target_res[2] / img.shape[1]
+
+      resize_img = zoom(img, (target_w_ratio, target_h_ratio, target_d_ratio))
+
+      resize_img = np.swapaxes(resize_img, 0, 2)
+      resize_img = np.swapaxes(resize_img, 1, 2)
+      resize_img = (resize_img - resize_img.min()) / (resize_img.max() - resize_img.min() + EPS)
+      if transforms != None:
+        resize_img = self.transforms(resize_img)
+      return resize_img, resize_img
+
+class acdc_gan_single_slice(Dataset):
+  def __init__(self, train_path = '/well/papiez/shared/ACDC/clean_training/images/'):
+    self.t_p = train_path
+    self.files = os.listdir(self.t_p)
+
+  def __len__(self):
+    return len(self.files)
+
+  def __getitem__(self, index):
+    img = self.files[index]
+    img = nib.load(self.t_p + img).get_fdata()
+    depth = img.shape[2]
+    mid_d = int(depth/2)
+    mid_slice = img[:,:,mid_d]
+    mid_slice = resize(mid_slice, (128, 128), anti_aliasing=True, preserve_range=True)
+    mid_slice = (mid_slice-mid_slice.min())/(mid_slice.max()-mid_slice.min()+EPS)
+    # print(mid_slice.max(),mid_slice.min())
+
+    return mid_slice, mid_slice
+
+
+
+

Dataloader/dataloader_tester.py

@@ -0,0 +1,39 @@
+from dataLoader import *
+import torchvision.transforms as tf
+import SimpleITK as sitk
+import os 
+
+
+transform = tf.Compose([
+    tf.ToTensor(),  # Convert image to tensor
+])
+
+mapping_files_bert = {
+    # 'TotalSegmentor': '/home/data/Github/data/data_gen_def/DATASETS_processed/TotalSegmentorCT_MRI/nifti_mappings.json',
+    # 'MSD': '/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/MSD_processed/nifti_mappings_updated.json',
+    'CancerImageArchive': '/home/data/Github/OmniMorph/Dataloader/nifty_mappings/CIA_mappings.json',
+}
+
+if __name__ == "__main__":
+    # dataset = OminiDataset_v1(transform=None)
+    # datasetp = OminiDataset_paired(transform=None)
+    # dataset = OminiDataset_paired_inf(transform=None)
+    # dataset = OminiDataset_inference_w_all(transform=None)
+    # dataset = OminiDataset_bertembd(transform=None,mapping_files=mapping_files_bert)
+    dataset = OminiDataset(transform=None)
+    
+    
+    
+    
+    # print(dataset.get_keys_dist())
+    # print(len(dataset))
+    # print(dataset.build_batch().shape)
+    # exit()
+    dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
+
+    for i, data in enumerate(dataloader):
+        print(data[1])
+        exit()
+
+    
+    # print(dataset.get_ALLdata())

Dataloader/dataloader_utils.py

@@ -0,0 +1,193 @@
+import os
+import torch
+# from torch import nn, optim
+# from torch.autograd.variable import Variable
+# from torchvision import transforms, datasets
+# from torchvision.utils import save_image
+# import torch.nn.functional as F
+# import scipy.ndimage as spimg
+# import pyquaternion as quater
+# import random
+import numpy as np
+from scipy.ndimage import gaussian_filter, binary_dilation, binary_erosion, generate_binary_structure
+import pydicom
+from scipy.ndimage import zoom
+from einops import rearrange, reduce, repeat
+
+def get_sizeRange_dict(roi=''):
+    """
+    Returns a dictionary with size ranges for different regions of interest (ROIs).
+    If a specific ROI is provided, returns the size range for that ROI.
+    If no ROI is provided, returns the entire dictionary.
+    Args:
+        roi (str): The region of interest for which to get the size range.
+    Returns:
+        dict or list: A dictionary with size ranges for all ROIs, or a list with the size range for the specified ROI.
+    """
+    # Define the size ranges for different ROIs
+    # The values are in the format [min_size, max_size]
+    # The sizes are in mm for the minimum and maximum dimensions
+    sizeRange_dict = {
+      'whole-body': [420, 2048],
+      'neck-thorax-abdomen-pelvis-leg': [400, 2048],
+      'neck-thorax-abdomen-pelvis': [380, 2048],
+      'thorax-abdomen-pelvis-leg': [360, 2048],
+      'neck-thorax-abdomen': [320, 1024],
+      'head-neck-thorax-abdomen': [360, 2048],
+      'head-neck-thorax': [340, 1024],
+      'thorax-abdomen-pelvis': [340, 1024],
+      'abdomen-pelvis-leg': [320, 1024],
+      'neck-thorax': [220, 1024],
+      'thorax-abdomen': [260, 1024],
+      'abdomen-pelvis': [260, 1024],
+      'pelvis-leg': [240, 1024],
+      'head-neck': [240, 1024],
+      'head': [150, 1024],
+      'brain': [128, 1024],
+      'neck': [140, 1024],
+      'abdomen': [240, 1024],
+      'pelvis': [220, 1024],
+      'thorax': [220, 1024],
+      'arm': [140, 1024],
+      'hand': [140, 1024],
+      'leg': [160, 1024],
+      'skeleton': [130, 1024],
+    }
+    if roi in sizeRange_dict:
+        return sizeRange_dict[roi]
+    else:
+        return sizeRange_dict
+
+
+def remove_background(img,replace_value=None,num_bin=256,dim_ch=0,sigma=None):
+    # common_value1,common_value2=[], []
+    # if replace_value is None:
+    if dim_ch is None:
+        dim_ch=0
+        img=np.expand_dims(img,axis=dim_ch)
+    ims = np.split(img,img.shape[dim_ch],axis=dim_ch)
+    # ims =[img]
+    ims = [np.squeeze(im,axis=dim_ch) for im in ims]
+    msk1 = np.ones_like(ims[0])
+    for im in ims:
+        if num_bin>0:
+            flatten_im=im.flatten()
+            hist, bins = np.histogram(flatten_im,bins=range(num_bin))
+            # common_value1.append(np.argmax(hist))
+            common_value1 = np.argmax(hist)
+            # hist[common_value1] = -10**5
+            msk1[im!=common_value1] = 0
+            # common_value2 = np.argmax(hist)
+    if sigma is not None and sigma > 0:
+        # struct=generate_binary_structure()
+        msk1 = binary_dilation(msk1,iterations=int(sigma*4)).astype(float)
+        msk0 = binary_erosion(1-msk1,iterations=int(sigma*4)).astype(float)
+        msk_blur = gaussian_filter(msk0, sigma=sigma*4,truncate=sigma//4, mode='nearest')
+        # msk_blur = msk0
+    for id, im in enumerate(ims):
+        if replace_value is None:
+            # a=im[np.logical_not(msk1)]
+            # replace_value[id] = np.min(im[np.logical_not(msk1)])
+            replace_v=np.min(im[np.logical_not(msk1)])
+        else:
+            replace_v=replace_value[id]
+        # im[msk1==1] = replace_v
+        if sigma is not None and sigma>0:
+            im_blur=im
+            im_blur[msk1==1]=replace_v
+            im_blur = gaussian_filter(im_blur, sigma=sigma*4,truncate=sigma//4, mode='nearest')
+            # im[msk1==1] = im_blur[msk1==1]
+            im=im*(msk_blur) + im_blur*(1-msk_blur)
+        else:
+            im[msk1 == 1] = replace_v
+        # print(im.shape)
+        ims[id]=im
+    return np.stack(ims,axis=dim_ch)
+
+def thresh_img(img,thresh = None,EPS = 10**-7):
+    
+    if isinstance(thresh,list):
+        threshold=np.random.uniform(thresh[0],thresh[1])
+        upbound=1-np.random.uniform(thresh[0],thresh[1])-threshold
+    else:
+        threshold=thresh
+    if threshold is not None:
+        # img=img-threshold
+        # img=np.where(img>=0,img,0)
+        # img = np.maximum(img-threshold,0)
+        # img = torch.maximum(img - threshold,torch.tensor(0.))
+        if isinstance(img,list):
+            device=img[0].device
+            for i in range(len(img)):
+                img[i] = torch.clamp(img[i]-threshold,min=torch.tensor(0.).to(device),max=torch.tensor(upbound).to(device))
+        else:
+            device=img.device
+            img = torch.clamp(img-threshold,min=torch.tensor(0.).to(device),max=torch.tensor(upbound).to(device))
+    # return (img - img.min()) / (img.max() - img.min() + EPS)
+    return img
+
+def clamp_img_tensor(img,clamp = [None,None]):
+    device=img.device
+    if clamp[0] is not None and clamp[1] is not None:
+        img = torch.clamp(img, min=torch.tensor(clamp[0]).to(device),max=torch.tensor(clamp[1]).to(device))
+    else:
+      if clamp[0] is not None:
+          img = torch.clamp(img, min=torch.tensor(clamp[0]).to(device))
+      if clamp[1] is not None:
+          img = torch.clamp(img, max=torch.tensor(clamp[1]).to(device))
+    return img
+
+def read_CT_volume(folder_path,target_res = 128):
+# read CT into a (128x128x128) cube and pad the insufficient dimension
+    
+  dicom_slices = []
+  # Iterate over each file in the folder
+  for filename in sorted(os.listdir(folder_path), reverse=True):
+    if filename.endswith(".dcm"):  # Check if the file is a DICOM file
+      file_path = os.path.join(folder_path, filename)
+        
+    # Read the DICOM file
+      dicom_data = pydicom.dcmread(file_path)
+        
+    # Append DICOM pixel data to the list
+      dicom_slices.append(dicom_data.pixel_array)
+
+  # Convert the list of slices to a numpy array
+  
+  dicom_slices = np.array(dicom_slices)
+  dicome_volume = rearrange(dicom_slices, 'z h w -> h w z')
+
+  # Get spatial information from the first DICOM file
+  first_dicom = pydicom.dcmread(os.path.join(folder_path, os.listdir(folder_path)[0]))
+  slice_thickness = first_dicom.SliceThickness
+  pixel_spacing = first_dicom.PixelSpacing
+  
+#   Get the scaling ratio for each dim
+  h_axis_ratio = pixel_spacing[0]
+  w_axis_ratio = pixel_spacing[1]
+  z_axis_ratio = slice_thickness
+
+# find the longest dim that need to rescale
+  longest_axis = max([h_axis_ratio*dicome_volume.shape[0], w_axis_ratio*dicome_volume.shape[1],z_axis_ratio*dicome_volume.shape[2]])
+  c_factor = longest_axis/target_res
+#   print((h_axis_ratio/c_factor, w_axis_ratio/c_factor ,z_axis_ratio/c_factor))
+  resized_volume = zoom(dicome_volume, (h_axis_ratio/c_factor, w_axis_ratio/c_factor ,z_axis_ratio/c_factor))
+#   print('resize', resized_volume.shape)
+
+  
+  max_dim_size = max(resized_volume.shape)
+
+ # Calculate padding for each dimension
+  padding_h = max_dim_size - resized_volume.shape[0]
+  padding_w = max_dim_size - resized_volume.shape[1]
+  padding_z = max_dim_size - resized_volume.shape[2]
+  
+  pad_depth = (padding_z // 2, padding_z - padding_z // 2)
+  pad_height = (padding_h // 2, padding_h - padding_h // 2)
+  pad_width = (padding_w // 2, padding_w - padding_w // 2)
+
+#   Pad the array symmetrically
+  padded_resized_volume = np.pad(resized_volume, (pad_height, pad_width, pad_depth), mode='constant')
+
+  return padded_resized_volume, slice_thickness, pixel_spacing
+  

Dataloader/embding_gen.py

@@ -0,0 +1,149 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+import json
+import SimpleITK as sitk
+import numpy as np
+from skimage.transform import rescale, resize, downscale_local_mean
+# from torchvision.transforms import v2
+import sys
+from bert_helper import *
+sys.path.append('./')
+from Dataloader.dataloader_utils import *
+import random
+
+
+
+mapping_files = {
+    # 'MSD': '/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/MSD_processed/nifti_mappings_updated.json',
+    # 'TotalSegmentor': '/home/data/Github/data/data_gen_def/DATASETS_processed/TotalSegmentorCT_MRI/nifti_mappings.json',
+    # 'Kaggle_osic': '/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/Kaggle_osic_new/nifti_mappings.json',
+    # 'CancerImageArchive': '/home/data/Github/data/data_gen_def/DATASETS_processed/CancerImageArchive_test/nifti_mappings.json',
+    # 'MnMs': '/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/MnMs/nifti_mappings.json',
+    # 'Brats2019': '/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/BRATS/BRATS2019/nifti_mappings.json',
+    # 'Brats2020': '/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/BRATS/BRATS2020/nifti_mappings.json',
+    # 'Brats2021': '/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/BRATS/BRATS2021/nifti_mappings.json',
+    # 'OASIS_1': '/home/data/Github/data/data_gen_def/DATASETS_processed/OASIS/OASIS_1/CS_SECTIONAL/nifti_mappings.json',
+    'OASIS_2': '/home/data/Github/data/data_gen_def/DATASETS_processed/OASIS/OASIS_2/RAW_V2/nifti_mappings.json',
+    # 'PSMA-FDG-PET-CT-LESION':'/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/PSMA/PSMA-FDG-PET-CT-LESION/V2/nifti_mappings.json',
+    # 'PSMA-CT':'/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/PSMA/Longitudinal-CT/nifti_mappings.json',
+    # 'AbdomenAtlas':'/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/AbdomenAtlas_v2/nifti_mappings.json',
+    # 'AbdomenCT1k':'/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/AbdomenCT1k/nifti_mappings.json',
+    
+}
+save_paths = {
+    'MSD': '/home/data/Github/OmniMorph/Dataloader/nifty_mappings/MSD_mappings.json',
+    'TotalSegmentor': '/home/data/Github/OmniMorph/Dataloader/nifty_mappings/TotalSegmentorCT_MRI_mappings.json',
+    'Kaggle_osic': '/home/data/Github/OmniMorph/Dataloader/nifty_mappings/Kaggle_osic_mappings.json',
+    'CancerImageArchive': '/home/data/Github/OmniMorph/Dataloader/nifty_mappings/CIA_mappings.json',
+    'MnMs': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/MnMs_mappings.json',
+    'Brats2019': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/Brats2019_mappings.json',
+    'Brats2020': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/Brats2020_mappings.json',
+    'Brats2021': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/Brats2021_mappings.json',
+    'OASIS_1': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/OASIS_1_mappings.json',
+    'OASIS_2': '/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/OASIS_2_mappings.json',
+    'PSMA-FDG-PET-CT-LESION':'/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/PSMA-FDG-PET-CT-LESION_mappings.json',
+    'PSMA-CT':'/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/PSMA-CT-Longitud_mappings.json',
+    'AbdomenAtlas':'/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/AbdomenAtlas_mappings.json',
+    'AbdomenCT1k':'/home/jachin/data/Github/OmniMorph/Dataloader/nifty_mappings/AbdomenCT1k_mappings.json',
+}
+query = {
+    'MSD': ['description'],
+    'TotalSegmentor': ['age','gender'],
+    'Kaggle_osic': ['Age','Sex','Smoke_Status','Weeks','FVC','Percent'],
+    'CancerImageArchive':['Series_Description', 'Study_Description', 'Manufacturer'],
+    'MnMs': ['Age','Sex','Height','Weight'],
+    'Brats2019': ['Age', 'Grade', 'Survival','ResectionStatus'],
+    'Brats2020': ['Age', 'Grade', 'Survival','ResectionStatus'],
+    'Brats2021': ['Age', 'Grade', 'Survival','ResectionStatus'],
+    'OASIS_1': ['Age', 'M/F','ASF','Educ','SES','MMSE','eTIV','CDR','nWBV'],
+    'OASIS_2': ['Age', 'Group','M/F','ASF','Educ','SES','MMSE','eTIV','CDR','nWBV'],
+    'PSMA-FDG-PET-CT-LESION':['Study Description', 'diagnosis','age','sex',"pet_radionuclide",'ct_contrast_agent'],
+    'PSMA-CT':[],
+    'AbdomenAtlas':[],
+    'AbdomenCT1k':[],
+}
+add_text = {
+    'MSD': {},
+    'TotalSegmentor': {},
+    'Kaggle_osic': {'description': 'pulmonary fibrosis progression'},
+    'CancerImageArchive': {},
+    'MnMs': {},
+    'Brats2019': {'description': 'could include brain tumor, glioma, glioblastoma, low grade glioma, high grade glioma'},
+    'Brats2020': {'description': 'could include brain tumor, glioma, glioblastoma, low grade glioma, high grade glioma'},
+    'Brats2021': {'description': 'could include brain tumor, glioma, glioblastoma, low grade glioma, high grade glioma'},
+    'OASIS_1': {},
+    'OASIS_2': {},
+    'PSMA-CT':{'description': 'melanoma patients'},
+    'PSMA-FDG-PET-CT-LESION':{'description': 'malignant melanoma, lymphoma, lung cancer, or healthy'},
+    'AbdomenAtlas':{},
+    'AbdomenCT1k':{},
+}
+
+
+# bert intialization
+model_name = '/home/jachin/data/Github/OmniMorph/External/Models/bert_large_uncased'
+reduce_method = 'mean'
+max_words_num = 32  # max number of words in the caption > 2
+# max_words_num = 64  # max number of words in the caption > 2
+
+embeder, tokenizer = get_frozen_embeder(model_name)
+def embed_str_filter(str_input, filter_words=['segmentation', 'registration']):
+    '''
+    Filter out specific words from the input string.
+    '''
+    for word in filter_words:
+        str_input = str_input.replace(word, '')
+    return str_input
+
+for dataset in mapping_files.keys():
+    jsn_path = mapping_files[dataset]
+    
+    with open(jsn_path, 'r') as f:
+        embd_json = json.load(f)
+        for key in embd_json.keys():
+            embd_json_temp = {}
+
+
+            embd_json_temp['Modality'] = embd_json[key]['Modality']
+            embd_json_temp['ROI'] = embd_json[key]['ROI']
+
+
+            query_key = query[dataset]
+
+            meta_data = embd_json[key]['Metadata']
+            for q in query_key:
+                if q in meta_data:
+                    embd_json_temp[q] = meta_data[q]
+                else:
+                    embd_json_temp[q] = 'N/A'
+            for q in add_text[dataset].keys():
+                if q in embd_json_temp:
+                    embd_json_temp[q] += ', ' + add_text[dataset][q]
+                else:
+                    embd_json_temp[q] = add_text[dataset][q]
+            emdb_str = str(embd_json_temp)[1:-1].lower()
+            embd_str = replace_text(emdb_str, get_synonyms_dict(None))
+            embd_str = embed_str_filter(embd_str)
+
+            print(f'embd_json_temp: {str(embd_json_temp)}')
+            print(f'embd_str: {embd_str}')
+            print(f'words_num: {len(embd_str.split())}')
+            assert(len(embd_str.split()) <= max_words_num), f'Too many words in the caption: {embd_str}'
+       
+            embd = str2emb(embd_str, max_words_num, embeder, tokenizer, reduce_method=reduce_method)
+            print(embd)
+            embd_json[key]['embd'] = embd.tolist()[0]
+            embd_json[key]['embd_key'] = embd_str
+        
+        # exit()
+
+        new_jsn_path = save_paths[dataset]
+        with open(new_jsn_path, 'w') as f:
+            json.dump(embd_json, f, indent=4)
+
+
+
+            
+        
+    
+

Dataloader/nifty_mappings/AbdomenAtlas_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:303c3fb7388e7b3b01cb6f494c3ac3f542da98487039e5b2415786ac4af58ba0
+size 179457573

Dataloader/nifty_mappings/AbdomenCT1k_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0abaaa1013fdafe3fae6d5544746a66d8b20892ceb3cf9141a125113984e8350
+size 37315918

Dataloader/nifty_mappings/Brats2019_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c5b80fc861484d36d8d6e0f97c404e2c321ee965cc1556a868205f5937d24fe
+size 12126490

Dataloader/nifty_mappings/Brats2020_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de345c6a66a4f33552aacbb961cd034ac488500ff5d48810579055f0543162dc
+size 17743015

Dataloader/nifty_mappings/Brats2021_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4990a7031d6ac91e1c33e6db046dddf234f67dd8edecd07691675945b9d00af5
+size 44722001

Dataloader/nifty_mappings/CIA_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98cbd21d3d5b7f5fb84091705fbbfcd0f8f26cb26ff4b34ffcf546cf1cedb48a
+size 32744567

Dataloader/nifty_mappings/MSD_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1ab13c61cd6829f088ee92bff4ce12a0f0e19fc9367682291fbd9717b149e83
+size 92620864

Dataloader/nifty_mappings/OASIS_1_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8784bff1bb5c9ba08fccc8ca9776f3f26c9b2993c1c446ef17d5ba1dd2bda490
+size 15609846

Dataloader/nifty_mappings/OASIS_2_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f88910a0846e056b0d4caacd6e6ebfebde52b537828756e217d9a6c6343177c
+size 13396017

Dataloader/nifty_mappings/PSMA-CT-Longitud_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3c8729df59b6e9771fa791c5fe1cd7636e83a3c17109613984cdce0d92eefdc
+size 11700732

Dataloader/nifty_mappings/PSMA-FDG-PET-CT-LESION_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:922363b739e1f14243731ea283ee730bc55724a27360d2f28f32b01b23ede5d9
+size 48425273

Dataloader/nifty_mappings/TotalSegmentorCT_MRI_mappings.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c36ba45053fea97244c259af0151ddb02e8281fce8c8f439cc88733bd71d668f
+size 67962146

Diffusion/__init__.py

@@ -0,0 +1,8 @@
+import Diffusion
+from . import diffuser
+from . import networks
+from . import losses
+
+import sys
+sys.path.append('./Diffusion')
+sys.path.append('./')

Diffusion/diffuser.py

@@ -0,0 +1,531 @@
+from torch import nn
+import torch
+import numpy as np
+from torch.nn.utils.stateless import functional_call
+
+import Diffusion.utils_diff as utils
+from Diffusion.networks import *
+# from networks import *
+
+import random
+
+EPS = 1e-8
+
+
+
+class DeformDDPM(nn.Module):
+    def __init__(
+        self, 
+        network, 
+        n_steps=50, 
+        beta_schedule_fn = None, 
+        device='cpu', 
+        image_chw=(1, 28, 28),
+        batch_size = 1,
+        img_pad_mode = "zeros",
+        ddf_pad_mode="border",
+        padding_mode="border",
+        v_scale = 0.008/256,
+        resample_mode=None,
+        ):
+        super(DeformDDPM, self).__init__()
+        self.rec_num=2
+        self.ndims=len(image_chw)-1
+        self.n_steps = n_steps
+        self.v_scale = v_scale
+        self.device = device
+        self.msk_noise_scale = torch.tensor(0)
+        
+        # print('================')
+        # print("device:",device)
+        # if device == 'cpu':
+        #     print("num_device: 1")
+        # else:
+        #     print("num_device:", torch.cuda.device_count())
+        # print('================')
+
+        self.num_device = torch.cuda.device_count()
+
+        self.batch_size = batch_size #//self.num_device
+        self.img_pad_mode = img_pad_mode 
+        self.ddf_pad_mode = ddf_pad_mode
+        self.padding_mode = padding_mode
+        self.resample_mode = resample_mode
+        self.image_chw = image_chw
+        self.network = network#.to(self.device)
+        self.ddf_stn_full = STN(
+                                    img_sz = self.image_chw[1],
+                                    ndims = self.ndims,
+                                    padding_mode = self.padding_mode,
+                                    device = self.device,
+                                )
+        self._DDF_Encoder_init()
+        self.copy_opt = nn.Identity()
+        return
+    
+    def get_stn(self):
+        return self.img_stn, self.ddf_stn_full
+
+    def _DDF_Encoder_init(self, ctl_ratio=4, ctl_sz=None, resample_mode=None):
+        if ctl_sz is None:
+            ctl_sz = self.image_chw[1] // ctl_ratio
+        self.ctl_sz=ctl_sz
+        self.img_sz=self.image_chw[1]
+        self.ddf_stn_rec=STN(img_sz=ctl_sz,ndims=self.ndims,device=self.device,padding_mode=self.ddf_pad_mode)
+        self.img_stn=STN(img_sz=self.img_sz,ndims=self.ndims,device=self.device,padding_mode=self.img_pad_mode,resample_mode=self.resample_mode)
+        self.msk_stn=STN(img_sz=self.img_sz,ndims=self.ndims,device=self.device,padding_mode=self.img_pad_mode,resample_mode='nearest')
+
+    def _get_ddf_scale(self,t,divide_num=1,max_ddf_num=200):   # 128
+        rec_num = 1
+        mul_num_ddf = torch.floor_divide(2*torch.pow(t,1.3), 3*divide_num).int()
+        mul_num_dvf = torch.floor_divide(torch.pow(t,0.6), divide_num).int()
+        # print("time_step:",t,"mul_num_ddf:",mul_num_ddf,"mul_num_dvf:",mul_num_dvf)
+        # mul_num_ddf = self._sample_random_uniform_multi_order(high=mul_num_ddf)
+        # mul_num_dvf = self._sample_random_uniform_multi_order(high=mul_num_dvf)
+        mul_num_ddf = torch.clamp(mul_num_ddf, min=1, max=max_ddf_num)
+        mul_num_dvf = torch.clamp(mul_num_dvf, min=0, max=max_ddf_num)
+        # print("time_step:",t,"mul_num_ddf:",mul_num_ddf,"mul_num_dvf:",mul_num_dvf)
+        return rec_num,mul_num_ddf,mul_num_dvf
+
+    # def _sample_random_uniform_multi_order(self, high=None, low=0, order_num=3):
+    #     # high: tensor of shape (...), low: int or tensor broadcastable to high
+    #     sample_num = torch.full_like(high, low) if not isinstance(low, torch.Tensor) else low.clone()
+    #     for _ in range(order_num):
+    #         # For each element, sample in [sample_num, high]
+    #         # torch.randint requires scalar low/high, so we use elementwise sampling
+    #         rand_shape = high.shape
+    #         # Clamp sample_num to be <= high
+    #         sample_num = torch.minimum(sample_num, high)
+    #         # Generate random numbers for each element
+    #         rand = torch.empty(rand_shape, dtype=high.dtype, device=high.device)
+    #         for idx in np.ndindex(rand_shape):
+    #             l = sample_num[idx].item()
+    #             h = high[idx].item()
+    #             if l >= h:
+    #                 rand[idx] = l
+    #             else:
+    #                 rand[idx] = torch.randint(l, h + 1, (1,), device=high.device)
+    #         sample_num = rand.to(high.dtype)
+    #     return sample_num
+    
+    def _get_random_ddf(self,img,t):
+        rec_num, mul_num_ddf, mul_num_dvf = self._get_ddf_scale(t=t)
+        ddf_forward,dvf_forward = self._random_ddf_generate(rec_num=rec_num, mul_num=[mul_num_ddf,mul_num_dvf])
+        warped_img = self.img_stn(img,ddf_forward)
+        return warped_img, dvf_forward,ddf_forward
+
+    def _multiscale_dvf_generate(self,v_scale,ctl_szs=[4,8,16,32,64], rand_v_scale=True):
+        dvf=0
+        if self.img_sz is None:
+            self.img_sz=max(ctl_szs)
+        if 1 in ctl_szs:
+            dvf_rot = utils.random_ddf(batch_size=self.batch_size, ndims=self.ndims, img_sz=[self.ctl_sz]*self.ndims, range_gauss=0, rot_range=np.pi/90)
+            dvf = dvf + dvf_rot
+        for ctl_sz in ctl_szs:
+            _v_scale = self._sample_random_uniform_multi_order(high=v_scale, low=1e-8, order_num=2) if rand_v_scale else v_scale
+            # temp>>
+            if ctl_sz <= 2:
+                _v_scale = _v_scale/2
+            # temp<<
+            dvf_comp = torch.randn([self.batch_size, self.ndims] + [ctl_sz]*self.ndims) * _v_scale
+            dvf_comp = F.interpolate(dvf_comp * self.ctl_sz / ctl_sz, [self.ctl_sz]*self.ndims, align_corners=False, mode='bilinear' if self.ndims == 2 else 'trilinear')
+            dvf=dvf+dvf_comp
+        return dvf
+
+    def _sample_random_uniform_multi_order(self, high=None, low=0., order_num=3):
+        sample_value = low
+        for _ in range(order_num):
+            sample_value = np.random.uniform(low=sample_value, high=high)
+        return sample_value
+
+    def _random_ddf_generate(self,rec_num=3,mul_num=[torch.tensor([5]),torch.tensor([5])],ddf0=None,keep_inverse=False,noise_ratio=0.08,select_num=4, flip_ratio=0.5): 
+        crop_rate=2
+        for _ in range(self.ndims+1):
+            mul_num=[torch.unsqueeze(n,-1) for n in mul_num]
+        # v_scale = v_scale *crop_rate
+        ctl_ddf_sz=[self.batch_size, self.ndims] + [self.ctl_sz] * self.ndims
+        if ddf0 is not None:
+            ddf=ddf0
+        else:
+            ddf = torch.zeros(ctl_ddf_sz) * 0
+        dddf = torch.zeros(ctl_ddf_sz) * 0
+        scale_num = min(8,int(math.log2(self.ctl_sz)))   # allow affine
+        # scale_num = min(5,int(math.log2(self.ctl_sz))-1)   # semi-allow affine
+        # scale_num = min(5,int(math.log2(self.ctl_sz))-2)   # avoid coupling between deformation and affine
+        ctl_szs_all = [self.ctl_sz // (2 ** i) for i in range(scale_num)]
+        
+        for i in range(rec_num):
+            # Randomly select 5 elements from ctl_szs (if there are at least 5)
+            if len(ctl_szs_all) > select_num:
+                ctl_szs = random.sample(ctl_szs_all, select_num)
+            dvf = self._multiscale_dvf_generate(self.v_scale, ctl_szs=ctl_szs).to(self.device)
+            # if True:
+            if noise_ratio==0:
+                dvf0=dvf
+            else:
+                dvf0=dvf+self.ddf_stn_rec(self._multiscale_dvf_generate(self.v_scale*noise_ratio,ctl_szs=ctl_szs, rand_v_scale=False).to(self.device),dvf)
+            # print([num.shape for num in mul_num])
+            for j in range(torch.max(mul_num[0]).item()):
+                flag = [(n>j).int().to(self.device) for n in mul_num]
+                ddf = dvf0*flag[0] + self.ddf_stn_rec(ddf, dvf0*flag[0])
+                dddf = dvf*flag[1] + self.ddf_stn_rec(dddf, dvf*flag[1])
+        
+        ddf = F.interpolate(ddf * self.img_sz/self.ctl_sz, self.img_sz*crop_rate, mode='bilinear' if self.ndims == 2 else 'trilinear')
+        # ddf = ddf[...,img_sz//2:img_sz*3//2,img_sz//2:img_sz*3//2]
+        if self.ndims==2:
+            ddf = ddf[..., self.img_sz // 2:self.img_sz * 3 // 2, self.img_sz // 2:self.img_sz * 3 // 2]
+        else:
+            ddf = ddf[..., self.img_sz // 2:self.img_sz * 3 // 2, self.img_sz // 2:self.img_sz * 3 // 2, self.img_sz // 2:self.img_sz * 3 // 2]
+        # if rec_num==1:
+        if True:
+            dddf = F.interpolate(dddf * self.img_sz/self.ctl_sz, self.img_sz*crop_rate, mode='bilinear' if self.ndims == 2 else 'trilinear')
+            # dddf = dddf[...,img_sz//2:img_sz*3//2,img_sz//2:img_sz*3//2]
+            if self.ndims == 2:
+                dddf = dddf[..., self.img_sz // 2:self.img_sz * 3 // 2, self.img_sz // 2:self.img_sz * 3 // 2]
+            else:
+                dddf = dddf[..., self.img_sz // 2:self.img_sz * 3 // 2, self.img_sz // 2:self.img_sz * 3 // 2, self.img_sz // 2:self.img_sz * 3 // 2]
+            return ddf,dddf
+        else:
+            return ddf
+
+    def create_noise_map(self, img, noise_type='gaussian', noise_ratio=0.2):
+        if noise_type == 'gaussian':
+            noise_map = torch.randn_like(img) * noise_ratio
+        elif noise_type == 'uniform':
+            noise_map = torch.rand_like(img) # 0-1
+        elif noise_type == 'binary':
+            noise_map = torch.bernoulli(torch.rand_like(img))
+        else:
+            noise_map = torch.zeros_like(img)
+        noise_map = noise_map.to(img.device)
+        return noise_map
+
+    def add_noise(self, img, noise_map=None, noise_ratio_range=[0.,1.]):
+        noise_ratio = np.random.uniform(noise_ratio_range[0], noise_ratio_range[1])
+        return img * (1-noise_ratio) + noise_map * noise_ratio, noise_ratio
+
+    def apply_noise(self, img, noise_map=None, apply_mask=None):
+        return img * apply_mask + noise_map * (1-apply_mask)
+
+    def downsample(self, img, down_ratio_range=[1./32,1]):
+        down_ratio = list(np.random.uniform(down_ratio_range[0], down_ratio_range[1],[self.ndims]))
+        # print(down_ratio)
+        down_img = F.interpolate(img, scale_factor=down_ratio, mode='bilinear' if self.ndims == 2 else 'trilinear')
+        # print(down_img)
+        # return F.interpolate(down_img, size=[self.image_chw[1]]*self.ndims, mode='bilinear' if self.ndims == 2 else 'trilinear', align_corners=False), np.prod(down_ratio)
+        return F.interpolate(down_img, size=[self.image_chw[1]]*self.ndims, mode='bilinear' if self.ndims == 2 else 'trilinear', align_corners=False), np.sqrt(np.prod(down_ratio)) # jzheng: cond weight based on entropy
+
+    def get_slice_mask(self, img, slice_num_range=[0,32]):
+        slice_num_range[1] = min(slice_num_range[1], self.image_chw[1])
+        mask = torch.zeros_like(img)
+        sample_ratio = 0
+        for i in range(self.ndims):
+            slice_num = random.randint(slice_num_range[0], slice_num_range[1])
+            slice_idx = random.sample(range(self.image_chw[1]), slice_num)
+            transpose_list = [0, 1, 1 + self.ndims] + list(range(2, 1 + self.ndims))
+            for idx in slice_idx:
+                mask[..., idx] = 1
+            mask = mask.permute(*transpose_list)
+            # sample_ratio += slice_num / self.image_chw[1] / self.ndims    
+            sample_ratio += np.sqrt(slice_num / self.image_chw[1]) / self.ndims     # jzheng: cond weight based on entropy
+
+        # print(mask)
+        # print("sample_ratio:", sample_ratio)
+        return mask, sample_ratio
+
+    def project(self, img):
+        proj_img = torch.zeros_like(img)
+        rand_bourn = np.random.randint(0, 2, size=[self.ndims])
+        proj_dim_num = np.sum(rand_bourn)
+        for i,pflag in zip(range(2, 2 + self.ndims), rand_bourn):
+            if pflag:
+                proj_img += torch.mean(img, dim=i, keepdim=True)
+                # print("projecting dim:", i)
+        return proj_img/(proj_dim_num+EPS), proj_dim_num
+
+    def proc_cond_img(self, img, proc_type=None):
+        # Remove torch.no_grad() since most operations are not differentiable anyway
+        proc_img = img.clone().detach()
+        if proc_type is None:
+            # Heavily bias towards 'uncon' for efficiency
+            proc_type = random.choices(
+                # ['adding', 'independ', 'downsample', 'slice', 'project', 'none', 'uncon'],
+                # weights=[1, 1, 1, 1, 1, 1, 3], k=1
+                ['adding', 'independ', 'downsample', 'slice', 'none', 'uncon'],
+                weights=[1, 1, 1, 1, 1, 3], k=1
+            )[0]
+        mask = torch.tensor(1, device=img.device)
+        cond_ratio = torch.tensor(1., device=img.device)
+        self.msk_noise_scale = torch.tensor(0, device=img.device)
+        noise_type = random.choice(['gaussian', 'uniform', 'none'])
+        # Precompute noise_map only if needed
+        noise_map = None
+        if proc_type not in ['none', None, '']:
+            if proc_type == 'uncon':
+                noise_map = self.create_noise_map(img, noise_type=noise_type)
+                proc_img = noise_map
+                mask = torch.tensor(0, device=img.device)
+                cond_ratio = torch.tensor(0, device=img.device)
+                return proc_img, mask, cond_ratio
+            if proc_type in ['adding', 'independ', 'slice']:
+                # self.msk_noise_scale = 0
+                noise_map = self.create_noise_map(img, noise_type=noise_type)
+            if proc_type == 'adding':
+                proc_img, noise_ratio = self.add_noise(proc_img, noise_map=noise_map, noise_ratio_range=[0., 1.])
+                cond_ratio = torch.tensor(1 - noise_ratio, device=img.device)
+            elif proc_type == 'independ':
+                mask = self.create_noise_map(img, noise_type='binary')
+                if self.msk_noise_scale == 0:
+                    proc_img = img * mask
+                else:
+                    proc_img = self.apply_noise(proc_img, noise_map=noise_map*self.msk_noise_scale, apply_mask=mask)
+                with torch.no_grad():
+                    cond_ratio = mask.float().mean()
+            elif proc_type == 'downsample':
+                # proc_img, down_ratio = self.downsample(proc_img, down_ratio_range=[1./32, 1])
+                proc_img, down_ratio = self.downsample(proc_img, down_ratio_range=[1./64, 1])
+                cond_ratio = torch.tensor(down_ratio, device=img.device)
+            elif proc_type == 'slice':
+                slice_num_max = random.randint(1, 64)
+                slice_num_max = random.randint(1, slice_num_max)
+                mask, sample_ratio = self.get_slice_mask(img, slice_num_range=[0, slice_num_max])
+                if self.msk_noise_scale == 0:
+                    proc_img = img * mask
+                else:
+                    proc_img = self.apply_noise(proc_img, noise_map=noise_map*self.msk_noise_scale, apply_mask=mask)
+                cond_ratio = torch.tensor(sample_ratio, device=img.device)
+            elif proc_type == 'project':
+                proc_img, proj_num = self.project(proc_img)
+                cond_ratio = torch.tensor(proj_num / (128 * self.ndims), device=img.device)
+                # cond_ratio = torch.tensor(proj_num / (32 * self.ndims), device=img.device)  # jzheng: cond weight based on entropy
+        return proc_img, mask, cond_ratio
+    
+    def diffuse(self, x_0, t):
+        t=torch.tensor(t)
+        # img_t, dvf_forward, ddf_forward, ddf_stn, img_stn = self.ddf_enc(img= x_0, t=t)
+        # return img_t, dvf_forward,ddf_forward,ddf_stn,img_stn
+        return self._get_random_ddf(img = x_0, t = t)
+    
+    
+    def recover(self, x, y, t,rec_num=2, text=None):
+        if isinstance(t, list):
+            t=[torch.tensor(t0) for t0 in t]
+            t=[t0.to(x.device) for t0 in t]
+        else:
+            t=torch.tensor(t)
+            t.to(x.device)
+        if rec_num is None:
+            rec_num = self.rec_num
+        return self.network(x=x, y=y, t=t, rec_num=rec_num, text=text)
+
+    def recover_frozen_params_but_grad_input(self, x, y, t,rec_num=2, text=None):
+        """
+        use detach to recover:
+        - but not include no_grad
+        """
+        if isinstance(t, list):
+            t = [torch.tensor(t0, device=x.device) for t0 in t]
+        else:
+            t = torch.tensor(t, device=x.device)
+        
+        if rec_num is None:
+            rec_num = self.rec_num
+
+        # params = {k: v.detach() for k, v in self.network.named_parameters()}
+        # buffers = dict(self.network.named_buffers())  # BN running stats etc. buffer
+        # # functional_call require position args，here kwargs doesnot work, so：
+        # def _forward(module, kw):
+        #     return module(**kw)
+        # # functional_call(module, ...) can only pass args/kwargs to module.forward
+        # # PyTorch 2.x support functional_call(module, (params, buffers), args, kwargs)
+        # return functional_call(
+        #     self.network,
+        #     (params, buffers),
+        #     args=(),
+        #     kwargs=dict(x=x, y=y, t=t, rec_num=rec_num, text=text),
+        # )
+
+        # 1) param detached
+        params = {k: v.detach() for k, v in self.network.named_parameters()}
+        # 2) buffers keeps unchanged
+        buffers = dict(self.network.named_buffers())
+
+        # 3) old version of PyTorch doesnot support passing params and buffers together
+        params_and_buffers = {}
+        params_and_buffers.update(params)
+        params_and_buffers.update(buffers)
+        return functional_call(
+            self.network,
+            params_and_buffers,
+            (),
+            kwargs=dict(x=x, y=y, t=t, rec_num=rec_num, text=text),
+        )
+        
+
+    def _single_step(self, x0, t, rec_num=2, proc_type=None,mask=None, cond_imgs=None, text=None):
+        if mask is None:
+            mask = 1
+        # org_imgs=self.copy_opt(x0)
+        if cond_imgs is None:
+            cond_imgs, mask_tgt, cond_ratio = self.proc_cond_img(x0,proc_type=proc_type)
+        noisy_imgs, dvf_I,_ = self.diffuse(x0, t)
+        if isinstance(self.network,DefRec_MutAttnNet):
+            t = [t] * 1
+        return self.recover(x=noisy_imgs*mask, y=cond_imgs, t=t, rec_num=rec_num, text=text), dvf_I
+
+    def forward(self, img_org, cond_imgs=None, proc_type=None, T=None, **kwargs):
+        if T is not None:
+            return self.diff_recover(img_org=img_org, T=T, proc_type=proc_type, cond_imgs=cond_imgs, **kwargs)
+        else:
+            return self._single_step(x0=img_org, proc_type=proc_type, cond_imgs=cond_imgs, **kwargs) 
+            # if mask is None:
+            #     mask = 1
+            # cond_imgs = self.proc_cond_img(x0, proc_type=proc_type, **kwargs)
+            # noisy_imgs, dvf_I, _ = self.diffuse(x0, t)
+            # if isinstance(self.network, DefRec_MutAttnNet):
+            #     t = [t] * 1
+            # return self.recover(x=noisy_imgs * mask, y=cond_imgs, t=t, rec_num=rec_num), dvf_I
+        
+    def diff_recover(self,
+                     img_org,
+                     msk_org=None,
+                     T=[None,None],
+                     ddf_rand=None,
+                     v_scale = None,
+                     t_save=None,
+                     cond_imgs=None,
+                     proc_type=None,
+                     text=None,
+                     ):
+        if cond_imgs is None:
+            cond_imgs = img_org.clone().detach()
+        # if proc_type is not None:
+        cond_imgs,mask_tgt,cond_ratio=self.proc_cond_img(cond_imgs, proc_type=proc_type)
+        if ddf_rand is None:
+            if v_scale is not None:
+                self.v_scale=v_scale
+                self._DDF_Encoder_init()
+            if T[0] is None or T[0] == 0:
+                img_diff = img_org.clone().detach()
+                ddf_rand = torch.zeros_like(img_diff)
+            else:
+                img_diff, _, ddf_rand = self._get_random_ddf(img= img_org, t=torch.tensor(np.array([T[0]])).to(self.device))
+        else:
+            img_diff = self.img_stn(img_org.clone().detach(), ddf_rand)
+        ddf_comp = ddf_rand.clone().detach()
+        img_rec = img_diff.clone().detach()
+        if msk_org is not None:
+            msk_diff = self.msk_stn(msk_org.clone().detach(), ddf_rand)
+        else:
+            msk_diff = None
+        msk_rec = msk_diff.clone().detach() if msk_org is not None else None
+        img_save=[]
+        msk_save=[]
+        
+        if isinstance(self.network,DefRec_MutAttnNet):
+            # Denosing image via list of t
+            t_list = list(range(T[1]-1, -1, -1))
+            pre_dvf_I = self.recover(x=img_rec, y=cond_imgs, t=t_list,rec_num=None, text=text)
+            ddf_comp = self.ddf_stn_full(ddf_comp, pre_dvf_I) + pre_dvf_I
+            img_rec = self.img_stn(img_org.clone().detach(), ddf_comp)
+            if msk_org is not None:
+                msk_rec = self.msk_stn(msk_org.clone().detach(), ddf_comp)
+        else:
+            # Denosing image
+            if isinstance(T[-1], int):
+                time_steps = range(T[-1] - 1, -1, -1)
+                trainable_iterations =[]
+            else:
+                time_steps = T[-1]
+                
+                # # Randomly select k iterations to make their parameters trainable
+                # win_len = 2  # Number of iterations to make trainable
+                # if len(time_steps) <= win_len:
+                #     win_start = 0
+                # else:
+                #     win_start = random.randint(len(time_steps)//2, len(time_steps) - win_len)
+                # win_end = win_start + win_len - 1
+                
+                k=2
+                # trainable_iterations = time_steps[win_start: win_start + win_len]
+                # trainable_iterations = random.sample(time_steps, k)
+                trainable_iterations = time_steps[-1:-k-1:-1]
+                # print(time_steps)
+                # print("trainable_iterations:", trainable_iterations)
+            for i in time_steps:
+                t = torch.tensor(np.array([i])).to(self.device)
+                
+                if i in trainable_iterations:
+                    # Make parameters trainable for this iteration
+                    pre_dvf_I = self.recover(x=img_rec, y=cond_imgs, t=t, rec_num=None, text=text)
+                else:
+                    # Freeze parameters for this iteration using torch.no_grad()
+                    with torch.no_grad():
+                        pre_dvf_I = self.recover(x=img_rec, y=cond_imgs, t=t, rec_num=None, text=text)
+            # for idx, i in enumerate(time_steps):
+            #     t = torch.tensor(np.array([i])).to(self.device)
+            #     if idx < win_start:
+            #         # just no_grad
+            #         with torch.no_grad():
+            #             pre_dvf_I = self.recover(x=img_rec, y=cond_imgs, t=t, rec_num=None, text=text)
+            #     elif win_start <= idx <= win_end:
+            #         # normal update
+            #         pre_dvf_I = self.recover(x=img_rec, y=cond_imgs, t=t, rec_num=None, text=text)
+            #     else:
+            #         # freeze params but keep grad for input
+            #         pre_dvf_I = self.recover_frozen_params_but_grad_input(
+            #             x=img_rec, y=cond_imgs, t=t, rec_num=None, text=text
+            #         )
+
+                ddf_comp = self.ddf_stn_full(ddf_comp, pre_dvf_I) + pre_dvf_I
+                # Apply to image
+                img_rec = self.img_stn(img_org.clone().detach(), ddf_comp)
+                if msk_org is not None:
+                    msk_rec = self.msk_stn(msk_org.clone().detach(), ddf_comp)
+                if t_save is not None:
+                    if i in t_save:
+                        img_save.append(img_rec)
+                        if msk_org is not None:
+                            msk_save.append(msk_rec)
+
+            # for i in time_steps:
+            #     t = torch.tensor(np.array([i])).to(self.device)
+            #     pre_dvf_I = self.recover(x=img_rec, y=cond_imgs, t=t,rec_num=None)
+            #     ddf_comp = self.ddf_stn_full(ddf_comp, pre_dvf_I) + pre_dvf_I
+            #     # apply to image
+            #     img_rec = self.img_stn(img_org.clone().detach(), ddf_comp)
+            #     if msk_org is not None:
+            #         msk_rec = self.img_stn(msk_org.clone().detach(), ddf_comp)
+            #     if t_save is not None:
+            #         if i in t_save:
+            #             img_save.append(img_rec)
+            #             if msk_org is not None:
+            #                 msk_save.append(msk_rec)
+        # print(torch.max(torch.abs(ddf_comp)))
+        # print(torch.max(torch.abs(ddf_rand)))
+
+        return [ddf_comp,ddf_rand],[img_rec,img_diff,img_save],[msk_rec,msk_diff,msk_save]
+    
+if __name__ == "__main__":
+    H, W = 8, 8
+    deformddpm = DeformDDPM(network=get_net(name="recmutattnnet")(n_steps=80, ndims=2, num_input_chn=1),image_chw=(1, H, W),device='cpu')
+    # img = torch.zeros([1, 1, H, W])
+    img = torch.randn([1, 1, H, W])
+    t = 1
+    rec_num = 2
+    # proc_type = 'adding'
+    # proc_type = 'independ'
+    # proc_type = 'downsample'
+    proc_type = 'slice'
+    # proc_type = 'project'
+    # proc_type = 'none'
+    print(img)
+    cond_imgs, mask_tgt = deformddpm.proc_cond_img(img, proc_type=proc_type)
+    print(cond_imgs)
+    # img_rec, dvf_I = deformddpm.forward(img, t, rec_num=rec_num, proc_type=proc_type)
+    # print(img_rec.shape, dvf_I.shape)
+    
+    # proc_type = 'adding'
+    # ddf_comp, ddf_rand = deformddpm.diff_recover(img, T=[1,1], proc_type=proc_type)
+
+    

Diffusion/losses.py

@@ -0,0 +1,534 @@
+"""
+losses for DRDM
+"""
+
+import numpy as np
+import sys
+import torch
+import torch.nn.functional as F
+
+
+EPS=1e-7
+
+# eps_scale = 10e-5
+# eps_scale = 10e-4
+# eps_scale = 1e-4
+eps_scale = 1e-5
+
+
+class LMSE(torch.nn.Module):
+    """
+    Labeled Mean Square Error (LMSE)
+    """
+
+    def __init__(self, eps=1e-7, relate_eps=5e-1, win=None, smooth=False):
+        super(LMSE, self).__init__()
+        self.eps = eps
+        self.relate_eps = relate_eps
+        self.ndims = 3
+        self.smooth = smooth
+        self.win = win
+        # Set window size
+        if self.win is None:
+            self.win = [5] * self.ndims
+        if smooth:
+            self.kernels = self._build_kernel(std=0.0)
+
+    def _build_kernel(self, std=0.0):
+        if std == 0.0:
+            return torch.ones([1, 1, *self.win])
+        else:
+            tail = int(np.ceil(std)) * 3
+            k = torch.exp(-0.5 * torch.arange(-tail, tail + 1, dtype=torch.float32) ** 2 / std ** 2)
+            kernel = k / torch.sum(k)
+            kernel = kernel.view(-1, 1, 1) * kernel.view(1, -1, 1) * kernel.view(1, 1, -1)
+            # print(kernel.item)
+            return kernel.unsqueeze(0).unsqueeze(0)
+        
+    def forward(self, I, J, label=None):
+        """
+        Computes the labeled mean squared error between I and J (ref).
+        If label is provided, computes the MSE only over the labeled regions.
+        """
+        padding = [(w-1) // 2 for w in self.win]
+        if self.smooth:
+            I = torch.nn.functional.conv3d(I, self.kernels, stride=1, padding=padding)
+            J = torch.nn.functional.conv3d(J, self.kernels, stride=1, padding=padding)
+        mse = (I - J) ** 2
+        if self.relate_eps is not None:
+            mse = mse/((J**2) + self.relate_eps)
+        if label is not None:
+            label = label.float()
+            mse = mse * label
+            mse_sum = torch.sum(mse, dim=(2, 3, 4))
+            label_sum = torch.sum(label, dim=(2, 3, 4)) + self.eps
+            loss = torch.mean(mse_sum / label_sum)
+        else:
+            loss = torch.mean(mse)
+        return loss
+    
+class LNCC(torch.nn.Module):
+    """
+    Local (over window) normalized cross-correlation (LNCC)
+    """
+
+    def __init__(self, win=None, num_ch=1, eps=1e-6, central=True, smooth=True):
+        super(LNCC, self).__init__()
+        self.scale = 2e0
+        self.win = win
+        self.eps = eps
+        self.central = central
+        self.ndims = 3
+        self.strides = [1] * (self.ndims + 2)
+        self.smooth = smooth
+
+        # Set window size
+        if self.win is None:
+            self.win = [9] * self.ndims
+        self.padding = [(w-1) // 2 for w in self.win]
+
+        if smooth:
+            self.kernels = self._build_kernel(std=0.45)
+        self.sum_filt = self._build_kernel(std=0.0)
+
+    def _build_kernel(self, std=0.0):
+        if std == 0.0:
+            return torch.ones([1, 1, *self.win])/np.prod(self.win)
+        else:
+            self.tail = int(np.ceil(std)) * 2
+            k = torch.exp(-0.5 * (torch.arange(-self.tail, self.tail + 1, dtype=torch.float32) ** 2) / std ** 2)
+            kernel = k / torch.sum(k)
+            # print(kernel)
+            kernel = kernel.view(-1, 1, 1) * kernel.view(1, -1, 1) * kernel.view(1, 1, -1)
+            # kernel = kernel * np.prod(self.win)
+            # print('Gaussian kernel created with std:', std)
+            # print('Kernel sum:', torch.sum(kernel))
+            
+            return kernel.unsqueeze(0).unsqueeze(0)
+
+    def lncc(self, I, J, label=None):
+        self.sum_filt = self.sum_filt.to(I.device)
+
+        if self.smooth:
+            self.kernels = self.kernels.to(I.device)
+            I = torch.nn.functional.conv3d(I, self.kernels, stride=1, padding=self.tail)
+            J = torch.nn.functional.conv3d(J, self.kernels, stride=1, padding=self.tail)
+
+        # if self.central:
+        #     I = I - torch.mean(I, dim=(2, 3, 4), keepdim=True)
+        #     J = J - torch.mean(J, dim=(2, 3, 4), keepdim=True)
+        # Compute CC squares
+        I2 = I * I
+        J2 = J * J
+        IJ = I * J
+
+        if self.central:
+            # Compute local sums via convolution
+            I_sum = torch.nn.functional.conv3d(I, self.sum_filt, stride=1, padding=self.padding)
+            J_sum = torch.nn.functional.conv3d(J, self.sum_filt, stride=1, padding=self.padding)
+            I2_sum = torch.nn.functional.conv3d(I2, self.sum_filt, stride=1, padding=self.padding)
+            J2_sum = torch.nn.functional.conv3d(J2, self.sum_filt, stride=1, padding=self.padding)
+            IJ_sum = torch.nn.functional.conv3d(IJ, self.sum_filt, stride=1, padding=self.padding)
+
+            # Compute cross-correlation
+            win_size = np.prod(self.win)
+            # print('Window size:', win_size)
+            # u_I = I_sum / win_size
+            # u_J = J_sum / win_size
+            # cross = IJ_sum - ((I_sum * J_sum) / win_size)
+            # I_var = I2_sum - ((I_sum * I_sum) / win_size)
+            # J_var = J2_sum - ((J_sum * J_sum) / win_size)
+            cross = IJ_sum - (I_sum * J_sum)
+            I_var = I2_sum - (I_sum * I_sum)
+            J_var = J2_sum - (J_sum * J_sum)
+        else:
+        # if 1:
+            # Compute local sums via convolution
+            I2_sum = torch.nn.functional.conv3d(I2, self.sum_filt, stride=1, padding=self.padding)
+            J2_sum = torch.nn.functional.conv3d(J2, self.sum_filt, stride=1, padding=self.padding)
+            IJ_sum = torch.nn.functional.conv3d(IJ, self.sum_filt, stride=1, padding=self.padding)
+
+            cross = IJ_sum
+            I_var = I2_sum
+            J_var = J2_sum
+
+        # cc = (cross * cross) / (I_var * J_var + self.eps)
+        cc = (cross * cross) / (I_var + self.eps) / (J_var + self.eps)
+        if label is not None:
+            label = label.float()
+            cc = torch.sum(cc * label, dim=(2, 3, 4)) / (torch.sum(label, dim=(2, 3, 4)) + self.eps)
+
+        return torch.mean(cc)
+
+    def forward(self, I, J, label=None):
+        return -self.lncc(I*self.scale, J*self.scale, label=label)
+
+
+
+class NCC(torch.nn.Module):
+    # def __init__(self, eps_scale=10e-7,img_sz=256):
+    def __init__(self, eps_scale=10e-5,img_sz=256):
+        super(NCC, self).__init__()
+        self.eps_scale=eps_scale#*img_sz/256
+        # self.scale=10e4
+        self.scale=1e2
+
+    def forward(self,pred,inv_lab=None,ddf_stn=None,mask=None):
+        if ddf_stn is None:
+            trm_pred=pred
+        else:
+            trm_pred=-ddf_stn(pred, inv_lab)
+        trm_pred = self.scale * trm_pred
+        inv_lab = self.scale * inv_lab
+        if mask is None:
+            loss_gen = torch.mean(torch.sum(trm_pred*inv_lab,dim=1)/(torch.sqrt(torch.sum(torch.square(trm_pred),dim=1)*torch.sum(torch.square(inv_lab),dim=1)+self.eps_scale)))
+        else:
+            batch_size = inv_lab.shape[0]
+            loss_gen = torch.sum(torch.sum(trm_pred*inv_lab,dim=1)*mask/(torch.sqrt(torch.sum(torch.square(trm_pred),dim=1)*torch.sum(torch.square(inv_lab),dim=1)+self.eps_scale)))/torch.sum(mask)/batch_size
+        return loss_gen
+
+class MRSE(torch.nn.Module):
+    def __init__(self, eps_scale=eps_scale,img_sz=256):
+        super(MRSE, self).__init__()
+        self.eps_scale=eps_scale#*img_sz/256
+        self.scale = 10e1
+
+    def forward(self,pred,inv_lab=None,ddf_stn=None,mask=None):
+        if ddf_stn is None:
+            trm_pred=pred
+        else:
+            trm_pred=-ddf_stn(pred, inv_lab)
+        trm_pred = self.scale * trm_pred
+        inv_lab = self.scale * inv_lab
+        if mask is None:
+            loss_gen = torch.mean(
+                torch.sum(torch.square(trm_pred + inv_lab), dim=1)
+                / (torch.sum(torch.square(inv_lab), dim=1) + self.eps_scale)
+            )
+        else:
+            batch_size = inv_lab.shape[0]
+            loss_gen = torch.sum(
+                torch.sum(torch.square(trm_pred + inv_lab), dim=1) * mask
+                / (torch.sum(torch.square(inv_lab), dim=1) + self.eps_scale)
+            )/torch.sum(mask)/batch_size
+        return loss_gen/1
+
+class RMSE(torch.nn.Module):
+    def __init__(self, eps_scale=eps_scale,img_sz=256,ndims=2):
+        super(RMSE, self).__init__()
+        self.eps_scale=eps_scale#*img_sz/256
+        self.ndims=ndims
+
+    def forward(self,pred,inv_lab=None,ddf_stn=None):
+        if ddf_stn is None:
+            trm_pred=pred
+        else:
+            trm_pred=-ddf_stn(pred, inv_lab)
+        loss_gen = torch.mean(torch.mean(torch.sum(torch.square(trm_pred - inv_lab), dim=1),
+                              dim=list(range(1, 1 + self.ndims))) / (
+                               torch.mean(torch.sum(torch.square(inv_lab), dim=1), dim=list(range(1, 1 + self.ndims))) + self.eps_scale))
+        return loss_gen
+# loss_gen = torch.mean(torch.mean(torch.sum(torch.square(ddf_stn(pre_dvf_I, dvf_I) + dvf_I), dim=1),dim=list(range(1,1+ndims))) / (torch.mean(torch.sum(torch.square(dvf_I), dim=1),dim=list(range(1,1+ndims))) + EPS))
+
+
+class Grad(torch.nn.Module):
+    """
+    N-D gradient loss
+    """
+
+    def __init__(self, penalty=['l1'],ndims=2, eps=1e-8, outrange_weight=1e4,outrange_thresh=0.5, detj_weight=2, apear_scale=4, dist=1, sign=1,waive_thresh=10**-5):
+        super(Grad, self).__init__()
+        self.penalty = penalty
+        self.eps = eps
+        self.outrange_weight = outrange_weight
+        self.detj_weight=detj_weight
+        self.apear_scale = apear_scale
+        self.ndims=ndims
+        self.max_sz = torch.reshape(torch.tensor([outrange_thresh]*ndims, dtype=torch.float32) , [1]+[ndims]+[1]*(ndims))
+        self.act = torch.nn.ReLU(inplace=False)
+        self.dist=dist
+        self.sign=sign
+        self.waive_thresh=waive_thresh
+
+    def _diffs(self, y,dist=None):
+        if dist is None:
+            dist=self.dist
+        # vol_shape = y.size()[2:]
+        # vol_shape = y.get_shape().as_list()[1:-1]
+        # ndims = len(vol_shape)
+
+        df = [None] * self.ndims
+        for i in range(self.ndims):
+            d = i + 2
+            # permute dimensions to put the ith dimension first
+            r = [d, *range(d), *range(d + 1, self.ndims + 2)]
+            yp = y.permute(r)
+            dfi = (yp[dist:, ...] - yp[:-dist, ...])/float(dist)
+
+            # permute back
+            # note: this might not be necessary for this loss specifically,
+            # since the results are just summed over anyway.
+            r = [*range(1, d + 1), 0, *range(d + 1, self.ndims + 2)]
+            df[i] = dfi.permute(r)
+        return df
+
+    def _eq_diffs(self, y,dist=None):
+        if dist is None:
+            dist=self.dist
+        # vol_shape = y.get_shape().as_list()[1:-1]
+        vol_shape = y.size()[2:]
+        ndims = len(vol_shape)
+        pad = [0, 0] * (ndims + 1) +[dist, 0]
+        pad1 = [0, 0] * (ndims + 1) +[0, dist]
+        # df = [None, None] * ndims
+        df = [None] * ndims
+        for i in range(ndims):
+            d = i + 2
+            r=[d, *range(d), *range(d + 1, ndims + 2)]
+            ri=[*range(1, d + 1), 0, *range(d + 1, ndims + 2)]
+            yt = y.permute(r)
+            dy=(yt[dist:, ...] - yt[:-dist, ...])/float(dist)
+            df[i] = (F.pad(dy, pad,mode='constant',value=0)).permute(ri)
+            # df[2*i] = (F.pad(dy, pad,mode='constant',value=0)).permute(ri)
+            # df[2*i+1] = (F.pad(dy, pad1, mode='constant', value=0)).permute(ri)
+            y.permute(ri)
+        return df
+
+    def _weighted_diffs_error(self, y,dist=None,w=None,expect=None,mean_dim=None):
+        if dist is None:
+            dist=self.dist
+        vol_shape = y.size()[2:]
+        ndims = len(vol_shape)
+        df = [None] * ndims
+
+        for i in range(ndims):
+            d = i + 2
+            r=[d, *range(d), *range(d + 1, ndims + 2)]
+            ri=[*range(1, d + 1), 0, *range(d + 1, ndims + 2)]
+            yt = y.permute(r)
+            wt = w.permute(r)
+            dy=(torch.abs(yt[dist:, ...] - yt[:-dist, ...])-expect.permute(r))*(wt[dist:, ...]*wt[:-dist, ...])
+            df[i] = torch.mean((dy).permute(ri),dim=mean_dim,keepdim=True)
+            y.permute(ri)
+            w.permute(ri)
+        return df
+
+    def _outl_dist(self, y,range_thresh=0.2):
+        self.device = y.device
+        vol_shape = y.size()[2:]
+        self.max_sz=self.max_sz.to(self.device)
+        act=torch.nn.ReLU(inplace=True)
+        loss=0.
+        for i in range(self.ndims):
+            d = i + 2
+            # permute dimensions to put the ith dimension first
+            r = [d, *range(d), *range(d + 1, self.ndims + 2)]
+            ri = [*range(1, d + 1), 0, *range(d + 1, self.ndims + 2)]
+            yt = y.permute(r)
+            loss += torch.mean(torch.square(act(-range_thresh-yt[0,:,i, ...])))+torch.mean(torch.square(act(yt[-1,:,i, ...]-range_thresh)))
+            # loss += torch.mean(torch.square(act(-range_thresh-yt[0,:,i, ...])+act(yt[-1,:,i, ...]-range_thresh)))
+            y.permute(ri)
+        return loss/self.ndims
+
+    def _center_dist(self, y):
+        self.device = y.device
+        vol_shape = y.size()[2:]
+        self.max_sz=self.max_sz.to(self.device)
+        select_loc = [s // 2 for s in vol_shape]
+        if self.ndims==3:
+            # return torch.mean(self.act(torch.abs(y[:,:, select_loc[0], select_loc[1], select_loc[2]]) - self.max_sz))
+            return torch.mean(torch.square(self.act(torch.abs(y[:, :, select_loc[0], select_loc[1], select_loc[2]]) - self.max_sz)))
+        elif self.ndims == 2:
+            # return torch.mean(self.act(torch.abs(y[:, :, select_loc[0], select_loc[1]]) - self.max_sz))
+            return torch.mean(torch.square(self.act(torch.abs(y[:, :, select_loc[0], select_loc[1]]) - self.max_sz)))
+
+        
+    # def _eval_detJ(self, disp=None, weight=None):
+    #     weight = 1
+    #     if self.ndims==3:
+    #         detj = (disp[0][:, 0, ...] * disp[1][:, 1, ...] * disp[2][:, 2, ...]) + (
+    #                 disp[0][:, 1, ...] * disp[1][:, 2, ...] * disp[2][:, 0, ...]) + (
+    #                        disp[0][:, 2, ...] * disp[1][:, 0, ...] * disp[2][:, 1, ...]) - (
+    #                        disp[0][:, 2, ...] * disp[1][:, 1, ...] * disp[2][:, 0, ...]) - (
+    #                        disp[0][:, 0, ...] * disp[1][:, 2, ...] * disp[2][:, 1, ...]) - (
+    #                        disp[0][:, 1, ...] * disp[1][:, 0, ...] * disp[2][:, 2, ...])
+    #     elif self.ndims==2:
+    #         detj = (disp[0][:, 0, ...] * disp[1][:, 1, ...]) - (disp[0][:, 1, ...] * disp[1][:, 0, ...])
+
+    #     return detj * weight
+    
+    def _eval_detJ(self, disp, add_identity=True, spacing=1.0):
+        """
+        disp: list length ndims
+            disp[i] is derivative wrt spatial dim i (forward diff),
+            tensor shape [B, C=ndims, ...]
+        add_identity: True if y_pred is displacement u and phi=x+u
+        spacing: voxel spacing (or 1.0). If you care about physical units,
+                divide derivatives by spacing (and dist). Sign won't change.
+        """
+        # Optional scaling (won't affect sign as long as spacing>0)
+        if spacing != 1.0:
+            disp = [d / spacing for d in disp]
+
+        if self.ndims == 2:
+            dux_dx = disp[0][:, 0, ...]
+            duy_dx = disp[0][:, 1, ...]
+            dux_dy = disp[1][:, 0, ...]
+            duy_dy = disp[1][:, 1, ...]
+
+            if add_identity:
+                j11 = 1.0 + dux_dx
+                j22 = 1.0 + duy_dy
+            else:
+                j11 = dux_dx
+                j22 = duy_dy
+
+            detj = j11 * j22 - dux_dy * duy_dx
+            return detj
+
+        elif self.ndims == 3:
+            dux_dx = disp[0][:, 0, ...]
+            duy_dx = disp[0][:, 1, ...]
+            duz_dx = disp[0][:, 2, ...]
+
+            dux_dy = disp[1][:, 0, ...]
+            duy_dy = disp[1][:, 1, ...]
+            duz_dy = disp[1][:, 2, ...]
+
+            dux_dz = disp[2][:, 0, ...]
+            duy_dz = disp[2][:, 1, ...]
+            duz_dz = disp[2][:, 2, ...]
+
+            if add_identity:
+                j11 = 1.0 + dux_dx 
+                j22 = 1.0 + duy_dy 
+                j33 = 1.0 + duz_dz
+            else:
+                j11 = dux_dx       
+                j22 = duy_dy
+                j33 = duz_dz
+
+            j12 = dux_dy; j13 = dux_dz
+            j21 = duy_dx; j23 = duy_dz
+            j31 = duz_dx; j32 = duz_dy
+
+            detj = (
+                j11 * (j22 * j33 - j23 * j32)
+                - j12 * (j21 * j33 - j23 * j31)
+                + j13 * (j21 * j32 - j22 * j31)
+            )
+            return detj
+
+        else:
+            raise ValueError(f"Unsupported ndims={self.ndims}")
+
+        
+    def forward(self,  y_pred=None,x_in=None, img=None, msk=None):
+        reg_loss = 0
+        act=torch.nn.ReLU(inplace=True)
+        
+        dg = 1
+        if img is not None:
+            dg = torch.exp(-self.apear_scale * sum([torch.sum(g * g, dim=1, keepdim=True) for g in self._eq_diffs(img)]) / torch.sum(torch.square(0.2 + img), dim=1, keepdim=True))   
+        if msk is not None:
+            dg = dg * msk
+
+        if 'l1' in self.penalty:
+            df = [torch.mean(dg*F.relu(torch.abs(f) - self.waive_thresh,inplace=True)) for f in self._eq_diffs(y_pred)]
+            reg_loss += sum(df) / len(df)
+            
+        if 'l2' in self.penalty:
+            df = [torch.mean(dg*F.relu(f * f - self.waive_thresh**2,inplace=True)) for f in self._eq_diffs(y_pred)]
+            reg_loss += torch.sqrt(sum(df) / len(df))
+
+        if 'negdetj' in self.penalty:
+            df = self.detj_weight*torch.mean(act(-self._eval_detJ(self._eq_diffs(y_pred,dist=1))))  # , dg[...,0])
+            reg_loss += 0.5*df
+        if 'range' in self.penalty:
+            reg_loss += self.outrange_weight * (self._center_dist(y_pred)) #self._outl_dist(y_pred))#+
+        if 'param' in self.penalty or 'detj' in self.penalty or 'std' in self.penalty:
+            mean_dim=list(range(1, self.ndims + 2))
+            dg = torch.sum(torch.abs(img),dim=1,keepdim=True)* torch.exp(-self.apear_scale * torch.nn.ReLU(inplace=True)(.1-sum([torch.sum(g * g, dim=1, keepdim=True) for g in self._eq_diffs(img,dist=3)]) / torch.sum(torch.square(.1 + img), dim=1, keepdim=True)))
+            dg = dg/(EPS+torch.mean(dg,dim=mean_dim,keepdim=True))
+
+            y_pred = torch.clamp(y_pred, min=-0.8, max=0.8)
+            x_in = x_in if isinstance(x_in,list) else [x_in]
+            if 'std' in self.penalty:
+                reg_loss += self.sign*torch.mean(torch.clamp(grad_std((y_pred-torch.mean(y_pred,dim=list(range(2,ndims+2)),keepdim=True))*dg), max=.2, min=0))
+            if 'param' in self.penalty:
+                for id, d in enumerate(self.dist):
+                    df = torch.mean(torch.abs(sum(self._weighted_diffs_error(y_pred, dist=d, w=dg, expect=torch.abs(x_in[-1][:, id:id + 1, ...]),mean_dim=mean_dim))))
+                    reg_loss += 1 * (df) / len(self.dist)
+
+            if 'detj' in self.penalty:
+                df = torch.mean(torch.abs(
+                    torch.mean((torch.abs(self._eval_detJ(self._eq_diffs(y_pred, dist=1))) - torch.abs(x_in[0])) * dg, dim=mean_dim)))
+                reg_loss += 0.5*df
+
+        return reg_loss
+
+
+def avg_std_skew_kurt(array,ndims=2):
+    dim = list(range(2, ndims + 2))
+    mean = torch.mean(array,dim=dim)
+    diffs = array - mean
+    var = torch.mean(torch.pow(diffs, 2.0),dim=dim)
+    std = torch.pow(var, 0.5)
+    zscores = diffs / std
+    skews = torch.mean(torch.pow(zscores, 3.0),dim=dim)
+    kurtoses = torch.mean(torch.pow(zscores, 4.0),dim=dim) - 3.0
+    return [mean,std,skews,kurtoses]
+
+def grad_std(array,ndims=2):
+    dim = list(range(2, ndims + 2))
+    array=torch.clamp(array,min=-0.8,max=0.8)
+    dim0=list(range(1,ndims+2))
+    std = torch.sqrt(torch.mean(torch.square(array - torch.mean(array, dim=dim, keepdim=True)), dim=dim0))
+    return std
+
+def avg_std(array,ndims=2):
+    dim = list(range(2, ndims + 2))
+    return [torch.mean(array,dim=dim),grad_std(array,dim=dim)]
+
+
+if __name__ == "__main__":
+    # ndims=2
+    # dist=[16,32]
+    # ddf = torch.rand(1,2,128,128)
+    # # ddf[:,:,0,:]=ddf[:,:,0,:]-1
+    # # ddf[:,:,1,:]=ddf[:,:,1,:]+1
+    # # ddf[:,:,0,0]=ddf[:,:,0,0] -1
+    # # ddf[:,:,1,1]=ddf[:,:,1,1] +1
+    # # ddf[:,0,0,1]=ddf[:,0,0,1] +1
+    # # ddf[:,1,0,1]=ddf[:,1,0,1] -1
+    # # ddf[:,0,0,1]=ddf[:,0,0,1] -1
+    # # ddf[:,1,0,1]=ddf[:,1,0,1] +1
+    # # ddf[:,1,1,0]=ddf[:,1,1,0] -1
+    # # ddf[:,0,1,0]=ddf[:,0,1,0] +1
+    # ddf=ddf
+    # img = torch.rand(1,1,128,128)
+    # x_in=np.reshape([0.2,0.3],newshape=[1,ndims]+[1]*ndims)
+    # x_in=[torch.tensor(x_in).type(torch.float32),0.]
+
+    # Loss_detj = Grad(penalty=['detj'],ndims=ndims,dist=dist)
+    # loss_detj = Loss_detj(ddf,x_in,img)
+    # print(loss_detj)
+
+    size = 128
+    smooth = True
+    # smooth = False
+    img3d = torch.empty(1,1,size,size,size).uniform_(0,1)
+    img3d_t = torch.empty(1,1,size,size,size).uniform_(0,1)#*-0.000001
+    # img3d_t = img3d.clone().detach()
+    # img3d_t = torch.zeros_like(img3d)
+    translation = 2
+    start = 0
+    end = 32
+    # img3d_t[:,:,translation:,translation:,translation:] = img3d[:,:,:size-translation,:size-translation,:size-translation]
+    # img3d_t[:,:,:,translation:,translation:] = img3d[:,:,:,:size-translation,:size-translation]
+    img3d_t[:,:,:,:,translation:] = img3d[:,:,:,:,:size-translation]
+    # img3d_t[:,:,start:end,start:end,start:end] = img3d[:,:,start+translation:end+translation,start+translation:end+translation,start+translation:end+translation]
+    img3d_t = img3d_t
+    loss_ncc = LNCC(smooth=smooth,central=True)
+    loss_sim = loss_ncc(img3d, img3d_t)
+    print(loss_sim)

Diffusion/losses_ncc0.py

@@ -0,0 +1,496 @@
+"""
+losses for DRDM
+"""
+
+import numpy as np
+import sys
+import torch
+import torch.nn.functional as F
+
+
+EPS=1e-7
+
+# eps_scale = 10e-5
+# eps_scale = 10e-4
+# eps_scale = 1e-4
+eps_scale = 1e-5
+
+
+
+class LMSE(torch.nn.Module):
+    """
+    Labeled Mean Square Error (LMSE)
+    """
+
+    def __init__(self, eps=1e-7, relate_eps=5e-1, win=None, smooth=False):
+        super(LMSE, self).__init__()
+        self.eps = eps
+        self.relate_eps = relate_eps
+        self.ndims = 3
+        self.smooth = smooth
+        self.win = win
+        # Set window size
+        if self.win is None:
+            self.win = [5] * self.ndims
+        if smooth:
+            self.kernels = self._build_kernel(std=0.0)
+
+    def _build_kernel(self, std=0.0):
+        if std == 0.0:
+            return torch.ones([1, 1, *self.win])
+        else:
+            tail = int(np.ceil(std)) * 3
+            k = torch.exp(-0.5 * torch.arange(-tail, tail + 1, dtype=torch.float32) ** 2 / std ** 2)
+            kernel = k / torch.sum(k)
+            kernel = kernel.view(-1, 1, 1) * kernel.view(1, -1, 1) * kernel.view(1, 1, -1)
+            return kernel.unsqueeze(0).unsqueeze(0)
+        
+    def forward(self, I, J, label=None):
+        """
+        Computes the labeled mean squared error between I and J (ref).
+        If label is provided, computes the MSE only over the labeled regions.
+        """
+        padding = [(w-1) // 2 for w in self.win]
+        if self.smooth:
+            I = torch.nn.functional.conv3d(I, self.kernels, stride=1, padding=padding)
+            J = torch.nn.functional.conv3d(J, self.kernels, stride=1, padding=padding)
+        mse = (I - J) ** 2
+        if self.relate_eps is not None:
+            mse = mse/((J**2) + self.relate_eps)
+        if label is not None:
+            label = label.float()
+            mse = mse * label
+            mse_sum = torch.sum(mse, dim=(2, 3, 4))
+            label_sum = torch.sum(label, dim=(2, 3, 4)) + self.eps
+            loss = torch.mean(mse_sum / label_sum)
+        else:
+            loss = torch.mean(mse)
+        return loss
+    
+class LNCC(torch.nn.Module):
+    """
+    Local (over window) normalized cross-correlation (LNCC)
+    """
+
+    def __init__(self, win=None, num_ch=1, eps=1e-7, central=True, smooth=False):
+        super(LNCC, self).__init__()
+        self.win = win
+        self.eps = eps
+        self.central = central
+        self.ndims = 3
+        self.strides = [1] * (self.ndims + 2)
+        self.smooth = smooth
+
+        # Set window size
+        if self.win is None:
+            self.win = [11] * self.ndims
+
+        if smooth:
+            self.kernels = self._build_kernel(std=0.5)
+        self.sum_filt = self._build_kernel(std=0.0)
+
+    def _build_kernel(self, std=0.0):
+        if std == 0.0:
+            return torch.ones([1, 1, *self.win])
+        else:
+            tail = int(np.ceil(std)) * 3
+            k = torch.exp(-0.5 * torch.arange(-tail, tail + 1, dtype=torch.float32) ** 2 / std ** 2)
+            kernel = k / torch.sum(k)
+            kernel = kernel.view(-1, 1, 1) * kernel.view(1, -1, 1) * kernel.view(1, 1, -1)
+            return kernel.unsqueeze(0).unsqueeze(0)
+
+    def lncc(self, I, J, label=None):
+        self.sum_filt = self.sum_filt.to(I.device)
+        padding = [(w-1) // 2 for w in self.win]
+
+        if self.smooth:
+            I = torch.nn.functional.conv3d(I, self.kernels, stride=1, padding=padding)
+            J = torch.nn.functional.conv3d(J, self.kernels, stride=1, padding=padding)
+
+        # Compute CC squares
+        I2 = I * I
+        J2 = J * J
+        IJ = I * J
+
+        if self.central:
+            # Compute local sums via convolution
+            I_sum = torch.nn.functional.conv3d(I, self.sum_filt, stride=1, padding=padding)
+            J_sum = torch.nn.functional.conv3d(J, self.sum_filt, stride=1, padding=padding)
+            I2_sum = torch.nn.functional.conv3d(I2, self.sum_filt, stride=1, padding=padding)
+            J2_sum = torch.nn.functional.conv3d(J2, self.sum_filt, stride=1, padding=padding)
+            IJ_sum = torch.nn.functional.conv3d(IJ, self.sum_filt, stride=1, padding=padding)
+
+            # Compute cross-correlation
+            win_size = np.prod(self.win)
+            cross = IJ_sum - (I_sum * J_sum) / win_size
+            I_var = I2_sum - (I_sum * I_sum) / win_size
+            J_var = J2_sum - (J_sum * J_sum) / win_size
+        else:
+            # Compute local sums via convolution
+            I2_sum = torch.nn.functional.conv3d(I2, self.sum_filt, stride=1, padding=padding)
+            J2_sum = torch.nn.functional.conv3d(J2, self.sum_filt, stride=1, padding=padding)
+            IJ_sum = torch.nn.functional.conv3d(IJ, self.sum_filt, stride=1, padding=padding)
+
+            cross = IJ_sum
+            I_var = I2_sum
+            J_var = J2_sum
+
+        cc = (cross * cross) / (I_var * J_var + self.eps)
+        if label is not None:
+            label = label.float()
+            cc = torch.sum(cc * label, dim=(2, 3, 4)) / (torch.sum(label, dim=(2, 3, 4)) + self.eps)
+
+        return torch.mean(cc)
+
+    def forward(self, I, J, label=None):
+        return -self.lncc(I, J, label=label)
+
+
+
+class NCC(torch.nn.Module):
+    # def __init__(self, eps_scale=10e-7,img_sz=256):
+    def __init__(self, eps_scale=10e-5,img_sz=256):
+        super(NCC, self).__init__()
+        self.eps_scale=eps_scale#*img_sz/256
+        # self.scale=10e4
+        self.scale=1e2
+
+    def forward(self,pred,inv_lab=None,ddf_stn=None,mask=None):
+        if ddf_stn is None:
+            trm_pred=pred
+        else:
+            trm_pred=-ddf_stn(pred, inv_lab)
+        trm_pred = self.scale * trm_pred
+        inv_lab = self.scale * inv_lab
+        if mask is None:
+            loss_gen = torch.mean(torch.sum(trm_pred*inv_lab,dim=1)/(torch.sqrt(torch.sum(torch.square(trm_pred),dim=1)*torch.sum(torch.square(inv_lab),dim=1)+self.eps_scale)))
+        else:
+            batch_size = inv_lab.shape[0]
+            loss_gen = torch.sum(torch.sum(trm_pred*inv_lab,dim=1)*mask/(torch.sqrt(torch.sum(torch.square(trm_pred),dim=1)*torch.sum(torch.square(inv_lab),dim=1)+self.eps_scale)))/torch.sum(mask)/batch_size
+        return loss_gen
+
+class MRSE(torch.nn.Module):
+    def __init__(self, eps_scale=eps_scale,img_sz=256):
+        super(MRSE, self).__init__()
+        self.eps_scale=eps_scale#*img_sz/256
+        self.scale = 10e1
+
+    def forward(self,pred,inv_lab=None,ddf_stn=None,mask=None):
+        if ddf_stn is None:
+            trm_pred=pred
+        else:
+            trm_pred=-ddf_stn(pred, inv_lab)
+        trm_pred = self.scale * trm_pred
+        inv_lab = self.scale * inv_lab
+        if mask is None:
+            loss_gen = torch.mean(
+                torch.sum(torch.square(trm_pred + inv_lab), dim=1)
+                / (torch.sum(torch.square(inv_lab), dim=1) + self.eps_scale)
+            )
+        else:
+            batch_size = inv_lab.shape[0]
+            loss_gen = torch.sum(
+                torch.sum(torch.square(trm_pred + inv_lab), dim=1) * mask
+                / (torch.sum(torch.square(inv_lab), dim=1) + self.eps_scale)
+            )/torch.sum(mask)/batch_size
+        return loss_gen/1
+
+class RMSE(torch.nn.Module):
+    def __init__(self, eps_scale=eps_scale,img_sz=256,ndims=2):
+        super(RMSE, self).__init__()
+        self.eps_scale=eps_scale#*img_sz/256
+        self.ndims=ndims
+
+    def forward(self,pred,inv_lab=None,ddf_stn=None):
+        if ddf_stn is None:
+            trm_pred=pred
+        else:
+            trm_pred=-ddf_stn(pred, inv_lab)
+        loss_gen = torch.mean(torch.mean(torch.sum(torch.square(trm_pred - inv_lab), dim=1),
+                              dim=list(range(1, 1 + self.ndims))) / (
+                               torch.mean(torch.sum(torch.square(inv_lab), dim=1), dim=list(range(1, 1 + self.ndims))) + self.eps_scale))
+        return loss_gen
+# loss_gen = torch.mean(torch.mean(torch.sum(torch.square(ddf_stn(pre_dvf_I, dvf_I) + dvf_I), dim=1),dim=list(range(1,1+ndims))) / (torch.mean(torch.sum(torch.square(dvf_I), dim=1),dim=list(range(1,1+ndims))) + EPS))
+
+class Grad(torch.nn.Module):
+    """
+    N-D gradient loss
+    """
+
+    def __init__(self, penalty=['l1'],ndims=3, eps=1e-8, outrange_weight=1e4,outrange_thresh=0.5, detj_weight=2, apear_scale=4, dist=1, sign=1,waive_thresh=10**-5):
+        super(Grad, self).__init__()
+        self.penalty = penalty
+        self.eps = eps
+        self.outrange_weight = outrange_weight
+        self.detj_weight=detj_weight
+        self.apear_scale = apear_scale
+        self.ndims=ndims
+        self.max_sz = torch.reshape(torch.tensor([outrange_thresh]*ndims, dtype=torch.float32) , [1]+[ndims]+[1]*(ndims))
+        self.act = torch.nn.ReLU(inplace=False)
+        self.dist=dist
+        self.sign=sign
+        self.waive_thresh=waive_thresh
+
+    def _diffs(self, y,dist=None):
+        if dist is None:
+            dist=self.dist
+        # vol_shape = y.size()[2:]
+        # vol_shape = y.get_shape().as_list()[1:-1]
+        # ndims = len(vol_shape)
+
+        df = [None] * self.ndims
+        for i in range(self.ndims):
+            d = i + 2
+            # permute dimensions to put the ith dimension first
+            r = [d, *range(d), *range(d + 1, self.ndims + 2)]
+            yp = y.permute(r)
+            dfi = (yp[dist:, ...] - yp[:-dist, ...])/float(dist)
+
+            # permute back
+            # note: this might not be necessary for this loss specifically,
+            # since the results are just summed over anyway.
+            r = [*range(1, d + 1), 0, *range(d + 1, self.ndims + 2)]
+            df[i] = dfi.permute(r)
+        return df
+
+    def _eq_diffs(self, y,dist=None):
+        if dist is None:
+            dist=self.dist
+        # vol_shape = y.get_shape().as_list()[1:-1]
+        vol_shape = y.size()[2:]
+        ndims = len(vol_shape)
+        pad = [0, 0] * (ndims + 1) +[dist, 0]
+        pad1 = [0, 0] * (ndims + 1) +[0, dist]
+        # df = [None, None] * ndims
+        df = [None] * ndims
+        for i in range(ndims):
+            d = i + 2
+            r=[d, *range(d), *range(d + 1, ndims + 2)]
+            ri=[*range(1, d + 1), 0, *range(d + 1, ndims + 2)]
+            yt = y.permute(r)
+            dy=(yt[dist:, ...] - yt[:-dist, ...])/float(dist)
+            df[i] = (F.pad(dy, pad,mode='constant',value=0)).permute(ri)
+            # df[2*i] = (F.pad(dy, pad,mode='constant',value=0)).permute(ri)
+            # df[2*i+1] = (F.pad(dy, pad1, mode='constant', value=0)).permute(ri)
+            y.permute(ri)
+        return df
+
+    def _weighted_diffs_error(self, y,dist=None,w=None,expect=None,mean_dim=None):
+        if dist is None:
+            dist=self.dist
+        vol_shape = y.size()[2:]
+        ndims = len(vol_shape)
+        df = [None] * ndims
+
+        for i in range(ndims):
+            d = i + 2
+            r=[d, *range(d), *range(d + 1, ndims + 2)]
+            ri=[*range(1, d + 1), 0, *range(d + 1, ndims + 2)]
+            yt = y.permute(r)
+            wt = w.permute(r)
+            dy=(torch.abs(yt[dist:, ...] - yt[:-dist, ...])-expect.permute(r))*(wt[dist:, ...]*wt[:-dist, ...])
+            df[i] = torch.mean((dy).permute(ri),dim=mean_dim,keepdim=True)
+            y.permute(ri)
+            w.permute(ri)
+        return df
+
+    def _outl_dist(self, y,range_thresh=0.2):
+        self.device = y.device
+        vol_shape = y.size()[2:]
+        self.max_sz=self.max_sz.to(self.device)
+        act=torch.nn.ReLU(inplace=True)
+        loss=0.
+        for i in range(self.ndims):
+            d = i + 2
+            # permute dimensions to put the ith dimension first
+            r = [d, *range(d), *range(d + 1, self.ndims + 2)]
+            ri = [*range(1, d + 1), 0, *range(d + 1, self.ndims + 2)]
+            yt = y.permute(r)
+            loss += torch.mean(torch.square(act(-range_thresh-yt[0,:,i, ...])))+torch.mean(torch.square(act(yt[-1,:,i, ...]-range_thresh)))
+            # loss += torch.mean(torch.square(act(-range_thresh-yt[0,:,i, ...])+act(yt[-1,:,i, ...]-range_thresh)))
+            y.permute(ri)
+        return loss/self.ndims
+
+    def _center_dist(self, y):
+        self.device = y.device
+        vol_shape = y.size()[2:]
+        self.max_sz=self.max_sz.to(self.device)
+        select_loc = [s // 2 for s in vol_shape]
+        if self.ndims==3:
+            # return torch.mean(self.act(torch.abs(y[:,:, select_loc[0], select_loc[1], select_loc[2]]) - self.max_sz))
+            return torch.mean(torch.square(self.act(torch.abs(y[:, :, select_loc[0], select_loc[1], select_loc[2]]) - self.max_sz)))
+        elif self.ndims == 2:
+            # return torch.mean(self.act(torch.abs(y[:, :, select_loc[0], select_loc[1]]) - self.max_sz))
+            return torch.mean(torch.square(self.act(torch.abs(y[:, :, select_loc[0], select_loc[1]]) - self.max_sz)))
+
+        
+    # def _eval_detJ(self, disp=None, weight=None):
+    #     weight = 1
+    #     if self.ndims==3:
+    #         detj = (disp[0][:, 0, ...] * disp[1][:, 1, ...] * disp[2][:, 2, ...]) + (
+    #                 disp[0][:, 1, ...] * disp[1][:, 2, ...] * disp[2][:, 0, ...]) + (
+    #                        disp[0][:, 2, ...] * disp[1][:, 0, ...] * disp[2][:, 1, ...]) - (
+    #                        disp[0][:, 2, ...] * disp[1][:, 1, ...] * disp[2][:, 0, ...]) - (
+    #                        disp[0][:, 0, ...] * disp[1][:, 2, ...] * disp[2][:, 1, ...]) - (
+    #                        disp[0][:, 1, ...] * disp[1][:, 0, ...] * disp[2][:, 2, ...])
+    #     elif self.ndims==2:
+    #         detj = (disp[0][:, 0, ...] * disp[1][:, 1, ...]) - (disp[0][:, 1, ...] * disp[1][:, 0, ...])
+
+    #     return detj * weight
+    
+    def _eval_detJ(self, disp, add_identity=True, spacing=1.0):
+        """
+        disp: list length ndims
+            disp[i] is derivative wrt spatial dim i (forward diff),
+            tensor shape [B, C=ndims, ...]
+        add_identity: True if y_pred is displacement u and phi=x+u
+        spacing: voxel spacing (or 1.0). If you care about physical units,
+                divide derivatives by spacing (and dist). Sign won't change.
+        """
+        # Optional scaling (won't affect sign as long as spacing>0)
+        if spacing != 1.0:
+            disp = [d / spacing for d in disp]
+
+        if self.ndims == 2:
+            dux_dx = disp[0][:, 0, ...]
+            duy_dx = disp[0][:, 1, ...]
+            dux_dy = disp[1][:, 0, ...]
+            duy_dy = disp[1][:, 1, ...]
+
+            if add_identity:
+                j11 = 1.0 + dux_dx
+                j22 = 1.0 + duy_dy
+            else:
+                j11 = dux_dx
+                j22 = duy_dy
+
+            detj = j11 * j22 - dux_dy * duy_dx
+            return detj
+
+        elif self.ndims == 3:
+            dux_dx = disp[0][:, 0, ...]
+            duy_dx = disp[0][:, 1, ...]
+            duz_dx = disp[0][:, 2, ...]
+
+            dux_dy = disp[1][:, 0, ...]
+            duy_dy = disp[1][:, 1, ...]
+            duz_dy = disp[1][:, 2, ...]
+
+            dux_dz = disp[2][:, 0, ...]
+            duy_dz = disp[2][:, 1, ...]
+            duz_dz = disp[2][:, 2, ...]
+
+            if add_identity:
+                j11 = 1.0 + dux_dx 
+                j22 = 1.0 + duy_dy 
+                j33 = 1.0 + duz_dz
+            else:
+                j11 = dux_dx       
+                j22 = duy_dy
+                j33 = duz_dz
+
+            j12 = dux_dy; j13 = dux_dz
+            j21 = duy_dx; j23 = duy_dz
+            j31 = duz_dx; j32 = duz_dy
+
+            detj = (
+                j11 * (j22 * j33 - j23 * j32)
+                - j12 * (j21 * j33 - j23 * j31)
+                + j13 * (j21 * j32 - j22 * j31)
+            )
+            return detj
+
+        else:
+            raise ValueError(f"Unsupported ndims={self.ndims}")
+
+        
+    def forward(self,  y_pred=None,x_in=None, img=None, msk=None):
+        reg_loss = 0
+        act=torch.nn.ReLU(inplace=True)
+        
+        dg = 1
+        if img is not None:
+            dg = torch.exp(-self.apear_scale * sum([torch.sum(g * g, dim=1, keepdim=True) for g in self._eq_diffs(img)]) / torch.sum(torch.square(0.2 + img), dim=1, keepdim=True))   
+        if msk is not None:
+            dg = dg * msk
+
+        if 'l1' in self.penalty:
+            df = [torch.mean(dg*F.relu(torch.abs(f) - self.waive_thresh,inplace=True)) for f in self._eq_diffs(y_pred)]
+            reg_loss += sum(df) / len(df)
+            
+        if 'l2' in self.penalty:
+            df = [torch.mean(dg*F.relu(f * f - self.waive_thresh**2,inplace=True)) for f in self._eq_diffs(y_pred)]
+            reg_loss += torch.sqrt(sum(df) / len(df))
+
+        if 'negdetj' in self.penalty:
+            df = self.detj_weight*torch.mean(act(-self._eval_detJ(self._eq_diffs(y_pred,dist=1))))  # , dg[...,0])
+            reg_loss += 0.5*df
+        if 'range' in self.penalty:
+            reg_loss += self.outrange_weight * (self._center_dist(y_pred)) #self._outl_dist(y_pred))#+
+        if 'param' in self.penalty or 'detj' in self.penalty or 'std' in self.penalty:
+            mean_dim=list(range(1, self.ndims + 2))
+            dg = torch.sum(torch.abs(img),dim=1,keepdim=True)* torch.exp(-self.apear_scale * torch.nn.ReLU(inplace=True)(.1-sum([torch.sum(g * g, dim=1, keepdim=True) for g in self._eq_diffs(img,dist=3)]) / torch.sum(torch.square(.1 + img), dim=1, keepdim=True)))
+            dg = dg/(EPS+torch.mean(dg,dim=mean_dim,keepdim=True))
+
+            y_pred = torch.clamp(y_pred, min=-0.8, max=0.8)
+            x_in = x_in if isinstance(x_in,list) else [x_in]
+            if 'std' in self.penalty:
+                reg_loss += self.sign*torch.mean(torch.clamp(grad_std((y_pred-torch.mean(y_pred,dim=list(range(2,ndims+2)),keepdim=True))*dg), max=.2, min=0))
+            if 'param' in self.penalty:
+                for id, d in enumerate(self.dist):
+                    df = torch.mean(torch.abs(sum(self._weighted_diffs_error(y_pred, dist=d, w=dg, expect=torch.abs(x_in[-1][:, id:id + 1, ...]),mean_dim=mean_dim))))
+                    reg_loss += 1 * (df) / len(self.dist)
+
+            if 'detj' in self.penalty:
+                df = torch.mean(torch.abs(
+                    torch.mean((torch.abs(self._eval_detJ(self._eq_diffs(y_pred, dist=1))) - torch.abs(x_in[0])) * dg, dim=mean_dim)))
+                reg_loss += 0.5*df
+
+        return reg_loss
+
+
+def avg_std_skew_kurt(array,ndims=2):
+    dim = list(range(2, ndims + 2))
+    mean = torch.mean(array,dim=dim)
+    diffs = array - mean
+    var = torch.mean(torch.pow(diffs, 2.0),dim=dim)
+    std = torch.pow(var, 0.5)
+    zscores = diffs / std
+    skews = torch.mean(torch.pow(zscores, 3.0),dim=dim)
+    kurtoses = torch.mean(torch.pow(zscores, 4.0),dim=dim) - 3.0
+    return [mean,std,skews,kurtoses]
+
+def grad_std(array,ndims=2):
+    dim = list(range(2, ndims + 2))
+    array=torch.clamp(array,min=-0.8,max=0.8)
+    dim0=list(range(1,ndims+2))
+    std = torch.sqrt(torch.mean(torch.square(array - torch.mean(array, dim=dim, keepdim=True)), dim=dim0))
+    return std
+
+def avg_std(array,ndims=2):
+    dim = list(range(2, ndims + 2))
+    return [torch.mean(array,dim=dim),grad_std(array,dim=dim)]
+
+
+if __name__ == "__main__":
+    ndims=2
+    dist=[16,32]
+    ddf = torch.rand(1,2,128,128)
+    # ddf[:,:,0,:]=ddf[:,:,0,:]-1
+    # ddf[:,:,1,:]=ddf[:,:,1,:]+1
+    # ddf[:,:,0,0]=ddf[:,:,0,0] -1
+    # ddf[:,:,1,1]=ddf[:,:,1,1] +1
+    # ddf[:,0,0,1]=ddf[:,0,0,1] +1
+    # ddf[:,1,0,1]=ddf[:,1,0,1] -1
+    # ddf[:,0,0,1]=ddf[:,0,0,1] -1
+    # ddf[:,1,0,1]=ddf[:,1,0,1] +1
+    # ddf[:,1,1,0]=ddf[:,1,1,0] -1
+    # ddf[:,0,1,0]=ddf[:,0,1,0] +1
+    ddf=ddf
+    img = torch.rand(1,1,128,128)
+    x_in=np.reshape([0.2,0.3],newshape=[1,ndims]+[1]*ndims)
+    x_in=[torch.tensor(x_in).type(torch.float32),0.]
+
+    Loss_detj = Grad(penalty=['detj'],ndims=ndims,dist=dist)
+    loss_detj = Loss_detj(ddf,x_in,img)
+    print(loss_detj)

Diffusion/networks.py

@@ -0,0 +1,1167 @@
+from torch import nn
+import torch
+import torch.nn.functional as F
+import numpy as np
+import math
+
+def get_net(name="recresnet"):
+    name = name.lower()
+    if name == "recresacnet":
+        net = RecResACNet
+    elif name == "recmutattnnet":
+        net = RecMutAttnNet
+    elif name == "recmutattnnet0":
+        net = RecMutAttnNet0
+    elif name == "recmutattnnet1":
+        net = RecMutAttnNet1
+    elif name == "defrecmutattnnet":
+        net = DefRec_MutAttnNet
+    elif name == "recmutattnnet_contrastive":
+        net = RecMutAttnNet_contrastive
+    else:
+        net = None
+    return net
+
+
+
+def sinusoidal_embedding(n, d):
+    # Returns the standard positional embedding
+    embedding = torch.zeros(n, d)
+    wk = torch.tensor([1 / 10_000 ** (2 * j / d) for j in range(d)])
+    wk = wk.reshape((1, d))
+    t = torch.arange(n).reshape((n, 1))
+    embedding[:,::2] = torch.sin(t * wk[:,::2])
+    embedding[:,1::2] = torch.cos(t * wk[:,::2])
+    return embedding
+
+class AtrousBlock(nn.Module):
+    def __init__(self, shape, in_c, out_c, kernel_size=3, stride=1, atrous_rates=[1,3], ndims=2, activation=None, normalize=True):
+        super(AtrousBlock, self).__init__()
+        # if 0 not in shape:
+        if normalize:
+            # print(shape)
+            # self.ln = nn.LayerNorm(shape)     # jzheng 15/03/2024
+            norm=getattr(nn, 'InstanceNorm%dd' % ndims)     # jzheng 15/03/2024
+            self.ln = norm(out_c,affine=True)
+        else:
+            self.ln = nn.Identity()
+        Conv=getattr(nn,'Conv%dd' % ndims)
+        if in_c!=out_c:
+            self.conv0 = Conv(in_c, out_c, kernel_size, 1, (kernel_size-1)//2*1) #if in_c!=out_c else None
+        else:
+            self.conv0 = None
+        self.convs = nn.ModuleList([
+            Conv(out_c, out_c, kernel_size, 1, (kernel_size-1)//2*ar, dilation=ar)
+            if ar>0 else Conv(out_c, out_c, 1, 1, 0)
+            for ar in atrous_rates
+        ])
+        # self.conv1 = Conv(out_c, out_c, kernel_size, stride, padding)
+        # self.conv2 = Conv(out_c, out_c, kernel_size, stride, padding)
+        self.activation = nn.LeakyReLU(1e-6) if activation is None else activation
+        # self.activation = nn.ReLU() if activation is None else activation
+        # self.activation = nn.ReLU()
+        self.normalize = normalize
+
+    def forward(self, x):
+        if self.conv0 is not None:
+            x = self.conv0(x) #if self.conv0 is not None else x
+        x = self.ln(x) if self.normalize else x     # jzheng 15/03/2024
+        out=nn.Identity()(x)
+        for conv in self.convs:
+            out = self.activation(out)
+            out = conv(out)
+        return self.activation(out+x)
+
+# ==============================================
+# Unconditional Network
+# ==============================================
+
+class RecResACNet(nn.Module):
+    def __init__(self, n_steps=1000, time_emb_dim=100, ndims=2, num_input_chn=1, res=0):
+        super(RecResACNet, self).__init__()
+
+        self.dimension = ndims
+        self.Conv = getattr(nn, 'Conv%dd' % self.dimension)
+        self.ConvT = getattr(nn, 'ConvTranspose%dd' % self.dimension)
+
+        # Sinusoidal embedding
+        self.time_embed = nn.Embedding(n_steps, time_emb_dim)
+        self.time_embed.weight.data = sinusoidal_embedding(n_steps, time_emb_dim)
+        self.time_embed.requires_grad_(False)
+
+        # First half
+        self.te1 = self._make_te(time_emb_dim, 1)
+        self.b1 = nn.Sequential(
+            AtrousBlock([num_input_chn] + [res] * ndims, num_input_chn, 10, ndims=ndims),
+            AtrousBlock([10] + [res] * ndims, 10, 10, ndims=ndims),
+            AtrousBlock([10] + [res] * ndims, 10, 10, ndims=ndims),
+
+        )
+        self.down1 = self.Conv(10, 10, 4, 2, 1)
+
+        self.te2 = self._make_te(time_emb_dim, 10)
+        self.b2 = nn.Sequential(
+            AtrousBlock([10] + [res // 2] * ndims, 10, 20, ndims=ndims),
+            AtrousBlock([20] + [res // 2] * ndims, 20, 20, ndims=ndims),
+            AtrousBlock([20] + [res // 2] * ndims, 20, 20, ndims=ndims)
+        )
+        self.down2 = self.Conv(20, 20, 4, 2, 1)
+
+        self.te3 = self._make_te(time_emb_dim, 20)
+        self.b3 = nn.Sequential(
+            AtrousBlock([20] + [res // 4] * ndims, 20, 40, ndims=ndims),
+            AtrousBlock([40] + [res // 4] * ndims, 40, 40, ndims=ndims),
+            AtrousBlock([40] + [res // 4] * ndims, 40, 40, ndims=ndims)
+        )
+        self.down3 = self.Conv(40, 40, 4, 2, 1)
+
+        # Bottleneck
+        self.te_mid = self._make_te(time_emb_dim, 40)
+        self.b_mid = nn.Sequential(
+            AtrousBlock([40] + [res // 8] * ndims, 40, 20, ndims=ndims),
+            AtrousBlock([20] + [res // 8] * ndims, 20, 20, ndims=ndims),
+            AtrousBlock([20] + [res // 8] * ndims, 20, 40, ndims=ndims)
+        )
+
+        # Second half
+        self.up1 = self.ConvT(40, 40, 4, 2, 1)
+
+        self.te4 = self._make_te(time_emb_dim, 80)
+        self.b4 = nn.Sequential(
+            AtrousBlock([80] + [res // 4] * ndims, 80, 40, ndims=ndims, normalize=False),
+            AtrousBlock([40] + [res // 4] * ndims, 40, 20, ndims=ndims, normalize=False),
+            AtrousBlock([20] + [res // 4] * ndims, 20, 20, ndims=ndims, normalize=False)
+        )
+
+        self.up2 = self.ConvT(20, 20, 4, 2, 1)
+        self.te5 = self._make_te(time_emb_dim, 40)
+        self.b5 = nn.Sequential(
+            AtrousBlock([40] + [res // 2] * ndims, 40, 20, ndims=ndims, normalize=False),
+            AtrousBlock([20] + [res // 2] * ndims, 20, 10, ndims=ndims, normalize=False),
+            AtrousBlock([10] + [res // 2] * ndims, 10, 10, ndims=ndims, normalize=False)
+        )
+
+        self.up3 = self.ConvT(10, 10, 4, 2, 1)
+        self.te_out = self._make_te(time_emb_dim, 20)
+        self.b_out = nn.Sequential(
+            AtrousBlock([20] + [res // 1] * ndims, 20, 10, ndims=ndims, normalize=False),
+            AtrousBlock([10] + [res // 1] * ndims, 10, 10, ndims=ndims, normalize=False),
+            AtrousBlock([10] + [res // 1] * ndims, 10, 10, ndims=ndims, normalize=False)
+        )
+
+        self.conv_out = self.Conv(10, ndims, 3, 1, 1)
+
+    def boundary_limit(self, sample_coords0, max_sz, plus=0., minus=1.):
+        sample_coords = torch.split(sample_coords0, split_size_or_sections=1, dim=1)
+        return torch.cat([torch.clamp(x * sz, min=minus - 1 * sz + plus, max=1 * sz - minus + plus) / sz for x, sz in
+                          zip(sample_coords, max_sz)], 1)
+
+    def resample(self, vol, ddf, ref=None, img_sz=None, padding_mode="zeros"):
+        ref = self.ref_grid if ref is None else ref
+        img_sz = self.max_sz if img_sz is None else img_sz
+        # resample_mode = 'bicubic'
+        resample_mode = 'bilinear' # if self.dimension==2 else 'trilinear'
+        # padding_mode = "border"
+
+        if True:
+            # return F.grid_sample(vol, torch.flip(torch.transpose(ddf * torch.Tensor(np.reshape(np.array(self.max_sz), [1, 1, 1, self.dimension])).cuda() + ref,[0, 2, 3, 1]) / img_sz - 1, dims=[-1]), mode=resample_mode, padding_mode=padding_mode,align_corners=True)
+            return F.grid_sample(vol, torch.flip((ddf * torch.Tensor(
+                np.reshape(np.array(self.max_sz), [1, self.dimension]+[1]*self.dimension)).to(self.device) + ref).permute(
+                [0]+list(range(2,2+self.dimension))+[1]) / img_sz - 1, dims=[-1]), mode=resample_mode, padding_mode=padding_mode,
+                                 align_corners=True)
+
+    def forward(self, x=None, t=None, y=None, rec_num=2, ndims=2):
+        #
+        self.device = x.device
+        # [h, w] = x.size()[2:]
+        img_sz = x.size()[2:]
+        n = x.size()[0]
+        self.max_sz = [img_sz[0]] * self.dimension
+        ts_emb_shape=[n,-1]+[1]*self.dimension
+        # [h,w]=img_sz
+        # self.img_sz = torch.reshape(torch.tensor([(h - 1) / 2., (w - 1) / 2.], device=self.device), [1, 1, 1, 2])
+        self.img_sz = torch.reshape(torch.tensor([(imsz - 1) / 2 for imsz in img_sz], device=self.device), [1]*(self.dimension+1)+[self.dimension])
+        # self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=h), torch.arange(end=w)]), 0),
+        #                               [1, 2, h, w]).to(self.device)
+        self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=imsz) for imsz in img_sz]), 0),
+                                      [1, self.dimension]+list(img_sz)).to(self.device)
+        img = x
+
+        # x is (N, 2, 28, 28) (image with positional embedding stacked on channel dimension)
+        t = self.time_embed(t)
+
+        for rec_id in range(rec_num):
+            out1 = self.b1(img + self.te1(t).reshape(ts_emb_shape))  # (N, 10, 28, 28)
+            out2 = self.b2(self.down1(out1) + self.te2(t).reshape(ts_emb_shape))  # (N, 20, 14, 14)
+            out3 = self.b3(self.down2(out2) + self.te3(t).reshape(ts_emb_shape))  # (N, 40, 7, 7)
+
+            out_mid = self.b_mid(self.down3(out3) * self.te_mid(t).reshape(ts_emb_shape))  # (N, 40, 3, 3)
+
+            out4 = torch.cat((out3, self.up1(out_mid)), dim=1)  # (N, 80, 7, 7)
+            out4 = self.b4(out4 + self.te4(t).reshape(ts_emb_shape))  # (N, 20, 7, 7)
+
+            out5 = torch.cat((out2, self.up2(out4)), dim=1)  # (N, 40, 14, 14)
+            out5 = self.b5(out5 + self.te5(t).reshape(ts_emb_shape))  # (N, 10, 14, 14)
+
+            out = torch.cat((out1, self.up3(out5)), dim=1)  # (N, 20, 28, 28)
+            out = self.b_out(out + self.te_out(t).reshape(ts_emb_shape))  # (N, 1, 28, 28)
+
+            out = self.conv_out(out)
+
+            ddf_one = self.boundary_limit(out, max_sz=1 * self.max_sz)
+            if rec_id == 0:
+                ddf = ddf_one
+            else:
+                ddf = ddf_one + self.resample(ddf, ddf=ddf_one, img_sz=self.img_sz, padding_mode="border")
+            img = self.resample(x, ddf=ddf, img_sz=self.img_sz)
+
+        return ddf
+
+    def _make_te(self, dim_in, dim_out):
+        # make time embedding
+
+        return nn.Sequential(
+            nn.Linear(dim_in, dim_out),
+            # nn.SiLU(),
+            nn.ReLU(),
+            nn.Linear(dim_out, dim_out)
+        )
+
+# ==============================================
+# Conditional Network
+# ==============================================
+
+class cross_attn(nn.Module):
+    def __init__(self, q, k, v, ndims=2):
+        self.q = q
+        self.k = k
+        self.v = v
+        self.ndims = ndims
+        self.Conv = getattr(nn, 'Conv%dd' % self.ndims)
+        self.ConvT = getattr(nn, 'ConvTranspose%dd' % self.ndims)
+        self.softmax = nn.Softmax(dim=-1)
+        self.gamma = nn.Parameter(torch.zeros(1))
+    
+    def forward(self, x, y):
+        q = self.q(x)
+        k = self.k(y)
+        v = self.v(y)
+        attn = self.softmax(torch.matmul(q, k.transpose(-2, -1)))
+        out = torch.matmul(attn, v)
+        return out
+
+class DefRec_MutAttnNet(nn.Module):
+    def __init__(self, n_steps=1000, time_emb_dim=100, ndims=2, num_input_chn=1, res=0, conditional_input=True,text_feat_chn=1024, num_heads=4):
+        super(DefRec_MutAttnNet, self).__init__()
+
+        # self.feat_channels = [num_input_chn, 8, 16, 32, 32, 64]
+        # self.feat_channels = [num_input_chn, 16, 32, 64, 128, 256]
+        self.feat_channels = [num_input_chn, 16, 32, 128, 256, 512]
+        self.conditional_input = conditional_input
+        self.num_heads = num_heads
+        self.text_feat_chn = text_feat_chn
+
+        self.dimension = ndims
+        self.Conv = getattr(nn, 'Conv%dd' % self.dimension)
+        self.ConvT = getattr(nn, 'ConvTranspose%dd' % self.dimension)
+        self.copy = nn.Identity()
+        # Sinusoidal embedding
+        self.time_embed = nn.Embedding(n_steps, time_emb_dim)
+        self.time_embed.weight.data = sinusoidal_embedding(n_steps, time_emb_dim)
+        self.time_embed.requires_grad_(False)
+        self.hier_num = len(self.feat_channels) - 1
+        self.down_layers = nn.ModuleList()
+        self.up_layers = nn.ModuleList()
+        self.ted_layers = nn.ModuleList()
+        self.teu_layers = nn.ModuleList()
+        self.block_down = nn.ModuleList()
+        self.block_up = nn.ModuleList()
+        if self.conditional_input:
+            self.block_down_cond = nn.ModuleList()
+            self.fuse_conv0 = nn.ModuleList()
+            # self.fuse_conv1 = nn.ModuleList()
+            self.attn_layer = nn.MultiheadAttention(self.feat_channels[-1], self.num_heads)
+            Global_Maxpool = getattr(nn, 'AdaptiveMaxPool%dd' % self.dimension)
+            self.global_maxpool = Global_Maxpool(1)
+            self.img2txt = self.Conv(self.feat_channels[-1], self.text_feat_chn, 1, 1, 0)
+            self.txt_proc = AtrousBlock([self.text_feat_chn] + [1] * ndims, self.text_feat_chn, self.text_feat_chn, ndims=ndims, normalize=False, atrous_rates=[0, 0])
+            self.txt2img = self.Conv(self.text_feat_chn, self.feat_channels[-1], 1, 1, 0)
+            self.text = torch.zeros(1, self.text_feat_chn, *([1]*self.dimension))
+        self.img_res = [res]*self.dimension
+        self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=imsz) for imsz in self.img_res]), 0),
+                                      [1, self.dimension]+list(self.img_res))
+        
+        for i in range(1, self.hier_num + 1):
+            j=-i
+            self.down_layers.append(self.Conv(self.feat_channels[i], self.feat_channels[i], 4, 2, 1))
+            self.up_layers.append(self.ConvT(self.feat_channels[j], self.feat_channels[j], 4, 2, 1))
+            self.ted_layers.append(self._make_te(time_emb_dim, self.feat_channels[i-1]))
+            self.teu_layers.append(self._make_te(time_emb_dim, 2*self.feat_channels[j]))
+            self.block_down.append(nn.Sequential(
+                AtrousBlock([self.feat_channels[i-1]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i-1], self.feat_channels[i], ndims=ndims),
+                AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims),
+                AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims)
+            ))
+            if self.conditional_input:
+                self.block_down_cond.append(nn.Sequential(
+                    AtrousBlock([self.feat_channels[i-1]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i-1], self.feat_channels[i], ndims=ndims),
+                    AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims),
+                    AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims)
+                ))
+                self.fuse_conv0.append(self.Conv(2*self.feat_channels[i], self.feat_channels[i], 1, 1, 0))
+                # self.fuse_conv1.append(self.Conv(2*self.feat_channels[i], self.feat_channels[i], 1, 1, 0))
+            if i==self.hier_num:
+                k=j
+            else:
+                k=j-1
+            self.block_up.append(nn.Sequential(
+                AtrousBlock([2*self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, 2*self.feat_channels[j], self.feat_channels[j], ndims=ndims, normalize=False),
+                AtrousBlock([self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, self.feat_channels[j], self.feat_channels[j], ndims=ndims, normalize=False),
+                AtrousBlock([self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, self.feat_channels[j], self.feat_channels[k], ndims=ndims, normalize=False)
+            ))
+
+        # Bottleneck
+        self.tmid = self._make_te(time_emb_dim, self.feat_channels[-1])
+        self.b_mid = nn.Sequential(
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims),
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims),
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims)
+        )
+
+        self.conv_out = self.Conv(self.feat_channels[1], ndims, 3, 1, 1)
+
+    def boundary_limit(self, sample_coords0, max_sz, plus=0., minus=1.):
+        sample_coords = torch.split(sample_coords0, split_size_or_sections=1, dim=1)
+        return torch.cat([torch.clamp(x * sz, min=minus - 1 * sz + plus, max=1 * sz - minus + plus) / sz for x, sz in
+                          zip(sample_coords, max_sz)], 1)
+
+    def resample(self, vol, ddf, ref=None, img_sz=None, padding_mode="zeros"):
+        ref = self.ref_grid if ref is None else ref
+        img_sz = self.max_sz if img_sz is None else img_sz
+        resample_mode = 'bilinear' # if self.dimension==2 else 'trilinear'
+
+        return F.grid_sample(vol, torch.flip((ddf * torch.Tensor(
+            np.reshape(np.array(self.max_sz), [1, self.dimension]+[1]*self.dimension)).to(self.device) + ref).permute(
+            [0]+list(range(2,2+self.dimension))+[1]) / img_sz - 1, dims=[-1]), mode=resample_mode, padding_mode=padding_mode,
+                             align_corners=True)
+
+    def forward(self, x=None, y=None, t=None, text=None, rec_num=2, ndims=2):
+        self.device = x.device
+        img_sz = x.size()[2:]
+        n = x.size()[0]
+        self.max_sz = [img_sz[0]] * self.dimension
+        ts_emb_shape=[n,-1]+[1]*self.dimension
+        
+        self.img_sz = torch.reshape(torch.tensor([(imsz - 1) / 2 for imsz in img_sz], device=self.device), [1]*(self.dimension+1)+[self.dimension])
+        if list(img_sz) != self.img_res:
+            # print ("Reinitialize the ref_grid to match the model's input image size.")
+            # print(img_sz, self.img_res)
+            self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=imsz) for imsz in img_sz]), 0),
+                                        [1, self.dimension]+list(img_sz))
+        self.ref_grid = self.ref_grid.to(self.device)
+
+        img = x
+        if self.conditional_input:
+            tgt = y
+        # encode the conditional input
+        tgt_down_list = []
+        for i in range(self.hier_num):
+            # out = self.block_down[i](out + self.ted_layers[i](t_emb).reshape(ts_emb_shape))
+            if self.conditional_input:
+                tgt = self.block_down_cond[i](tgt)
+                tgt_down_list.append(self.copy(tgt))
+                tgt = self.down_layers[i](tgt)
+        tgt_mid = self.copy(tgt)
+        tgt_shape = tgt_mid.shape
+        # out = out.view(out_shape[0], out_shape[1], -1).permute(2, 0, 1)  # (N, C, H, W) -> (H*W, N, C)
+        tgt_mid = tgt_mid.view(tgt_shape[0], tgt_shape[1], -1).permute(2, 0, 1)  # (N, C, H, W) -> (H*W, N, C)
+
+        t = [t0.to(self.device) for t0 in t]
+        t = [t0 for _ in range(rec_num) for t0 in t]
+        for rec_id,time in enumerate(t):
+            t_emb = self.time_embed(time)
+
+        # for rec_id in range(rec_num):
+            # if self.conditional_input:
+            #     tgt = y
+            enc_list = []
+            out = img
+            for i in range(self.hier_num):
+                out = self.block_down[i](out + self.ted_layers[i](t_emb).reshape(ts_emb_shape))
+                if self.conditional_input:
+                    # tgt = self.block_down_cond[i](tgt)
+                    out = self.fuse_conv0[i](torch.cat([out, tgt_down_list[i]], axis=1))
+                    # tgt = self.fuse_conv1[i](torch.cat([tgt, out], axis=1))
+                enc_list.append(out)
+                out = self.down_layers[i](out)
+                # if self.conditional_input:
+                #     tgt = self.down_layers[i](tgt)
+            
+
+            out = self.b_mid(out + self.tmid(t_emb).reshape(ts_emb_shape))
+            if self.conditional_input:
+                # out += self.attn_layer(out, tgt, tgt)[0]
+                out_shape = out.shape
+                # tgt_shape = tgt.shape
+                # # out = out.view(out_shape[0], out_shape[1], -1).permute(2, 0, 1)  # (N, C, H, W) -> (H*W, N, C)
+                # tgt = tgt.view(tgt_shape[0], tgt_shape[1], -1).permute(2, 0, 1)  # (N, C, H, W) -> (H*W, N, C)
+                out_attn, _ = self.attn_layer(out.view(out_shape[0], out_shape[1], -1).permute(2, 0, 1), tgt_mid, tgt_mid)
+                out_attn = out_attn.permute(1, 2, 0).contiguous().view(out_shape)  # (H*W, N, C) -> (N, C, H, W)
+                out = out + out_attn
+
+            if self.conditional_input:
+                if text is None:
+                    text = self.text
+                    text = text.to(self.device)
+                out_txt = self.img2txt(out) + text
+                out_txt = self.txt_proc(out_txt)
+                out_txt = self.txt2img(out_txt)
+                out = out + out_txt
+
+            for i in range(self.hier_num):
+                out = torch.cat((self.up_layers[i](out),enc_list[-i-1]), dim=1)
+                out = self.block_up[i](out + self.teu_layers[i](t_emb).reshape(ts_emb_shape))
+
+            out = self.conv_out(out)/128
+
+            ddf_one = self.boundary_limit(out, max_sz=1 * self.max_sz)
+            if rec_id == 0:
+                ddf = ddf_one
+            else:
+                ddf = ddf_one + self.resample(ddf, ddf=ddf_one, img_sz=self.img_sz, padding_mode="border")
+            img = self.resample(x, ddf=ddf, img_sz=self.img_sz)
+
+        return ddf
+
+    def _make_te(self, dim_in, dim_out):
+        return nn.Sequential(
+            nn.Linear(dim_in, dim_out),
+            nn.ReLU(),
+            nn.Linear(dim_out, dim_out)
+        )
+
+class RecMutAttnNet1(nn.Module):
+    def __init__(self, n_steps=1000, time_emb_dim=100, ndims=2, num_input_chn=1, res=0, conditional_input=True,text_feat_chn=1024, num_heads=4):
+        super(RecMutAttnNet1, self).__init__()
+
+        # self.feat_channels = [num_input_chn, 8, 16, 32, 32, 64]
+        self.feat_channels = [num_input_chn, 16, 32, 64, 128, 256]
+        self.conditional_input = conditional_input
+        self.num_heads = num_heads
+        self.text_feat_chn = text_feat_chn
+
+        self.dimension = ndims
+        self.Conv = getattr(nn, 'Conv%dd' % self.dimension)
+        self.ConvT = getattr(nn, 'ConvTranspose%dd' % self.dimension)
+
+        # Sinusoidal embedding
+        self.time_embed = nn.Embedding(n_steps, time_emb_dim)
+        self.time_embed.weight.data = sinusoidal_embedding(n_steps, time_emb_dim)
+        self.time_embed.requires_grad_(False)
+        self.hier_num = len(self.feat_channels) - 1
+        self.down_layers = nn.ModuleList()
+        self.up_layers = nn.ModuleList()
+        self.ted_layers = nn.ModuleList()
+        self.teu_layers = nn.ModuleList()
+        self.block_down = nn.ModuleList()
+        if self.conditional_input:
+            self.block_down_cond = nn.ModuleList()
+            self.fuse_conv0 = nn.ModuleList()
+            self.fuse_conv1 = nn.ModuleList()
+            self.attn_layer = nn.MultiheadAttention(self.feat_channels[-1], self.num_heads)
+
+        self.block_up = nn.ModuleList()
+
+        for i in range(1, self.hier_num + 1):
+            j=-i
+            self.down_layers.append(self.Conv(self.feat_channels[i], self.feat_channels[i], 4, 2, 1))
+            self.up_layers.append(self.ConvT(self.feat_channels[j], self.feat_channels[j], 4, 2, 1))
+            self.ted_layers.append(self._make_te(time_emb_dim, self.feat_channels[i-1]))
+            self.teu_layers.append(self._make_te(time_emb_dim, 2*self.feat_channels[j]))
+            self.block_down.append(nn.Sequential(
+                AtrousBlock([self.feat_channels[i-1]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i-1], self.feat_channels[i], ndims=ndims),
+                AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims),
+                AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims)
+            ))
+            if self.conditional_input:
+                self.block_down_cond.append(nn.Sequential(
+                    AtrousBlock([self.feat_channels[i-1]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i-1], self.feat_channels[i], ndims=ndims),
+                    AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims),
+                    AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims)
+                ))
+                self.fuse_conv0.append(self.Conv(2*self.feat_channels[i], self.feat_channels[i], 1, 1, 0))
+                self.fuse_conv1.append(self.Conv(2*self.feat_channels[i], self.feat_channels[i], 1, 1, 0))
+            if i==self.hier_num:
+                k=j
+            else:
+                k=j-1
+            self.block_up.append(nn.Sequential(
+                AtrousBlock([2*self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, 2*self.feat_channels[j], self.feat_channels[j], ndims=ndims, normalize=False),
+                AtrousBlock([self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, self.feat_channels[j], self.feat_channels[j], ndims=ndims, normalize=False),
+                AtrousBlock([self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, self.feat_channels[j], self.feat_channels[k], ndims=ndims, normalize=False)
+            ))
+
+        # Bottleneck
+        self.tmid = self._make_te(time_emb_dim, self.feat_channels[-1])
+        self.b_mid = nn.Sequential(
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims),
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims),
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims)
+        )
+
+        self.conv_out = self.Conv(self.feat_channels[1], ndims, 3, 1, 1)
+
+    def boundary_limit(self, sample_coords0, max_sz, plus=0., minus=1.):
+        sample_coords = torch.split(sample_coords0, split_size_or_sections=1, dim=1)
+        return torch.cat([torch.clamp(x * sz, min=minus - 1 * sz + plus, max=1 * sz - minus + plus) / sz for x, sz in
+                          zip(sample_coords, max_sz)], 1)
+
+    def resample(self, vol, ddf, ref=None, img_sz=None, padding_mode="zeros"):
+        ref = self.ref_grid if ref is None else ref
+        img_sz = self.max_sz if img_sz is None else img_sz
+        resample_mode = 'bilinear' # if self.dimension==2 else 'trilinear'
+
+        return F.grid_sample(vol, torch.flip((ddf * torch.Tensor(
+            np.reshape(np.array(self.max_sz), [1, self.dimension]+[1]*self.dimension)).to(self.device) + ref).permute(
+            [0]+list(range(2,2+self.dimension))+[1]) / img_sz - 1, dims=[-1]), mode=resample_mode, padding_mode=padding_mode,
+                             align_corners=True)
+
+    def forward(self, x=None, y=None, t=None, rec_num=2, ndims=2):
+        self.device = x.device
+        img_sz = x.size()[2:]
+        n = x.size()[0]
+        self.max_sz = [img_sz[0]] * self.dimension
+        ts_emb_shape=[n,-1]+[1]*self.dimension
+        self.img_sz = torch.reshape(torch.tensor([(imsz - 1) / 2 for imsz in img_sz], device=self.device), [1]*(self.dimension+1)+[self.dimension])
+        self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=imsz) for imsz in img_sz]), 0),
+                                      [1, self.dimension]+list(img_sz)).to(self.device)
+        img = x
+        t = self.time_embed(t)
+
+        for rec_id in range(rec_num):
+            if self.conditional_input:
+                tgt = y
+            enc_list = []
+            out = img
+            for i in range(self.hier_num):
+                out = self.block_down[i](out + self.ted_layers[i](t).reshape(ts_emb_shape))
+                if self.conditional_input:
+                    tgt = self.block_down_cond[i](tgt)
+                    out = self.fuse_conv0[i](torch.cat([out, tgt], axis=1))
+                    tgt = self.fuse_conv1[i](torch.cat([tgt, out], axis=1))
+                enc_list.append(out)
+                out = self.down_layers[i](out)
+                if self.conditional_input:
+                    tgt = self.down_layers[i](tgt)
+
+            out = self.b_mid(out + self.tmid(t).reshape(ts_emb_shape))
+            if self.conditional_input:
+                # out += self.attn_layer(out, tgt, tgt)[0]
+                out_shape = out.shape
+                tgt_shape = tgt.shape
+                # out = out.view(out_shape[0], out_shape[1], -1).permute(2, 0, 1)  # (N, C, H, W) -> (H*W, N, C)
+                tgt = tgt.view(tgt_shape[0], tgt_shape[1], -1).permute(2, 0, 1)  # (N, C, H, W) -> (H*W, N, C)
+                out_attn, _ = self.attn_layer(out.view(out_shape[0], out_shape[1], -1).permute(2, 0, 1), tgt, tgt)
+                out_attn = out_attn.permute(1, 2, 0).contiguous().view(out_shape)  # (H*W, N, C) -> (N, C, H, W)
+                out = out + out_attn
+
+            for i in range(self.hier_num):
+                out = torch.cat((self.up_layers[i](out),enc_list[-i-1]), dim=1)
+                out = self.block_up[i](out + self.teu_layers[i](t).reshape(ts_emb_shape))
+
+            out = self.conv_out(out)/128
+
+            ddf_one = self.boundary_limit(out, max_sz=1 * self.max_sz)
+            if rec_id == 0:
+                ddf = ddf_one
+            else:
+                ddf = ddf_one + self.resample(ddf, ddf=ddf_one, img_sz=self.img_sz, padding_mode="border")
+            img = self.resample(x, ddf=ddf, img_sz=self.img_sz)
+
+        return ddf
+
+    def _make_te(self, dim_in, dim_out):
+        return nn.Sequential(
+            nn.Linear(dim_in, dim_out),
+            nn.ReLU(),
+            nn.Linear(dim_out, dim_out)
+        )
+
+class RecMutAttnNet(nn.Module):
+    def __init__(self, n_steps=1000, time_emb_dim=100, ndims=2, num_input_chn=1, res=0, conditional_input=True,text_feat_chn=1024, num_heads=4):
+        super(RecMutAttnNet, self).__init__()
+
+        # self.feat_channels = [num_input_chn, 8, 16, 32, 32, 64]
+        self.feat_channels = [num_input_chn, 16, 32, 64, 128, 256]
+        self.conditional_input = conditional_input
+        self.num_heads = num_heads
+        self.text_feat_chn = text_feat_chn
+
+        self.dimension = ndims
+        self.Conv = getattr(nn, 'Conv%dd' % self.dimension)
+        self.ConvT = getattr(nn, 'ConvTranspose%dd' % self.dimension)
+
+        # Sinusoidal embedding
+        self.time_embed = nn.Embedding(n_steps, time_emb_dim)
+        self.time_embed.weight.data = sinusoidal_embedding(n_steps, time_emb_dim)
+        self.time_embed.requires_grad_(False)
+        self.hier_num = len(self.feat_channels) - 1
+        self.down_layers = nn.ModuleList()
+        self.up_layers = nn.ModuleList()
+        self.ted_layers = nn.ModuleList()
+        self.teu_layers = nn.ModuleList()
+        self.block_down = nn.ModuleList()
+        self.block_up = nn.ModuleList()
+        if self.conditional_input:
+            self.block_down_cond = nn.ModuleList()
+            self.fuse_conv0 = nn.ModuleList()
+            self.fuse_conv1 = nn.ModuleList()
+            self.attn_layer = nn.MultiheadAttention(self.feat_channels[-1], self.num_heads)
+            Global_Maxpool = getattr(nn, 'AdaptiveMaxPool%dd' % self.dimension)
+            self.global_maxpool = Global_Maxpool(1)
+            self.img2txt = self.Conv(self.feat_channels[-1], self.text_feat_chn, 1, 1, 0)
+            self.txt_proc = AtrousBlock([self.text_feat_chn] + [1] * ndims, self.text_feat_chn, self.text_feat_chn, ndims=ndims, normalize=False, atrous_rates=[0, 0])
+            self.txt2img = self.Conv(self.text_feat_chn, self.feat_channels[-1], 1, 1, 0)
+            self.text = torch.zeros(1, self.text_feat_chn, *([1]*self.dimension))
+        self.img_res = [res]*self.dimension
+        self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=imsz) for imsz in self.img_res]), 0),
+                                      [1, self.dimension]+list(self.img_res))
+        
+        for i in range(1, self.hier_num + 1):
+            j=-i
+            self.down_layers.append(self.Conv(self.feat_channels[i], self.feat_channels[i], 4, 2, 1))
+            self.up_layers.append(self.ConvT(self.feat_channels[j], self.feat_channels[j], 4, 2, 1))
+            self.ted_layers.append(self._make_te(time_emb_dim, self.feat_channels[i-1]))
+            self.teu_layers.append(self._make_te(time_emb_dim, 2*self.feat_channels[j]))
+            self.block_down.append(nn.Sequential(
+                AtrousBlock([self.feat_channels[i-1]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i-1], self.feat_channels[i], ndims=ndims),
+                AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims),
+                AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims)
+            ))
+            if self.conditional_input:
+                self.block_down_cond.append(nn.Sequential(
+                    AtrousBlock([self.feat_channels[i-1]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i-1], self.feat_channels[i], ndims=ndims),
+                    AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims),
+                    AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims)
+                ))
+                self.fuse_conv0.append(self.Conv(2*self.feat_channels[i], self.feat_channels[i], 1, 1, 0))
+                self.fuse_conv1.append(self.Conv(2*self.feat_channels[i], self.feat_channels[i], 1, 1, 0))
+            if i==self.hier_num:
+                k=j
+            else:
+                k=j-1
+            self.block_up.append(nn.Sequential(
+                AtrousBlock([2*self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, 2*self.feat_channels[j], self.feat_channels[j], ndims=ndims, normalize=False),
+                AtrousBlock([self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, self.feat_channels[j], self.feat_channels[j], ndims=ndims, normalize=False),
+                AtrousBlock([self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, self.feat_channels[j], self.feat_channels[k], ndims=ndims, normalize=False)
+            ))
+
+        # Bottleneck
+        self.tmid = self._make_te(time_emb_dim, self.feat_channels[-1])
+        self.b_mid = nn.Sequential(
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims),
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims),
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims)
+        )
+
+        self.conv_out = self.Conv(self.feat_channels[1], ndims, 3, 1, 1)
+
+    def boundary_limit(self, sample_coords0, max_sz, plus=0., minus=1.):
+        sample_coords = torch.split(sample_coords0, split_size_or_sections=1, dim=1)
+        return torch.cat([torch.clamp(x * sz, min=minus - 1 * sz + plus, max=1 * sz - minus + plus) / sz for x, sz in
+                          zip(sample_coords, max_sz)], 1)
+
+    def resample(self, vol, ddf, ref=None, img_sz=None, padding_mode="zeros"):
+        ref = self.ref_grid if ref is None else ref
+        img_sz = self.max_sz if img_sz is None else img_sz
+        resample_mode = 'bilinear' # if self.dimension==2 else 'trilinear'
+
+        return F.grid_sample(vol, torch.flip((ddf * torch.Tensor(
+            np.reshape(np.array(self.max_sz), [1, self.dimension]+[1]*self.dimension)).to(self.device) + ref).permute(
+            [0]+list(range(2,2+self.dimension))+[1]) / img_sz - 1, dims=[-1]), mode=resample_mode, padding_mode=padding_mode,
+                             align_corners=True)
+
+    def forward(self, x=None, y=None, t=None, text=None, rec_num=2, ndims=2):
+        self.device = x.device
+        img_sz = x.size()[2:]
+        n = x.size()[0]
+        self.max_sz = [img_sz[0]] * self.dimension
+        ts_emb_shape=[n,-1]+[1]*self.dimension
+        
+        self.img_sz = torch.reshape(torch.tensor([(imsz - 1) / 2 for imsz in img_sz], device=self.device), [1]*(self.dimension+1)+[self.dimension])
+        if list(img_sz) != self.img_res:
+            # print ("Reinitialize the ref_grid to match the model's input image size.")
+            # print(img_sz, self.img_res)
+            self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=imsz) for imsz in img_sz]), 0),
+                                        [1, self.dimension]+list(img_sz))
+        self.ref_grid = self.ref_grid.to(self.device)
+
+        img = x
+        t = self.time_embed(t)
+
+        for rec_id in range(rec_num):
+            if self.conditional_input:
+                tgt = y
+            enc_list = []
+            out = img
+            for i in range(self.hier_num):
+                out = self.block_down[i](out + self.ted_layers[i](t).reshape(ts_emb_shape))
+                if self.conditional_input:
+                    tgt = self.block_down_cond[i](tgt)
+                    out = self.fuse_conv0[i](torch.cat([out, tgt], axis=1))
+                    tgt = self.fuse_conv1[i](torch.cat([tgt, out], axis=1))
+                enc_list.append(out)
+                out = self.down_layers[i](out)
+                if self.conditional_input:
+                    tgt = self.down_layers[i](tgt)
+            
+
+            out = self.b_mid(out + self.tmid(t).reshape(ts_emb_shape))
+            if self.conditional_input:
+                # out += self.attn_layer(out, tgt, tgt)[0]
+                out_shape = out.shape
+                tgt_shape = tgt.shape
+                # out = out.view(out_shape[0], out_shape[1], -1).permute(2, 0, 1)  # (N, C, H, W) -> (H*W, N, C)
+                tgt = tgt.view(tgt_shape[0], tgt_shape[1], -1).permute(2, 0, 1)  # (N, C, H, W) -> (H*W, N, C)
+                out_attn, _ = self.attn_layer(out.view(out_shape[0], out_shape[1], -1).permute(2, 0, 1), tgt, tgt)
+                out_attn = out_attn.permute(1, 2, 0).contiguous().view(out_shape)  # (H*W, N, C) -> (N, C, H, W)
+                out = out + out_attn
+
+            if self.conditional_input:
+                if text is None:
+                    text = self.text
+                    text = text.to(self.device)
+                text = text.view(-1, self.text_feat_chn, *([1]*self.dimension))
+                out_txt = self.img2txt(out) + text
+                out_txt = self.txt_proc(out_txt)
+                out_txt = self.txt2img(out_txt)
+                out = out + out_txt
+
+            for i in range(self.hier_num):
+                out = torch.cat((self.up_layers[i](out),enc_list[-i-1]), dim=1)
+                out = self.block_up[i](out + self.teu_layers[i](t).reshape(ts_emb_shape))
+
+            out = self.conv_out(out)/128
+
+            ddf_one = self.boundary_limit(out, max_sz=1 * self.max_sz)
+            if rec_id == 0:
+                ddf = ddf_one
+            else:
+                ddf = ddf_one + self.resample(ddf, ddf=ddf_one, img_sz=self.img_sz, padding_mode="border")
+            img = self.resample(x, ddf=ddf, img_sz=self.img_sz)
+
+        return ddf
+
+    def _make_te(self, dim_in, dim_out):
+        return nn.Sequential(
+            nn.Linear(dim_in, dim_out),
+            nn.ReLU(),
+            nn.Linear(dim_out, dim_out)
+        )
+
+class RecMutAttnNet_contrastive(nn.Module):
+    def __init__(self, n_steps=1000, time_emb_dim=100, ndims=2, num_input_chn=1, res=0, conditional_input=True,text_feat_chn=1024, num_heads=4):
+        super(RecMutAttnNet_contrastive, self).__init__()
+
+        # self.feat_channels = [num_input_chn, 8, 16, 32, 32, 64]
+        self.feat_channels = [num_input_chn, 16, 32, 64, 128, 256]
+        self.conditional_input = conditional_input
+        self.num_heads = num_heads
+        self.text_feat_chn = text_feat_chn
+
+        self.dimension = ndims
+        self.Conv = getattr(nn, 'Conv%dd' % self.dimension)
+        self.ConvT = getattr(nn, 'ConvTranspose%dd' % self.dimension)
+
+        # Sinusoidal embedding
+        self.time_embed = nn.Embedding(n_steps, time_emb_dim)
+        self.time_embed.weight.data = sinusoidal_embedding(n_steps, time_emb_dim)
+        self.time_embed.requires_grad_(False)
+        self.hier_num = len(self.feat_channels) - 1
+        self.down_layers = nn.ModuleList()
+        self.up_layers = nn.ModuleList()
+        self.ted_layers = nn.ModuleList()
+        self.teu_layers = nn.ModuleList()
+        self.block_down = nn.ModuleList()
+        self.block_up = nn.ModuleList()
+        if self.conditional_input:
+            self.block_down_cond = nn.ModuleList()
+            self.fuse_conv0 = nn.ModuleList()
+            self.fuse_conv1 = nn.ModuleList()
+            self.attn_layer = nn.MultiheadAttention(self.feat_channels[-1], self.num_heads)
+            Global_Maxpool = getattr(nn, 'AdaptiveMaxPool%dd' % self.dimension)
+            self.global_maxpool = Global_Maxpool(1)
+            self.img2txt = self.Conv(self.feat_channels[-1], self.text_feat_chn, 1, 1, 0)
+            self.txt_proc = AtrousBlock([self.text_feat_chn] + [1] * ndims, self.text_feat_chn, self.text_feat_chn, ndims=ndims, normalize=False, atrous_rates=[0, 0])
+            self.txt2img = self.Conv(self.text_feat_chn, self.feat_channels[-1], 1, 1, 0)
+            self.text = torch.zeros(1, self.text_feat_chn, *([1]*self.dimension))
+        self.img_res = [res]*self.dimension
+        self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=imsz) for imsz in self.img_res]), 0),
+                                      [1, self.dimension]+list(self.img_res))
+        
+        for i in range(1, self.hier_num + 1):
+            j=-i
+            self.down_layers.append(self.Conv(self.feat_channels[i], self.feat_channels[i], 4, 2, 1))
+            self.up_layers.append(self.ConvT(self.feat_channels[j], self.feat_channels[j], 4, 2, 1))
+            self.ted_layers.append(self._make_te(time_emb_dim, self.feat_channels[i-1]))
+            self.teu_layers.append(self._make_te(time_emb_dim, 2*self.feat_channels[j]))
+            self.block_down.append(nn.Sequential(
+                AtrousBlock([self.feat_channels[i-1]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i-1], self.feat_channels[i], ndims=ndims),
+                AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims),
+                AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims)
+            ))
+            if self.conditional_input:
+                self.block_down_cond.append(nn.Sequential(
+                    AtrousBlock([self.feat_channels[i-1]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i-1], self.feat_channels[i], ndims=ndims),
+                    AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims),
+                    AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims)
+                ))
+                self.fuse_conv0.append(self.Conv(2*self.feat_channels[i], self.feat_channels[i], 1, 1, 0))
+                self.fuse_conv1.append(self.Conv(2*self.feat_channels[i], self.feat_channels[i], 1, 1, 0))
+            if i==self.hier_num:
+                k=j
+            else:
+                k=j-1
+            self.block_up.append(nn.Sequential(
+                AtrousBlock([2*self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, 2*self.feat_channels[j], self.feat_channels[j], ndims=ndims, normalize=False),
+                AtrousBlock([self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, self.feat_channels[j], self.feat_channels[j], ndims=ndims, normalize=False),
+                AtrousBlock([self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, self.feat_channels[j], self.feat_channels[k], ndims=ndims, normalize=False)
+            ))
+
+        # Bottleneck
+        self.tmid = self._make_te(time_emb_dim, self.feat_channels[-1])
+        self.b_mid = nn.Sequential(
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims),
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims),
+            AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims)
+        )
+
+        self.conv_out = self.Conv(self.feat_channels[1], ndims, 3, 1, 1)
+
+    def boundary_limit(self, sample_coords0, max_sz, plus=0., minus=1.):
+        sample_coords = torch.split(sample_coords0, split_size_or_sections=1, dim=1)
+        return torch.cat([torch.clamp(x * sz, min=minus - 1 * sz + plus, max=1 * sz - minus + plus) / sz for x, sz in
+                          zip(sample_coords, max_sz)], 1)
+
+    def resample(self, vol, ddf, ref=None, img_sz=None, padding_mode="zeros"):
+        ref = self.ref_grid if ref is None else ref
+        img_sz = self.max_sz if img_sz is None else img_sz
+        resample_mode = 'bilinear' # if self.dimension==2 else 'trilinear'
+
+        return F.grid_sample(vol, torch.flip((ddf * torch.Tensor(
+            np.reshape(np.array(self.max_sz), [1, self.dimension]+[1]*self.dimension)).to(self.device) + ref).permute(
+            [0]+list(range(2,2+self.dimension))+[1]) / img_sz - 1, dims=[-1]), mode=resample_mode, padding_mode=padding_mode,
+                             align_corners=True)
+
+    def forward(self, x=None, y=None, t=None, text=None, rec_num=2, ndims=2):
+        self.device = x.device
+        img_sz = x.size()[2:]
+        n = x.size()[0]
+        self.max_sz = [img_sz[0]] * self.dimension
+        ts_emb_shape=[n,-1]+[1]*self.dimension
+        
+        self.img_sz = torch.reshape(torch.tensor([(imsz - 1) / 2 for imsz in img_sz], device=self.device), [1]*(self.dimension+1)+[self.dimension])
+        if list(img_sz) != self.img_res:
+            # print ("Reinitialize the ref_grid to match the model's input image size.")
+            # print(img_sz, self.img_res)
+            self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=imsz) for imsz in img_sz]), 0),
+                                        [1, self.dimension]+list(img_sz))
+        self.ref_grid = self.ref_grid.to(self.device)
+
+        img = x
+        t = self.time_embed(t)
+
+        for rec_id in range(rec_num):
+            if self.conditional_input:
+                tgt = y
+            enc_list = []
+            out = img
+            for i in range(self.hier_num):
+                out = self.block_down[i](out + self.ted_layers[i](t).reshape(ts_emb_shape))
+                if self.conditional_input:
+                    tgt = self.block_down_cond[i](tgt)
+                    out = self.fuse_conv0[i](torch.cat([out, tgt], axis=1))
+                    tgt = self.fuse_conv1[i](torch.cat([tgt, out], axis=1))
+                enc_list.append(out)
+                out = self.down_layers[i](out)
+                if self.conditional_input:
+                    tgt = self.down_layers[i](tgt)
+            
+
+            out = self.b_mid(out + self.tmid(t).reshape(ts_emb_shape))
+            if self.conditional_input:
+                # out += self.attn_layer(out, tgt, tgt)[0]
+                out_shape = out.shape
+                tgt_shape = tgt.shape
+                # out = out.view(out_shape[0], out_shape[1], -1).permute(2, 0, 1)  # (N, C, H, W) -> (H*W, N, C)
+                tgt = tgt.view(tgt_shape[0], tgt_shape[1], -1).permute(2, 0, 1)  # (N, C, H, W) -> (H*W, N, C)
+                out_attn, _ = self.attn_layer(out.view(out_shape[0], out_shape[1], -1).permute(2, 0, 1), tgt, tgt)
+                out_attn = out_attn.permute(1, 2, 0).contiguous().view(out_shape)  # (H*W, N, C) -> (N, C, H, W)
+                out = out + out_attn
+
+            if self.conditional_input:
+                if text is None:
+                    text = self.text
+                    text = text.to(self.device)
+                text = text.view(-1, self.text_feat_chn, *([1]*self.dimension))
+                img_embd = self.global_maxpool(self.img2txt(out)).view(n, -1)  # [B, 1024]
+                out_txt = self.img2txt(out) + text
+                out_txt = self.txt_proc(out_txt)
+                out_txt = self.txt2img(out_txt)
+                out = out + out_txt
+
+            for i in range(self.hier_num):
+                out = torch.cat((self.up_layers[i](out),enc_list[-i-1]), dim=1)
+                out = self.block_up[i](out + self.teu_layers[i](t).reshape(ts_emb_shape))
+
+            out = self.conv_out(out)/128
+
+            ddf_one = self.boundary_limit(out, max_sz=1 * self.max_sz)
+            if rec_id == 0:
+                ddf = ddf_one
+            else:
+                ddf = ddf_one + self.resample(ddf, ddf=ddf_one, img_sz=self.img_sz, padding_mode="border")
+            img = self.resample(x, ddf=ddf, img_sz=self.img_sz)
+
+        return ddf, img_embd
+
+    def _make_te(self, dim_in, dim_out):
+        return nn.Sequential(
+            nn.Linear(dim_in, dim_out),
+            nn.ReLU(),
+            nn.Linear(dim_out, dim_out)
+        )
+# class RecMutAttnNet(nn.Module):
+#     def __init__(self, n_steps=1000, time_emb_dim=100, ndims=2, num_input_chn=1, res=0, conditional_input=True):
+#         super(RecMutAttnNet, self).__init__()
+
+#         self.feat_channels = [num_input_chn, 8, 16, 32, 32, 64]
+#         self.conditional_input = conditional_input
+        
+#         self.dimension = ndims
+#         self.Conv = getattr(nn, 'Conv%dd' % self.dimension)
+#         self.ConvT = getattr(nn, 'ConvTranspose%dd' % self.dimension)
+
+#         # Sinusoidal embedding
+#         self.time_embed = nn.Embedding(n_steps, time_emb_dim)
+#         self.time_embed.weight.data = sinusoidal_embedding(n_steps, time_emb_dim)
+#         self.time_embed.requires_grad_(False)
+#         self.hier_num = len(self.feat_channels) - 1
+#         self.down_layers = nn.ModuleList()
+#         self.up_layers = nn.ModuleList()
+#         self.ted_layers = nn.ModuleList()
+#         self.teu_layers = nn.ModuleList()
+#         self.block_down = nn.ModuleList()
+#         if self.conditional_input:
+#             self.block_down_cond = nn.ModuleList()
+#             self.fuse_conv0 = nn.ModuleList()
+#             self.fuse_conv1 = nn.ModuleList()
+#         self.block_up = nn.ModuleList()
+
+#         for i in range(1, self.hier_num + 1):
+#             j=-i
+#             self.down_layers.append(self.Conv(self.feat_channels[i], self.feat_channels[i], 4, 2, 1))
+#             self.up_layers.append(self.ConvT(self.feat_channels[j], self.feat_channels[j], 4, 2, 1))
+#             self.ted_layers.append(self._make_te(time_emb_dim, self.feat_channels[i-1]))
+#             self.teu_layers.append(self._make_te(time_emb_dim, 2*self.feat_channels[j]))
+#             self.block_down.append(nn.Sequential(
+#                 AtrousBlock([self.feat_channels[i-1]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i-1], self.feat_channels[i], ndims=ndims),
+#                 AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims),
+#                 AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims)
+#             ))
+#             if self.conditional_input:
+#                 self.block_down_cond.append(nn.Sequential(
+#                     AtrousBlock([self.feat_channels[i-1]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i-1], self.feat_channels[i], ndims=ndims),
+#                     AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims),
+#                     AtrousBlock([self.feat_channels[i]] + [res // (2 ** (i-1))] * ndims, self.feat_channels[i], self.feat_channels[i], ndims=ndims)
+#                 ))
+#                 self.fuse_conv0.append(self.Conv(2*self.feat_channels[i], self.feat_channels[i], 1, 1, 0))
+#                 self.fuse_conv1.append(self.Conv(2*self.feat_channels[i], self.feat_channels[i], 1, 1, 0))
+#             if i==self.hier_num:
+#                 k=j
+#             else:
+#                 k=j-1
+#             self.block_up.append(nn.Sequential(
+#                 AtrousBlock([2*self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, 2*self.feat_channels[j], self.feat_channels[j], ndims=ndims, normalize=False),
+#                 AtrousBlock([self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, self.feat_channels[j], self.feat_channels[j], ndims=ndims, normalize=False),
+#                 AtrousBlock([self.feat_channels[j]] + [res // (2 ** (self.hier_num-i-1))] * ndims, self.feat_channels[j], self.feat_channels[k], ndims=ndims, normalize=False)
+#             ))
+
+#         # Bottleneck
+#         self.tmid = self._make_te(time_emb_dim, self.feat_channels[-1])
+#         self.b_mid = nn.Sequential(
+#             AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims),
+#             AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims),
+#             AtrousBlock([self.feat_channels[self.hier_num]] + [res // (2**self.hier_num)] * ndims, self.feat_channels[self.hier_num], self.feat_channels[self.hier_num], ndims=ndims)
+#         )
+
+#         self.conv_out = self.Conv(self.feat_channels[1], ndims, 3, 1, 1)
+
+#     def boundary_limit(self, sample_coords0, max_sz, plus=0., minus=1.):
+#         sample_coords = torch.split(sample_coords0, split_size_or_sections=1, dim=1)
+#         return torch.cat([torch.clamp(x * sz, min=minus - 1 * sz + plus, max=1 * sz - minus + plus) / sz for x, sz in
+#                           zip(sample_coords, max_sz)], 1)
+
+#     def resample(self, vol, ddf, ref=None, img_sz=None, padding_mode="zeros"):
+#         ref = self.ref_grid if ref is None else ref
+#         img_sz = self.max_sz if img_sz is None else img_sz
+#         resample_mode = 'bilinear' # if self.dimension==2 else 'trilinear'
+
+#         return F.grid_sample(vol, torch.flip((ddf * torch.Tensor(
+#             np.reshape(np.array(self.max_sz), [1, self.dimension]+[1]*self.dimension)).to(self.device) + ref).permute(
+#             [0]+list(range(2,2+self.dimension))+[1]) / img_sz - 1, dims=[-1]), mode=resample_mode, padding_mode=padding_mode,
+#                              align_corners=True)
+
+#     def forward(self, x=None, y=None, t=None, rec_num=2, ndims=2):
+#         self.device = x.device
+#         img_sz = x.size()[2:]
+#         n = x.size()[0]
+#         self.max_sz = [img_sz[0]] * self.dimension
+#         ts_emb_shape=[n,-1]+[1]*self.dimension
+#         self.img_sz = torch.reshape(torch.tensor([(imsz - 1) / 2 for imsz in img_sz], device=self.device), [1]*(self.dimension+1)+[self.dimension])
+#         self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=imsz) for imsz in img_sz]), 0),
+#                                       [1, self.dimension]+list(img_sz)).to(self.device)
+#         img = x
+#         t = self.time_embed(t)
+
+#         for rec_id in range(rec_num):
+#             if self.conditional_input:
+#                 tgt = y
+#             enc_list = []
+#             out = img
+#             for i in range(self.hier_num):
+#                 out = self.block_down[i](out + self.ted_layers[i](t).reshape(ts_emb_shape))
+#                 if self.conditional_input:
+#                     tgt = self.block_down_cond[i](tgt)
+#                     out = self.fuse_conv0[i](torch.cat([out, tgt], axis=1))
+#                     tgt = self.fuse_conv1[i](torch.cat([tgt, out], axis=1))
+#                 enc_list.append(out)
+#                 out = self.down_layers[i](out)
+#                 if self.conditional_input:
+#                     tgt = self.down_layers[i](tgt)
+
+#             out = self.b_mid(out + self.tmid(t).reshape(ts_emb_shape))
+#             if self.conditional_input:
+#                 out = out + tgt
+            
+#             for i in range(self.hier_num):
+#                 out = torch.cat((self.up_layers[i](out),enc_list[-i-1]), dim=1)
+#                 out = self.block_up[i](out + self.teu_layers[i](t).reshape(ts_emb_shape))
+
+#             out = self.conv_out(out)/128
+
+#             ddf_one = self.boundary_limit(out, max_sz=1 * self.max_sz)
+#             if rec_id == 0:
+#                 ddf = ddf_one
+#             else:
+#                 ddf = ddf_one + self.resample(ddf, ddf=ddf_one, img_sz=self.img_sz, padding_mode="border")
+#             img = self.resample(x, ddf=ddf, img_sz=self.img_sz)
+
+#         return ddf
+
+#     def _make_te(self, dim_in, dim_out):
+#         return nn.Sequential(
+#             nn.Linear(dim_in, dim_out),
+#             nn.ReLU(),
+#             nn.Linear(dim_out, dim_out)
+#         )
+# ==============================================
+# Layers
+# ==============================================
+
+
+def ddf_multiplier(dvf,mul_num=10,stn=None):
+    ddf=dvf
+    for i in range(mul_num):
+        ddf = dvf + stn(ddf, dvf)
+    return ddf
+
+
+def composite(ddfs,stn=None):
+    if stn is None:
+        stn = STN(device=ddfs[0].device,padding_mode="border")
+    comp_ddf=ddfs[0]
+    for i in range(1,len(ddfs)):
+        comp_ddf = ddfs[i] + stn(comp_ddf,ddfs[i])
+    return comp_ddf
+
+class STN(nn.Module):
+    def __init__(self,ndims=2,img_sz=None,max_sz=None,device=None,padding_mode="border",resample_mode=None):
+        super(STN, self).__init__()
+        self.ndims=ndims
+        self.img_sz=[img_sz]*ndims
+        # self.img_sz=img_sz
+        self.device = device
+        self.padding_mode = padding_mode
+        # max_sz=[128]*self.ndims
+        max_sz=[img_sz]*self.ndims
+        # max_sz=img_sz
+        # max_sz=img_sz if max_sz is None else ([128,128] if img_sz is None else img_sz)
+        # self.max_sz=torch.Tensor(np.reshape(np.array(max_sz), [1, self.ndims, 1, 1])).to(self.device)
+        self.max_sz=torch.Tensor(np.reshape(np.array(max_sz), [1, self.ndims]+[1]*self.ndims)).to(self.device)
+        self.resample_mode=resample_mode
+        if self.img_sz is not None:
+            self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=s) for s in self.img_sz]), 0),
+                                        [1, self.ndims] + self.img_sz).to(self.device)
+        return
+    def max_limit(self, sample_coords0, plus=0., minus=1.):
+        sample_coords = torch.split(sample_coords0, split_size_or_sections=1, dim=1)
+        # return tf.stack([tf.maximum(tf.minimum(x, sz - minus + plus), 0 + plus) for x, sz in zip(sample_coords, input_size0)],-1)
+        return torch.cat([torch.clamp(x * sz, min=minus - 1 * sz + plus, max=1 * sz - minus + plus) / sz for x, sz in
+                        zip(sample_coords, self.max_sz)], 1)
+
+    def boundary_limit(self, sample_coords0, plus=0., minus=1.):
+
+        sample_coords = torch.split(sample_coords0, split_size_or_sections=1, dim=1)
+        # return tf.stack([tf.maximum(tf.minimum(x, sz - minus + plus), 0 + plus) for x, sz in zip(sample_coords, input_size0)],-1)
+        return torch.cat([(torch.clamp(x * sz+ref, min=minus - 1 * sz + plus, max=1 * sz - minus + plus)-ref) / sz for x, sz,ref in
+                        zip(sample_coords, self.max_sz, self.ref_grid)], 1)
+
+    def resample(self, vol, ddf, ref=None, img_sz=None,padding_mode = "zeros"):
+        # print(vol.device, ddf.device)
+        # print(self.device)
+        # print('===================')
+        device = ddf.device
+
+        ref = self.ref_grid if ref is None else ref
+        if img_sz is None:
+            img_sz = self.max_sz
+        else:
+            img_sz = torch.reshape(torch.tensor([(s - 1) / 2. for s in img_sz], device=device), [1]+[1]*self.ndims+[self.ndims])
+        # resample_mode = 'bicubic'
+        if self.resample_mode is None:
+            resample_mode = 'bilinear' # if self.ndims==2 else 'trilinear'
+        else:
+            resample_mode=self.resample_mode
+        # padding_mode = "border"
+        # print(ddf.shape, ref.shape)
+        return F.grid_sample(vol.to(device), torch.flip((ddf * self.max_sz.to(device) + ref.to(device)).permute(
+            [0] + list(range(2, 2 + self.ndims)) + [1]) / img_sz - 1, dims=[-1]), mode=resample_mode,
+                            padding_mode=padding_mode,
+                            align_corners=True)
+
+    def forward(self,x,ddf):
+        self.device = x.device if self.device is None else self.device
+        if self.img_sz is None:
+            self.img_sz = list(x.size()[2:]).to(self.device)
+            self.ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=s) for s in self.img_sz]), 0),[1, self.ndims]+self.img_sz).to(self.device)
+        resampled_x = self.resample(x, ddf=ddf, img_sz=self.img_sz, padding_mode=self.padding_mode)
+        return resampled_x
+
+if __name__ == '__main__':
+    ndims = 3
+    res = 128
+    x = torch.rand([1, 1] + [res]*ndims)
+    t = torch.randint(0, 1000, (1,))
+    text = torch.rand([1, 1024] + [1]*ndims)
+    model = RecMutAttnNet(n_steps=1000, time_emb_dim=100, ndims=ndims, num_input_chn=1, res=res, conditional_input=True)
+    y = model(x, x, t, text=text)
+    print("Ouput shape", y.shape)
+
+    # Total parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    # Trainable parameters only
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    print(f"Total parameters: {total_params}")
+    print(f"Trainable parameters: {trainable_params}")

Diffusion/utils_diff.py

@@ -0,0 +1,477 @@
+import os
+import torch
+import torchvision
+from torch import nn, optim
+from torch.autograd.variable import Variable
+from torchvision import transforms, datasets
+from torchvision.utils import save_image
+import torch.nn.functional as F
+import scipy.ndimage as spimg
+import pyquaternion as quater
+import random
+import numpy as np
+import math
+from typing import Optional, Tuple, List
+# from data_loader.acdc_dataloader import acdc_gan
+
+# from Adaptive_Motion_Generator.Dataloader.Archive.acdc_dataloader import *
+
+def get_barcode(index=[],header=['Patient','Slice','AugImg','NoiseStep'],digit=[4,6,4,4],split='_'):
+    # Patient0001_Slice0001_NosieImg0001_NoiseStep0070
+    barcode_str=''
+    header=header.copy()
+    digit=digit.copy()
+    if len(index)<3:
+        header[2] = 'ORG'
+        header[3] = 'NA'
+        digit[2] = 0
+        digit[3] = 0
+        index +=['','']
+
+    for id, h in enumerate(header):
+        barcode_str+=h+str(index[id]).zfill(digit[id])+split
+    return barcode_str[:-1]
+
+class RandomResizedCrop3D(nn.Module):
+    """Crop a random portion of a 3D volume and resize it to a given size.
+
+    Args:
+        size (tuple of int): Expected output size of the crop, for each dimension (D, H, W).
+        scale (tuple of float): Specifies the lower and upper bounds for the random area of the crop,
+                                before resizing. The scale is defined with respect to the volume of the original image.
+        ratio (tuple of float): Lower and upper bounds for the random aspect ratio of the crop, before resizing.
+        interpolation (str): Desired interpolation mode ('trilinear' or 'nearest').
+    """
+
+    def __init__(
+            self,
+            size: Tuple[int, int, int],
+            scale=(0.6, 1.0),
+            ratio=(0.5, 1.5),
+            interpolation='trilinear'
+    ):
+        super().__init__()
+        self.size = size
+        self.scale = scale
+        self.ratio = ratio
+        self.interpolation = interpolation
+
+    @staticmethod
+    def get_params(img: torch.Tensor, rand_scale: float, scale: List[float], ratio: List[float]) -> Tuple[int, int, int, int, int, int]:
+        """Get parameters for `crop` for a random sized crop.
+
+        Args:
+            img (Tensor): Input image.
+            scale (list): Range of scale of the origin size cropped.
+            ratio (list): Range of aspect ratio of the origin aspect ratio cropped.
+
+        Returns:
+            tuple: params (i, j, k, d, h, w) to be passed to `crop` for a random sized crop.
+        """
+        img_sz = np.array(list(img.size())[2:])
+        crop_sz = (img_sz * rand_scale).astype(np.int32)  #[int(s*rand_scale) for s in img_sz]
+        start_id = np.random.randint(0, img_sz - crop_sz + 1, size=(img_sz.size,))
+        return start_id.tolist()+crop_sz.tolist()
+
+        # volume = depth * height * width
+        #
+        # log_ratio = torch.log(torch.tensor(ratio))
+        # for _ in range(10):
+        #     target_volume = volume * torch.empty(1).uniform_(*scale).item()
+        #     aspect_ratio = torch.exp(torch.empty(1).uniform_(log_ratio[0], log_ratio[1])).item()
+        #
+        #     w = int(round(math.sqrt(target_volume * aspect_ratio)))
+        #     h = int(round(math.sqrt(target_volume / aspect_ratio)))
+        #     d = int(round(math.sqrt(target_volume / (w * h))))
+        #
+        #     if 0 < w <= width and 0 < h <= height and 0 < d <= depth:
+        #         i = torch.randint(0, depth - d + 1, size=(1,)).item()
+        #         j = torch.randint(0, height - h + 1, size=(1,)).item()
+        #         k = torch.randint(0, width - w + 1, size=(1,)).item()
+        #         return i, j, k, d, h, w
+        #
+        # # Fallback to central crop
+        # return (depth - d) // 2, (height - h) // 2, (width - w) // 2, d, h, w
+
+    def forward(self, img: torch.Tensor) -> torch.Tensor:
+        """Apply the RandomResizedCrop transformation.
+
+        Args:
+            img (Tensor): Input 3D image.
+
+        Returns:
+            Tensor: Cropped and resized image.
+        """
+        rand_scale = np.random.uniform(self.scale[0], self.scale[1])
+        [i, j, k, d, h, w] = self.get_params(img,rand_scale, self.scale, self.ratio)
+        # print(i, j, k, d, h, w)
+        img_cropped = img[:, :, i:i + d, j:j + h, k:k + w]
+        # print(img_cropped.shape)
+        img_resized = F.interpolate(img_cropped, size=self.size, mode=self.interpolation,
+                                    align_corners=False if self.interpolation == 'trilinear' else None)
+        return img_resized#.squeeze(0)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, scale={self.scale}, ratio={self.ratio}, interpolation={self.interpolation})"
+
+def random_permute(X, select_dims=[-1,-2],include_flip=True):
+    axes=list(range(X[0].ndim))
+    selected_axes = [axes[i] for i in select_dims]
+    random.shuffle(selected_axes)
+    for i, dim in enumerate(select_dims):
+        axes[dim] = selected_axes[i]
+        if include_flip and random.choice([True,False]):
+            # X = [np.flip(x, axis=dim) for x in X]
+            X = [torch.flip(x, [dim]) for x in X]
+    # return [np.transpose(x,axes) for x in X]
+    return [x.permute(axes) for x in X]
+
+# def thresh_img(img,thresh = None,EPS = 10**-7):
+#     threshold0 = np.random.uniform(thresh[0], thresh[1])
+#     threshold1 = np.random.uniform(thresh[0], thresh[1])
+#     scale =
+#     if threshold is not None:
+#         # img=img-threshold
+#         # img=np.where(img>=0,img,0)
+#         # img = np.maximum(img-threshold,0)
+#         img = torch.maximum(img - threshold,torch.tensor(0.))
+#     # return (img - img.min()) / (img.max() - img.min() + EPS)
+#     return img
+
+def get_transformer(degrees=180,translate=0.125,ndims=2,prob=0.8,fill=0.,img_sz=None):
+    prob_crop=0. if img_sz==None else 0.8
+    # prob_crop=0. if len(img_sz)==2 else 0.8
+
+    if img_sz==None or len(img_sz)==2:
+        return torchvision.transforms.Compose([
+            torchvision.transforms.RandomApply([
+                torchvision.transforms.RandomAffine(degrees=degrees, translate=[translate] * ndims, fill=fill,
+                                                    interpolation=torchvision.transforms.InterpolationMode.BILINEAR),
+            ],prob),
+            # torchvision.transforms.RandomApply([
+            #     torchvision.transforms.RandomResizedCrop(size=img_sz),
+            # ], prob_crop),
+            torchvision.transforms.RandomVerticalFlip(p=0.5),
+            torchvision.transforms.RandomAutocontrast(p=0.5),
+        ])
+    else:
+        return torchvision.transforms.Compose([
+            torchvision.transforms.RandomApply([
+                torchvision.transforms.RandomResizedCrop(size=img_sz) if len(img_sz) == 2 else RandomResizedCrop3D(
+                    size=img_sz),
+            ], prob_crop),
+        ])
+
+
+def get_random_affine_transformer(degrees=180,translate=0.125,ndims=2):
+    return torchvision.transforms.RandomAffine(degrees=degrees, translate=[translate] * ndims,interpolation=torchvision.transforms.InterpolationMode.BILINEAR)
+
+def channel_merge_acdc(img):
+#   input: a torch tensor (C,H,W)
+  ch = img.shape[0]
+  output = np.zeros((img.shape[1], img.shape[2]))
+  # output[img[2,:,:] == 1] = 1
+  for i in range(ch):
+    output= output + img[i]
+  return output
+    
+def img_crop(img, crop_rate=2, img_sz=[256,256]):
+    ndims=len(img_sz)
+    crop = [np.random.randint(0.*imgs, 1. * imgs)//crop_rate for imgs in img_sz]
+    crop = [crop, [1 * imgs//crop_rate - c for imgs, c in zip(img_sz, crop)]]
+    if ndims==2:
+        return img[..., crop[0][0]: img_sz[0] - crop[1][0], crop[0][1]: img_sz[1] - crop[1][1]]
+    else:
+        return img[..., crop[0][0]: img_sz[0] - crop[1][0], crop[0][1]:img_sz[1] - crop[1][1], crop[0][2]: img_sz[2] - crop[1][2]]
+
+
+def boundary_limit(sample_coords0, max_sz, plus=0., minus=1.):
+    sample_coords = torch.split(sample_coords0, split_size_or_sections=1, dim=1)
+    # return tf.stack([tf.maximum(tf.minimum(x, sz - minus + plus), 0 + plus) for x, sz in zip(sample_coords, input_size0)],-1)
+    return torch.cat([torch.clamp(x * sz, min=minus - 1 * sz + plus, max=1 * sz - minus + plus) for x, sz in
+                      zip(sample_coords, max_sz)], 1)
+
+
+def resample(vol, ddf, ref=None, img_sz=None,max_sz=[128,128],ndims=2):
+    device = vol.device
+    img_sz = vol.size()[2:]
+    ndims=len(img_sz)
+    if ndims==2:
+        [h,w]=img_sz
+        img_shape = torch.reshape(torch.tensor([(h - 1) / 2., (w - 1) / 2.], device=device), [1, 1, 1, ndims])
+        ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=h), torch.arange(end=w)]), 0), [1, ndims,h, w ])
+    elif ndims==3:
+        [h, w, d] = img_sz
+        img_shape = torch.reshape(torch.tensor([(h - 1) / 2., (w - 1) / 2., (d-1)/2], device=device), [1, 1, 1, 1, ndims])
+        ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=h), torch.arange(end=w), torch.arange(end=d)]), 0), [1, ndims,h, w, d])
+    # ref_grid.to(device)
+    # img_shape.to(device)
+    # ddf.to(device)
+    # ref = self.ref_grid if ref is None else ref
+    # img_sz = self.img_sz if img_sz is None else img_sz
+    resample_mode = 'bilinear'
+    # padding_mode = "border"
+    padding_mode = "zeros"
+
+    # img_sz = np.reshape(img_sz, [1] *(ndims+1)+[ndims])
+    # if ndims==2:
+    if True:
+        re=[0]+list(range(2,ndims+2))+[1]
+        # re=list(range(ndims+2))
+        # print((torch.flip((ddf.to(device) + ref_grid.permute(re))/ img_shape - 1, dims=[-1])).tolist())
+        return F.grid_sample(vol, torch.flip((ddf + ref_grid.permute(re).to(device))/ img_shape - 1, dims=[-1]).type(torch.float32).to(device), mode=resample_mode, padding_mode=padding_mode,align_corners=True)
+        #
+        # return F.grid_sample(vol, torch.flip(
+        #     torch.permute(ddf * torch.Tensor(np.reshape(np.array(max_sz), [1, 1, 1, ndims])) + ref_grid,
+        #                   [0, 2, 3, 1]) / img_shape - 1, dims=[-1]), mode=resample_mode, padding_mode=padding_mode,
+        #                      align_corners=True)
+
+def random_resample(vol,deform_scale=32.):
+    vol_size=vol.size()
+    device=vol.device
+    ndims = len(vol_size)-2
+    img_size=[s for s in vol_size[2:]]
+    if ndims==2:
+        img_size=img_size+[16]
+    # ddf,_,_=random_ddf(vol_size[0],img_size)
+    _,_,ddf=random_ddf(vol_size[0],img_size,ndims=ndims,range_gauss=deform_scale)
+    ddf=Variable(torch.tensor(ddf,dtype=torch.float32)).to(device)
+    if ndims==2:
+        return resample(vol,ddf[...,8,:ndims])
+    else:
+        return resample(vol, ddf[..., :ndims])
+
+def get_random_deformed_mask(msk_shape, deform_scale=32.,apply_possibility=0.75):
+    msk = torch.ones([1, 1]+list(msk_shape),dtype=torch.float32)
+    if random.uniform(0,1) < apply_possibility:
+        return random_resample(msk, deform_scale=deform_scale)
+    else:
+        return msk
+
+# grid option
+def get_tranf_mat(grid_size, vec=[[0., 0., 1.]], ang=[[0.]],transl=[[0,0,0]]):
+    return np.concatenate([get_rot_mat(grid_size, vec=vec, ang=ang),transl],-1)
+
+
+def get_rot_mat(grid_size, vec=[[0., 0., 1.]], ang=[[0.]],ndims=3):
+    vec = np.array(vec)
+    ang = np.array(ang)
+    batch_num = ang.shape[0]
+    return np.reshape(vecang2rotmats(vec, ang), [batch_num] + [ndims*(ndims)])
+
+def random_mat(batch_sz, img_sz, num_class=2,pn_spline=20, pn_gauss=10, range_spline=2., range_gauss=48, spread_range=[5., 24.],
+               transl_range=32., rot_range=np.pi / 2):
+    scale=4
+    ndims=3
+    vec=np.reshape(np.random.uniform(-1., 1., [batch_sz,1, ndims])+np.random.uniform(-.1, .1, [batch_sz,num_class, ndims]),[batch_sz*num_class, ndims])
+    ang=np.reshape(np.random.uniform(-rot_range, rot_range, [batch_sz,1])+np.random.uniform(-rot_range/scale, rot_range/scale, [batch_sz,num_class]),[batch_sz*num_class])
+    transl=np.reshape(np.random.uniform(-transl_range, transl_range, [batch_sz,1,ndims])+np.random.uniform(-transl_range/scale, transl_range/scale, [batch_sz,num_class,ndims]),[batch_sz*num_class,ndims])
+    return np.reshape(np.concatenate([get_rot_mat(img_sz, vec=vec, ang=ang),transl],-1),[batch_sz,num_class,4,3])
+
+    # return np.reshape(get_tranf_mat(img_sz, vec=np.random.uniform(-1., 1., [batch_sz*num_class, 3]), ang=np.random.uniform(-rot_range, rot_range, [batch_sz*num_class]),transl=np.random.uniform(-transl_range, transl_range, [batch_sz*num_class,3])),[batch_sz,num_class,4,3])
+
+def random_ddf(batch_sz, img_sz, pn_spline=20, pn_gauss=10, range_spline=1., range_gauss=16., spread_range=[16., 64.],
+               transl_range=0., rot_range=np.pi / 1,ndims=3):
+    rand_ang=np.random.uniform(-rot_range, rot_range, [batch_sz])
+    # rand_ang = np.random.randint(-4, 4, [batch_sz])*rot_range
+
+    if ndims==3:
+        rot_df = get_rot_ddf(img_sz, vec=np.random.uniform(-1., 1., [batch_sz, 3]),
+                             ang=rand_ang)
+    else:
+        rot_df = get_rot_ddf(img_sz, vec=np.concatenate([np.zeros([batch_sz, 2]),np.ones([batch_sz, 1])],-1),
+                             ang=rand_ang)
+    ndims = 3
+    # rot_df = +np.random.uniform(-1., 1., [batch_sz, ndims,ndims])
+    # ddf0=np.stack([generate_random_gaussian_ddf(img_sz, pn_gauss, range_sz=range_gauss, spread_std=spread_range)\
+    #                +generate_random_spline_ddf(img_sz, pn_spline, range_sz=range_spline)\
+    #                +np.random.uniform(-transl_range,transl_range,[3]) for i in range(batch_sz)],axis=0)\
+    #      +rot_df
+    if range_gauss>0:
+        ddf0 = np.tile([generate_random_gaussian_ddf(img_sz, pn_gauss, range_sz=range_gauss, spread_std=spread_range) \
+                        # + generate_random_spline_ddf(img_sz, pn_spline, range_sz=range_spline) \
+                        + np.random.uniform(-transl_range, transl_range, [ndims])], [batch_sz, 1, 1, 1, 1]) \
+               + rot_df
+    else:
+        ddf0 = rot_df
+
+    def boundary_replicate(sample_coords, input_size, padding=5):
+        return np.stack(
+            [np.maximum(np.minimum(sample_coords[..., i], input_size[i] - 1 + padding), 0 - padding) for i in
+             range(len(input_size))], axis=-1), \
+               np.prod([((sample_coords[..., i] < input_size[i]) * (sample_coords[..., i] >= 0)) for i in
+                        range(len(input_size))], axis=0)
+
+    ref = get_reference_grid(img_sz)
+    cf1, ind = boundary_replicate(ddf0 + ref, img_sz)
+    return cf1 - ref, np.expand_dims(ind, -1), rot_df
+
+
+def generate_random_gaussian_ddf(img_sz, pn=30, range_sz=5, spread_std=[0.1, 1.]):
+    x = np.floor(np.random.uniform(range_sz / 2., img_sz[0] - range_sz / 2., [1, pn])).astype('int')
+    y = np.floor(np.random.uniform(range_sz / 2., img_sz[1] - range_sz / 2., [1, pn])).astype('int')
+    z = np.floor(np.random.uniform(range_sz / 2., img_sz[2] - range_sz / 2., [1, pn])).astype('int')
+
+    odf = np.random.uniform(-range_sz, range_sz, [pn, 3])
+    vol = np.zeros([img_sz[0], img_sz[1], img_sz[2], 3])
+    vol[x, y, z] = odf
+
+    return spimg.gaussian_filter(vol, np.random.uniform(spread_std[0], spread_std[1]))
+
+
+def get_rot_ddf(grid_size, vec=[[0., 0., 1.]], ang=[[0.]]):
+    vec = np.array(vec)
+    ang = np.array(ang)
+    batch_num = ang.shape[0]
+    ref_grids = get_reference_grid(grid_size,
+                                   bias_scale=1.)
+    # a=vecang2rotmats(vec, ang)
+    return np.reshape(np.matmul(np.reshape(np.tile(ref_grids, [batch_num, 1, 1, 1, 1]), [batch_num, -1, 3]),
+                                vecang2rotmats(vec, ang)), [batch_num] + grid_size + [3]) - ref_grids
+
+
+def get_reference_grid(grid_size, bias_scale=0.):
+    return np.stack(np.meshgrid(
+        [i for i in range(grid_size[0])],
+        [j for j in range(grid_size[1])],
+        [k for k in range(grid_size[2])],
+        indexing='ij'), axis=-1).astype('float') - bias_scale * (np.array(grid_size) - 1) / 2.
+
+
+def resample_linear(inputs, ddf=None, sample_coords=None,random_boundary=True):
+    if random_boundary:
+        random_factor = np.random.uniform(0., 1.)
+        min_val = np.min(inputs)
+        inputs[:, 0, :, :] = min_val * random_factor + (1 - random_factor) * inputs[:, 0, :, :]
+        inputs[:, -1, :, :] = min_val * random_factor + (1 - random_factor) * inputs[:, -1, :, :]
+        inputs[:, :, 0, :] = min_val * random_factor + (1 - random_factor) * inputs[:, :, 0, :]
+        inputs[:, :, -1, :] = min_val * random_factor + (1 - random_factor) * inputs[:, :, -1, :]
+        inputs[:, :, :, 0] = min_val * random_factor + (1 - random_factor) * inputs[:, :, :, 0]
+        inputs[:, :, :, -1] = min_val * random_factor + (1 - random_factor) * inputs[:, :, :, -1]
+
+    input_size = inputs.shape[1:4]
+    sample_coords = get_reference_grid(input_size) + ddf if sample_coords is None else sample_coords
+    spatial_rank = 3  # inputs.ndim - 2
+    xy = [sample_coords[..., i] for i in
+          range(sample_coords.shape[-1])]  # tf.unstack(sample_coords, axis=len(sample_coords.shape)-1)
+    index_voxel_coords = [np.floor(x) for x in xy]
+
+    def boundary_replicate(sample_coords0, input_size0, plus=0):
+        return np.maximum(np.minimum(sample_coords0, input_size0 - 2 + plus), 0 + plus)
+
+    def boundary_replicate_float(sample_coords0, input_size0, plus=0.):
+        return np.maximum(np.minimum(sample_coords0, input_size0 - 1 + plus), 0 + plus)
+
+    xy = [boundary_replicate_float(x.astype('float32'), input_size[idx]) for idx, x in enumerate(xy)]
+    spatial_coords = [boundary_replicate(x.astype('int32'), input_size[idx])
+                      for idx, x in enumerate(index_voxel_coords)]
+    spatial_coords_plus1 = [boundary_replicate((x + 1).astype('int32'), input_size[idx], 1)
+                            for idx, x in enumerate(index_voxel_coords)]
+
+    weight = [np.expand_dims(x - i.astype('float32'), -1) for x, i in zip(xy, spatial_coords)]
+    weight_c = [np.expand_dims(i.astype('float32') - x, -1) for x, i in zip(xy, spatial_coords_plus1)]
+
+    sz = list(spatial_coords[0].shape)
+    batch_coords = np.tile(np.reshape(range(sz[0]), [sz[0]] + [1] * (len(sz) - 1)), [1] + sz[1:])
+    sc = (spatial_coords, spatial_coords_plus1)
+    binary_codes = [[int(c) for c in format(i, '0%ib' % spatial_rank)] for i in range(2 ** spatial_rank)]
+
+    make_sample = lambda bc: inputs[batch_coords, sc[bc[0]][0], sc[bc[1]][1], sc[bc[2]][
+        2], ...]  # tf.gather_nd(inputs, np.stack([batch_coords] + [sc[c][i] for i, c in enumerate(bc)], -1))
+    samples = [make_sample(bc) for bc in binary_codes]
+
+    def pyramid_combination(samples0, weight0, weight_c0):
+        if len(weight0) == 1:
+            return samples0[0] * weight_c0[0] + samples0[1] * weight0[0]
+        else:
+            return pyramid_combination(samples0[::2], weight0[:-1], weight_c0[:-1]) * weight_c0[-1] + \
+                   pyramid_combination(samples0[1::2], weight0[:-1], weight_c0[:-1]) * weight0[-1]
+
+    return pyramid_combination(samples, weight, weight_c)
+
+
+def vecang2rotmats(vec, ang):
+    return np.stack([np.reshape(vecang2rotmat(vec[i, ...], ang[i, ...]), [3, 3]) for i in range(len(vec))], 0)
+
+
+def vecang2rotmat(vec, ang):
+    q = quater.Quaternion(axis=vec, angle=ang)
+    return q.rotation_matrix
+
+
+def images_to_vectors(images):
+  return images.view(images.size(0), 16384).to(device)
+
+def vectors_to_images(vectors):
+  return vectors.view(vectors.size(0), 1, 128, 128).to(device)
+
+def noise(size):
+  n = Variable(torch.randn(size, 100)).to(device)
+  return n
+
+def ones_target(size):
+  data = Variable(torch.ones(size, 1)).to(device)
+  return data
+
+def zeros_target(size):
+  data = Variable(torch.zeros(size, 1)).to(device)
+  return data
+
+
+def eval_detJ_lab(disp=None,vol1=None,vol2=None,thresh=0.5):
+    ndims=disp.ndim-2
+    if vol1 ==None or thresh==None:
+        label=1
+    else:
+        label=vol1>thresh
+        label=label*(spimg.laplace(label) < 0.1)
+        rescale_factor=2
+        label=label[...,::rescale_factor,::rescale_factor,::rescale_factor]
+
+    # disp = disp.permute([0, *range(2,ndims+2), 1])
+    # print(disp.shape)
+    disp = np.transpose(disp, [0, *range(2,ndims+2), 1])
+    # Jacob=np.stack(np.gradient(disp,axis=[-4,-3,-2]),-1)
+    Jacob=np.stack(np.gradient(disp,axis=[*range(1,ndims+1)]),-1)
+    for ii in range(ndims):
+        Jacob[..., ii, ii] = Jacob[..., ii, ii] + 1
+    # Jacob[..., 0, 0] = Jacob[..., 0, 0] + 1
+    # Jacob[..., 1, 1] = Jacob[..., 1, 1] + 1
+    # Jacob[..., 2, 2] = Jacob[..., 2, 2] + 1
+    return np.sum((np.linalg.det(Jacob)<0)*label)
+
+def eval_def_mag(disp=None,vol1=None,vol2=None,thresh=0.5):
+    ndims=3
+    # if vol1 ==None or thresh==None:
+    #     label=1
+    # else:
+    #     label=vol1>thresh
+    #     label=label*(spimg.laplace(label) < 0.1)
+    #     rescale_factor=2
+    #     label=label[...,::rescale_factor,::rescale_factor,::rescale_factor]
+    mag=np.sqrt(np.sum(np.square(disp),axis=1))
+    sz=mag.shape
+    max_mag=np.mean(np.max(np.reshape(mag,[sz[0],-1]),axis=-1))
+    avg_mag=np.mean(mag)
+    return [avg_mag,max_mag]
+
+
+
+def print_memory_usage(tag=""):
+    print(f"[{tag}] Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB | Cached: {torch.cuda.memory_reserved() / 1e9:.2f} GB")
+
+
+if __name__ == "__main__":
+    vol_shape=[4,1,64,64]
+
+    vol=np.random.uniform(-1,1,vol_shape)
+    vol=Variable(torch.tensor(vol,dtype=torch.float32))
+    vol_res=random_resample(vol)
+    vol_crop=img_crop(vol_res)
+
+    mask = get_random_deformed_mask(vol.shape[2:])
+
+    print(mask)
+    
+    # print(vol.tolist())
+    # print(vol_res.tolist())

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

OM_aug.py

@@ -0,0 +1,254 @@
+import torch
+import torchvision
+from torch import nn
+from torchvision.utils import save_image
+from torch.utils.data import DataLoader
+from torch.optim import Adam
+from torchvision.utils import make_grid
+from Diffusion.diffuser import DeformDDPM
+from Diffusion.networks import get_net, STN
+from torchvision.transforms import Lambda
+import random
+import os
+import utils
+from Dataloader.dataloader0 import get_dataloader
+from Dataloader.dataLoader import *
+
+from torchvision.utils import save_image
+from einops import rearrange, reduce, repeat
+# import matplotlib.image
+import numpy as np
+import nibabel as nib
+from tqdm import tqdm 
+import yaml
+import argparse
+
+EPS = 10e-8
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument(
+        "--config",
+        "-C",
+        help="Path for the config file",
+        type=str,
+        default="Config/config_cmr.yaml",
+        # default="Config/config_lct.yaml",
+        required=False,
+    )
+args = parser.parse_args()
+#=======================================================================================================================
+
+# config_path = 'Config/config_cmr.yaml'
+# config_path = 'Config/config_lct.yaml'
+
+# Load the YAML file into a dictionary
+with open(args.config, 'r') as file:
+    hyp_parameters = yaml.safe_load(file)
+    print(hyp_parameters)
+# hyp_parameters["aug_img_savepath"] = os.path.join(hyp_parameters["aug_img_savepath"],hyp_parameters["data_name"],'')
+if not os.path.exists(hyp_parameters["aug_img_savepath"]):
+    os.makedirs(hyp_parameters["aug_img_savepath"])
+if not os.path.exists(hyp_parameters["aug_msk_savepath"]):
+    os.makedirs(hyp_parameters["aug_msk_savepath"])
+if not os.path.exists(hyp_parameters["aug_ddf_savepath"]):
+    os.makedirs(hyp_parameters["aug_ddf_savepath"])
+print(hyp_parameters["aug_img_savepath"])
+
+hyp_parameters['batchsize'] = 1
+
+
+# =======================================================================================================================
+select_channels_dict={}
+# min_crop_ratio = 0.5
+min_crop_ratio = 0.9
+
+# label_keys = ['heart']
+# label_keys = ['brain']
+# label_keys = ['pancreas']
+# label_keys = ['spleen']
+# label_keys = ['liver']
+# database = ['MSD']
+label_keys = ['heart']
+database = ['MnMs']
+# subtype = "ed"   # 'ed' or 'es' for MnMs
+subtype = "es"   # 'ed' or 'es' for MnMs
+hyp_parameters["aug_img_savepath"]=f"Data/Aug_data/mnms_{subtype}/img/"
+hyp_parameters["aug_msk_savepath"]=f"Data/Aug_data/mnms_{subtype}/msk/"
+hyp_parameters["aug_ddf_savepath"]=f"Data/Aug_data/mnms_{subtype}/ddf/"
+select_channels_dict={
+  "ImgDict":[subtype]
+}
+
+# dataset = OminiDataset_v1(transform=None,min_crop_ratio=min_crop_ratio)
+dataset = OminiDataset_inference_w_all(transform=None,min_crop_ratio=min_crop_ratio,label_key = label_keys, database=database, select_channels_dict=select_channels_dict)
+Infer_Loader = DataLoader(
+        dataset,
+        batch_size=hyp_parameters['batchsize'],
+        shuffle=False
+    )
+# =======================================================================================================================
+
+# Data_Loader=get_dataloader(hyp_parameters['data_name'],mode='aug')
+# transformer = utils.get_transformer(img_sz=hyp_parameters["ndims"]*[hyp_parameters['img_size']])
+# dataset = Data_Loader(patient_index = hyp_parameters["patients_list"])
+# train_loader = DataLoader(dataset, batch_size = hyp_parameters['batchsize'], shuffle = False) 
+
+
+
+epoch=f'{hyp_parameters["model_id_str"]}_{hyp_parameters["data_name"]}_{hyp_parameters["net_name"]}'
+model_save_path = f'Models/{hyp_parameters["data_name"]}_{hyp_parameters["net_name"]}/'
+model_save_path = os.path.join(model_save_path, str(epoch)+'.pth')
+
+
+
+Net = get_net(hyp_parameters["net_name"])
+
+Deformddpm = DeformDDPM(
+    network=Net(n_steps = hyp_parameters["timesteps"],
+                ndims = hyp_parameters["ndims"],
+                num_input_chn = hyp_parameters["num_input_chn"],
+                res = hyp_parameters['img_size']
+                ),
+    n_steps = hyp_parameters["timesteps"],
+    image_chw = [hyp_parameters["num_input_chn"]] + [hyp_parameters["img_size"]]*hyp_parameters["ndims"],
+    device = hyp_parameters["device"],
+    batch_size = hyp_parameters["batchsize"],
+    img_pad_mode = hyp_parameters["img_pad_mode"],
+    ddf_pad_mode = hyp_parameters["ddf_pad_mode"],
+    padding_mode = hyp_parameters["padding_mode"],
+    v_scale = hyp_parameters["v_scale"],
+    resample_mode = hyp_parameters["resample_mode"],
+)
+Deformddpm.to(hyp_parameters["device"])
+
+ddf_stn = STN(
+    img_sz = hyp_parameters["img_size"],
+    ndims = hyp_parameters["ndims"],
+    padding_mode = hyp_parameters['padding_mode'],
+    device = hyp_parameters["device"],
+)
+ddf_stn.to(hyp_parameters["device"])
+
+print("Loading model from:", model_save_path)
+# Deformddpm.load_state_dict(torch.load(model_save_path))
+checkpoint = torch.load(model_save_path)
+Deformddpm.load_state_dict(checkpoint['model_state_dict'])
+Deformddpm.eval()
+
+os.makedirs(hyp_parameters['aug_img_savepath'], exist_ok=True)
+os.makedirs(hyp_parameters['aug_msk_savepath'], exist_ok=True)
+os.makedirs(hyp_parameters['aug_ddf_savepath'], exist_ok=True)
+
+print("total num of image:", len(Infer_Loader))
+for e, d in tqdm(enumerate(Infer_Loader)):
+  # if e<1:
+  #   continue
+  # img, mask, pid = d
+  # img = d
+  # mask = d
+  img = d['img']
+  mask = d['labels']
+  label_str = str(d['label_channels'])
+  # mask = np.concatenate([v for v in d['labels'].values()], axis=1)
+  # print('img shape:', img.shape, 'mask shape:', mask.shape)
+
+  # pid = pid.cpu().detach().numpy()
+  # pid = pid[0] 
+  pid = e
+
+  print('Processing to patient:', pid, ' image:',e)
+  
+
+  img = img.type(torch.float32)
+  img = img.to(hyp_parameters["device"]) 
+  image_original = img.cpu().detach().numpy()
+
+  mask = mask.type(torch.float32)
+  mask = mask.to(hyp_parameters["device"]) 
+  mask_original = mask.cpu().detach().numpy()
+  # print(pid, image_original.shape, mask_original.max())
+
+
+  # if hyp_parameters["ndims"] == 2:
+  #   nifti_img = nib.Nifti1Image(image_original[0,0,:,:], np.eye(4))  
+  #   nifti_mask = nib.Nifti1Image(mask_original[0,:,:,:], np.eye(4))  
+  # elif hyp_parameters["ndims"] == 3:
+  #   nifti_img = nib.Nifti1Image(image_original[0,0,:,:,:], np.eye(4))  
+  #   nifti_mask = nib.Nifti1Image(mask_original[0,0,:,:,:], np.eye(4))  
+  nifti_img = utils.converet_to_nibabel(image_original,ndims=hyp_parameters["ndims"])
+  nifti_mask = utils.converet_to_nibabel(mask_original,ndims=hyp_parameters["ndims"])
+
+  # Saving original (undeformed image)
+  # CMR: format: Patient0001_Slice0001_ORG_NA.nii.gz
+  # Lung CT: Patient0001_Slice0001_ORG_NA.nii.gz
+  nib.save(nifti_img, os.path.join(hyp_parameters['aug_img_savepath'],utils.get_barcode([pid,e])+'.nii.gz'))
+
+  # Saving original (undeformed image)
+  # CMR: format: Patient0001_Slice0001_ORG_NA_GT.nii.gz
+  # Lung CT: ...
+  nib.save(nifti_mask, os.path.join(hyp_parameters['aug_msk_savepath'],utils.get_barcode([pid,e])+'_GT.nii.gz'))
+
+ 
+  noise_step = hyp_parameters["start_noise_step"]
+  with torch.no_grad():
+    for im in range(hyp_parameters["aug_coe"]):
+      # # Permute 
+      # if hyp_parameters["ndims"] == 2:
+      #   [img, mask] = utils.random_permute([img, mask], select_dims=[-1, -2])          # add random rotation to image
+      # elif hyp_parameters["ndims"] == 3:
+      #   [img, mask] = utils.random_permute([img, mask], select_dims=[-1, -2, -3])  # add random rotation to image
+
+      print('Generating - >', 'Subject-',pid,', Scan-',e,' (',im,'/',hyp_parameters["aug_coe"],')', end='\r')
+      
+      [ddf_comp,ddf_rand],[img_rec,img_diff,img_save],[msk_rec,msk_diff,msk_save] = Deformddpm.diff_recover(img_org=img,msk_org=mask,T=[noise_step,hyp_parameters["timesteps"]],v_scale=hyp_parameters["v_scale"],t_save=None,proc_type=hyp_parameters["condition_type"])
+      
+      denoise_imgs = img_rec.cpu().detach().numpy()
+      denoise_msks = msk_rec.cpu().detach().numpy()
+      noisy_imgs_np = img_diff.cpu().detach().numpy()
+      noisy_msks_np = msk_diff.cpu().detach().numpy()
+
+      # if hyp_parameters["ndims"] == 2:
+      #   nifti_img_aug = nib.Nifti1Image(denoise_imgs[0,0,:,:], np.eye(4))  
+      #   nifti_mask_aug = nib.Nifti1Image(denoise_msks[0,:,:,:], np.eye(4))  
+      #   nifti_img = nib.Nifti1Image(noisy_imgs_np[0,0,:,:], np.eye(4))
+      #   nifti_mask = nib.Nifti1Image(noisy_msks_np[0, :, :, :], np.eye(4))
+      # elif hyp_parameters["ndims"] == 3:
+      #   nifti_img_aug = nib.Nifti1Image(denoise_imgs[0,0,:,:,:], np.eye(4))  
+      #   nifti_mask_aug = nib.Nifti1Image(denoise_msks[0,0,:,:,:], np.eye(4))  
+      #   nifti_img = nib.Nifti1Image(noisy_imgs_np[0,0,:,:,:], np.eye(4))
+      #   nifti_mask = nib.Nifti1Image(noisy_msks_np[0, 0, :, :], np.eye(4))    ###
+      nifti_img_aug = utils.converet_to_nibabel(denoise_imgs,ndims=hyp_parameters["ndims"])
+      nifti_mask_aug = utils.converet_to_nibabel(denoise_msks,ndims=hyp_parameters["ndims"])
+      nifti_img = utils.converet_to_nibabel(noisy_imgs_np,ndims=hyp_parameters["ndims"])
+      nifti_mask = utils.converet_to_nibabel(noisy_msks_np,ndims=hyp_parameters["ndims"])
+      
+      nib.save(nifti_img_aug, os.path.join(hyp_parameters['aug_img_savepath'],utils.get_barcode([pid,e,im,noise_step])+'.nii.gz'))
+      nib.save(nifti_mask_aug, os.path.join(hyp_parameters['aug_msk_savepath'],utils.get_barcode([pid,e,im,noise_step])+'_GT.nii.gz'))
+      
+      # Saving noisy image to nifti
+      # CMR: format: Patient0001_Slice0001_NosieImg0001_NoiseStep0070.nii.gz
+      # Lung CT: ...
+      nib.save(nifti_img, os.path.join(hyp_parameters['aug_img_savepath'],utils.get_barcode([pid,e,im,noise_step],header=['Patient','Slice','NoiseImg','NoiseStep'])+'.nii.gz'))
+      nib.save(nifti_mask, os.path.join(hyp_parameters['aug_msk_savepath'],utils.get_barcode([pid,e,im,noise_step],header=['Patient','Slice','NoiseImg','NoiseStep'])+'_GT.nii.gz'))
+      
+          
+      if (im - hyp_parameters["start_noise_step"])%2 == 0:
+        noise_step = noise_step + hyp_parameters["noise_step"]
+      # break   # for testing
+  if e >= 0:
+    exit()
+          
+
+
+
+
+
+
+
+
+
+
+
+
+

OM_aug_highres.py

@@ -0,0 +1,233 @@
+import torch
+import torchvision
+from torch import nn
+from torchvision.utils import save_image
+from torch.utils.data import DataLoader
+from torch.optim import Adam
+from torchvision.utils import make_grid
+from Diffusion.diffuser import DeformDDPM
+from Diffusion.networks import get_net, STN
+from torchvision.transforms import Lambda
+import random
+import os
+import utils
+from Dataloader.dataloader0 import get_dataloader
+from Dataloader.dataLoader import *
+
+from torchvision.utils import save_image
+from einops import rearrange, reduce, repeat
+# import matplotlib.image
+import numpy as np
+import nibabel as nib
+from tqdm import tqdm 
+import yaml
+import argparse
+
+EPS = 10e-8
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument(
+        "--config",
+        "-C",
+        help="Path for the config file",
+        type=str,
+        default="Config/config_cmr.yaml",
+        # default="Config/config_lct.yaml",
+        required=False,
+    )
+args = parser.parse_args()
+#=======================================================================================================================
+
+# config_path = 'Config/config_cmr.yaml'
+# config_path = 'Config/config_lct.yaml'
+
+# Load the YAML file into a dictionary
+with open(args.config, 'r') as file:
+    hyp_parameters = yaml.safe_load(file)
+    print(hyp_parameters)
+# hyp_parameters["aug_img_savepath"] = os.path.join(hyp_parameters["aug_img_savepath"],hyp_parameters["data_name"],'')
+if not os.path.exists(hyp_parameters["aug_img_savepath"]):
+    os.makedirs(hyp_parameters["aug_img_savepath"])
+if not os.path.exists(hyp_parameters["aug_msk_savepath"]):
+    os.makedirs(hyp_parameters["aug_msk_savepath"])
+if not os.path.exists(hyp_parameters["aug_ddf_savepath"]):
+    os.makedirs(hyp_parameters["aug_ddf_savepath"])
+print(hyp_parameters["aug_img_savepath"])
+
+hyp_parameters['batchsize'] = 1
+
+
+# =======================================================================================================================
+# min_crop_ratio = 0.5
+min_crop_ratio = 0.9
+
+# label_keys = ['heart']
+# label_keys = ['brain']
+label_keys = ['pancreas']
+database = ['MSD']
+
+# dataset = OminiDataset_v1(transform=None,min_crop_ratio=min_crop_ratio)
+dataset = OminiDataset_inference_w_all(transform=None,min_crop_ratio=min_crop_ratio,label_key = label_keys, database=database)
+Infer_Loader = DataLoader(
+        dataset,
+        batch_size=hyp_parameters['batchsize'],
+        shuffle=False
+    )
+# =======================================================================================================================
+
+# Data_Loader=get_dataloader(hyp_parameters['data_name'],mode='aug')
+# transformer = utils.get_transformer(img_sz=hyp_parameters["ndims"]*[hyp_parameters['img_size']])
+# dataset = Data_Loader(patient_index = hyp_parameters["patients_list"])
+# train_loader = DataLoader(dataset, batch_size = hyp_parameters['batchsize'], shuffle = False) 
+
+
+
+epoch=f'{hyp_parameters["model_id_str"]}_{hyp_parameters["data_name"]}_{hyp_parameters["net_name"]}'
+model_save_path = f'Models/{hyp_parameters["data_name"]}_{hyp_parameters["net_name"]}/'
+model_save_path = os.path.join(model_save_path, str(epoch)+'.pth')
+
+
+
+Net = get_net(hyp_parameters["net_name"])
+
+Deformddpm = DeformDDPM(
+    network=Net(n_steps = hyp_parameters["timesteps"],
+                ndims = hyp_parameters["ndims"],
+                num_input_chn = hyp_parameters["num_input_chn"],
+                res = hyp_parameters['img_size']
+                ),
+    n_steps = hyp_parameters["timesteps"],
+    image_chw = [hyp_parameters["num_input_chn"]] + [hyp_parameters["img_size"]]*hyp_parameters["ndims"],
+    device = hyp_parameters["device"],
+    batch_size = hyp_parameters["batchsize"],
+    img_pad_mode = hyp_parameters["img_pad_mode"],
+    ddf_pad_mode = hyp_parameters["ddf_pad_mode"],
+    padding_mode = hyp_parameters["padding_mode"],
+    v_scale = hyp_parameters["v_scale"],
+    resample_mode = hyp_parameters["resample_mode"],
+)
+Deformddpm.to(hyp_parameters["device"])
+
+ddf_stn = STN(
+    img_sz = hyp_parameters["img_size"],
+    ndims = hyp_parameters["ndims"],
+    padding_mode = hyp_parameters['padding_mode'],
+    device = hyp_parameters["device"],
+)
+ddf_stn.to(hyp_parameters["device"])
+
+print("Loading model from:", model_save_path)
+# Deformddpm.load_state_dict(torch.load(model_save_path))
+checkpoint = torch.load(model_save_path)
+Deformddpm.load_state_dict(checkpoint['model_state_dict'])
+Deformddpm.eval()
+
+os.makedirs(hyp_parameters['aug_img_savepath'], exist_ok=True)
+os.makedirs(hyp_parameters['aug_msk_savepath'], exist_ok=True)
+os.makedirs(hyp_parameters['aug_ddf_savepath'], exist_ok=True)
+
+print("total num of image:", len(Infer_Loader))
+for e, d in tqdm(enumerate(Infer_Loader)):
+  
+  # img, mask, pid = d
+  # img = d
+  # mask = d
+  img = d['img']
+  mask = d['labels']
+  # mask = np.concatenate([v for v in d['labels'].values()], axis=1)
+  # print('img shape:', img.shape, 'mask shape:', mask.shape)
+
+  # pid = pid.cpu().detach().numpy()
+  # pid = pid[0] 
+  pid = e
+
+  print('Processing to patient:', pid, ' image:',e)
+  
+
+  img = img.type(torch.float32)
+  img = img.to(hyp_parameters["device"]) 
+  image_original = img.cpu().detach().numpy()
+
+  mask = mask.type(torch.float32)
+  mask = mask.to(hyp_parameters["device"]) 
+  mask_original = mask.cpu().detach().numpy()
+  # print(pid, image_original.shape, mask_original.max())
+
+
+  if hyp_parameters["ndims"] == 2:
+    nifti_img = nib.Nifti1Image(image_original[0,0,:,:], np.eye(4))  
+    nifti_mask = nib.Nifti1Image(mask_original[0,:,:,:], np.eye(4))  
+  elif hyp_parameters["ndims"] == 3:
+    nifti_img = nib.Nifti1Image(image_original[0,0,:,:,:], np.eye(4))  
+    nifti_mask = nib.Nifti1Image(mask_original[0,0,:,:,:], np.eye(4))  
+
+  # Saving original (undeformed image)
+  # CMR: format: Patient0001_Slice0001_ORG_NA.nii.gz
+  # Lung CT: Patient0001_Slice0001_ORG_NA.nii.gz
+  nib.save(nifti_img, os.path.join(hyp_parameters['aug_img_savepath'],utils.get_barcode([pid,e])+'.nii.gz'))
+
+  # Saving original (undeformed image)
+  # CMR: format: Patient0001_Slice0001_ORG_NA_GT.nii.gz
+  # Lung CT: ...
+  nib.save(nifti_mask, os.path.join(hyp_parameters['aug_msk_savepath'],utils.get_barcode([pid,e])+'_GT.nii.gz'))
+
+ 
+  noise_step = hyp_parameters["start_noise_step"]
+  with torch.no_grad():
+    for im in range(hyp_parameters["aug_coe"]):
+      # # Permute 
+      # if hyp_parameters["ndims"] == 2:
+      #   [img, mask] = utils.random_permute([img, mask], select_dims=[-1, -2])          # add random rotation to image
+      # elif hyp_parameters["ndims"] == 3:
+      #   [img, mask] = utils.random_permute([img, mask], select_dims=[-1, -2, -3])  # add random rotation to image
+
+      print('Generating - >', 'Subject-',pid,', Scan-',e,' (',im,'/',hyp_parameters["aug_coe"],')', end='\r')
+      
+      [ddf_comp,ddf_rand],[img_rec,img_diff,img_save],[msk_rec,msk_diff,msk_save] = Deformddpm.diff_recover(img_org=img,msk_org=mask,T=[noise_step,hyp_parameters["timesteps"]],v_scale=hyp_parameters["v_scale"],t_save=None,proc_type=hyp_parameters["condition_type"])
+      
+      denoise_imgs = img_rec.cpu().detach().numpy()
+      denoise_msks = msk_rec.cpu().detach().numpy()
+      noisy_imgs_np = img_diff.cpu().detach().numpy()
+      noisy_msks_np = msk_diff.cpu().detach().numpy()
+
+      if hyp_parameters["ndims"] == 2:
+        nifti_img_aug = nib.Nifti1Image(denoise_imgs[0,0,:,:], np.eye(4))  
+        nifti_mask_aug = nib.Nifti1Image(denoise_msks[0,:,:,:], np.eye(4))  
+        nifti_img = nib.Nifti1Image(noisy_imgs_np[0,0,:,:], np.eye(4))
+        nifti_mask = nib.Nifti1Image(noisy_msks_np[0, :, :, :], np.eye(4))
+      elif hyp_parameters["ndims"] == 3:
+        nifti_img_aug = nib.Nifti1Image(denoise_imgs[0,0,:,:,:], np.eye(4))  
+        nifti_mask_aug = nib.Nifti1Image(denoise_msks[0,0,:,:,:], np.eye(4))  
+        nifti_img = nib.Nifti1Image(noisy_imgs_np[0,0,:,:,:], np.eye(4))
+        nifti_mask = nib.Nifti1Image(noisy_msks_np[0, 0, :, :], np.eye(4))
+      
+      nib.save(nifti_img_aug, os.path.join(hyp_parameters['aug_img_savepath'],utils.get_barcode([pid,e,im,noise_step])+'.nii.gz'))
+      nib.save(nifti_mask_aug, os.path.join(hyp_parameters['aug_msk_savepath'],utils.get_barcode([pid,e,im,noise_step])+'_GT.nii.gz'))
+      
+      # Saving noisy image to nifti
+      # CMR: format: Patient0001_Slice0001_NosieImg0001_NoiseStep0070.nii.gz
+      # Lung CT: ...
+      nib.save(nifti_img, os.path.join(hyp_parameters['aug_img_savepath'],utils.get_barcode([pid,e,im,noise_step],header=['Patient','Slice','NoiseImg','NoiseStep'])+'.nii.gz'))
+      nib.save(nifti_mask, os.path.join(hyp_parameters['aug_msk_savepath'],utils.get_barcode([pid,e,im,noise_step],header=['Patient','Slice','NoiseImg','NoiseStep'])+'_GT.nii.gz'))
+      
+          
+      if (im - hyp_parameters["start_noise_step"])%2 == 0:
+        noise_step = noise_step + hyp_parameters["noise_step"]
+      # break   # for testing
+  # if e > 5:
+  #   break
+          
+
+
+
+
+
+
+
+
+
+
+
+
+

OM_contrastive.py

@@ -0,0 +1,72 @@
+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 *
+import argparse
+import yaml
+import os
+import time
+import swanlab
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--config", "-C", type=str, default="Config/config_om_contrastive.yaml")
+args = parser.parse_args()
+
+with open(args.config, 'r') as file:
+    hyp = yaml.safe_load(file)
+
+# Setup
+device = torch.device(hyp['device'] if torch.cuda.is_available() else 'cpu')
+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)
+
+# SwanLab
+swanlab.init(project="OM", config=hyp)
+
+# 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
+dataset = OMDataset_indiv(out_sz=img_size, transform=None)
+train_loader = DataLoader(dataset, batch_size=hyp['batchsize'], shuffle=True, drop_last=True)
+
+# Training
+print('start training...')
+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()
+        swanlab.log({"loss": loss.item()})
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        epoch_loss += loss.item()
+        t1 = time.time()
+        dt = t1 - t0
+        swanlab.log({"Time(mins)/batch": dt/60})
+    avg_loss = epoch_loss / max(len(train_loader), 1)
+    print(f"Epoch {epoch:04d} | Loss: {avg_loss:.6f}")
+    swanlab.log({"Avg Loss/epoch": avg_loss})
+    
+    # if epoch % hyp['epoch_per_save'] == 0:
+    #     save_path = model_save_path + str(epoch).rjust(6, '0') + f'_{data_name}_{net_name}.pth'
+    #     torch.save({'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict()}, save_path)
+    #     print(f"Saved: {save_path}")

OM_reg.py

@@ -0,0 +1,240 @@
+import torch
+import torchvision
+from torch import nn
+from torchvision.utils import save_image
+from torch.utils.data import DataLoader
+from torch.optim import Adam
+from torchvision.utils import make_grid
+from Diffusion.diffuser import DeformDDPM
+from Diffusion.networks import get_net, STN
+from torchvision.transforms import Lambda
+import random
+import os
+import utils
+from Dataloader.dataloader0 import get_dataloader
+from Dataloader.dataLoader import *
+
+from torchvision.utils import save_image
+from einops import rearrange, reduce, repeat
+# import matplotlib.image
+import numpy as np
+import nibabel as nib
+from tqdm import tqdm 
+import yaml
+import argparse
+
+EPS = 10e-8
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument(
+        "--config",
+        "-C",
+        help="Path for the config file",
+        type=str,
+        default="Config/config_cmr.yaml",
+        # default="Config/config_lct.yaml",
+        required=False,
+    )
+args = parser.parse_args()
+#=======================================================================================================================
+
+# config_path = 'Config/config_cmr.yaml'
+# config_path = 'Config/config_lct.yaml'
+
+# Load the YAML file into a dictionary
+with open(args.config, 'r') as file:
+    hyp_parameters = yaml.safe_load(file)
+    print(hyp_parameters)
+# hyp_parameters["aug_img_savepath"] = os.path.join(hyp_parameters["aug_img_savepath"],hyp_parameters["data_name"],'')
+if not os.path.exists(hyp_parameters["aug_img_savepath"]):
+    os.makedirs(hyp_parameters["aug_img_savepath"])
+if not os.path.exists(hyp_parameters["aug_msk_savepath"]):
+    os.makedirs(hyp_parameters["aug_msk_savepath"])
+if not os.path.exists(hyp_parameters["aug_ddf_savepath"]):
+    os.makedirs(hyp_parameters["aug_ddf_savepath"])
+print(hyp_parameters["aug_img_savepath"])
+
+hyp_parameters['batchsize'] = 1
+
+
+# =======================================================================================================================
+# min_crop_ratio = 0.5
+min_crop_ratio = 0.9
+
+# dataset = OminiDataset_v1(transform=None,min_crop_ratio=min_crop_ratio)
+# Infer_Loader = DataLoader(
+#         dataset,
+#         batch_size=hyp_parameters['batchsize'],
+#         shuffle=False
+#     )
+
+# label_keys = ['heart']
+label_keys = ['brain']
+# label_keys = ['pancreas']
+database = ['MSD']
+
+dataset = OminiDataset_inference_w_all(transform=None,min_crop_ratio=min_crop_ratio,label_key = label_keys, database=database)
+Infer_Loader = DataLoader(
+        dataset,
+        batch_size=hyp_parameters['batchsize'],
+        shuffle=False
+    )
+# =======================================================================================================================
+
+# Data_Loader=get_dataloader(hyp_parameters['data_name'],mode='aug')
+# transformer = utils.get_transformer(img_sz=hyp_parameters["ndims"]*[hyp_parameters['img_size']])
+# dataset = Data_Loader(patient_index = hyp_parameters["patients_list"])
+# train_loader = DataLoader(dataset, batch_size = hyp_parameters['batchsize'], shuffle = False) 
+
+
+
+epoch=f'{hyp_parameters["model_id_str"]}_{hyp_parameters["data_name"]}_{hyp_parameters["net_name"]}'
+model_save_path = f'Models/{hyp_parameters["data_name"]}_{hyp_parameters["net_name"]}/'
+model_save_path = os.path.join(model_save_path, str(epoch)+'.pth')
+
+
+
+Net = get_net(hyp_parameters["net_name"])
+
+Deformddpm = DeformDDPM(
+    network=Net(n_steps = hyp_parameters["timesteps"],
+                ndims = hyp_parameters["ndims"],
+                num_input_chn = hyp_parameters["num_input_chn"],
+                res = hyp_parameters['img_size']
+                ),
+    n_steps = hyp_parameters["timesteps"],
+    image_chw = [hyp_parameters["num_input_chn"]] + [hyp_parameters["img_size"]]*hyp_parameters["ndims"],
+    device = hyp_parameters["device"],
+    batch_size = hyp_parameters["batchsize"],
+    img_pad_mode = hyp_parameters["img_pad_mode"],
+    ddf_pad_mode = hyp_parameters["ddf_pad_mode"],
+    padding_mode = hyp_parameters["padding_mode"],
+    v_scale = hyp_parameters["v_scale"],
+    resample_mode = hyp_parameters["resample_mode"],
+)
+Deformddpm.to(hyp_parameters["device"])
+
+ddf_stn = STN(
+    img_sz = hyp_parameters["img_size"],
+    ndims = hyp_parameters["ndims"],
+    padding_mode = hyp_parameters['padding_mode'],
+    device = hyp_parameters["device"],
+)
+ddf_stn.to(hyp_parameters["device"])
+
+print("Loading model from:", model_save_path)
+# Deformddpm.load_state_dict(torch.load(model_save_path))
+checkpoint = torch.load(model_save_path)
+Deformddpm.load_state_dict(checkpoint['model_state_dict'])
+Deformddpm.eval()
+
+os.makedirs(hyp_parameters['reg_img_savepath'], exist_ok=True)
+os.makedirs(hyp_parameters['reg_msk_savepath'], exist_ok=True)
+os.makedirs(hyp_parameters['reg_ddf_savepath'], exist_ok=True)
+
+print("total num of image:", len(Infer_Loader))
+for e, d in tqdm(enumerate(Infer_Loader)):
+# for e, d in enumerate(Infer_Loader):
+  # img, mask, pid = d
+  # img = d
+  # mask = d
+  img = d['img']
+  mask = d['labels']
+  
+  # pid = pid.cpu().detach().numpy()
+  # pid = pid[0] 
+  pid = e
+
+  print('Processing to patient:', pid, ' image:',e)
+  
+  img = img.to(hyp_parameters["device"]) 
+  img = img.type(torch.float32)
+  image_original = img.cpu().detach().numpy()
+  # 
+  # 
+  if e <= 0:
+    target_img = img.clone().detach()  # save the first image as target image for conditioning
+
+  mask = mask.to(hyp_parameters["device"]) 
+  mask = mask.type(torch.float32)
+  mask_original = mask.cpu().detach().numpy()
+  # print(pid, image_original.shape, mask_original.max())
+
+
+  if hyp_parameters["ndims"] == 2:
+    nifti_img = nib.Nifti1Image(image_original[0,0,:,:], np.eye(4))  
+    nifti_mask = nib.Nifti1Image(mask_original[0,:,:,:], np.eye(4))  
+  elif hyp_parameters["ndims"] == 3:
+    nifti_img = nib.Nifti1Image(image_original[0,0,:,:,:], np.eye(4))  
+    nifti_mask = nib.Nifti1Image(mask_original[0,0,:,:,:], np.eye(4))  
+
+  # Saving original (undeformed image)
+  # CMR: format: Patient0001_Slice0001_ORG_NA.nii.gz
+  # Lung CT: Patient0001_Slice0001_ORG_NA.nii.gz
+  nib.save(nifti_img, os.path.join(hyp_parameters['reg_img_savepath'],utils.get_barcode([pid,e])+'.nii.gz'))
+
+  # Saving original (undeformed image)
+  # CMR: format: Patient0001_Slice0001_ORG_NA_GT.nii.gz
+  # Lung CT: ...
+  nib.save(nifti_img, os.path.join(hyp_parameters['reg_msk_savepath'],utils.get_barcode([pid,e])+'_GT.nii.gz'))
+
+ 
+  noise_step = hyp_parameters["start_noise_step"]
+  with torch.no_grad():
+    for im in range(1):
+      # # Permute 
+      # if hyp_parameters["ndims"] == 2:
+      #   [img, mask] = utils.random_permute([img, mask], select_dims=[-1, -2])          # add random rotation to image
+      # elif hyp_parameters["ndims"] == 3:
+      #   [img, mask] = utils.random_permute([img, mask], select_dims=[-1, -2, -3])  # add random rotation to image
+
+      print('Generating - >', 'Subject-',pid,', Scan-',e,' (',im,'/',hyp_parameters["aug_coe"],')', end='\r')
+      
+      [ddf_comp,ddf_rand],[img_rec,img_diff,img_save],[msk_rec,msk_diff,msk_save] = Deformddpm.diff_recover(img_org=img,cond_imgs=target_img.clone().detach(),msk_org=mask,T=[None,hyp_parameters["timesteps"]],v_scale=hyp_parameters["v_scale"],t_save=None,proc_type=hyp_parameters["condition_type"])
+      
+      denoise_imgs = img_rec.cpu().detach().numpy()
+      denoise_msks = msk_rec.cpu().detach().numpy()
+      noisy_imgs_np = img_diff.cpu().detach().numpy()
+      noisy_msks_np = msk_diff.cpu().detach().numpy()
+
+      if hyp_parameters["ndims"] == 2:
+        nifti_img_aug = nib.Nifti1Image(denoise_imgs[0,0,:,:], np.eye(4))  
+        nifti_mask_aug = nib.Nifti1Image(denoise_msks[0,:,:,:], np.eye(4))  
+        nifti_img = nib.Nifti1Image(noisy_imgs_np[0,0,:,:], np.eye(4))
+        nifti_mask = nib.Nifti1Image(noisy_msks_np[0, :, :, :], np.eye(4))
+      elif hyp_parameters["ndims"] == 3:
+        nifti_img_aug = nib.Nifti1Image(denoise_imgs[0,0,:,:,:], np.eye(4))  
+        nifti_mask_aug = nib.Nifti1Image(denoise_msks[0,0,:,:,:], np.eye(4))  
+        nifti_img = nib.Nifti1Image(noisy_imgs_np[0,0,:,:,:], np.eye(4))
+        nifti_mask = nib.Nifti1Image(noisy_msks_np[0, 0, :, :], np.eye(4))
+      
+      nib.save(nifti_img_aug, os.path.join(hyp_parameters['reg_img_savepath'],utils.get_barcode([pid,e,im,noise_step])+'.nii.gz'))
+      nib.save(nifti_mask_aug, os.path.join(hyp_parameters['reg_msk_savepath'],utils.get_barcode([pid,e,im,noise_step])+'_GT.nii.gz'))
+      
+      # Saving noisy image to nifti
+      # CMR: format: Patient0001_Slice0001_NosieImg0001_NoiseStep0070.nii.gz
+      # Lung CT: ...
+      nib.save(nifti_img, os.path.join(hyp_parameters['reg_img_savepath'],utils.get_barcode([pid,e,im,noise_step],header=['Patient','Slice','NoiseImg','NoiseStep'])+'.nii.gz'))
+      nib.save(nifti_mask, os.path.join(hyp_parameters['reg_msk_savepath'],utils.get_barcode([pid,e,im,noise_step],header=['Patient','Slice','NoiseImg','NoiseStep'])+'_GT.nii.gz'))
+      
+          
+      if (im - hyp_parameters["start_noise_step"])%2 == 0:
+        noise_step = noise_step + hyp_parameters["noise_step"]
+      # break   # for testing
+  if e > 5:
+    break
+          
+
+
+
+
+
+
+
+
+
+
+
+
+

OM_train.py

@@ -0,0 +1,309 @@
+import os
+import gc
+import torch
+import torchvision
+from torch import nn
+from torchvision.utils import save_image
+from torch.utils.data import DataLoader
+
+from torch.optim import Adam, SGD
+from Diffusion.diffuser import DeformDDPM
+from Diffusion.networks import get_net, STN
+from torchvision.transforms import Lambda
+import Diffusion.losses as losses
+import random
+import glob
+import numpy as np
+import utils
+
+from Dataloader.dataloader0 import get_dataloader
+from Dataloader.dataLoader import *
+
+from Dataloader.dataloader_utils import thresh_img
+import yaml
+import argparse
+
+####################
+import torch.multiprocessing as mp
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
+# from torch.distributed import init_process_group
+###############
+def ddp_setup(rank, world_size):
+    """
+    Args:
+        rank: Unique identifier of each process
+        world_size: Total number of processes
+    """
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12355"
+    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    torch.cuda.set_device(rank)
+
+use_distributed = True
+# use_distributed = False
+
+EPS = 1e-5
+
+parser = argparse.ArgumentParser()
+
+# config_file_path = 'Config/config_cmr.yaml'
+parser.add_argument(
+        "--config",
+        "-C",
+        help="Path for the config file",
+        type=str,
+        # default="Config/config_cmr.yaml",
+        # default="Config/config_lct.yaml",
+        default="Config/config_all.yaml",
+        required=False,
+    )
+args = parser.parse_args()
+#=======================================================================================================================
+
+
+
+def main_train(rank=0,world_size=1):
+    if use_distributed:
+        ddp_setup(rank,world_size)
+    gpu_id = rank
+
+    # Load the YAML file into a dictionary
+    with open(args.config, 'r') as file:
+        hyp_parameters = yaml.safe_load(file)
+        print(hyp_parameters)
+
+    # epoch_per_save=10
+    epoch_per_save=hyp_parameters['epoch_per_save']
+
+    data_name=hyp_parameters['data_name']
+    net_name = hyp_parameters['net_name']
+
+    Net=get_net(net_name)
+
+    suffix_pth=f'_{data_name}_{net_name}.pth'
+    model_save_path = os.path.join('Models',f'{data_name}_{net_name}/')
+    model_dir=model_save_path
+    transformer=utils.get_transformer(img_sz=hyp_parameters["ndims"]*[hyp_parameters['img_size']])
+    
+    # Data_Loader=get_dataloader(data_name=hyp_parameters['data_name'], mode='train')
+
+    # tsfm = torchvision.transforms.Compose([
+    #             torchvision.transforms.ToTensor(),
+    #             ])
+
+    # dataset = Data_Loader(target_res = [hyp_parameters["img_size"]]*hyp_parameters["ndims"], transforms=None, noise_scale=hyp_parameters['noise_scale'])
+    # train_loader = DataLoader(
+    #     dataset,
+    #     batch_size=hyp_parameters['batchsize'],
+    #     # shuffle=False,
+    #     shuffle=True,
+    #     drop_last=True,
+    # )
+
+    dataset = OminiDataset_v1(transform=None)
+    train_loader = DataLoader(
+        dataset,
+        batch_size=hyp_parameters['batchsize'],
+        shuffle=True,
+        drop_last=True,
+    )
+
+
+
+    Deformddpm = DeformDDPM(
+        network=Net(
+            n_steps=hyp_parameters["timesteps"], 
+            ndims=hyp_parameters["ndims"], 
+            num_input_chn = hyp_parameters["num_input_chn"],
+            res = hyp_parameters['img_size']
+            ),
+        n_steps=hyp_parameters["timesteps"],
+        image_chw=[1] + [hyp_parameters["img_size"]]*hyp_parameters["ndims"],
+        device=hyp_parameters["device"],
+        batch_size=hyp_parameters["batchsize"],
+        img_pad_mode=hyp_parameters["img_pad_mode"],
+        v_scale=hyp_parameters["v_scale"],
+    )
+
+
+    ddf_stn = STN(
+        img_sz=hyp_parameters["img_size"],
+        ndims=hyp_parameters["ndims"],
+        # padding_mode="zeros",
+        padding_mode=hyp_parameters["padding_mode"],
+        device=hyp_parameters["device"],
+    )
+
+
+    if use_distributed:
+        Deformddpm.to(rank)
+        Deformddpm = DDP(Deformddpm, device_ids=[rank])
+        ddf_stn.to(rank)
+    else:
+        Deformddpm.to(hyp_parameters["device"])
+        ddf_stn.to(hyp_parameters["device"])
+    # ddf_stn = DDP(ddf_stn, device_ids=[rank])
+
+
+    # mse = nn.MSELoss()
+    loss_reg = losses.Grad(penalty=['l1', 'negdetj'], ndims=hyp_parameters["ndims"])
+    loss_dist = losses.MRSE(img_sz=hyp_parameters["img_size"])
+    # loss_ang = losses.MRSE(img_sz=hyp_parameters["img_size"])
+    loss_ang = losses.NCC(img_sz=hyp_parameters["img_size"])
+
+    optimizer = Adam(Deformddpm.parameters(), lr=hyp_parameters["lr"])
+    # hyp_parameters["lr"]=0.00000001
+    # # optimizer = SGD(Deformddpm.parameters(), lr=hyp_parameters["lr"], momentum=0.95)
+    # optimizer = SGD(Deformddpm.parameters(), lr=hyp_parameters["lr"], momentum=0.9)
+
+    # # LR scheduler ----- YHM
+    # scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, hyp_parameters["lr"], hyp_parameters["lr"]*10, step_size_up=500, step_size_down=500, mode='triangular', gamma=1.0, scale_fn=None, scale_mode='cycle', cycle_momentum=True, base_momentum=0.8, max_momentum=0.9, last_epoch=-1)
+
+    # Deformddpm.network.load_state_dict(torch.load('/home/data/jzheng/Adaptive_Motion_Generator-master/models/1000.pth'))
+
+    # check for existing models
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir, exist_ok=True)
+    model_files = glob.glob(os.path.join(model_dir, "*.pth"))
+    model_files.sort()
+    if model_files:
+        if gpu_id == 0:
+            print(model_files)
+        initial_epoch, Deformddpm, optimizer = ddp_load_dict(gpu_id, Deformddpm, optimizer, model_files[-1])
+    else:
+        initial_epoch = 0
+
+    if gpu_id == 0:
+        print('len_train_data: ',len(dataset))
+    for epoch in range(initial_epoch,hyp_parameters["epoch"]):
+
+        epoch_loss_tot = 0.0
+        epoch_loss_gen_d = 0.0
+        epoch_loss_gen_a = 0.0
+        epoch_loss_reg = 0.0
+        # Set model inside to train model
+        Deformddpm.train()
+
+        for step, batch in enumerate(train_loader):
+        # for step, batch in enumerate(train_loader_omni):
+            # x0, _ = batch
+            x0 = batch # for omni dataset
+            x0 = x0.to(hyp_parameters["device"]).type(torch.float32)
+        
+            n = x0.size()[0]  # batch_size -> n
+            x0 = x0.to(hyp_parameters["device"])
+            
+            blind_mask = utils.get_random_deformed_mask(x0.shape[2:],apply_possibility=0.6).to(hyp_parameters["device"])
+
+            # random deformation + rotation
+            if hyp_parameters["ndims"]>2:
+                if np.random.uniform(0,1)<0.6:
+                    x0 = utils.random_resample(x0, deform_scale=0)
+            x0 = transformer(x0)
+            if hyp_parameters['noise_scale']>0:
+                x0 = thresh_img(x0, [0, 2*hyp_parameters['noise_scale']])
+                x0 = x0 * (np.random.normal(1, hyp_parameters['noise_scale'] * 1)) + np.random.normal(0, hyp_parameters['noise_scale'] * 1)
+
+            # Picking some noise for each of the images in the batch, a timestep and the respective alpha_bars
+            t = torch.randint(0, hyp_parameters["timesteps"], (n,)).to(
+                hyp_parameters["device"]
+            )  # pick up a seq of rand number from 0 to 'timestep'
+
+            
+            pre_dvf_I,dvf_I = Deformddpm(img_org=x0, t=t, mask=blind_mask)  # forward diffusion process
+
+            loss_tot=0
+
+            loss_ddf = loss_reg(pre_dvf_I)
+            trm_pred = ddf_stn(pre_dvf_I, dvf_I)
+            loss_gen_d = loss_dist(pred=trm_pred,inv_lab=dvf_I,ddf_stn=None,mask=blind_mask)
+            loss_gen_a = loss_ang(pred=trm_pred,inv_lab=dvf_I,ddf_stn=None,mask=blind_mask)
+
+            loss_tot += 1. * loss_gen_d + 1. * loss_gen_a
+            loss_tot += 1.0 * loss_ddf
+            optimizer.zero_grad()
+            loss_tot.backward()
+            optimizer.step()
+
+            epoch_loss_tot += loss_tot.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_gen_d += loss_gen_d.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_gen_a += loss_gen_a.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_reg += loss_ddf.item() * len(x0) / len(train_loader.dataset)
+            # print('step:',step,':', loss_tot.item(),'=',loss_gen_a.item(),'+', loss_gen_d.item(),'+',loss_ddf.item())
+
+            # break   # FOR TESTING
+
+        if gpu_id == 0:
+            print(epoch,':', epoch_loss_tot,'=',epoch_loss_gen_a,'+', epoch_loss_gen_d,'+',epoch_loss_reg, ' (ang+dist+regul)')
+
+        # # LR schedular step ----- YHM
+        # scheduler.step()
+
+        if 0 == epoch % epoch_per_save:
+            save_dir=model_save_path + str(epoch).rjust(6, '0') + suffix_pth    
+            os.makedirs(os.path.dirname(model_save_path), exist_ok=True)
+            # break   # FOR TESTING
+            if not use_distributed:
+                print(f"saved in {save_dir}")
+                # torch.save(Deformddpm.state_dict(), save_dir)
+                torch.save({
+                    'model_state_dict': Deformddpm.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'epoch': epoch
+                }, save_dir)
+            elif gpu_id == 0:
+                print(f"saved in {save_dir}")
+                # torch.save(Deformddpm.module.state_dict(), save_dir)
+                torch.save({
+                    'model_state_dict': Deformddpm.module.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'epoch': epoch
+                }, save_dir)
+
+def ddp_load_dict(gpu_id, Deformddpm, optimizer, model_file,use_distributed=True):
+    
+    if gpu_id == 0:
+    # if 0:
+        utils.print_memory_usage("Before Loading Model")
+        if 1:
+            gc.collect()
+            torch.cuda.empty_cache()
+        # Deformddpm.network.load_state_dict(torch.load(latest_model_file))
+        # Deformddpm.load_state_dict(torch.load(latest_model_file), strict=False)
+        checkpoint = torch.load(model_file)
+        # checkpoint = torch.load(latest_model_file, map_location=f"cuda:{rank}")
+        if use_distributed:
+            Deformddpm.module.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            Deformddpm.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        utils.print_memory_usage("After Loading Checkpoint on GPU")
+
+    if use_distributed:
+        # Broadcast model weights from rank 0 to all other GPUs
+        dist.barrier()
+        for param in Deformddpm.parameters():
+            dist.broadcast(param.data, src=0)  # Synchronize model across ranks
+        dist.barrier()
+        for param_group in optimizer.param_groups:
+            for param in param_group['params']:
+                if param.grad is not None:
+                    dist.broadcast(param.grad, src=0)  # Sync optimizer gradients
+        
+    # initial_epoch = checkpoint['epoch'] + 1
+    # get the epoch number from the filename and add 1 to set as initial_epoch
+    initial_epoch = int(os.path.basename(model_file).split('.')[0][:6]) + 1
+
+    return initial_epoch, Deformddpm, optimizer
+
+            
+
+if __name__ == "__main__":
+    if use_distributed:
+        world_size = torch.cuda.device_count()
+        print(f"Distributed GPU number = {world_size}")
+        mp.spawn(main_train,args = (world_size,),nprocs = world_size)
+    else:
+        main_train(0,1)

OM_train_2modes.py

@@ -0,0 +1,528 @@
+import os
+import gc
+import torch
+import torchvision
+from torch import nn
+from torchvision.utils import save_image
+from torch.utils.data import DataLoader
+
+from torch.optim import Adam, SGD
+from Diffusion.diffuser import DeformDDPM
+from Diffusion.networks import get_net, STN
+from torchvision.transforms import Lambda
+import Diffusion.losses as losses
+import random
+import glob
+import numpy as np
+import utils
+from tqdm import tqdm 
+
+from Dataloader.dataloader0 import get_dataloader
+from Dataloader.dataLoader import *
+
+from Dataloader.dataloader_utils import thresh_img
+import yaml
+import argparse
+
+####################
+import torch.multiprocessing as mp
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
+# from torch.distributed import init_process_group
+###############
+def ddp_setup(rank, world_size):
+    """
+    Args:
+        rank: Unique identifier of each process
+        world_size: Total number of processes
+    """
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12355"
+    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    torch.cuda.set_device(rank)
+
+use_distributed = True
+# use_distributed = False
+
+EPS = 1e-5
+MSK_EPS = 0.01
+TEXT_EMBED_PROB = 0.7
+AUG_RESAMPLE_PROB = 0.6
+LOSS_WEIGHTS_DIFF = [2.0, 1.0, 30]  # [ang, dist, reg]
+# LOSS_WEIGHTS_REGIST = [9.0, 1.0, 16.0]  # [imgsim, imgmse, ddf]
+# LOSS_WEIGHTS_REGIST = [10.0, 1.0, 1.0]  # [imgsim, imgmse, ddf]
+# LOSS_WEIGHTS_REGIST = [2.0, 0.1, 1e3]  # [imgsim, imgmse, ddf]
+LOSS_WEIGHTS_REGIST = [2.0, 0.1, 256]  # [imgsim, imgmse, ddf]
+
+# AUG_PERMUTE_PROB = 0.35
+
+parser = argparse.ArgumentParser()
+
+# config_file_path = 'Config/config_cmr.yaml'
+parser.add_argument(
+        "--config",
+        "-C",
+        help="Path for the config file",
+        type=str,
+        # default="Config/config_cmr.yaml",
+        # default="Config/config_lct.yaml",
+        default="Config/config_all.yaml",
+        required=False,
+    )
+args = parser.parse_args()
+#=======================================================================================================================
+
+
+
+def main_train(rank=0,world_size=1,train_mode_ratio=1,thresh_imgsim=0.01):
+    if use_distributed:
+        ddp_setup(rank,world_size)
+
+        if torch.distributed.is_initialized():
+            print(f"World size: {torch.distributed.get_world_size()}")
+            print(f"Communication backend: {torch.distributed.get_backend()}")
+    gpu_id = rank
+    
+    # Load the YAML file into a dictionary
+    with open(args.config, 'r') as file:
+        hyp_parameters = yaml.safe_load(file)
+        print(hyp_parameters)
+
+    # epoch_per_save=10
+    epoch_per_save=hyp_parameters['epoch_per_save']
+
+    data_name=hyp_parameters['data_name']
+    net_name = hyp_parameters['net_name']
+
+    Net=get_net(net_name)
+
+    suffix_pth=f'_{data_name}_{net_name}.pth'
+    model_save_path = os.path.join('Models',f'{data_name}_{net_name}/')
+    model_dir=model_save_path
+    transformer=utils.get_transformer(img_sz=hyp_parameters["ndims"]*[hyp_parameters['img_size']])
+    
+    # Data_Loader=get_dataloader(data_name=hyp_parameters['data_name'], mode='train')
+
+    # tsfm = torchvision.transforms.Compose([
+    #             torchvision.transforms.ToTensor(),
+    #             ])
+
+    # dataset = Data_Loader(target_res = [hyp_parameters["img_size"]]*hyp_parameters["ndims"], transforms=None, noise_scale=hyp_parameters['noise_scale'])
+    # train_loader = DataLoader(
+    #     dataset,
+    #     batch_size=hyp_parameters['batchsize'],
+    #     # shuffle=False,
+    #     shuffle=True,
+    #     drop_last=True,
+    # )
+
+    # dataset = OminiDataset_v1(transform=None)
+    dataset = OMDataset_indiv(transform=None)
+    train_loader = DataLoader(
+        dataset,
+        batch_size=hyp_parameters['batchsize'],
+        shuffle=True,
+        drop_last=True,
+    )
+
+    # datasetp = OminiDataset_paired(transform=None)
+    datasetp = OMDataset_pair(transform=None)
+    train_loader_p = DataLoader(
+        datasetp,
+        batch_size=hyp_parameters['batchsize']//2,
+        shuffle=True,
+        drop_last=True,
+    )
+
+
+
+    Deformddpm = DeformDDPM(
+        network=Net(
+            n_steps=hyp_parameters["timesteps"], 
+            ndims=hyp_parameters["ndims"], 
+            num_input_chn = hyp_parameters["num_input_chn"],
+            res = hyp_parameters['img_size']
+            ),
+        n_steps=hyp_parameters["timesteps"],
+        image_chw=[1] + [hyp_parameters["img_size"]]*hyp_parameters["ndims"],
+        device=hyp_parameters["device"],
+        batch_size=hyp_parameters["batchsize"],
+        img_pad_mode=hyp_parameters["img_pad_mode"],
+        v_scale=hyp_parameters["v_scale"],
+    )
+
+
+    ddf_stn = STN(
+        img_sz=hyp_parameters["img_size"],
+        ndims=hyp_parameters["ndims"],
+        # padding_mode="zeros",
+        padding_mode=hyp_parameters["padding_mode"],
+        device=hyp_parameters["device"],
+    )
+
+
+    if use_distributed:
+        Deformddpm.to(rank)
+        Deformddpm = DDP(Deformddpm, device_ids=[rank])
+        ddf_stn.to(rank)
+    else:
+        Deformddpm.to(hyp_parameters["device"])
+        ddf_stn.to(hyp_parameters["device"])
+    # ddf_stn = DDP(ddf_stn, device_ids=[rank])
+
+
+    # mse = nn.MSELoss()
+    # loss_reg = losses.Grad(penalty=['l1', 'negdetj'], ndims=hyp_parameters["ndims"])
+    loss_reg = losses.Grad(penalty=['l1', 'negdetj', 'range'], ndims=hyp_parameters["ndims"],outrange_thresh=0.2,outrange_weight=1e3)
+    loss_reg1 = losses.Grad(penalty=['l1', 'negdetj', 'range'], ndims=hyp_parameters["ndims"],outrange_thresh=0.6,outrange_weight=1e3)
+    loss_dist = losses.MRSE(img_sz=hyp_parameters["img_size"])
+    # loss_ang = losses.MRSE(img_sz=hyp_parameters["img_size"])
+    loss_ang = losses.NCC(img_sz=hyp_parameters["img_size"])
+    loss_imgsim = losses.LNCC()
+    loss_imgmse = losses.LMSE()
+
+    optimizer = Adam(Deformddpm.parameters(), lr=hyp_parameters["lr"])
+    # hyp_parameters["lr"]=0.00000001
+    # optimizer_regist = Adam(Deformddpm.parameters(), lr=hyp_parameters["lr"]*0.01)
+    # optimizer_regist = SGD(Deformddpm.parameters(), lr=hyp_parameters["lr"]*0.01, momentum=0.98)
+    # optimizer = SGD(Deformddpm.parameters(), lr=hyp_parameters["lr"], momentum=0.9)
+
+    # # LR scheduler ----- YHM
+    # scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, hyp_parameters["lr"], hyp_parameters["lr"]*10, step_size_up=500, step_size_down=500, mode='triangular', gamma=1.0, scale_fn=None, scale_mode='cycle', cycle_momentum=True, base_momentum=0.8, max_momentum=0.9, last_epoch=-1)
+
+    # Deformddpm.network.load_state_dict(torch.load('/home/data/jzheng/Adaptive_Motion_Generator-master/models/1000.pth'))
+
+    # check for existing models
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir, exist_ok=True)
+    model_files = glob.glob(os.path.join(model_dir, "*.pth"))
+    model_files.sort()
+    if model_files:
+        if gpu_id == 0:
+            print(model_files)
+        initial_epoch, Deformddpm, optimizer = ddp_load_dict(gpu_id, Deformddpm, optimizer, model_files[-1])
+    else:
+        initial_epoch = 0
+
+    if gpu_id == 0:
+        print('len_train_data: ',len(dataset))
+    # Training loop
+    for epoch in range(initial_epoch,hyp_parameters["epoch"]):
+
+        epoch_loss_tot = 0.0
+        epoch_loss_gen_d = 0.0
+        epoch_loss_gen_a = 0.0
+        epoch_loss_reg = 0.0
+        epoch_loss_regist = 0.0
+        epoch_loss_imgsim = 0.0
+        epoch_loss_imgmse = 0.0
+        epoch_loss_ddfreg = 0.0
+        # Set model inside to train model
+        Deformddpm.train()
+        
+        loss_nan_step = 0  # yu: count the number of nan loss steps
+
+        total = min(len(train_loader), len(train_loader_p))
+        for step, (batch, batch_p) in tqdm(enumerate(zip(train_loader, train_loader_p)), total=total):
+        # for step, batch in tqdm(enumerate(train_loader)):
+        # for step, batch in tqdm(enumerate(train_loader)):
+
+        # for step, batch in enumerate(train_loader_omni):
+            # x0, _ = batch
+            
+            
+            # ==========================================================================
+            # diffusion train on single image
+
+            # x0 = batch # for omni dataset
+            [x0,embd] = batch # for om dataset
+            x0 = x0.to(hyp_parameters["device"]).type(torch.float32)
+            # print('embd:', embd.shape)
+            if np.random.uniform(0,1)<TEXT_EMBED_PROB:
+                embd = embd.to(hyp_parameters["device"]).type(torch.float32)
+            else:
+                embd = None
+
+            
+
+            n = x0.size()[0]  # batch_size -> n
+            x0 = x0.to(hyp_parameters["device"])
+            
+            blind_mask = utils.get_random_deformed_mask(x0.shape[2:],apply_possibility=0.6).to(hyp_parameters["device"])
+
+            # random deformation + rotation
+            if hyp_parameters["ndims"]>2:
+                if np.random.uniform(0,1)<AUG_RESAMPLE_PROB:
+                    x0 = utils.random_resample(x0, deform_scale=0)
+                # elif np.random.uniform(0,1)<AUG_RESAMPLE_PROB+AUG_PERMUTE_PROB:
+                else:
+                    [x0] = utils.random_permute([x0], select_dims=[-1,-2,-3])
+            x0 = transformer(x0)
+            if hyp_parameters['noise_scale']>0:
+                if np.random.uniform(0,1)<AUG_RESAMPLE_PROB:
+                    x0 = thresh_img(x0, [0, 1*hyp_parameters['noise_scale']])
+                x0 = x0 * (np.random.normal(1, hyp_parameters['noise_scale'] * 1)) + np.random.normal(0, hyp_parameters['noise_scale'] * 1)
+
+            # Picking some noise for each of the images in the batch, a timestep and the respective alpha_bars
+            t = torch.randint(0, hyp_parameters["timesteps"], (n,)).to(
+                hyp_parameters["device"]
+            )  # pick up a seq of rand number from 0 to 'timestep'
+
+            # proc_type = random.choice(['adding', 'independ', 'downsample', 'slice', 'project', 'none', 'uncon', 'uncon', 'uncon'])
+            proc_type = random.choice(['adding', 'downsample', 'slice', 'none', 'uncon', 'uncon', 'uncon'])
+            # print('proc_type:', proc_type)
+            cond_img, _, cond_ratio = Deformddpm.module.proc_cond_img(x0,proc_type=proc_type)
+
+            pre_dvf_I,dvf_I = Deformddpm(img_org=x0, t=t, cond_imgs=cond_img, mask=blind_mask,proc_type=[],text=embd)  # forward diffusion process
+
+            loss_tot=0
+
+            loss_ddf = loss_reg(pre_dvf_I,img=x0)
+            trm_pred = ddf_stn(pre_dvf_I, dvf_I)
+            loss_gen_d = loss_dist(pred=trm_pred,inv_lab=dvf_I,ddf_stn=None,mask=blind_mask)
+            loss_gen_a = loss_ang(pred=trm_pred,inv_lab=dvf_I,ddf_stn=None,mask=blind_mask)
+
+            loss_tot += LOSS_WEIGHTS_DIFF[0] * loss_gen_a + LOSS_WEIGHTS_DIFF[1] * loss_gen_d
+            loss_tot += LOSS_WEIGHTS_DIFF[2] * loss_ddf
+            loss_tot = torch.sqrt(1.+MSK_EPS-cond_ratio) * loss_tot
+
+            # >> JZ: print nan in x0
+            if torch.isnan(x0).any():
+                print(f"*** Encountered NaN in input image x0 at epoch {epoch}, step {step}.")
+            # >> JZ: print loss of ddf
+            if loss_ddf>0.001:
+                print(f"*** High diffusion DDF loss at epoch {epoch}, step {step}: {loss_ddf.item()}.")
+            # yu: check if loss_tot==nan or inf
+            if torch.isnan(loss_tot) or torch.isinf(loss_tot):
+                print(f"*** Encountered NaN or Inf loss at epoch {epoch}, step {step}. Skipping this batch.")
+                loss_nan_step += 1
+                continue
+            if loss_nan_step > 5:
+                print(f"*** Too many NaN or Inf losses ({loss_nan_step} times) at epoch {epoch}, step {step}. Stopping training.")
+                raise ValueError("Too many NaN losses detected in loss_tot. Code terminated.")
+
+
+            optimizer.zero_grad()
+            loss_tot.backward()
+            optimizer.step()
+
+            epoch_loss_tot += loss_tot.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_gen_d += loss_gen_d.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_gen_a += loss_gen_a.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_reg += loss_ddf.item() * len(x0) / len(train_loader.dataset)
+
+            # print(loss_gen_a.item())
+            # if 0:
+            # if loss_gen_a.item() < -0.3 and step%train_mode_ratio == 0:
+            if step%train_mode_ratio == 0:
+                # ==========================================================================
+                # registration train on paired images
+                # x1, y1 = next(iter(train_loader_p))
+                # [x1, y1, _, embd_y] = next(iter(train_loader_p))
+                [x1, y1, _, embd_y] = batch_p
+                if np.random.uniform(0,1)<TEXT_EMBED_PROB:
+                    # embd_x = embd_x.to(hyp_parameters["device"]).type(torch.float32)
+                    embd_y = embd_y.to(hyp_parameters["device"]).type(torch.float32)
+                else:
+                    # embd_x = None
+                    embd_y = None
+
+                x1 = x1.to(hyp_parameters["device"]).type(torch.float32)
+                y1 = y1.to(hyp_parameters["device"]).type(torch.float32)
+                n = x1.size()[0]  # batch_size -> n
+                # random deformation + rotation
+                # if hyp_parameters["ndims"]>2:
+                #     if np.random.uniform(0,1)<0.6:
+                #         x1 = utils.random_resample(x1, deform_scale=0)
+                #         y1 = utils.random_resample(y1, deform_scale=0)
+                x1 = transformer(x1)
+                y1 = transformer(y1)
+                [x1, y1] = utils.random_permute([x1, y1], select_dims=[-1,-2,-3])
+                if hyp_parameters['noise_scale']>0:
+                    [x1, y1] = thresh_img([x1, y1], [0, 2*hyp_parameters['noise_scale']])
+                    random_scale = np.random.normal(1, hyp_parameters['noise_scale'] * 1)
+                    random_shift = np.random.normal(0, hyp_parameters['noise_scale'] * 1)
+                    x1 = x1 * random_scale + random_shift
+                    y1 = y1 * random_scale + random_shift
+                    # x1 = thresh_img(x1, [0, 2*hyp_parameters['noise_scale']])
+                    # x1 = x1 * (np.random.normal(1, hyp_parameters['noise_scale'] * 1)) + np.random.normal(0, hyp_parameters['noise_scale'] * 1)
+                    # y1 = thresh_img(y1, [0, 2*hyp_parameters['noise_scale']])
+                    # y1 = y1 * (np.random.normal(1, hyp_parameters['noise_scale'] * 1)) + np.random.normal(0, hyp_parameters['noise_scale'] * 1)
+                # # Picking some noise for each of the images in the batch, a timestep and the respective alpha_bars
+                # t = torch.randint(0, hyp_parameters["timesteps"], (n,)).to(
+                #     hyp_parameters["device"]
+                # )  # pick up a seq of rand number from 0 to 'timestep'
+
+
+                # scale_regist = np.random.uniform(0.2,0.25)
+                # T_regist = sorted(random.sample(range(0, int(hyp_parameters["timesteps"] * scale_regist) + 1), 16), reverse=True)
+                scale_regist = np.random.uniform(0.05,0.7)
+                T_regist = sorted(random.sample(range(int(hyp_parameters["timesteps"] * scale_regist),hyp_parameters["timesteps"]), 16), reverse=True)
+                # scale_regist = np.random.uniform(0.4,1.)
+                # T_regist = [int(hyp_parameters["timesteps"]*scale_regist)]
+                # scale_regist = np.random.uniform(0.6,1.)
+                # init_T = int(hyp_parameters["timesteps"] * scale_regist)
+                # T_regist = sorted(random.sample(range(0, int(hyp_parameters["timesteps"] * scale_regist)), 2)+list(range(init_T,hyp_parameters["timesteps"]+1)), reverse=True)
+
+                T_regist = [[t for _ in range(hyp_parameters["batchsize"]//2)] for t in T_regist]
+
+                # print('T_regist:', T_regist)
+                # proc_type = random.choice(['adding', 'independ', 'downsample', 'slice', 'project', 'none', 'none'])
+                proc_type = random.choice(['adding', 'downsample', 'slice', 'none', 'none'])
+                # proc_type = random.choice(['project'])
+                y1_proc, msk_tgt, cond_ratio = Deformddpm.module.proc_cond_img(y1,proc_type=proc_type)
+                # msk_tgt = msk_tgt + MSK_EPS
+                [ddf_comp,ddf_rand],[img_rec,img_diff,img_save],_ = Deformddpm(img_org=x1, cond_imgs=y1_proc, T=[None, T_regist], proc_type=[],text=embd_y)  # forward diffusion process
+                # loss_ddf1 = loss_reg1(ddf_comp,img=y1,msk=msk_tgt)  # calculate loss for the registration process
+                # loss_sim = loss_imgsim(img_rec, y1, label=msk_tgt*(y1>thresh_imgsim))  # calculate loss for the registration process
+                # loss_mse = loss_imgmse(img_rec, y1, label=msk_tgt*(y1>0.0))  # calculate loss for the registration process
+                loss_sim = loss_imgsim(img_rec, y1, label=(y1>thresh_imgsim))  # calculate loss for the registration process
+                loss_mse = loss_imgmse(img_rec, y1, label=(y1>0.0))  # calculate loss for the registration process
+                loss_ddf1 = loss_reg1(ddf_comp,img=y1)  # calculate loss for the registration process
+
+                loss_regist = 0
+                loss_regist += LOSS_WEIGHTS_REGIST[0] * loss_sim
+                loss_regist += LOSS_WEIGHTS_REGIST[1] * loss_mse
+                loss_regist += LOSS_WEIGHTS_REGIST[2] * loss_ddf1
+                # print('proc_type:', proc_type, 'cond_ratio:', cond_ratio.item())
+                # print('loss_regist:', loss_regist.item(), 'loss_sim:', loss_sim.item(), 'loss_ddf1:', loss_ddf1.item())
+
+                # >> JZ: print nan in x0
+                if torch.isnan(x0).any():
+                    print(f"*** Encountered NaN in input image x0 at epoch {epoch}, step {step}.")
+                
+                
+
+                # >> JZ: print loss of ddf
+                if loss_ddf1>0.001:
+                    print(f"*** High registration DDF loss at epoch {epoch}, step {step}: {loss_ddf1.item()}.")
+                    # # Print gradients for each parameter
+                    # for name, param in Deformddpm.named_parameters():
+                    #     if param.grad is not None:
+                    #         print(f"Gradient for {name}: {param.grad.norm()}")
+                    #     else:
+                    #         print(f"Gradient for {name}: None")
+                
+                loss_regist = torch.sqrt(cond_ratio+MSK_EPS) *loss_regist
+                optimizer.zero_grad()
+                loss_regist.backward()
+
+
+
+                torch.nn.utils.clip_grad_norm_(Deformddpm.parameters(), max_norm=0.1)
+                optimizer.step()
+
+                epoch_loss_regist += loss_regist.item() * len(x0) / len(train_loader.dataset)
+                epoch_loss_imgsim += loss_sim.item() * len(x0) / len(train_loader.dataset)
+                epoch_loss_imgmse += loss_mse.item() * len(x0) / len(train_loader.dataset)
+                epoch_loss_ddfreg += loss_ddf1.item() * len(x0) / len(train_loader.dataset)
+
+            
+            print('step:',step,':', loss_tot.item(),'=',loss_gen_a.item(),'+', loss_gen_d.item(),'+',loss_ddf.item())
+            print(f'     loss_regist: {loss_regist} = {loss_sim} (imgsim) + {loss_mse} (imgmse) + {loss_ddf1} (ddf)')
+            # >> JZ: if loss_imgsim is zero
+            if loss_sim.item()>-0.001:
+                print(f"*** Zero image similarity loss at epoch {epoch}, step {step}.")
+                def save_niftiimage(tensor, filename):
+                    import nibabel as nib
+                    import numpy as np
+                    array = tensor.squeeze().cpu().detach().numpy()
+                    nifti_img = nib.Nifti1Image(array, affine=np.eye(4))
+                    nib.save(nifti_img, filename)
+                # save the x1 and y1 images for debugging
+                save_path = os.path.join('/home/data/Github/OmniMorph/Log/error_files',f"debug_images_epoch{epoch}_step{step}/")
+                os.makedirs(save_path, exist_ok=True)
+                save_niftiimage(img_rec, os.path.join(save_path, 'img_rec.nii.gz'))
+                save_niftiimage(x1, os.path.join(save_path, 'x1.nii.gz'))
+                save_niftiimage(y1, os.path.join(save_path, 'y1.nii.gz'))
+                save_niftiimage(y1_proc, os.path.join(save_path, 'y1_proc.nii.gz'))
+                exit()
+            # print('step:',step,':', loss_tot.item(),'=',loss_gen_a.item(),'+', loss_gen_d.item(),'+',loss_ddf.item())
+
+            # break   # FOR TESTING
+            # else:
+            #     print('loss_gen_a:',loss_gen_a.item())     # FOR TESTING
+            #     pass
+
+        if 1:
+        # if gpu_id == 0:
+            print('==================')
+            print(epoch,':', epoch_loss_tot,'=',epoch_loss_gen_a,'+', epoch_loss_gen_d,'+',epoch_loss_reg, ' (ang+dist+regul)')
+            print(f'     loss_regist: {epoch_loss_regist} = {epoch_loss_imgsim} (imgsim) + {epoch_loss_imgmse} (imgmse) + {epoch_loss_ddfreg} (ddf)')
+            print('==================')
+        # # LR schedular step ----- YHM
+        # scheduler.step()
+
+        if 0 == epoch % epoch_per_save:
+            save_dir=model_save_path + str(epoch).rjust(6, '0') + suffix_pth    
+            os.makedirs(os.path.dirname(model_save_path), exist_ok=True)
+            # break   # FOR TESTING
+            if not use_distributed:
+                print(f"saved in {save_dir}")
+                # torch.save(Deformddpm.state_dict(), save_dir)
+                torch.save({
+                    'model_state_dict': Deformddpm.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'epoch': epoch
+                }, save_dir)
+            elif gpu_id == 0:
+                print(f"saved in {save_dir}")
+                # torch.save(Deformddpm.module.state_dict(), save_dir)
+                torch.save({
+                    'model_state_dict': Deformddpm.module.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'epoch': epoch
+                }, save_dir)
+
+    # Resource cleanup at the end of training
+    torch.cuda.empty_cache()
+    gc.collect()
+    if use_distributed and dist.is_initialized():
+        dist.destroy_process_group()
+
+def ddp_load_dict(gpu_id, Deformddpm, optimizer, model_file,use_distributed=True):
+    
+    if gpu_id == 0:
+    # if 0:
+        utils.print_memory_usage("Before Loading Model")
+        if 1:
+            gc.collect()
+            torch.cuda.empty_cache()
+        # Deformddpm.network.load_state_dict(torch.load(latest_model_file))
+        # Deformddpm.load_state_dict(torch.load(latest_model_file), strict=False)
+        checkpoint = torch.load(model_file)
+        # checkpoint = torch.load(latest_model_file, map_location=f"cuda:{rank}")
+        if use_distributed:
+            Deformddpm.module.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            Deformddpm.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        utils.print_memory_usage("After Loading Checkpoint on GPU")
+
+    if use_distributed:
+        # Broadcast model weights from rank 0 to all other GPUs
+        dist.barrier()
+        for param in Deformddpm.parameters():
+            dist.broadcast(param.data, src=0)  # Synchronize model across ranks
+        dist.barrier()
+        for param_group in optimizer.param_groups:
+            for param in param_group['params']:
+                if param.grad is not None:
+                    dist.broadcast(param.grad, src=0)  # Sync optimizer gradients
+        
+    # initial_epoch = checkpoint['epoch'] + 1
+    # get the epoch number from the filename and add 1 to set as initial_epoch
+    initial_epoch = int(os.path.basename(model_file).split('.')[0][:6]) + 1
+
+    return initial_epoch, Deformddpm, optimizer
+
+            
+
+if __name__ == "__main__":
+    if use_distributed:
+        world_size = torch.cuda.device_count()
+        print(f"Distributed GPU number = {world_size}")
+        mp.spawn(main_train,args = (world_size,),nprocs = world_size)
+    else:
+        main_train(0,1)

OM_train_3modes.py

@@ -0,0 +1,490 @@
+import os
+import gc
+import torch
+import torchvision
+from torch import nn
+from torchvision.utils import save_image
+from torch.utils.data import DataLoader
+
+from torch.optim import Adam, SGD
+from Diffusion.diffuser import DeformDDPM
+from Diffusion.networks import get_net, STN
+from torchvision.transforms import Lambda
+import Diffusion.losses as losses
+import random
+import glob
+import numpy as np
+import utils
+from tqdm import tqdm 
+
+from Dataloader.dataloader0 import get_dataloader
+from Dataloader.dataLoader import *
+
+from Dataloader.dataloader_utils import thresh_img
+import yaml
+import argparse
+
+####################
+import torch.multiprocessing as mp
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
+# from torch.distributed import init_process_group
+###############
+def ddp_setup(rank, world_size):
+    """
+    Args:
+        rank: Unique identifier of each process
+        world_size: Total number of processes
+    """
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12355"
+    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    torch.cuda.set_device(rank)
+
+use_distributed = True
+# use_distributed = False
+
+EPS = 1e-5
+MSK_EPS = 0.01
+TEXT_EMBED_PROB = 0.7
+AUG_RESAMPLE_PROB = 0.6
+LOSS_WEIGHTS_DIFF = [2.0, 1.0, 3.0]  # [ang, dist, reg]
+# LOSS_WEIGHTS_REGIST = [9.0, 1.0, 16.0]  # [imgsim, imgmse, ddf]
+LOSS_WEIGHTS_REGIST = [1.0, 0.2, 1e3]  # [imgsim, imgmse, ddf]
+
+# AUG_PERMUTE_PROB = 0.35
+
+parser = argparse.ArgumentParser()
+
+# config_file_path = 'Config/config_cmr.yaml'
+parser.add_argument(
+        "--config",
+        "-C",
+        help="Path for the config file",
+        type=str,
+        # default="Config/config_cmr.yaml",
+        # default="Config/config_lct.yaml",
+        default="Config/config_all.yaml",
+        required=False,
+    )
+args = parser.parse_args()
+#=======================================================================================================================
+
+
+
+def main_train(rank=0,world_size=1,train_mode_ratio=1,thresh_imgsim=0.01):
+    if use_distributed:
+        ddp_setup(rank,world_size)
+
+        if torch.distributed.is_initialized():
+            print(f"World size: {torch.distributed.get_world_size()}")
+            print(f"Communication backend: {torch.distributed.get_backend()}")
+    gpu_id = rank
+    
+    # Load the YAML file into a dictionary
+    with open(args.config, 'r') as file:
+        hyp_parameters = yaml.safe_load(file)
+        print(hyp_parameters)
+
+    # epoch_per_save=10
+    epoch_per_save=hyp_parameters['epoch_per_save']
+
+    data_name=hyp_parameters['data_name']
+    net_name = hyp_parameters['net_name']
+
+    Net=get_net(net_name)
+
+    suffix_pth=f'_{data_name}_{net_name}.pth'
+    model_save_path = os.path.join('Models',f'{data_name}_{net_name}/')
+    model_dir=model_save_path
+    transformer=utils.get_transformer(img_sz=hyp_parameters["ndims"]*[hyp_parameters['img_size']])
+    
+    # Data_Loader=get_dataloader(data_name=hyp_parameters['data_name'], mode='train')
+
+    # tsfm = torchvision.transforms.Compose([
+    #             torchvision.transforms.ToTensor(),
+    #             ])
+
+    # dataset = Data_Loader(target_res = [hyp_parameters["img_size"]]*hyp_parameters["ndims"], transforms=None, noise_scale=hyp_parameters['noise_scale'])
+    # train_loader = DataLoader(
+    #     dataset,
+    #     batch_size=hyp_parameters['batchsize'],
+    #     # shuffle=False,
+    #     shuffle=True,
+    #     drop_last=True,
+    # )
+
+    # dataset = OminiDataset_v1(transform=None)
+    dataset = OMDataset_indiv(transform=None)
+    train_loader = DataLoader(
+        dataset,
+        batch_size=hyp_parameters['batchsize'],
+        shuffle=True,
+        drop_last=True,
+    )
+
+    # datasetp = OminiDataset_paired(transform=None)
+    datasetp = OMDataset_pair(transform=None)
+    train_loader_p = DataLoader(
+        datasetp,
+        batch_size=hyp_parameters['batchsize']//2,
+        shuffle=True,
+        drop_last=True,
+    )
+
+
+
+    Deformddpm = DeformDDPM(
+        network=Net(
+            n_steps=hyp_parameters["timesteps"], 
+            ndims=hyp_parameters["ndims"], 
+            num_input_chn = hyp_parameters["num_input_chn"],
+            res = hyp_parameters['img_size']
+            ),
+        n_steps=hyp_parameters["timesteps"],
+        image_chw=[1] + [hyp_parameters["img_size"]]*hyp_parameters["ndims"],
+        device=hyp_parameters["device"],
+        batch_size=hyp_parameters["batchsize"],
+        img_pad_mode=hyp_parameters["img_pad_mode"],
+        v_scale=hyp_parameters["v_scale"],
+    )
+
+
+    ddf_stn = STN(
+        img_sz=hyp_parameters["img_size"],
+        ndims=hyp_parameters["ndims"],
+        # padding_mode="zeros",
+        padding_mode=hyp_parameters["padding_mode"],
+        device=hyp_parameters["device"],
+    )
+
+
+    if use_distributed:
+        Deformddpm.to(rank)
+        Deformddpm = DDP(Deformddpm, device_ids=[rank])
+        ddf_stn.to(rank)
+    else:
+        Deformddpm.to(hyp_parameters["device"])
+        ddf_stn.to(hyp_parameters["device"])
+    # ddf_stn = DDP(ddf_stn, device_ids=[rank])
+
+
+    # mse = nn.MSELoss()
+    # loss_reg = losses.Grad(penalty=['l1', 'negdetj'], ndims=hyp_parameters["ndims"])
+    loss_reg = losses.Grad(penalty=['l1', 'negdetj', 'range'], ndims=hyp_parameters["ndims"],outrange_thresh=0.2,outrange_weight=1e2)
+    loss_reg1 = losses.Grad(penalty=['l1', 'negdetj', 'range'], ndims=hyp_parameters["ndims"],outrange_thresh=0.6,outrange_weight=1e2)
+    loss_dist = losses.MRSE(img_sz=hyp_parameters["img_size"])
+    # loss_ang = losses.MRSE(img_sz=hyp_parameters["img_size"])
+    loss_ang = losses.NCC(img_sz=hyp_parameters["img_size"])
+    loss_imgsim = losses.LNCC()
+    loss_imgmse = losses.LMSE()
+
+    optimizer = Adam(Deformddpm.parameters(), lr=hyp_parameters["lr"])
+    # hyp_parameters["lr"]=0.00000001
+    # optimizer_regist = Adam(Deformddpm.parameters(), lr=hyp_parameters["lr"]*0.01)
+    # optimizer_regist = SGD(Deformddpm.parameters(), lr=hyp_parameters["lr"]*0.01, momentum=0.98)
+    # optimizer = SGD(Deformddpm.parameters(), lr=hyp_parameters["lr"], momentum=0.9)
+
+    # # LR scheduler ----- YHM
+    # scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, hyp_parameters["lr"], hyp_parameters["lr"]*10, step_size_up=500, step_size_down=500, mode='triangular', gamma=1.0, scale_fn=None, scale_mode='cycle', cycle_momentum=True, base_momentum=0.8, max_momentum=0.9, last_epoch=-1)
+
+    # Deformddpm.network.load_state_dict(torch.load('/home/data/jzheng/Adaptive_Motion_Generator-master/models/1000.pth'))
+
+    # check for existing models
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir, exist_ok=True)
+    model_files = glob.glob(os.path.join(model_dir, "*.pth"))
+    model_files.sort()
+    if model_files:
+        if gpu_id == 0:
+            print(model_files)
+        initial_epoch, Deformddpm, optimizer = ddp_load_dict(gpu_id, Deformddpm, optimizer, model_files[-1])
+    else:
+        initial_epoch = 0
+
+    if gpu_id == 0:
+        print('len_train_data: ',len(dataset))
+    # Training loop
+    for epoch in range(initial_epoch,hyp_parameters["epoch"]):
+
+        epoch_loss_tot = 0.0
+        epoch_loss_gen_d = 0.0
+        epoch_loss_gen_a = 0.0
+        epoch_loss_reg = 0.0
+        epoch_loss_regist = 0.0
+        epoch_loss_imgsim = 0.0
+        epoch_loss_imgmse = 0.0
+        epoch_loss_ddfreg = 0.0
+        # Set model inside to train model
+        Deformddpm.train()
+        
+        loss_nan_step = 0  # yu: count the number of nan loss steps
+
+        for step, batch in tqdm(enumerate(train_loader)):
+        # for step, batch in tqdm(enumerate(train_loader)):
+
+        # for step, batch in enumerate(train_loader_omni):
+            # x0, _ = batch
+            
+            
+            # ==========================================================================
+            # diffusion train on single image
+
+            # x0 = batch # for omni dataset
+            [x0,embd] = batch # for om dataset
+            x0 = x0.to(hyp_parameters["device"]).type(torch.float32)
+            # print('embd:', embd.shape)
+            if np.random.uniform(0,1)<TEXT_EMBED_PROB:
+                embd = embd.to(hyp_parameters["device"]).type(torch.float32)
+            else:
+                embd = None
+
+            
+
+            n = x0.size()[0]  # batch_size -> n
+            x0 = x0.to(hyp_parameters["device"])
+            
+            blind_mask = utils.get_random_deformed_mask(x0.shape[2:],apply_possibility=0.6).to(hyp_parameters["device"])
+
+            # random deformation + rotation
+            if hyp_parameters["ndims"]>2:
+                if np.random.uniform(0,1)<AUG_RESAMPLE_PROB:
+                    x0 = utils.random_resample(x0, deform_scale=0)
+                # elif np.random.uniform(0,1)<AUG_RESAMPLE_PROB+AUG_PERMUTE_PROB:
+                else:
+                    [x0] = utils.random_permute([x0], select_dims=[-1,-2,-3])
+            x0 = transformer(x0)
+            if hyp_parameters['noise_scale']>0:
+                if np.random.uniform(0,1)<AUG_RESAMPLE_PROB:
+                    x0 = thresh_img(x0, [0, 1*hyp_parameters['noise_scale']])
+                x0 = x0 * (np.random.normal(1, hyp_parameters['noise_scale'] * 1)) + np.random.normal(0, hyp_parameters['noise_scale'] * 1)
+
+            # Picking some noise for each of the images in the batch, a timestep and the respective alpha_bars
+            t = torch.randint(0, hyp_parameters["timesteps"], (n,)).to(
+                hyp_parameters["device"]
+            )  # pick up a seq of rand number from 0 to 'timestep'
+
+            # proc_type = random.choice(['adding', 'independ', 'downsample', 'slice', 'project', 'none', 'uncon', 'uncon', 'uncon'])
+            proc_type = random.choice(['adding', 'independ', 'downsample', 'slice', 'none', 'uncon', 'uncon', 'uncon'])
+            # print('proc_type:', proc_type)
+            cond_img, _, cond_ratio = Deformddpm.module.proc_cond_img(x0,proc_type=proc_type)
+
+            pre_dvf_I,dvf_I = Deformddpm(img_org=x0, t=t, cond_imgs=cond_img, mask=blind_mask,proc_type=[],text=embd)  # forward diffusion process
+
+            loss_tot=0
+
+            loss_ddf = loss_reg(pre_dvf_I,img=x0)
+            trm_pred = ddf_stn(pre_dvf_I, dvf_I)
+            loss_gen_d = loss_dist(pred=trm_pred,inv_lab=dvf_I,ddf_stn=None,mask=blind_mask)
+            loss_gen_a = loss_ang(pred=trm_pred,inv_lab=dvf_I,ddf_stn=None,mask=blind_mask)
+
+            loss_tot += LOSS_WEIGHTS_DIFF[0] * loss_gen_a + LOSS_WEIGHTS_DIFF[1] * loss_gen_d
+            loss_tot += LOSS_WEIGHTS_DIFF[2] * loss_ddf
+            loss_tot = torch.sqrt(1.+MSK_EPS-cond_ratio) * loss_tot
+
+            # >> JZ: print nan in x0
+            if torch.isnan(x0).any():
+                print(f"*** Encountered NaN in input image x0 at epoch {epoch}, step {step}.")
+            # >> JZ: print loss of ddf
+            if loss_ddf>0.001:
+                print(f"*** High diffusion DDF loss at epoch {epoch}, step {step}: {loss_ddf.item()}.")
+            # yu: check if loss_tot==nan or inf
+            if torch.isnan(loss_tot) or torch.isinf(loss_tot):
+                print(f"*** Encountered NaN or Inf loss at epoch {epoch}, step {step}. Skipping this batch.")
+                loss_nan_step += 1
+                continue
+            if loss_nan_step > 5:
+                print(f"*** Too many NaN or Inf losses ({loss_nan_step} times) at epoch {epoch}, step {step}. Stopping training.")
+                raise ValueError("Too many NaN losses detected in loss_tot. Code terminated.")
+
+
+            optimizer.zero_grad()
+            loss_tot.backward()
+            optimizer.step()
+
+            epoch_loss_tot += loss_tot.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_gen_d += loss_gen_d.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_gen_a += loss_gen_a.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_reg += loss_ddf.item() * len(x0) / len(train_loader.dataset)
+
+            # print(loss_gen_a.item())
+            # if 0:
+            # if loss_gen_a.item() < -0.3 and step%train_mode_ratio == 0:
+            if step%train_mode_ratio == 0:
+                # ==========================================================================
+                # registration train on paired images
+                # x1, y1 = next(iter(train_loader_p))
+                [x1, y1, _, embd_y] = next(iter(train_loader_p))
+                if np.random.uniform(0,1)<TEXT_EMBED_PROB:
+                    # embd_x = embd_x.to(hyp_parameters["device"]).type(torch.float32)
+                    embd_y = embd_y.to(hyp_parameters["device"]).type(torch.float32)
+                else:
+                    # embd_x = None
+                    embd_y = None
+
+                x1 = x1.to(hyp_parameters["device"]).type(torch.float32)
+                y1 = y1.to(hyp_parameters["device"]).type(torch.float32)
+                n = x1.size()[0]  # batch_size -> n
+                # random deformation + rotation
+                # if hyp_parameters["ndims"]>2:
+                #     if np.random.uniform(0,1)<0.6:
+                #         x1 = utils.random_resample(x1, deform_scale=0)
+                #         y1 = utils.random_resample(y1, deform_scale=0)
+                x1 = transformer(x1)
+                y1 = transformer(y1)
+                [x1, y1] = utils.random_permute([x1, y1], select_dims=[-1,-2,-3])
+                if hyp_parameters['noise_scale']>0:
+                    x1 = thresh_img(x1, [0, 2*hyp_parameters['noise_scale']])
+                    x1 = x1 * (np.random.normal(1, hyp_parameters['noise_scale'] * 1)) + np.random.normal(0, hyp_parameters['noise_scale'] * 1)
+                    y1 = thresh_img(y1, [0, 2*hyp_parameters['noise_scale']])
+                    y1 = y1 * (np.random.normal(1, hyp_parameters['noise_scale'] * 1)) + np.random.normal(0, hyp_parameters['noise_scale'] * 1)
+                # Picking some noise for each of the images in the batch, a timestep and the respective alpha_bars
+                t = torch.randint(0, hyp_parameters["timesteps"], (n,)).to(
+                    hyp_parameters["device"]
+                )  # pick up a seq of rand number from 0 to 'timestep'
+
+
+                scale_regist = np.random.uniform(0.6,1.)
+                T_regist = sorted(random.sample(range(0, int(hyp_parameters["timesteps"] * scale_regist) + 1), 16), reverse=True)
+                # scale_regist = np.random.uniform(0.4,1.)
+                # T_regist = [int(hyp_parameters["timesteps"]*scale_regist)]
+                # scale_regist = np.random.uniform(0.6,1.)
+                # init_T = int(hyp_parameters["timesteps"] * scale_regist)
+                # T_regist = sorted(random.sample(range(0, int(hyp_parameters["timesteps"] * scale_regist)), 2)+list(range(init_T,hyp_parameters["timesteps"]+1)), reverse=True)
+
+                T_regist = [[t for _ in range(hyp_parameters["batchsize"]//2)] for t in T_regist]
+
+                # print('T_regist:', T_regist)
+                # proc_type = random.choice(['adding', 'independ', 'downsample', 'slice', 'project', 'none', 'none'])
+                proc_type = random.choice(['adding', 'independ', 'downsample', 'slice', 'none', 'none'])
+                # proc_type = random.choice(['project'])
+                y1, msk_tgt, cond_ratio = Deformddpm.module.proc_cond_img(y1,proc_type=proc_type)
+                msk_tgt = msk_tgt + MSK_EPS
+                [ddf_comp,ddf_rand],[img_rec,img_diff,img_save],_ = Deformddpm(img_org=x1, cond_imgs=y1, T=[None, T_regist], proc_type=[],text=embd_y)  # forward diffusion process
+                loss_ddf1 = loss_reg1(ddf_comp,img=y1,msk=msk_tgt)  # calculate loss for the registration process
+                loss_sim = loss_imgsim(img_rec, y1, label=msk_tgt*(y1>thresh_imgsim))  # calculate loss for the registration process
+                loss_mse = loss_imgmse(img_rec, y1, label=msk_tgt*(y1>0.0))  # calculate loss for the registration process
+
+                loss_regist = 0
+                loss_regist += LOSS_WEIGHTS_REGIST[0] * loss_sim
+                loss_regist += LOSS_WEIGHTS_REGIST[1] * loss_mse
+                loss_regist += LOSS_WEIGHTS_REGIST[2] * loss_ddf1
+                # print('proc_type:', proc_type, 'cond_ratio:', cond_ratio.item())
+                # print('loss_regist:', loss_regist.item(), 'loss_sim:', loss_sim.item(), 'loss_ddf1:', loss_ddf1.item())
+
+                # >> JZ: print nan in x0
+                if torch.isnan(x0).any():
+                    print(f"*** Encountered NaN in input image x0 at epoch {epoch}, step {step}.")
+                # >> JZ: print loss of ddf
+                if loss_ddf1>0.001:
+                    print(f"*** High registration DDF loss at epoch {epoch}, step {step}: {loss_ddf1.item()}.")
+                
+                loss_regist = torch.sqrt(cond_ratio+MSK_EPS) *loss_regist
+                optimizer.zero_grad()
+                loss_regist.backward()
+
+                # # Print gradients for each parameter
+                # for name, param in Deformddpm.named_parameters():
+                #     if param.grad is not None:
+                #         print(f"Gradient for {name}: {param.grad.norm()}")
+                #     else:
+                #         print(f"Gradient for {name}: None")
+
+                torch.nn.utils.clip_grad_norm_(Deformddpm.parameters(), max_norm=0.1)
+                optimizer.step()
+
+                epoch_loss_regist += loss_regist.item() * len(x0) / len(train_loader.dataset)
+                epoch_loss_imgsim += loss_sim.item() * len(x0) / len(train_loader.dataset)
+                epoch_loss_imgmse += loss_mse.item() * len(x0) / len(train_loader.dataset)
+                epoch_loss_ddfreg += loss_ddf1.item() * len(x0) / len(train_loader.dataset)
+
+            
+            
+            # print('step:',step,':', loss_tot.item(),'=',loss_gen_a.item(),'+', loss_gen_d.item(),'+',loss_ddf.item())
+
+            # break   # FOR TESTING
+            # else:
+            #     print('loss_gen_a:',loss_gen_a.item())     # FOR TESTING
+            #     pass
+
+        if 1:
+        # if gpu_id == 0:
+            print(epoch,':', epoch_loss_tot,'=',epoch_loss_gen_a,'+', epoch_loss_gen_d,'+',epoch_loss_reg, ' (ang+dist+regul)')
+            print(f'     loss_regist: {epoch_loss_regist} = {epoch_loss_imgsim} (imgsim) + {epoch_loss_imgmse} (imgmse) + {epoch_loss_ddfreg} (ddf)')
+
+        # # LR schedular step ----- YHM
+        # scheduler.step()
+
+        if 0 == epoch % epoch_per_save:
+            save_dir=model_save_path + str(epoch).rjust(6, '0') + suffix_pth    
+            os.makedirs(os.path.dirname(model_save_path), exist_ok=True)
+            # break   # FOR TESTING
+            if not use_distributed:
+                print(f"saved in {save_dir}")
+                # torch.save(Deformddpm.state_dict(), save_dir)
+                torch.save({
+                    'model_state_dict': Deformddpm.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'epoch': epoch
+                }, save_dir)
+            elif gpu_id == 0:
+                print(f"saved in {save_dir}")
+                # torch.save(Deformddpm.module.state_dict(), save_dir)
+                torch.save({
+                    'model_state_dict': Deformddpm.module.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'epoch': epoch
+                }, save_dir)
+
+    # Resource cleanup at the end of training
+    torch.cuda.empty_cache()
+    gc.collect()
+    if use_distributed and dist.is_initialized():
+        dist.destroy_process_group()
+
+def ddp_load_dict(gpu_id, Deformddpm, optimizer, model_file,use_distributed=True):
+    
+    if gpu_id == 0:
+    # if 0:
+        utils.print_memory_usage("Before Loading Model")
+        if 1:
+            gc.collect()
+            torch.cuda.empty_cache()
+        # Deformddpm.network.load_state_dict(torch.load(latest_model_file))
+        # Deformddpm.load_state_dict(torch.load(latest_model_file), strict=False)
+        checkpoint = torch.load(model_file)
+        # checkpoint = torch.load(latest_model_file, map_location=f"cuda:{rank}")
+        if use_distributed:
+            Deformddpm.module.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            Deformddpm.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        utils.print_memory_usage("After Loading Checkpoint on GPU")
+
+    if use_distributed:
+        # Broadcast model weights from rank 0 to all other GPUs
+        dist.barrier()
+        for param in Deformddpm.parameters():
+            dist.broadcast(param.data, src=0)  # Synchronize model across ranks
+        dist.barrier()
+        for param_group in optimizer.param_groups:
+            for param in param_group['params']:
+                if param.grad is not None:
+                    dist.broadcast(param.grad, src=0)  # Sync optimizer gradients
+        
+    # initial_epoch = checkpoint['epoch'] + 1
+    # get the epoch number from the filename and add 1 to set as initial_epoch
+    initial_epoch = int(os.path.basename(model_file).split('.')[0][:6]) + 1
+
+    return initial_epoch, Deformddpm, optimizer
+
+            
+
+if __name__ == "__main__":
+    if use_distributed:
+        world_size = torch.cuda.device_count()
+        print(f"Distributed GPU number = {world_size}")
+        mp.spawn(main_train,args = (world_size,),nprocs = world_size)
+    else:
+        main_train(0,1)

OM_train_uncon.py

@@ -0,0 +1,258 @@
+import os
+import torch
+import torchvision
+from torch import nn
+from torchvision.utils import save_image
+from torch.utils.data import DataLoader
+from torch.optim import Adam, SGD
+from Diffusion.diffuser import DeformDDPM
+from Diffusion.networks import get_net, STN
+from torchvision.transforms import Lambda
+import Diffusion.losses as losses
+import random
+import glob
+import numpy as np
+import utils
+
+from Dataloader.dataloader0 import get_dataloader
+
+from Dataloader.dataloader_utils import thresh_img
+import yaml
+import argparse
+
+####################
+import torch.multiprocessing as mp
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.distributed import init_process_group, destroy_process_group
+###############
+def ddp_setup(rank, world_size):
+    """
+    Args:
+        rank: Unique identifier of each process
+        world_size: Total number of processes
+    """
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = "12355"
+    init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    torch.cuda.set_device(rank)
+
+use_parallel=False
+use_distributed = False
+
+EPS = 1e-5
+
+parser = argparse.ArgumentParser()
+
+# config_file_path = 'Config/config_cmr.yaml'
+parser.add_argument(
+        "--config",
+        "-C",
+        help="Path for the config file",
+        type=str,
+        default="Config/config_cmr.yaml",
+        # default="Config/config_lct.yaml",
+        required=False,
+    )
+args = parser.parse_args()
+#=======================================================================================================================
+
+
+
+def main_train(rank,world_size):
+    
+    ddp_setup(rank,world_size)
+    gpu_id = rank
+
+    # Load the YAML file into a dictionary
+    with open(args.config, 'r') as file:
+        hyp_parameters = yaml.safe_load(file)
+        print(hyp_parameters)
+
+
+
+    # epoch_per_save=10
+    epoch_per_save=hyp_parameters['epoch_per_save']
+
+    data_name=hyp_parameters['data_name']
+    net_name = hyp_parameters['net_name']
+
+    Net=get_net(net_name)
+
+    suffix_pth=f'_{data_name}_{net_name}.pth'
+    model_save_path = os.path.join('Models',f'{data_name}_{net_name}/')
+    model_dir=model_save_path
+    transformer=utils.get_transformer(img_sz=hyp_parameters["ndims"]*[hyp_parameters['img_size']])
+    Data_Loader=get_dataloader(data_name=hyp_parameters['data_name'], mode='train')
+
+    tsfm = torchvision.transforms.Compose([
+                torchvision.transforms.ToTensor(),
+                ])
+
+
+    dataset = Data_Loader(target_res = [hyp_parameters["img_size"]]*hyp_parameters["ndims"], transforms=None, noise_scale=hyp_parameters['noise_scale'])
+    train_loader = DataLoader(
+        dataset,
+        batch_size=hyp_parameters['batchsize'],
+        # shuffle=False,
+        shuffle=True,
+        drop_last=True,
+    )
+
+
+
+    Deformddpm = DeformDDPM(
+        network=Net(n_steps=hyp_parameters["timesteps"], ndims=hyp_parameters["ndims"], num_input_chn=1),
+        n_steps=hyp_parameters["timesteps"],
+        image_chw=[1] + [hyp_parameters["img_size"]]*hyp_parameters["ndims"],
+        device=hyp_parameters["device"],
+        batch_size=hyp_parameters["batchsize"],
+        img_pad_mode=hyp_parameters["img_pad_mode"],
+        v_scale=hyp_parameters["v_scale"],
+    )
+
+
+    ddf_stn = STN(
+        img_sz=hyp_parameters["img_size"],
+        ndims=hyp_parameters["ndims"],
+        # padding_mode="zeros",
+        padding_mode=hyp_parameters["padding_mode"],
+        device=hyp_parameters["device"],
+    )
+
+    # Deformddpm.to(hyp_parameters["device"])
+    # ddf_stn.to(hyp_parameters["device"])
+
+    # if use_distributed:
+    #     torch.distributed.init_process_group(backend='nccl')
+    #     Deformddpm = nn.parallel.DistributedDataParallel(Deformddpm, device_ids=[torch.cuda.current_device()])
+    #     ddf_stn = nn.parallel.DistributedDataParallel(ddf_stn, device_ids=[torch.cuda.current_device()])
+    # elif use_parallel:
+    #     Deformddpm = nn.DataParallel(Deformddpm)
+    #     ddf_stn = nn.DataParallel(ddf_stn)
+
+    Deformddpm.to(rank)
+    Deformddpm = DDP(Deformddpm, device_ids=[rank])
+    ddf_stn.to(rank)
+    # ddf_stn = DDP(ddf_stn, device_ids=[rank])
+
+
+    # mse = nn.MSELoss()
+    loss_reg = losses.Grad(penalty=['l1', 'negdetj'], ndims=hyp_parameters["ndims"])
+    loss_dist = losses.MRSE(img_sz=hyp_parameters["img_size"])
+    # loss_ang = losses.MRSE(img_sz=hyp_parameters["img_size"])
+    loss_ang = losses.NCC(img_sz=hyp_parameters["img_size"])
+
+    optimizer = Adam(Deformddpm.parameters(), lr=hyp_parameters["lr"])
+    # hyp_parameters["lr"]=0.00000001
+    # # optimizer = SGD(Deformddpm.parameters(), lr=hyp_parameters["lr"], momentum=0.95)
+    # optimizer = SGD(Deformddpm.parameters(), lr=hyp_parameters["lr"], momentum=0.9)
+
+    # # LR scheduler ----- YHM
+    # scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, hyp_parameters["lr"], hyp_parameters["lr"]*10, step_size_up=500, step_size_down=500, mode='triangular', gamma=1.0, scale_fn=None, scale_mode='cycle', cycle_momentum=True, base_momentum=0.8, max_momentum=0.9, last_epoch=-1)
+
+    # Deformddpm.network.load_state_dict(torch.load('/home/data/jzheng/Adaptive_Motion_Generator-master/models/1000.pth'))
+
+    # check for existing models
+    if not os.path.exists(model_dir):
+        os.makedirs(model_dir, exist_ok=True)
+    model_files = glob.glob(os.path.join(model_dir, "*.pth"))
+    model_files.sort()
+    print(model_files)
+    if model_files:
+        # if there are any model files, load the most recent one
+        latest_model_file = model_files[-1]
+        # Deformddpm.network.load_state_dict(torch.load(latest_model_file))
+        if use_parallel:
+            Deformddpm.module.load_state_dict(torch.load(latest_model_file), strict=False)
+        else:
+            Deformddpm.load_state_dict(torch.load(latest_model_file), strict=False)
+        # get the epoch number from the filename and add 1 to set as initial_epoch
+        initial_epoch = int(os.path.basename(latest_model_file).split('.')[0][:6]) + 1
+    else:
+        initial_epoch = 0
+    print('len_train_data: ',len(dataset))
+    for epoch in range(initial_epoch,hyp_parameters["epoch"]):
+
+        epoch_loss_tot = 0.0
+        epoch_loss_gen_d = 0.0
+        epoch_loss_gen_a = 0.0
+        epoch_loss_reg = 0.0
+        # Set model inside to train model
+        Deformddpm.train()
+
+        for step, batch in enumerate(train_loader):
+            # x0, _ = batch
+            x0, _, _ = batch
+            x0 = x0.to(hyp_parameters["device"]).type(torch.float32)
+        
+            n = x0.size()[0]  # batch_size -> n
+            x0 = x0.to(hyp_parameters["device"])
+            # random deformation + rotation
+            if hyp_parameters["ndims"]>2:
+                if np.random.uniform(0,1)<0.6:
+                    x0 = utils.random_resample(x0, deform_scale=0)
+            x0 = transformer(x0)
+            if hyp_parameters['noise_scale']>0:
+                x0 = thresh_img(x0, [0, 2*hyp_parameters['noise_scale']])
+                x0 = x0 * (np.random.normal(1, hyp_parameters['noise_scale'] * 1)) + np.random.normal(0, hyp_parameters['noise_scale'] * 1)
+
+            # Picking some noise for each of the images in the batch, a timestep and the respective alpha_bars
+            t = torch.randint(0, hyp_parameters["timesteps"], (n,)).to(
+                hyp_parameters["device"]
+            )  # pick up a seq of rand number from 0 to 'timestep'
+
+            
+            if use_parallel:
+                # # noisy_imgs, dvf_I = ddf_enc(img= x0, t)
+                # noisy_imgs, dvf_I,_ = Deformddpm.module.diffuse(x0, t)
+                # pre_dvf_I = Deformddpm.backward(noisy_imgs, t.reshape(16, -1))
+                pre_dvf_I, _ = Deformddpm.module(x0, t)
+            else:
+                # # noisy_imgs, dvf_I = ddf_enc(img= x0, t)
+                # noisy_imgs, dvf_I,_ = Deformddpm.diffuse(x0, t)
+                # pre_dvf_I = Deformddpm.backward(noisy_imgs, t.reshape(16, -1))
+                pre_dvf_I,dvf_I = Deformddpm(x0, t)
+
+            loss_tot=0
+
+            loss_ddf = loss_reg(pre_dvf_I)
+            trm_pred = ddf_stn(pre_dvf_I, dvf_I)
+            loss_gen_d = loss_dist(pred=trm_pred,inv_lab=dvf_I,ddf_stn=None)
+            loss_gen_a = loss_ang(pred=trm_pred,inv_lab=dvf_I,ddf_stn=None)
+
+            loss_tot += 1.0 * loss_gen_d + 1.0 * loss_gen_a
+            loss_tot +=10 * loss_ddf
+            optimizer.zero_grad()
+            loss_tot.backward()
+            optimizer.step()
+
+            epoch_loss_tot += loss_tot.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_gen_d += loss_gen_d.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_gen_a += loss_gen_a.item() * len(x0) / len(train_loader.dataset)
+            epoch_loss_reg += loss_ddf.item() * len(x0) / len(train_loader.dataset)
+            # print('step:',step,':', loss_tot.item(),'=',loss_gen_a.item(),'+', loss_gen_d.item(),'+',loss_ddf.item())
+
+        if gpu_id == 0:
+            print(epoch,':', epoch_loss_tot,'=',epoch_loss_gen_a,'+', epoch_loss_gen_d,'+',epoch_loss_reg, ' (ang+dist+regul)')
+
+        # # LR schedular step ----- YHM
+        # scheduler.step()
+
+        if 0 == epoch % epoch_per_save:
+            save_dir=model_save_path + str(epoch).rjust(6, '0') + suffix_pth
+            if os.path.exists(model_save_path):
+                print(f"saved in {save_dir}")
+            else:
+                os.makedirs(os.path.dirname(model_save_path))
+            # break   # FOR TESTING
+            if use_parallel:
+                torch.save(Deformddpm.module.state_dict(), save_dir)
+            elif gpu_id == 0:
+                torch.save(Deformddpm.module.state_dict(), save_dir)
+            
+
+if __name__ == "__main__":
+    world_size = torch.cuda.device_count()
+    print(f"world size = {world_size}")
+    mp.spawn(main_train,args = (world_size,),nprocs = world_size)

README.md

@@ -0,0 +1,11 @@
+# OmniMorph: Deform All-in-One Framework for Medical Image Generation, Restoration and Registration based on conditional Deformation-Recovery Diffusion Model
+
+## Environment
+```
+conda activate torch
+conda deactivate
+```
+source  /home/data/Github/OmniMorph/ominenv/bin/activate
+
+## Masking CUDA
+CUDA_VISIBLE_DEVICES=0,1,3 python ...

bash_infer.sh

@@ -0,0 +1,9 @@
+
+source /home/data/jzheng/Adaptive_Motion_Generator-master/pipenv/bin/activate
+
+export CUDA_VISIBLE_DEVICES=2
+# export CUDA_VISIBLE_DEVICES=0
+
+# python -u OM_aug.py -C Config/config_om.yaml
+# python -u OM_reg.py -C Config/config_om.yaml
+nohup python -u OM_aug.py -C Config/config_om.yaml > aug_log.txt 2>&1 &

bash_train.sh

@@ -0,0 +1,12 @@
+
+source /home/data/jzheng/Adaptive_Motion_Generator-master/pipenv/bin/activate
+
+export CUDA_VISIBLE_DEVICES=3
+# export CUDA_VISIBLE_DEVICES=1,3
+# export CUDA_VISIBLE_DEVICES=1,2,3 
+# # python -u OM_train.py -C Config/config_lct.yaml
+# nohup python -u OM_train.py -C Config/config_lct.yaml > train_log.txt 2>&1 &
+
+# python -u OM_train_2modes.py -C Config/config_om.yaml
+nohup python -u OM_train_2modes.py -C Config/config_om.yaml > train_log.txt 2>&1 &
+# nohup python -u OM_train.py -C Config/config_om.yaml > train_log.txt 2>&1 &

dataloader_tester.py

@@ -0,0 +1,65 @@
+import os
+import torch
+import torchvision
+from torch import nn
+from torchvision.utils import save_image
+from torch.utils.data import DataLoader
+from torch.optim import Adam, SGD
+from Diffusion.diffuser import DeformDDPM
+from Diffusion.networks import get_net, STN
+from torchvision.transforms import Lambda
+import Diffusion.losses as losses
+import random
+import glob
+import numpy as np
+import utils
+
+from Dataloader.dataloader0 import get_dataloader
+from Dataloader.dataLoader import *
+from Dataloader.dataloader_utils import thresh_img
+import yaml
+import argparse
+
+tsfm = torchvision.transforms.Compose(
+    [
+        torchvision.transforms.ToTensor(),
+    ]
+)
+Data_Loader=get_dataloader(data_name = 'lct', mode='train')
+
+dataset = Data_Loader(
+    target_res=[128] * 3,
+    transforms=None,
+    noise_scale=4.0e-05,
+)
+train_loader = DataLoader(
+    dataset,
+    batch_size=32,
+    # shuffle=False,
+    shuffle=True,
+    drop_last=True,
+)
+
+
+dataset2 = OminiDataset_v1(transform=None)
+train_loader2 = DataLoader(dataset2, batch_size=32, shuffle=True)
+
+
+dataset = OminiDataset_paired(transform=None, ROIs = ['leg'])
+train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
+# print(dataset.get_all_ROI())    
+# print(dataset.getitem())
+# print(dataset.get_ALLdata())
+# print(dataset.getitem(idx=11))
+# exit()
+
+
+
+for i, batch in enumerate(train_loader):
+    x0, x1 = batch
+    print(x0.shape,x1.shape)
+    print(x0.dtype,x1.dtype)
+    print(x0.min(),x0.max())
+    break
+exit()
+

requirements.txt

@@ -0,0 +1,57 @@
+certifi==2022.12.7
+charset-normalizer==2.1.1
+contourpy==1.1.1
+cycler==0.12.1
+einops==0.3.2
+elasticdeform==0.5.0
+filelock==3.16.1
+fonttools==4.49.0
+fsspec==2025.3.0
+hausdorff==0.2.6
+huggingface-hub==0.29.3
+idna==3.4
+imageio==2.34.0
+importlib_metadata==7.1.0
+importlib_resources==6.1.2
+joblib==1.4.0
+kiwisolver==1.4.5
+lazy_loader==0.3
+llvmlite==0.41.1
+matplotlib==3.7.5
+networkx==3.1
+nibabel==5.1.0
+nptyping==2.5.0
+numba==0.58.1
+numpy==1.24.1
+opencv-python==4.9.0.80
+packaging==23.2
+pandas==2.0.3
+pillow==10.2.0
+pydicom==2.4.4
+pynrrd==1.0.0
+pyparsing==3.1.1
+pyquaternion==0.9.9
+python-dateutil==2.8.2
+pytz==2025.2
+PyWavelets==1.4.1
+PyYAML==6.0.2
+regex==2024.11.6
+requests==2.28.1
+safetensors==0.5.3
+scikit-image==0.21.0
+scikit-learn==1.3.2
+scipy==1.9.3
+SimpleITK==2.3.1
+six==1.16.0
+threadpoolctl==3.5.0
+tifffile==2023.7.10
+tokenizers==0.20.3
+torch==1.12.1+cu113
+torchaudio==0.12.1+cu113
+torchvision==0.13.1+cu113
+tqdm==4.66.2
+transformers==4.46.3
+typing_extensions==4.8.0
+tzdata==2025.2
+urllib3==1.26.13
+zipp==3.17.0

utils.py

@@ -0,0 +1,498 @@
+import os
+import torch
+import torchvision
+from torch import nn, optim
+from torch.autograd.variable import Variable
+from torchvision import transforms, datasets
+from torchvision.utils import save_image
+import torch.nn.functional as F
+import scipy.ndimage as spimg
+import pyquaternion as quater
+import random
+import numpy as np
+import math
+from typing import Optional, Tuple, List
+import nibabel as nib
+# from data_loader.acdc_dataloader import acdc_gan
+
+# from Adaptive_Motion_Generator.Dataloader.Archive.acdc_dataloader import *
+
+def get_barcode(index=[],header=['Patient','Slice','AugImg','NoiseStep'],digit=[4,6,4,4],split='_'):
+    # Patient0001_Slice0001_NosieImg0001_NoiseStep0070
+    barcode_str=''
+    header=header.copy()
+    digit=digit.copy()
+    if len(index)<3:
+        header[2] = 'ORG'
+        header[3] = 'NA'
+        digit[2] = 0
+        digit[3] = 0
+        index +=['','']
+
+    for id, h in enumerate(header):
+        barcode_str+=h+str(index[id]).zfill(digit[id])+split
+    return barcode_str[:-1]
+
+class RandomResizedCrop3D(nn.Module):
+    """Crop a random portion of a 3D volume and resize it to a given size.
+
+    Args:
+        size (tuple of int): Expected output size of the crop, for each dimension (D, H, W).
+        scale (tuple of float): Specifies the lower and upper bounds for the random area of the crop,
+                                before resizing. The scale is defined with respect to the volume of the original image.
+        ratio (tuple of float): Lower and upper bounds for the random aspect ratio of the crop, before resizing.
+        interpolation (str): Desired interpolation mode ('trilinear' or 'nearest').
+    """
+
+    def __init__(
+            self,
+            size: Tuple[int, int, int],
+            scale=(0.6, 1.0),
+            ratio=(0.5, 1.5),
+            interpolation='trilinear'
+    ):
+        super().__init__()
+        self.size = size
+        self.scale = scale
+        self.ratio = ratio
+        self.interpolation = interpolation
+
+    @staticmethod
+    def get_params(img: torch.Tensor, rand_scale: float, scale: List[float], ratio: List[float]) -> Tuple[int, int, int, int, int, int]:
+        """Get parameters for `crop` for a random sized crop.
+
+        Args:
+            img (Tensor): Input image.
+            scale (list): Range of scale of the origin size cropped.
+            ratio (list): Range of aspect ratio of the origin aspect ratio cropped.
+
+        Returns:
+            tuple: params (i, j, k, d, h, w) to be passed to `crop` for a random sized crop.
+        """
+        img_sz = np.array(list(img.size())[2:])
+        crop_sz = (img_sz * rand_scale).astype(np.int32)  #[int(s*rand_scale) for s in img_sz]
+        start_id = np.random.randint(0, img_sz - crop_sz + 1, size=(img_sz.size,))
+        return start_id.tolist()+crop_sz.tolist()
+
+        # volume = depth * height * width
+        #
+        # log_ratio = torch.log(torch.tensor(ratio))
+        # for _ in range(10):
+        #     target_volume = volume * torch.empty(1).uniform_(*scale).item()
+        #     aspect_ratio = torch.exp(torch.empty(1).uniform_(log_ratio[0], log_ratio[1])).item()
+        #
+        #     w = int(round(math.sqrt(target_volume * aspect_ratio)))
+        #     h = int(round(math.sqrt(target_volume / aspect_ratio)))
+        #     d = int(round(math.sqrt(target_volume / (w * h))))
+        #
+        #     if 0 < w <= width and 0 < h <= height and 0 < d <= depth:
+        #         i = torch.randint(0, depth - d + 1, size=(1,)).item()
+        #         j = torch.randint(0, height - h + 1, size=(1,)).item()
+        #         k = torch.randint(0, width - w + 1, size=(1,)).item()
+        #         return i, j, k, d, h, w
+        #
+        # # Fallback to central crop
+        # return (depth - d) // 2, (height - h) // 2, (width - w) // 2, d, h, w
+
+    def forward(self, img: torch.Tensor) -> torch.Tensor:
+        """Apply the RandomResizedCrop transformation.
+
+        Args:
+            img (Tensor): Input 3D image.
+
+        Returns:
+            Tensor: Cropped and resized image.
+        """
+        rand_scale = np.random.uniform(self.scale[0], self.scale[1])
+        [i, j, k, d, h, w] = self.get_params(img,rand_scale, self.scale, self.ratio)
+        # print(i, j, k, d, h, w)
+        img_cropped = img[:, :, i:i + d, j:j + h, k:k + w]
+        # print(img_cropped.shape)
+        img_resized = F.interpolate(img_cropped, size=self.size, mode=self.interpolation,
+                                    align_corners=False if self.interpolation == 'trilinear' else None)
+        return img_resized#.squeeze(0)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, scale={self.scale}, ratio={self.ratio}, interpolation={self.interpolation})"
+
+def random_permute(X, select_dims=[-1,-2],include_flip=True):
+    axes=list(range(X[0].ndim))
+    selected_axes = [axes[i] for i in select_dims]
+    random.shuffle(selected_axes)
+    for i, dim in enumerate(select_dims):
+        axes[dim] = selected_axes[i]
+        if include_flip and random.choice([True,False]):
+            # X = [np.flip(x, axis=dim) for x in X]
+            X = [torch.flip(x, [dim]) for x in X]
+    # return [np.transpose(x,axes) for x in X]
+    return [x.permute(axes) for x in X]
+
+# def thresh_img(img,thresh = None,EPS = 10**-7):
+#     threshold0 = np.random.uniform(thresh[0], thresh[1])
+#     threshold1 = np.random.uniform(thresh[0], thresh[1])
+#     scale =
+#     if threshold is not None:
+#         # img=img-threshold
+#         # img=np.where(img>=0,img,0)
+#         # img = np.maximum(img-threshold,0)
+#         img = torch.maximum(img - threshold,torch.tensor(0.))
+#     # return (img - img.min()) / (img.max() - img.min() + EPS)
+#     return img
+
+def get_transformer(degrees=180,translate=0.125,ndims=2,prob=0.8,fill=0.,img_sz=None):
+    prob_crop=0. if img_sz==None else 0.8
+    # prob_crop=0. if len(img_sz)==2 else 0.8
+
+    if img_sz==None or len(img_sz)==2:
+        return torchvision.transforms.Compose([
+            torchvision.transforms.RandomApply([
+                torchvision.transforms.RandomAffine(degrees=degrees, translate=[translate] * ndims, fill=fill,
+                                                    interpolation=torchvision.transforms.InterpolationMode.BILINEAR),
+            ],prob),
+            # torchvision.transforms.RandomApply([
+            #     torchvision.transforms.RandomResizedCrop(size=img_sz),
+            # ], prob_crop),
+            torchvision.transforms.RandomVerticalFlip(p=0.5),
+            torchvision.transforms.RandomAutocontrast(p=0.5),
+        ])
+    else:
+        return torchvision.transforms.Compose([
+            torchvision.transforms.RandomApply([
+                torchvision.transforms.RandomResizedCrop(size=img_sz) if len(img_sz) == 2 else RandomResizedCrop3D(
+                    size=img_sz),
+            ], prob_crop),
+        ])
+
+
+def get_random_affine_transformer(degrees=180,translate=0.125,ndims=2):
+    return torchvision.transforms.RandomAffine(degrees=degrees, translate=[translate] * ndims,interpolation=torchvision.transforms.InterpolationMode.BILINEAR)
+
+def channel_merge_acdc(img):
+#   input: a torch tensor (C,H,W)
+  ch = img.shape[0]
+  output = np.zeros((img.shape[1], img.shape[2]))
+  # output[img[2,:,:] == 1] = 1
+  for i in range(ch):
+    output= output + img[i]
+  return output
+    
+def img_crop(img, crop_rate=2, img_sz=[256,256]):
+    ndims=len(img_sz)
+    crop = [np.random.randint(0.*imgs, 1. * imgs)//crop_rate for imgs in img_sz]
+    crop = [crop, [1 * imgs//crop_rate - c for imgs, c in zip(img_sz, crop)]]
+    if ndims==2:
+        return img[..., crop[0][0]: img_sz[0] - crop[1][0], crop[0][1]: img_sz[1] - crop[1][1]]
+    else:
+        return img[..., crop[0][0]: img_sz[0] - crop[1][0], crop[0][1]:img_sz[1] - crop[1][1], crop[0][2]: img_sz[2] - crop[1][2]]
+
+
+def boundary_limit(sample_coords0, max_sz, plus=0., minus=1.):
+    sample_coords = torch.split(sample_coords0, split_size_or_sections=1, dim=1)
+    # return tf.stack([tf.maximum(tf.minimum(x, sz - minus + plus), 0 + plus) for x, sz in zip(sample_coords, input_size0)],-1)
+    return torch.cat([torch.clamp(x * sz, min=minus - 1 * sz + plus, max=1 * sz - minus + plus) for x, sz in
+                      zip(sample_coords, max_sz)], 1)
+
+
+def resample(vol, ddf, ref=None, img_sz=None,max_sz=[128,128],ndims=2):
+    device = vol.device
+    img_sz = vol.size()[2:]
+    ndims=len(img_sz)
+    if ndims==2:
+        [h,w]=img_sz
+        img_shape = torch.reshape(torch.tensor([(h - 1) / 2., (w - 1) / 2.], device=device), [1, 1, 1, ndims])
+        ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=h), torch.arange(end=w)]), 0), [1, ndims,h, w ])
+    elif ndims==3:
+        [h, w, d] = img_sz
+        img_shape = torch.reshape(torch.tensor([(h - 1) / 2., (w - 1) / 2., (d-1)/2], device=device), [1, 1, 1, 1, ndims])
+        ref_grid = torch.reshape(torch.stack(torch.meshgrid([torch.arange(end=h), torch.arange(end=w), torch.arange(end=d)]), 0), [1, ndims,h, w, d])
+    # ref_grid.to(device)
+    # img_shape.to(device)
+    # ddf.to(device)
+    # ref = self.ref_grid if ref is None else ref
+    # img_sz = self.img_sz if img_sz is None else img_sz
+    resample_mode = 'bilinear'
+    # padding_mode = "border"
+    padding_mode = "zeros"
+
+    # img_sz = np.reshape(img_sz, [1] *(ndims+1)+[ndims])
+    # if ndims==2:
+    if True:
+        re=[0]+list(range(2,ndims+2))+[1]
+        # re=list(range(ndims+2))
+        # print((torch.flip((ddf.to(device) + ref_grid.permute(re))/ img_shape - 1, dims=[-1])).tolist())
+        return F.grid_sample(vol, torch.flip((ddf + ref_grid.permute(re).to(device))/ img_shape - 1, dims=[-1]).type(torch.float32).to(device), mode=resample_mode, padding_mode=padding_mode,align_corners=True)
+        #
+        # return F.grid_sample(vol, torch.flip(
+        #     torch.permute(ddf * torch.Tensor(np.reshape(np.array(max_sz), [1, 1, 1, ndims])) + ref_grid,
+        #                   [0, 2, 3, 1]) / img_shape - 1, dims=[-1]), mode=resample_mode, padding_mode=padding_mode,
+        #                      align_corners=True)
+
+def random_resample(vol,deform_scale=32.):
+    vol_size=vol.size()
+    device=vol.device
+    ndims = len(vol_size)-2
+    img_size=[s for s in vol_size[2:]]
+    if ndims==2:
+        img_size=img_size+[16]
+    # ddf,_,_=random_ddf(vol_size[0],img_size)
+    _,_,ddf=random_ddf(vol_size[0],img_size,ndims=ndims,range_gauss=deform_scale)
+    ddf=Variable(torch.tensor(ddf,dtype=torch.float32)).to(device)
+    if ndims==2:
+        return resample(vol,ddf[...,8,:ndims])
+    else:
+        return resample(vol, ddf[..., :ndims])
+
+def get_random_deformed_mask(msk_shape, deform_scale=32.,apply_possibility=0.75):
+    msk = torch.ones([1, 1]+list(msk_shape),dtype=torch.float32)
+    if random.uniform(0,1) < apply_possibility:
+        return random_resample(msk, deform_scale=deform_scale)
+    else:
+        return msk
+
+# grid option
+def get_tranf_mat(grid_size, vec=[[0., 0., 1.]], ang=[[0.]],transl=[[0,0,0]]):
+    return np.concatenate([get_rot_mat(grid_size, vec=vec, ang=ang),transl],-1)
+
+
+def get_rot_mat(grid_size, vec=[[0., 0., 1.]], ang=[[0.]],ndims=3):
+    vec = np.array(vec)
+    ang = np.array(ang)
+    batch_num = ang.shape[0]
+    return np.reshape(vecang2rotmats(vec, ang), [batch_num] + [ndims*(ndims)])
+
+def random_mat(batch_sz, img_sz, num_class=2,pn_spline=20, pn_gauss=10, range_spline=2., range_gauss=48, spread_range=[5., 24.],
+               transl_range=32., rot_range=np.pi / 2):
+    scale=4
+    ndims=3
+    vec=np.reshape(np.random.uniform(-1., 1., [batch_sz,1, ndims])+np.random.uniform(-.1, .1, [batch_sz,num_class, ndims]),[batch_sz*num_class, ndims])
+    ang=np.reshape(np.random.uniform(-rot_range, rot_range, [batch_sz,1])+np.random.uniform(-rot_range/scale, rot_range/scale, [batch_sz,num_class]),[batch_sz*num_class])
+    transl=np.reshape(np.random.uniform(-transl_range, transl_range, [batch_sz,1,ndims])+np.random.uniform(-transl_range/scale, transl_range/scale, [batch_sz,num_class,ndims]),[batch_sz*num_class,ndims])
+    return np.reshape(np.concatenate([get_rot_mat(img_sz, vec=vec, ang=ang),transl],-1),[batch_sz,num_class,4,3])
+
+    # return np.reshape(get_tranf_mat(img_sz, vec=np.random.uniform(-1., 1., [batch_sz*num_class, 3]), ang=np.random.uniform(-rot_range, rot_range, [batch_sz*num_class]),transl=np.random.uniform(-transl_range, transl_range, [batch_sz*num_class,3])),[batch_sz,num_class,4,3])
+
+def random_ddf(batch_sz, img_sz, pn_spline=20, pn_gauss=10, range_spline=1., range_gauss=16., spread_range=[16., 64.],
+               transl_range=0., rot_range=np.pi / 1,ndims=3):
+    rand_ang=np.random.uniform(-rot_range, rot_range, [batch_sz])
+    # rand_ang = np.random.randint(-4, 4, [batch_sz])*rot_range
+
+    if ndims==3:
+        rot_df = get_rot_ddf(img_sz, vec=np.random.uniform(-1., 1., [batch_sz, 3]),
+                             ang=rand_ang)
+    else:
+        rot_df = get_rot_ddf(img_sz, vec=np.concatenate([np.zeros([batch_sz, 2]),np.ones([batch_sz, 1])],-1),
+                             ang=rand_ang)
+    ndims = 3
+    # rot_df = +np.random.uniform(-1., 1., [batch_sz, ndims,ndims])
+    # ddf0=np.stack([generate_random_gaussian_ddf(img_sz, pn_gauss, range_sz=range_gauss, spread_std=spread_range)\
+    #                +generate_random_spline_ddf(img_sz, pn_spline, range_sz=range_spline)\
+    #                +np.random.uniform(-transl_range,transl_range,[3]) for i in range(batch_sz)],axis=0)\
+    #      +rot_df
+    if range_gauss>0:
+        ddf0 = np.tile([generate_random_gaussian_ddf(img_sz, pn_gauss, range_sz=range_gauss, spread_std=spread_range) \
+                        # + generate_random_spline_ddf(img_sz, pn_spline, range_sz=range_spline) \
+                        + np.random.uniform(-transl_range, transl_range, [ndims])], [batch_sz, 1, 1, 1, 1]) \
+               + rot_df
+    else:
+        ddf0 = rot_df
+
+    def boundary_replicate(sample_coords, input_size, padding=5):
+        return np.stack(
+            [np.maximum(np.minimum(sample_coords[..., i], input_size[i] - 1 + padding), 0 - padding) for i in
+             range(len(input_size))], axis=-1), \
+               np.prod([((sample_coords[..., i] < input_size[i]) * (sample_coords[..., i] >= 0)) for i in
+                        range(len(input_size))], axis=0)
+
+    ref = get_reference_grid(img_sz)
+    cf1, ind = boundary_replicate(ddf0 + ref, img_sz)
+    return cf1 - ref, np.expand_dims(ind, -1), rot_df
+
+
+def generate_random_gaussian_ddf(img_sz, pn=30, range_sz=5, spread_std=[0.1, 1.]):
+    x = np.floor(np.random.uniform(range_sz / 2., img_sz[0] - range_sz / 2., [1, pn])).astype('int')
+    y = np.floor(np.random.uniform(range_sz / 2., img_sz[1] - range_sz / 2., [1, pn])).astype('int')
+    z = np.floor(np.random.uniform(range_sz / 2., img_sz[2] - range_sz / 2., [1, pn])).astype('int')
+
+    odf = np.random.uniform(-range_sz, range_sz, [pn, 3])
+    vol = np.zeros([img_sz[0], img_sz[1], img_sz[2], 3])
+    vol[x, y, z] = odf
+
+    return spimg.gaussian_filter(vol, np.random.uniform(spread_std[0], spread_std[1]))
+
+
+def get_rot_ddf(grid_size, vec=[[0., 0., 1.]], ang=[[0.]]):
+    vec = np.array(vec)
+    ang = np.array(ang)
+    batch_num = ang.shape[0]
+    ref_grids = get_reference_grid(grid_size,
+                                   bias_scale=1.)
+    # a=vecang2rotmats(vec, ang)
+    return np.reshape(np.matmul(np.reshape(np.tile(ref_grids, [batch_num, 1, 1, 1, 1]), [batch_num, -1, 3]),
+                                vecang2rotmats(vec, ang)), [batch_num] + grid_size + [3]) - ref_grids
+
+
+def get_reference_grid(grid_size, bias_scale=0.):
+    return np.stack(np.meshgrid(
+        [i for i in range(grid_size[0])],
+        [j for j in range(grid_size[1])],
+        [k for k in range(grid_size[2])],
+        indexing='ij'), axis=-1).astype('float') - bias_scale * (np.array(grid_size) - 1) / 2.
+
+
+def resample_linear(inputs, ddf=None, sample_coords=None,random_boundary=True):
+    if random_boundary:
+        random_factor = np.random.uniform(0., 1.)
+        min_val = np.min(inputs)
+        inputs[:, 0, :, :] = min_val * random_factor + (1 - random_factor) * inputs[:, 0, :, :]
+        inputs[:, -1, :, :] = min_val * random_factor + (1 - random_factor) * inputs[:, -1, :, :]
+        inputs[:, :, 0, :] = min_val * random_factor + (1 - random_factor) * inputs[:, :, 0, :]
+        inputs[:, :, -1, :] = min_val * random_factor + (1 - random_factor) * inputs[:, :, -1, :]
+        inputs[:, :, :, 0] = min_val * random_factor + (1 - random_factor) * inputs[:, :, :, 0]
+        inputs[:, :, :, -1] = min_val * random_factor + (1 - random_factor) * inputs[:, :, :, -1]
+
+    input_size = inputs.shape[1:4]
+    sample_coords = get_reference_grid(input_size) + ddf if sample_coords is None else sample_coords
+    spatial_rank = 3  # inputs.ndim - 2
+    xy = [sample_coords[..., i] for i in
+          range(sample_coords.shape[-1])]  # tf.unstack(sample_coords, axis=len(sample_coords.shape)-1)
+    index_voxel_coords = [np.floor(x) for x in xy]
+
+    def boundary_replicate(sample_coords0, input_size0, plus=0):
+        return np.maximum(np.minimum(sample_coords0, input_size0 - 2 + plus), 0 + plus)
+
+    def boundary_replicate_float(sample_coords0, input_size0, plus=0.):
+        return np.maximum(np.minimum(sample_coords0, input_size0 - 1 + plus), 0 + plus)
+
+    xy = [boundary_replicate_float(x.astype('float32'), input_size[idx]) for idx, x in enumerate(xy)]
+    spatial_coords = [boundary_replicate(x.astype('int32'), input_size[idx])
+                      for idx, x in enumerate(index_voxel_coords)]
+    spatial_coords_plus1 = [boundary_replicate((x + 1).astype('int32'), input_size[idx], 1)
+                            for idx, x in enumerate(index_voxel_coords)]
+
+    weight = [np.expand_dims(x - i.astype('float32'), -1) for x, i in zip(xy, spatial_coords)]
+    weight_c = [np.expand_dims(i.astype('float32') - x, -1) for x, i in zip(xy, spatial_coords_plus1)]
+
+    sz = list(spatial_coords[0].shape)
+    batch_coords = np.tile(np.reshape(range(sz[0]), [sz[0]] + [1] * (len(sz) - 1)), [1] + sz[1:])
+    sc = (spatial_coords, spatial_coords_plus1)
+    binary_codes = [[int(c) for c in format(i, '0%ib' % spatial_rank)] for i in range(2 ** spatial_rank)]
+
+    make_sample = lambda bc: inputs[batch_coords, sc[bc[0]][0], sc[bc[1]][1], sc[bc[2]][
+        2], ...]  # tf.gather_nd(inputs, np.stack([batch_coords] + [sc[c][i] for i, c in enumerate(bc)], -1))
+    samples = [make_sample(bc) for bc in binary_codes]
+
+    def pyramid_combination(samples0, weight0, weight_c0):
+        if len(weight0) == 1:
+            return samples0[0] * weight_c0[0] + samples0[1] * weight0[0]
+        else:
+            return pyramid_combination(samples0[::2], weight0[:-1], weight_c0[:-1]) * weight_c0[-1] + \
+                   pyramid_combination(samples0[1::2], weight0[:-1], weight_c0[:-1]) * weight0[-1]
+
+    return pyramid_combination(samples, weight, weight_c)
+
+
+def vecang2rotmats(vec, ang):
+    return np.stack([np.reshape(vecang2rotmat(vec[i, ...], ang[i, ...]), [3, 3]) for i in range(len(vec))], 0)
+
+
+def vecang2rotmat(vec, ang):
+    q = quater.Quaternion(axis=vec, angle=ang)
+    return q.rotation_matrix
+
+
+def images_to_vectors(images):
+  return images.view(images.size(0), 16384).to(device)
+
+def vectors_to_images(vectors):
+  return vectors.view(vectors.size(0), 1, 128, 128).to(device)
+
+def noise(size):
+  n = Variable(torch.randn(size, 100)).to(device)
+  return n
+
+def ones_target(size):
+  data = Variable(torch.ones(size, 1)).to(device)
+  return data
+
+def zeros_target(size):
+  data = Variable(torch.zeros(size, 1)).to(device)
+  return data
+
+
+def eval_detJ_lab(disp=None,vol1=None,vol2=None,thresh=0.5):
+    ndims=disp.ndim-2
+    if vol1 ==None or thresh==None:
+        label=1
+    else:
+        label=vol1>thresh
+        label=label*(spimg.laplace(label) < 0.1)
+        rescale_factor=2
+        label=label[...,::rescale_factor,::rescale_factor,::rescale_factor]
+
+    # disp = disp.permute([0, *range(2,ndims+2), 1])
+    # print(disp.shape)
+    disp = np.transpose(disp, [0, *range(2,ndims+2), 1])
+    # Jacob=np.stack(np.gradient(disp,axis=[-4,-3,-2]),-1)
+    Jacob=np.stack(np.gradient(disp,axis=[*range(1,ndims+1)]),-1)
+    for ii in range(ndims):
+        Jacob[..., ii, ii] = Jacob[..., ii, ii] + 1
+    # Jacob[..., 0, 0] = Jacob[..., 0, 0] + 1
+    # Jacob[..., 1, 1] = Jacob[..., 1, 1] + 1
+    # Jacob[..., 2, 2] = Jacob[..., 2, 2] + 1
+    return np.sum((np.linalg.det(Jacob)<0)*label)
+
+def eval_def_mag(disp=None,vol1=None,vol2=None,thresh=0.5):
+    ndims=3
+    # if vol1 ==None or thresh==None:
+    #     label=1
+    # else:
+    #     label=vol1>thresh
+    #     label=label*(spimg.laplace(label) < 0.1)
+    #     rescale_factor=2
+    #     label=label[...,::rescale_factor,::rescale_factor,::rescale_factor]
+    mag=np.sqrt(np.sum(np.square(disp),axis=1))
+    sz=mag.shape
+    max_mag=np.mean(np.max(np.reshape(mag,[sz[0],-1]),axis=-1))
+    avg_mag=np.mean(mag)
+    return [avg_mag,max_mag]
+
+
+def converet_to_nibabel(vol_tensor,ndims=3):
+    if isinstance(vol_tensor, np.ndarray):
+        vol_np=vol_tensor
+    else:
+        vol_np=vol_tensor.cpu().numpy()
+    vol_np=vol_np.squeeze(0)
+    if ndims==3:
+        map_eyes = np.eye(4)
+    elif ndims==2:
+        map_eyes = np.eye(4)
+        map_eyes[2,2]=0
+
+    if vol_np.shape[0]==1:
+        vol_np=vol_np.squeeze(0)
+    elif vol_np.shape[0]>1:
+        # save as 4D volumes
+        # print(vol_np.shape)
+        vol_np=np.transpose(vol_np,[1,2,3,0])
+        
+    return nib.Nifti1Image(vol_np, affine=map_eyes)
+
+def print_memory_usage(tag=""):
+    print(f"[{tag}] Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB | Cached: {torch.cuda.memory_reserved() / 1e9:.2f} GB")
+
+
+if __name__ == "__main__":
+    vol_shape=[4,1,64,64]
+
+    vol=np.random.uniform(-1,1,vol_shape)
+    vol=Variable(torch.tensor(vol,dtype=torch.float32))
+    vol_res=random_resample(vol)
+    vol_crop=img_crop(vol_res)
+
+    mask = get_random_deformed_mask(vol.shape[2:])
+
+    print(mask)
+    
+    # print(vol.tolist())
+    # print(vol_res.tolist())