diff --git a/models/maisi_ct_generative/configs/inference.json b/models/maisi_ct_generative/configs/inference.json index 2ad6aa08..283305ed 100644 --- a/models/maisi_ct_generative/configs/inference.json +++ b/models/maisi_ct_generative/configs/inference.json @@ -221,21 +221,21 @@ "cross_attention_dim": 10 }, "autoencoder": "$@autoencoder_def.to(@device)", - "checkpoint_autoencoder": "$torch.load(@trained_autoencoder_path)", + "checkpoint_autoencoder": "$torch.load(@trained_autoencoder_path, weights_only=True)", "load_autoencoder": "$@autoencoder.load_state_dict(@checkpoint_autoencoder)", "diffusion_unet": "$@diffusion_unet_def.to(@device)", - "checkpoint_diffusion_unet": "$torch.load(@trained_diffusion_path)", + "checkpoint_diffusion_unet": "$torch.load(@trained_diffusion_path, weights_only=False)", "load_diffusion": "$@diffusion_unet.load_state_dict(@checkpoint_diffusion_unet['unet_state_dict'])", "controlnet": "$@controlnet_def.to(@device)", "copy_controlnet_state": "$monai.networks.utils.copy_model_state(@controlnet, @diffusion_unet.state_dict())", - "checkpoint_controlnet": "$torch.load(@trained_controlnet_path)", + "checkpoint_controlnet": "$torch.load(@trained_controlnet_path, weights_only=False)", "load_controlnet": "$@controlnet.load_state_dict(@checkpoint_controlnet['controlnet_state_dict'], strict=True)", "scale_factor": "$@checkpoint_diffusion_unet['scale_factor'].to(@device)", "mask_generation_autoencoder": "$@mask_generation_autoencoder_def.to(@device)", - "checkpoint_mask_generation_autoencoder": "$torch.load(@trained_mask_generation_autoencoder_path)", + "checkpoint_mask_generation_autoencoder": "$torch.load(@trained_mask_generation_autoencoder_path, weights_only=True)", "load_mask_generation_autoencoder": "$@mask_generation_autoencoder.load_state_dict(@checkpoint_mask_generation_autoencoder, strict=True)", "mask_generation_diffusion_unet": "$@mask_generation_diffusion_def.to(@device)", - "checkpoint_mask_generation_diffusion_unet": "$torch.load(@trained_mask_generation_diffusion_path)", + "checkpoint_mask_generation_diffusion_unet": "$torch.load(@trained_mask_generation_diffusion_path, weights_only=True)", "load_mask_generation_diffusion": "$@mask_generation_diffusion_unet.load_state_dict(@checkpoint_mask_generation_diffusion_unet['unet_state_dict'], strict=True)", "mask_generation_scale_factor": "$@checkpoint_mask_generation_diffusion_unet['scale_factor']", "noise_scheduler": { diff --git a/models/maisi_ct_generative/configs/label_dict.json b/models/maisi_ct_generative/configs/label_dict.json deleted file mode 100644 index 44bbb645..00000000 --- a/models/maisi_ct_generative/configs/label_dict.json +++ /dev/null @@ -1,134 +0,0 @@ -{ - "liver": 1, - "dummy1": 2, - "spleen": 3, - "pancreas": 4, - "right kidney": 5, - "aorta": 6, - "inferior vena cava": 7, - "right adrenal gland": 8, - "left adrenal gland": 9, - "gallbladder": 10, - "esophagus": 11, - "stomach": 12, - "duodenum": 13, - "left kidney": 14, - "bladder": 15, - "dummy2": 16, - "portal vein and splenic vein": 17, - "dummy3": 18, - "small bowel": 19, - "dummy4": 20, - "dummy5": 21, - "brain": 22, - "lung tumor": 23, - "pancreatic tumor": 24, - "hepatic vessel": 25, - "hepatic tumor": 26, - "colon cancer primaries": 27, - "left lung upper lobe": 28, - "left lung lower lobe": 29, - "right lung upper lobe": 30, - "right lung middle lobe": 31, - "right lung lower lobe": 32, - "vertebrae L5": 33, - "vertebrae L4": 34, - "vertebrae L3": 35, - "vertebrae L2": 36, - "vertebrae L1": 37, - "vertebrae T12": 38, - "vertebrae T11": 39, - "vertebrae T10": 40, - "vertebrae T9": 41, - "vertebrae T8": 42, - "vertebrae T7": 43, - "vertebrae T6": 44, - "vertebrae T5": 45, - "vertebrae T4": 46, - "vertebrae T3": 47, - "vertebrae T2": 48, - "vertebrae T1": 49, - "vertebrae C7": 50, - "vertebrae C6": 51, - "vertebrae C5": 52, - "vertebrae C4": 53, - "vertebrae C3": 54, - "vertebrae C2": 55, - "vertebrae C1": 56, - "trachea": 57, - "left iliac artery": 58, - "right iliac artery": 59, - "left iliac vena": 60, - "right iliac vena": 61, - "colon": 62, - "left rib 1": 63, - "left rib 2": 64, - "left rib 3": 65, - "left rib 4": 66, - "left rib 5": 67, - "left rib 6": 68, - "left rib 7": 69, - "left rib 8": 70, - "left rib 9": 71, - "left rib 10": 72, - "left rib 11": 73, - "left rib 12": 74, - "right rib 1": 75, - "right rib 2": 76, - "right rib 3": 77, - "right rib 4": 78, - "right rib 5": 79, - "right rib 6": 80, - "right rib 7": 81, - "right rib 8": 82, - "right rib 9": 83, - "right rib 10": 84, - "right rib 11": 85, - "right rib 12": 86, - "left humerus": 87, - "right humerus": 88, - "left scapula": 89, - "right scapula": 90, - "left clavicula": 91, - "right clavicula": 92, - "left femur": 93, - "right femur": 94, - "left hip": 95, - "right hip": 96, - "sacrum": 97, - "left gluteus maximus": 98, - "right gluteus maximus": 99, - "left gluteus medius": 100, - "right gluteus medius": 101, - "left gluteus minimus": 102, - "right gluteus minimus": 103, - "left autochthon": 104, - "right autochthon": 105, - "left iliopsoas": 106, - "right iliopsoas": 107, - "left atrial appendage": 108, - "brachiocephalic trunk": 109, - "left brachiocephalic vein": 110, - "right brachiocephalic vein": 111, - "left common carotid artery": 112, - "right common carotid artery": 113, - "costal cartilages": 114, - "heart": 115, - "left kidney cyst": 116, - "right kidney cyst": 117, - "prostate": 118, - "pulmonary vein": 119, - "skull": 120, - "spinal cord": 121, - "sternum": 122, - "left subclavian artery": 123, - "right subclavian artery": 124, - "superior vena cava": 125, - "thyroid gland": 126, - "vertebrae S1": 127, - "bone lesion": 128, - "dummy6": 129, - "dummy7": 130, - "dummy8": 131, - "airway": 132 -} diff --git a/models/maisi_ct_generative/configs/label_dict_124_to_132.json b/models/maisi_ct_generative/configs/label_dict_124_to_132.json deleted file mode 100644 index 96ff311e..00000000 --- a/models/maisi_ct_generative/configs/label_dict_124_to_132.json +++ /dev/null @@ -1,502 +0,0 @@ -{ - "background": [ - 0, - 0 - ], - "liver": [ - 1, - 1 - ], - "spleen": [ - 2, - 3 - ], - "pancreas": [ - 3, - 4 - ], - "right kidney": [ - 4, - 5 - ], - "aorta": [ - 5, - 6 - ], - "inferior vena cava": [ - 6, - 7 - ], - "right adrenal gland": [ - 7, - 8 - ], - "left adrenal gland": [ - 8, - 9 - ], - "gallbladder": [ - 9, - 10 - ], - "esophagus": [ - 10, - 11 - ], - "stomach": [ - 11, - 12 - ], - "duodenum": [ - 12, - 13 - ], - "left kidney": [ - 13, - 14 - ], - "bladder": [ - 14, - 15 - ], - "portal vein and splenic vein": [ - 15, - 17 - ], - "small bowel": [ - 16, - 19 - ], - "brain": [ - 17, - 22 - ], - "lung tumor": [ - 18, - 23 - ], - "pancreatic tumor": [ - 19, - 24 - ], - "hepatic vessel": [ - 20, - 25 - ], - "hepatic tumor": [ - 21, - 26 - ], - "colon cancer primaries": [ - 22, - 27 - ], - "left lung upper lobe": [ - 23, - 28 - ], - "left lung lower lobe": [ - 24, - 29 - ], - "right lung upper lobe": [ - 25, - 30 - ], - "right lung middle lobe": [ - 26, - 31 - ], - "right lung lower lobe": [ - 27, - 32 - ], - "vertebrae L5": [ - 28, - 33 - ], - "vertebrae L4": [ - 29, - 34 - ], - "vertebrae L3": [ - 30, - 35 - ], - "vertebrae L2": [ - 31, - 36 - ], - "vertebrae L1": [ - 32, - 37 - ], - "vertebrae T12": [ - 33, - 38 - ], - "vertebrae T11": [ - 34, - 39 - ], - "vertebrae T10": [ - 35, - 40 - ], - "vertebrae T9": [ - 36, - 41 - ], - "vertebrae T8": [ - 37, - 42 - ], - "vertebrae T7": [ - 38, - 43 - ], - "vertebrae T6": [ - 39, - 44 - ], - "vertebrae T5": [ - 40, - 45 - ], - "vertebrae T4": [ - 41, - 46 - ], - "vertebrae T3": [ - 42, - 47 - ], - "vertebrae T2": [ - 43, - 48 - ], - "vertebrae T1": [ - 44, - 49 - ], - "vertebrae C7": [ - 45, - 50 - ], - "vertebrae C6": [ - 46, - 51 - ], - "vertebrae C5": [ - 47, - 52 - ], - "vertebrae C4": [ - 48, - 53 - ], - "vertebrae C3": [ - 49, - 54 - ], - "vertebrae C2": [ - 50, - 55 - ], - "vertebrae C1": [ - 51, - 56 - ], - "trachea": [ - 52, - 57 - ], - "left iliac artery": [ - 53, - 58 - ], - "right iliac artery": [ - 54, - 59 - ], - "left iliac vena": [ - 55, - 60 - ], - "right iliac vena": [ - 56, - 61 - ], - "colon": [ - 57, - 62 - ], - "left rib 1": [ - 58, - 63 - ], - "left rib 2": [ - 59, - 64 - ], - "left rib 3": [ - 60, - 65 - ], - "left rib 4": [ - 61, - 66 - ], - "left rib 5": [ - 62, - 67 - ], - "left rib 6": [ - 63, - 68 - ], - "left rib 7": [ - 64, - 69 - ], - "left rib 8": [ - 65, - 70 - ], - "left rib 9": [ - 66, - 71 - ], - "left rib 10": [ - 67, - 72 - ], - "left rib 11": [ - 68, - 73 - ], - "left rib 12": [ - 69, - 74 - ], - "right rib 1": [ - 70, - 75 - ], - "right rib 2": [ - 71, - 76 - ], - "right rib 3": [ - 72, - 77 - ], - "right rib 4": [ - 73, - 78 - ], - "right rib 5": [ - 74, - 79 - ], - "right rib 6": [ - 75, - 80 - ], - "right rib 7": [ - 76, - 81 - ], - "right rib 8": [ - 77, - 82 - ], - "right rib 9": [ - 78, - 83 - ], - "right rib 10": [ - 79, - 84 - ], - "right rib 11": [ - 80, - 85 - ], - "right rib 12": [ - 81, - 86 - ], - "left humerus": [ - 82, - 87 - ], - "right humerus": [ - 83, - 88 - ], - "left scapula": [ - 84, - 89 - ], - "right scapula": [ - 85, - 90 - ], - "left clavicula": [ - 86, - 91 - ], - "right clavicula": [ - 87, - 92 - ], - "left femur": [ - 88, - 93 - ], - "right femur": [ - 89, - 94 - ], - "left hip": [ - 90, - 95 - ], - "right hip": [ - 91, - 96 - ], - "sacrum": [ - 92, - 97 - ], - "left gluteus maximus": [ - 93, - 98 - ], - "right gluteus maximus": [ - 94, - 99 - ], - "left gluteus medius": [ - 95, - 100 - ], - "right gluteus medius": [ - 96, - 101 - ], - "left gluteus minimus": [ - 97, - 102 - ], - "right gluteus minimus": [ - 98, - 103 - ], - "left autochthon": [ - 99, - 104 - ], - "right autochthon": [ - 100, - 105 - ], - "left iliopsoas": [ - 101, - 106 - ], - "right iliopsoas": [ - 102, - 107 - ], - "left atrial appendage": [ - 103, - 108 - ], - "brachiocephalic trunk": [ - 104, - 109 - ], - "left brachiocephalic vein": [ - 105, - 110 - ], - "right brachiocephalic vein": [ - 106, - 111 - ], - "left common carotid artery": [ - 107, - 112 - ], - "right common carotid artery": [ - 108, - 113 - ], - "costal cartilages": [ - 109, - 114 - ], - "heart": [ - 110, - 115 - ], - "prostate": [ - 111, - 118 - ], - "pulmonary vein": [ - 112, - 119 - ], - "skull": [ - 113, - 120 - ], - "spinal cord": [ - 114, - 121 - ], - "sternum": [ - 115, - 122 - ], - "left subclavian artery": [ - 116, - 123 - ], - "right subclavian artery": [ - 117, - 124 - ], - "superior vena cava": [ - 118, - 125 - ], - "thyroid gland": [ - 119, - 126 - ], - "vertebrae S1": [ - 120, - 127 - ], - "bone lesion": [ - 121, - 128 - ], - "kidney mass": [ - 122, - 129 - ], - "airway": [ - 123, - 132 - ], - "body": [ - 124, - 200 - ] -} diff --git a/models/maisi_ct_generative/configs/metadata.json b/models/maisi_ct_generative/configs/metadata.json index 2838ca92..a0d2c53e 100644 --- a/models/maisi_ct_generative/configs/metadata.json +++ b/models/maisi_ct_generative/configs/metadata.json @@ -1,7 +1,8 @@ { "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_ldm_20240318.json", - "version": "0.4.2", + "version": "0.4.3", "changelog": { + "0.4.3": "remove download large files, add weights_only when loading weights and add label_dict to large files", "0.4.2": "update train.json to fix finetune ckpt bug", "0.4.1": "update large files", "0.4.0": "update to use monai 1.4, model ckpt updated, rm GenerativeAI repo, add quality check", diff --git a/models/maisi_ct_generative/configs/train.json b/models/maisi_ct_generative/configs/train.json index f2c91fab..7138d616 100644 --- a/models/maisi_ct_generative/configs/train.json +++ b/models/maisi_ct_generative/configs/train.json @@ -98,11 +98,11 @@ }, "unzip_dataset": "$scripts.utils.unzip_dataset(@dataset_dir)", "diffusion_unet": "$@diffusion_unet_def.to(@device)", - "checkpoint_diffusion_unet": "$torch.load(@trained_diffusion_path)", + "checkpoint_diffusion_unet": "$torch.load(@trained_diffusion_path, weights_only=False)", "load_diffusion": "$@diffusion_unet.load_state_dict(@checkpoint_diffusion_unet['unet_state_dict'])", "controlnet": "$@controlnet_def.to(@device)", "copy_controlnet_state": "$monai.networks.utils.copy_model_state(@controlnet, @diffusion_unet.state_dict())", - "checkpoint_controlnet": "$torch.load(@trained_controlnet_path)", + "checkpoint_controlnet": "$torch.load(@trained_controlnet_path, weights_only=False)", "load_controlnet": "$@controlnet.load_state_dict(@checkpoint_controlnet['controlnet_state_dict'], strict=True)", "scale_factor": "$@checkpoint_diffusion_unet['scale_factor'].to(@device)", "loss": { diff --git a/models/maisi_ct_generative/docs/README.md b/models/maisi_ct_generative/docs/README.md index 20a87c19..bcb7c05f 100644 --- a/models/maisi_ct_generative/docs/README.md +++ b/models/maisi_ct_generative/docs/README.md @@ -23,10 +23,6 @@ The inference requires: - Disk Memory: at least 21GB disk memory ### Execute inference -The model weights can be downloaded with -``` -python -m scripts.download_files -``` The following code generates a synthetic image from a random sampled noise. ``` python -m monai.bundle run --config_file configs/inference.json diff --git a/models/maisi_ct_generative/large_files.yml b/models/maisi_ct_generative/large_files.yml index 58ca4f2d..97b26885 100644 --- a/models/maisi_ct_generative/large_files.yml +++ b/models/maisi_ct_generative/large_files.yml @@ -43,3 +43,11 @@ large_files: url: "https://developer.download.nvidia.com/assets/Clara/monai/tutorials/model_zoo/model_maisi_C4KC-KiTS_subset.json" hash_val: "2b8a0aebc31f5dcae5fc496d92e2beb9" hash_type: "md5" +- path: "configs/label_dict.json" + url: "https://raw.githubusercontent.com/Project-MONAI/tutorials/1c942f49693131c53365359a7b11153a563a97a6/generation/maisi/configs/label_dict.json" + hash_val: "c5fc6270c0fa2b4af8e67ba95a7dd5bd" + hash_type: "md5" +- path: "configs/label_dict_124_to_132.json" + url: "https://raw.githubusercontent.com/Project-MONAI/tutorials/1c942f49693131c53365359a7b11153a563a97a6/generation/maisi/configs/label_dict_124_to_132.json" + hash_val: "85ab16bd93c70c5514143a031872297a" + hash_type: "md5" diff --git a/models/maisi_ct_generative/scripts/download_files.py b/models/maisi_ct_generative/scripts/download_files.py deleted file mode 100644 index 50a7ef4b..00000000 --- a/models/maisi_ct_generative/scripts/download_files.py +++ /dev/null @@ -1,10 +0,0 @@ -import yaml -from monai.apps import download_url - -# Load YAML file -with open("large_files.yml", "r") as file: - data = yaml.safe_load(file) - -# Iterate over each file in the YAML and download it -for file in data["large_files"]: - download_url(url=file["url"], filepath=file["path"]) diff --git a/models/maisi_ct_generative/scripts/sample.py b/models/maisi_ct_generative/scripts/sample.py index 30127b80..d161597f 100644 --- a/models/maisi_ct_generative/scripts/sample.py +++ b/models/maisi_ct_generative/scripts/sample.py @@ -108,7 +108,7 @@ def ldm_conditional_sample_one_mask( """ recon_model = ReconModel(autoencoder=autoencoder, scale_factor=scale_factor).to(device) - with torch.no_grad(), torch.cuda.amp.autocast(): + with torch.no_grad(), torch.amp.autocast("cuda"): # Generate random noise latents = initialize_noise_latents(latent_shape, device) anatomy_size = torch.FloatTensor(anatomy_size).unsqueeze(0).unsqueeze(0).half().to(device) @@ -217,7 +217,7 @@ def ldm_conditional_sample_one_image( recon_model = ReconModel(autoencoder=autoencoder, scale_factor=scale_factor).to(device) - with torch.no_grad(), torch.cuda.amp.autocast(): + with torch.no_grad(), torch.amp.autocast("cuda"): logging.info("---- Start generating latent features... ----") start_time = time.time() # generate segmentation mask @@ -457,7 +457,7 @@ def check_input(body_region, anatomy_list, label_dict_json, output_size, spacing raise ValueError( f"The components in anatomy_list have to be chosen from {label_dict.keys()}, yet got {anatomy}." ) - logging.info(f"The generate results will have voxel size to be {spacing}mm, volume size to be {output_size}.") + logging.info(f"The generate results will have voxel size to be {spacing} mm, volume size to be {output_size}.") return @@ -923,7 +923,7 @@ def ensure_output_size_and_spacing(self, labels, check_contains_target_labels=Tr raise ValueError( ( f"Resampled mask does not contain required class labels {anatomy_label}. " - "Please tune spacing and output size." + "Please consider increasing the output spacing or specifying a larger output size." ) ) return labels diff --git a/models/maisi_ct_generative/scripts/trainer.py b/models/maisi_ct_generative/scripts/trainer.py index 7c779ef3..7b8eae9f 100644 --- a/models/maisi_ct_generative/scripts/trainer.py +++ b/models/maisi_ct_generative/scripts/trainer.py @@ -232,7 +232,7 @@ def _compute_pred_loss(): engine.optimizer.zero_grad(set_to_none=engine.optim_set_to_none) if engine.amp and engine.scaler is not None: - with torch.cuda.amp.autocast(**engine.amp_kwargs): + with torch.amp.autocast("cuda", **engine.amp_kwargs): _compute_pred_loss() engine.scaler.scale(engine.state.output[Keys.LOSS]).backward() engine.fire_event(IterationEvents.BACKWARD_COMPLETED)