# Copyright (c) MONAI Consortium # 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. import json import logging import math import os import random import time from datetime import datetime import warnings import gc import monai import torch from monai.data import MetaTensor from monai.inferers.inferer import DiffusionInferer from monai.transforms import Compose, SaveImage from monai.utils import set_determinism from tqdm import tqdm from monai.inferers.inferer import SlidingWindowInferer from monai.networks.schedulers import RFlowScheduler, DDPMScheduler from .augmentation import augmentation from .find_masks import find_masks from .quality_check import is_outlier from .utils import ( binarize_labels, general_mask_generation_post_process, get_body_region_index_from_mask, remap_labels, dynamic_infer, ) class ReconModel(torch.nn.Module): """ A PyTorch module for reconstructing images from latent representations. Attributes: autoencoder: The autoencoder model used for decoding. scale_factor: Scaling factor applied to the input before decoding. """ def __init__(self, autoencoder, scale_factor): super().__init__() self.autoencoder = autoencoder self.scale_factor = scale_factor def forward(self, z): """ Decode the input latent representation to an image. Args: z (torch.Tensor): The input latent representation. Returns: torch.Tensor: The reconstructed image. """ recon_pt_nda = self.autoencoder.decode_stage_2_outputs(z / self.scale_factor) return recon_pt_nda def initialize_noise_latents(latent_shape, device): """ Initialize random noise latents for image generation with float16. Args: latent_shape (tuple): The shape of the latent space. device (torch.device): The device to create the tensor on. Returns: torch.Tensor: Initialized noise latents. """ return ( torch.randn( [ 1, ] + list(latent_shape) ) .half() .to(device) ) def ldm_conditional_sample_one_mask( autoencoder, diffusion_unet, noise_scheduler, scale_factor, anatomy_size, device, latent_shape, label_dict_remap_json, num_inference_steps=1000, autoencoder_sliding_window_infer_size=[96, 96, 96], autoencoder_sliding_window_infer_overlap=0.6667, ): """ Generate a single synthetic mask using a latent diffusion model. Args: autoencoder (nn.Module): The autoencoder model. diffusion_unet (nn.Module): The diffusion U-Net model. noise_scheduler: The noise scheduler for the diffusion process. scale_factor (float): Scaling factor for the latent space. anatomy_size (torch.Tensor): Tensor specifying the desired anatomy sizes. device (torch.device): The device to run the computation on. latent_shape (tuple): The shape of the latent space. label_dict_remap_json (str): Path to the JSON file for label remapping. num_inference_steps (int): Number of inference steps for the diffusion process. autoencoder_sliding_window_infer_size (list, optional): Size of the sliding window for inference. Defaults to [96, 96, 96]. autoencoder_sliding_window_infer_overlap (float, optional): Overlap ratio for sliding window inference. Defaults to 0.6667. Returns: torch.Tensor: The generated synthetic mask. """ recon_model = ReconModel(autoencoder=autoencoder, scale_factor=scale_factor).to(device) 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) # synthesize latents if isinstance(noise_scheduler, DDPMScheduler) and num_inference_steps < noise_scheduler.num_train_timesteps: warnings.warn( "**************************************************************\n" "* WARNING: Mask noise_scheduler is a DDPMScheduler.\n" "* We expect num_inference_steps = noise_scheduler.num_train_timesteps" f" = {noise_scheduler.num_train_timesteps}.\n" f"* Yet got num_inference_steps = {num_inference_steps}.\n" "* The generated image quality is not guaranteed.\n" "**************************************************************" ) noise_scheduler.set_timesteps(num_inference_steps=num_inference_steps) # mask generator is DDPM inferer_ddpm = DiffusionInferer(noise_scheduler) latents = inferer_ddpm.sample( input_noise=latents, diffusion_model=diffusion_unet, scheduler=noise_scheduler, verbose=True, conditioning=anatomy_size.to(device), ) inferer = SlidingWindowInferer( roi_size=autoencoder_sliding_window_infer_size, sw_batch_size=1, progress=True, mode="gaussian", overlap=autoencoder_sliding_window_infer_overlap, sw_device=device, device=torch.device("cpu"), ) synthetic_mask = dynamic_infer(inferer, recon_model, latents) synthetic_mask = torch.softmax(synthetic_mask, dim=1) synthetic_mask = torch.argmax(synthetic_mask, dim=1, keepdim=True) # mapping raw index to 132 labels synthetic_mask = remap_labels(synthetic_mask, label_dict_remap_json) ###### post process ##### data = synthetic_mask.squeeze().cpu().detach().numpy() labels = [23, 24, 26, 27, 128] target_tumor_label = None for index, size in enumerate(anatomy_size[0, 0, 5:10]): if size.item() != -1.0: target_tumor_label = labels[index] logging.info(f"target_tumor_label for postprocess:{target_tumor_label}") data = general_mask_generation_post_process(data, target_tumor_label=target_tumor_label, device=device) synthetic_mask = torch.from_numpy(data).unsqueeze(0).unsqueeze(0).to(device) return synthetic_mask def ldm_conditional_sample_one_image( autoencoder, diffusion_unet, controlnet, noise_scheduler, scale_factor, device, combine_label_or, spacing_tensor, latent_shape, output_size, noise_factor, top_region_index_tensor=None, bottom_region_index_tensor=None, modality_tensor=None, num_inference_steps=1000, autoencoder_sliding_window_infer_size=[96, 96, 96], autoencoder_sliding_window_infer_overlap=0.6667, ): """ Generate a single synthetic image using a latent diffusion model with controlnet. Args: autoencoder (nn.Module): The autoencoder model. diffusion_unet (nn.Module): The diffusion U-Net model. controlnet (nn.Module): The controlnet model. noise_scheduler: The noise scheduler for the diffusion process. scale_factor (float): Scaling factor for the latent space. device (torch.device): The device to run the computation on. combine_label_or (torch.Tensor): The combined label tensor. spacing_tensor (torch.Tensor): Tensor specifying the spacing. latent_shape (tuple): The shape of the latent space. output_size (tuple): The desired output size of the image. noise_factor (float): Factor to scale the initial noise. top_region_index_tensor (torch.Tensor): Tensor specifying the top region index. Defaults to None. bottom_region_index_tensor (torch.Tensor): Tensor specifying the bottom region index. Defaults to None. modality_tensor (torch.Tensor): Int Tensor specifying the modality. num_inference_steps (int): Number of inference steps for the diffusion process. autoencoder_sliding_window_infer_size (list, optional): Size of the sliding window for inference. Defaults to [96, 96, 96]. autoencoder_sliding_window_infer_overlap (float, optional): Overlap ratio for sliding window inference. Defaults to 0.6667. Returns: tuple: A tuple containing the synthetic image and its corresponding label. """ # CT image intensity range a_min = -1000 a_max = 1000 # autoencoder output intensity range b_min = 0.0 b_max = 1 include_body_region = diffusion_unet.include_top_region_index_input include_modality = diffusion_unet.num_class_embeds is not None recon_model = ReconModel(autoencoder=autoencoder, scale_factor=scale_factor).to(device) with torch.no_grad(), torch.amp.autocast("cuda"): logging.info("---- Start generating latent features... ----") start_time = time.time() # generate segmentation mask combine_label = combine_label_or.to(device) if ( output_size[0] != combine_label.shape[2] or output_size[1] != combine_label.shape[3] or output_size[2] != combine_label.shape[4] ): logging.info( "output_size is not a desired value. Need to interpolate the mask to match with output_size. The result image will be very low quality." ) combine_label = torch.nn.functional.interpolate(combine_label, size=output_size, mode="nearest") controlnet_cond_vis = binarize_labels(combine_label.as_tensor().long()).half() # Generate random noise latents = initialize_noise_latents(latent_shape, device) * noise_factor # synthesize latents if isinstance(noise_scheduler, RFlowScheduler): noise_scheduler.set_timesteps( num_inference_steps=num_inference_steps, input_img_size_numel=torch.prod(torch.tensor(latents.shape[2:])), ) else: noise_scheduler.set_timesteps(num_inference_steps=num_inference_steps) if isinstance(noise_scheduler, DDPMScheduler) and num_inference_steps < noise_scheduler.num_train_timesteps: warnings.warn( "**************************************************************\n" "* WARNING: Image noise_scheduler is a DDPMScheduler.\n" "* We expect num_inference_steps = noise_scheduler.num_train_timesteps" f" = {noise_scheduler.num_train_timesteps}.\n" f"* Yet got num_inference_steps = {num_inference_steps}.\n" "* The generated image quality is not guaranteed.\n" "**************************************************************" ) all_timesteps = noise_scheduler.timesteps all_next_timesteps = torch.cat((all_timesteps[1:], torch.tensor([0], dtype=all_timesteps.dtype))) progress_bar = tqdm( zip(all_timesteps, all_next_timesteps), total=min(len(all_timesteps), len(all_next_timesteps)), ) for t, next_t in progress_bar: # get controlnet output # Create a dictionary to store the inputs controlnet_inputs = { "x": latents, "timesteps": torch.Tensor((t,)).to(device), "controlnet_cond": controlnet_cond_vis, } if include_modality: controlnet_inputs.update( { "class_labels": modality_tensor, } ) down_block_res_samples, mid_block_res_sample = controlnet(**controlnet_inputs) # get diffusion network output # Create a dictionary to store the inputs unet_inputs = { "x": latents, "timesteps": torch.Tensor((t,)).to(device), "spacing_tensor": spacing_tensor, "down_block_additional_residuals": down_block_res_samples, "mid_block_additional_residual": mid_block_res_sample, } # Add extra arguments if include_body_region is True if include_body_region: unet_inputs.update( { "top_region_index_tensor": top_region_index_tensor, "bottom_region_index_tensor": bottom_region_index_tensor, } ) if include_modality: unet_inputs.update( { "class_labels": modality_tensor, } ) model_output = diffusion_unet(**unet_inputs) if not isinstance(noise_scheduler, RFlowScheduler): latents, _ = noise_scheduler.step(model_output, t, latents) # type: ignore else: latents, _ = noise_scheduler.step(model_output, t, latents, next_t) # type: ignore end_time = time.time() logging.info(f"---- DM/ControlNet Latent features generation time: {end_time - start_time} seconds ----") del ( unet_inputs, controlnet_inputs, model_output, controlnet_cond_vis, down_block_res_samples, mid_block_res_sample, ) gc.collect() torch.cuda.empty_cache() # decode latents to synthesized images logging.info("---- Start decoding latent features into images... ----") start_time = time.time() inferer = SlidingWindowInferer( roi_size=autoencoder_sliding_window_infer_size, sw_batch_size=1, progress=True, mode="gaussian", overlap=autoencoder_sliding_window_infer_overlap, sw_device=device, device=torch.device("cpu"), ) synthetic_images = dynamic_infer(inferer, recon_model, latents) synthetic_images = torch.clip(synthetic_images, b_min, b_max).cpu() end_time = time.time() logging.info(f"---- Image VAE decoding time: {end_time - start_time} seconds ----") ## post processing: # project output to [0, 1] synthetic_images = (synthetic_images - b_min) / (b_max - b_min) # project output to [-1000, 1000] synthetic_images = synthetic_images * (a_max - a_min) + a_min # regularize background intensities synthetic_images = crop_img_body_mask(synthetic_images, combine_label) torch.cuda.empty_cache() return synthetic_images, combine_label def filter_mask_with_organs(combine_label, anatomy_list): """ Filter a mask to only include specified organs. Args: combine_label (torch.Tensor): The input mask. anatomy_list (list): List of organ labels to keep. Returns: torch.Tensor: The filtered mask. """ # final output mask file has shape of output_size, contains labels in anatomy_list # it is already interpolated to target size combine_label = combine_label.long() # filter out the organs that are not in anatomy_list for i in range(len(anatomy_list)): organ = anatomy_list[i] # replace it with a negative value so it will get mixed combine_label[combine_label == organ] = -(i + 1) # zero-out voxels with value not in anatomy_list combine_label[combine_label > 0] = 0 # output positive values combine_label = -combine_label return combine_label def crop_img_body_mask(synthetic_images, combine_label): """ Crop the synthetic image using a body mask. Args: synthetic_images (torch.Tensor): The synthetic images. combine_label (torch.Tensor): The body mask. Returns: torch.Tensor: The cropped synthetic images. """ synthetic_images[combine_label == 0] = -1000 return synthetic_images def check_input( body_region, anatomy_list, label_dict_json, output_size, spacing, controllable_anatomy_size=[("pancreas", 0.5)], ): """ Validate input parameters for image generation. Args: body_region (list): List of body regions. anatomy_list (list): List of anatomical structures. label_dict_json (str): Path to the label dictionary JSON file. output_size (tuple): Desired output size of the image. spacing (tuple): Desired voxel spacing. controllable_anatomy_size (list): List of tuples specifying controllable anatomy sizes. Raises: ValueError: If any input parameter is invalid. """ # check output_size and spacing format if output_size[0] != output_size[1]: raise ValueError(f"The first two components of output_size need to be equal, yet got {output_size}.") if (output_size[0] not in [256, 384, 512]) or (output_size[2] not in [128, 256, 384, 512, 640, 768]): raise ValueError( f"The output_size[0] have to be chosen from [256, 384, 512], and output_size[2] have to be chosen from [128, 256, 384, 512, 640, 768], yet got {output_size}." ) if spacing[0] != spacing[1]: raise ValueError(f"The first two components of spacing need to be equal, yet got {spacing}.") if spacing[0] < 0.5 or spacing[0] > 3.0 or spacing[2] < 0.5 or spacing[2] > 5.0: raise ValueError( f"spacing[0] have to be between 0.5 and 3.0 mm, spacing[2] have to be between 0.5 and 5.0 mm, yet got {spacing}." ) if output_size[0] * spacing[0] < 256: FOV = [output_size[axis] * spacing[axis] for axis in range(3)] raise ValueError( f"`'spacing'({spacing}mm) and 'output_size'({output_size}) together decide the output field of view (FOV). The FOV will be {FOV}mm. We recommend the FOV in x and y axis to be at least 256mm for head, and at least 384mm for other body regions like abdomen. There is no such restriction for z-axis." ) if controllable_anatomy_size == None: logging.info(f"`controllable_anatomy_size` is not provided.") return # check controllable_anatomy_size format if len(controllable_anatomy_size) > 10: raise ValueError( f"The length of list controllable_anatomy_size has to be less than 10. Yet got length equal to {len(controllable_anatomy_size)}." ) available_controllable_organ = [ "liver", "gallbladder", "stomach", "pancreas", "colon", ] available_controllable_tumor = [ "hepatic tumor", "bone lesion", "lung tumor", "colon cancer primaries", "pancreatic tumor", ] available_controllable_anatomy = available_controllable_organ + available_controllable_tumor controllable_tumor = [] controllable_organ = [] for controllable_anatomy_size_pair in controllable_anatomy_size: if controllable_anatomy_size_pair[0] not in available_controllable_anatomy: raise ValueError( f"The controllable_anatomy have to be chosen from {available_controllable_anatomy}, yet got {controllable_anatomy_size_pair[0]}." ) if controllable_anatomy_size_pair[0] in available_controllable_tumor: controllable_tumor += [controllable_anatomy_size_pair[0]] if controllable_anatomy_size_pair[0] in available_controllable_organ: controllable_organ += [controllable_anatomy_size_pair[0]] if controllable_anatomy_size_pair[1] == -1: continue if controllable_anatomy_size_pair[1] < 0 or controllable_anatomy_size_pair[1] > 1.0: raise ValueError( f"The controllable size scale have to be between 0 and 1,0, or equal to -1, yet got {controllable_anatomy_size_pair[1]}." ) if len(controllable_tumor + controllable_organ) != len(list(set(controllable_tumor + controllable_organ))): raise ValueError(f"Please do not repeat controllable_anatomy. Got {controllable_tumor + controllable_organ}.") if len(controllable_tumor) > 1: raise ValueError(f"Only one controllable tumor is supported. Yet got {controllable_tumor}.") if len(controllable_anatomy_size) > 0: logging.info( f"`controllable_anatomy_size` is not empty.\nWe will ignore `body_region` and `anatomy_list` and synthesize based on `controllable_anatomy_size`: ({controllable_anatomy_size})." ) else: logging.info( f"`controllable_anatomy_size` is empty.\nWe will synthesize based on `body_region`: ({body_region}) and `anatomy_list`: ({anatomy_list})." ) # check body_region format available_body_region = [ "head", "chest", "thorax", "abdomen", "pelvis", "lower", ] for region in body_region: if region not in available_body_region: raise ValueError( f"The components in body_region have to be chosen from {available_body_region}, yet got {region}." ) # check anatomy_list format with open(label_dict_json) as f: label_dict = json.load(f) for anatomy in anatomy_list: if anatomy not in label_dict.keys(): 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}.") return class LDMSampler: """ A sampler class for generating synthetic medical images and masks using latent diffusion models. Attributes: Various attributes related to model configuration, input parameters, and generation settings. """ def __init__( self, body_region, anatomy_list, all_mask_files_json, all_anatomy_size_condtions_json, all_mask_files_base_dir, label_dict_json, label_dict_remap_json, autoencoder, diffusion_unet, controlnet, noise_scheduler, scale_factor, mask_generation_autoencoder, mask_generation_diffusion_unet, mask_generation_scale_factor, mask_generation_noise_scheduler, device, latent_shape, mask_generation_latent_shape, output_size, output_dir, controllable_anatomy_size, image_output_ext=".nii.gz", label_output_ext=".nii.gz", real_img_median_statistics="./configs/image_median_statistics.json", spacing=[1, 1, 1], modality=1, num_inference_steps=None, mask_generation_num_inference_steps=None, random_seed=None, autoencoder_sliding_window_infer_size=[96, 96, 96], autoencoder_sliding_window_infer_overlap=0.6667, ) -> None: """ Initialize the LDMSampler with various parameters and models. Args: Various parameters related to model configuration, input settings, and output specifications. """ self.random_seed = random_seed if random_seed is not None: set_determinism(seed=random_seed) with open(label_dict_json, "r") as f: label_dict = json.load(f) self.all_anatomy_size_condtions_json = all_anatomy_size_condtions_json # intialize variables self.body_region = body_region self.anatomy_list = [label_dict[organ] for organ in anatomy_list] self.all_mask_files_json = all_mask_files_json self.data_root = all_mask_files_base_dir self.label_dict_remap_json = label_dict_remap_json self.autoencoder = autoencoder self.diffusion_unet = diffusion_unet self.controlnet = controlnet self.noise_scheduler = noise_scheduler self.scale_factor = scale_factor self.mask_generation_autoencoder = mask_generation_autoencoder self.mask_generation_diffusion_unet = mask_generation_diffusion_unet self.mask_generation_scale_factor = mask_generation_scale_factor self.mask_generation_noise_scheduler = mask_generation_noise_scheduler self.device = device self.latent_shape = latent_shape self.mask_generation_latent_shape = mask_generation_latent_shape self.output_size = output_size self.output_dir = output_dir self.noise_factor = 1.0 self.controllable_anatomy_size = controllable_anatomy_size if len(self.controllable_anatomy_size): logging.info("controllable_anatomy_size is given, mask generation is triggered!") # overwrite the anatomy_list by given organs in self.controllable_anatomy_size self.anatomy_list = [label_dict[organ_and_size[0]] for organ_and_size in self.controllable_anatomy_size] self.image_output_ext = image_output_ext self.label_output_ext = label_output_ext # Set the default value for number of inference steps to 1000 self.num_inference_steps = num_inference_steps if num_inference_steps is not None else 1000 self.mask_generation_num_inference_steps = ( mask_generation_num_inference_steps if mask_generation_num_inference_steps is not None else 1000 ) if any(size % 16 != 0 for size in autoencoder_sliding_window_infer_size): raise ValueError( f"autoencoder_sliding_window_infer_size must be divisible by 16.\n Got {autoencoder_sliding_window_infer_size}" ) if not (0 <= autoencoder_sliding_window_infer_overlap <= 1): raise ValueError( f"Value of autoencoder_sliding_window_infer_overlap must be between 0 and 1.\n Got {autoencoder_sliding_window_infer_overlap}" ) self.autoencoder_sliding_window_infer_size = autoencoder_sliding_window_infer_size self.autoencoder_sliding_window_infer_overlap = autoencoder_sliding_window_infer_overlap # quality check args self.max_try_time = 2 # if not pass quality check, will try self.max_try_time times with open(real_img_median_statistics, "r") as json_file: self.median_statistics = json.load(json_file) self.label_int_dict = { "liver": [1], "spleen": [3], "pancreas": [4], "kidney": [5, 14], "lung": [28, 29, 30, 31, 31], "brain": [22], "hepatic tumor": [26], "bone lesion": [128], "lung tumor": [23], "colon cancer primaries": [27], "pancreatic tumor": [24], "bone": list(range(33, 57)) + list(range(63, 98)) + [120, 122, 127], } # networks self.autoencoder.eval() self.diffusion_unet.eval() self.controlnet.eval() self.mask_generation_autoencoder.eval() self.mask_generation_diffusion_unet.eval() self.spacing = spacing self.modality_tensor = modality * torch.ones((1,), dtype=torch.long).to(device) self.include_body_region = self.diffusion_unet.include_top_region_index_input self.include_modality = self.diffusion_unet.num_class_embeds is not None val_transforms_list = [ monai.transforms.LoadImaged(keys=["pseudo_label"]), monai.transforms.EnsureChannelFirstd(keys=["pseudo_label"]), monai.transforms.Orientationd(keys=["pseudo_label"], axcodes="RAS"), monai.transforms.EnsureTyped(keys=["pseudo_label"], dtype=torch.uint8), monai.transforms.Lambdad(keys="spacing", func=lambda x: torch.FloatTensor(x)), monai.transforms.Lambdad(keys="spacing", func=lambda x: x * 1e2), ] if self.include_body_region: val_transforms_list += [ monai.transforms.Lambdad(keys="top_region_index", func=lambda x: torch.FloatTensor(x)), monai.transforms.Lambdad(keys="bottom_region_index", func=lambda x: torch.FloatTensor(x)), monai.transforms.Lambdad(keys="top_region_index", func=lambda x: x * 1e2), monai.transforms.Lambdad(keys="bottom_region_index", func=lambda x: x * 1e2), ] self.val_transforms = Compose(val_transforms_list) logging.info("LDM sampler initialized.") def sample_multiple_images(self, num_img): """ Generate multiple synthetic images and masks. Args: num_img (int): Number of images to generate. """ modality_tensor = self.modality_tensor output_filenames = [] if len(self.controllable_anatomy_size) > 0: # we will use mask generation instead of finding candidate masks # create a dummy selected_mask_files for placeholder selected_mask_files = list(range(num_img)) # prerpare organ size conditions anatomy_size_condtion = self.prepare_anatomy_size_condtion(self.controllable_anatomy_size) else: need_resample = False # find candidate mask and save to candidate_mask_files candidate_mask_files = find_masks( self.body_region, self.anatomy_list, self.spacing, self.output_size, True, self.all_mask_files_json, self.data_root, ) if len(candidate_mask_files) < num_img: # if we cannot find enough masks based on the exact match of anatomy list, spacing, and output size, # then we will try to find the closest mask in terms of spacing, and output size. logging.info("Resample mask file to get desired output size and spacing") candidate_mask_files = self.find_closest_masks(num_img) need_resample = True selected_mask_files = self.select_mask(candidate_mask_files, num_img) logging.info(f"Images will be generated based on {selected_mask_files}.") if len(selected_mask_files) < num_img: raise ValueError( ( f"len(selected_mask_files) ({len(selected_mask_files)}) < num_img ({num_img}). " "This should not happen. Please revisit function select_mask(self, candidate_mask_files, num_img)." ) ) num_generated_img = 0 for index_s in range(len(selected_mask_files)): item = selected_mask_files[index_s] if num_generated_img >= num_img: break logging.info("---- Start preparing masks... ----") start_time = time.time() if len(self.controllable_anatomy_size) > 0: # generate a synthetic mask ( combine_label_or, top_region_index_tensor, bottom_region_index_tensor, spacing_tensor, ) = self.prepare_one_mask_and_meta_info(anatomy_size_condtion) else: # read in mask file mask_file = item["mask_file"] if_aug = item["if_aug"] ( combine_label_or, top_region_index_tensor, bottom_region_index_tensor, spacing_tensor, ) = self.read_mask_information(mask_file) if need_resample: combine_label_or = self.ensure_output_size_and_spacing(combine_label_or) # mask augmentation if if_aug: combine_label_or = augmentation(combine_label_or, self.output_size, self.random_seed) end_time = time.time() logging.info(f"---- Mask preparation time: {end_time - start_time} seconds ----") torch.cuda.empty_cache() # generate image/label pairs to_generate = True try_time = 0 # start generation synthetic_images, synthetic_labels = self.sample_one_pair( combine_label_or, top_region_index_tensor, bottom_region_index_tensor, spacing_tensor, modality_tensor, ) # synthetic image quality check pass_quality_check = self.quality_check( synthetic_images.cpu().detach().numpy(), combine_label_or.cpu().detach().numpy() ) print(num_img - num_generated_img, (len(selected_mask_files) - index_s)) if pass_quality_check or (num_img - num_generated_img) >= (len(selected_mask_files) - index_s): if not pass_quality_check: logging.info( "Generated image/label pair did not pass quality check, but will still save them. " "Please consider changing spacing and output_size to facilitate a more realistic setting." ) num_generated_img = num_generated_img + 1 # save image/label pairs output_postfix = datetime.now().strftime("%Y%m%d_%H%M%S_%f") synthetic_labels.meta["filename_or_obj"] = "sample.nii.gz" synthetic_images = MetaTensor(synthetic_images, meta=synthetic_labels.meta) img_saver = SaveImage( output_dir=self.output_dir, output_postfix=output_postfix + "_image", output_ext=self.image_output_ext, separate_folder=False, ) img_saver(synthetic_images[0]) synthetic_images_filename = os.path.join( self.output_dir, "sample_" + output_postfix + "_image" + self.image_output_ext ) # filter out the organs that are not in anatomy_list synthetic_labels = filter_mask_with_organs(synthetic_labels, self.anatomy_list) label_saver = SaveImage( output_dir=self.output_dir, output_postfix=output_postfix + "_label", output_ext=self.label_output_ext, separate_folder=False, ) label_saver(synthetic_labels[0]) synthetic_labels_filename = os.path.join( self.output_dir, "sample_" + output_postfix + "_label" + self.label_output_ext ) output_filenames.append([synthetic_images_filename, synthetic_labels_filename]) to_generate = False else: logging.info("Generated image/label pair did not pass quality check, will re-generate another pair.") return output_filenames def select_mask(self, candidate_mask_files, num_img): """ Select mask files for image generation. Args: candidate_mask_files (list): List of candidate mask files. num_img (int): Number of images to generate. Returns: list: Selected mask files with augmentation flags. """ selected_mask_files = [] random.shuffle(candidate_mask_files) for n in range(len(candidate_mask_files)): mask_file = candidate_mask_files[n % len(candidate_mask_files)] selected_mask_files.append({"mask_file": mask_file, "if_aug": True}) return selected_mask_files def sample_one_pair( self, combine_label_or_aug, top_region_index_tensor, bottom_region_index_tensor, spacing_tensor, modality_tensor, ): """ Generate a single pair of synthetic image and mask. Args: combine_label_or_aug (torch.Tensor): Combined label tensor or augmented label. top_region_index_tensor (torch.Tensor): Tensor specifying the top region index. bottom_region_index_tensor (torch.Tensor): Tensor specifying the bottom region index. spacing_tensor (torch.Tensor): Tensor specifying the spacing. modality_tensor (torch.Tensor): Int Tensor specifying the modality. Returns: tuple: A tuple containing the synthetic image and its corresponding label. """ # generate image/label pairs synthetic_images, synthetic_labels = ldm_conditional_sample_one_image( autoencoder=self.autoencoder, diffusion_unet=self.diffusion_unet, controlnet=self.controlnet, noise_scheduler=self.noise_scheduler, scale_factor=self.scale_factor, device=self.device, combine_label_or=combine_label_or_aug, top_region_index_tensor=top_region_index_tensor, bottom_region_index_tensor=bottom_region_index_tensor, spacing_tensor=spacing_tensor, modality_tensor=modality_tensor, latent_shape=self.latent_shape, output_size=self.output_size, noise_factor=self.noise_factor, num_inference_steps=self.num_inference_steps, autoencoder_sliding_window_infer_size=self.autoencoder_sliding_window_infer_size, autoencoder_sliding_window_infer_overlap=self.autoencoder_sliding_window_infer_overlap, ) return synthetic_images, synthetic_labels def prepare_anatomy_size_condtion( self, controllable_anatomy_size, ): """ Prepare anatomy size conditions for mask generation. Args: controllable_anatomy_size (list): List of tuples specifying controllable anatomy sizes. Returns: list: Prepared anatomy size conditions. """ anatomy_size_idx = { "gallbladder": 0, "liver": 1, "stomach": 2, "pancreas": 3, "colon": 4, "lung tumor": 5, "pancreatic tumor": 6, "hepatic tumor": 7, "colon cancer primaries": 8, "bone lesion": 9, } provide_anatomy_size = [None for _ in range(10)] logging.info(f"controllable_anatomy_size: {controllable_anatomy_size}") for element in controllable_anatomy_size: anatomy_name, anatomy_size = element provide_anatomy_size[anatomy_size_idx[anatomy_name]] = anatomy_size with open(self.all_anatomy_size_condtions_json, "r") as f: all_anatomy_size_condtions = json.load(f) # loop through the database and find closest combinations candidate_list = [] for anatomy_size in all_anatomy_size_condtions: size = anatomy_size["organ_size"] diff = 0 for db_size, provide_size in zip(size, provide_anatomy_size): if provide_size is None: continue diff += abs(provide_size - db_size) candidate_list.append((size, diff)) candidate_condition = sorted(candidate_list, key=lambda x: x[1])[0][0] # overwrite the anatomy size provided by users for element in controllable_anatomy_size: anatomy_name, anatomy_size = element candidate_condition[anatomy_size_idx[anatomy_name]] = anatomy_size return candidate_condition def prepare_one_mask_and_meta_info(self, anatomy_size_condtion): """ Prepare a single mask and its associated meta information. Args: anatomy_size_condtion (list): Anatomy size conditions. Returns: tuple: A tuple containing the prepared mask and associated tensors. """ combine_label_or = self.sample_one_mask(anatomy_size=anatomy_size_condtion) # TODO: current mask generation model only can generate 256^3 volumes with 1.5 mm spacing. affine = torch.zeros((4, 4)) affine[0, 0] = 1.5 affine[1, 1] = 1.5 affine[2, 2] = 1.5 affine[3, 3] = 1.0 # dummy combine_label_or = MetaTensor(combine_label_or, affine=affine) combine_label_or = self.ensure_output_size_and_spacing(combine_label_or) top_region_index, bottom_region_index = get_body_region_index_from_mask(combine_label_or) spacing_tensor = torch.FloatTensor(self.spacing).unsqueeze(0).half().to(self.device) * 1e2 top_region_index_tensor = torch.FloatTensor(top_region_index).unsqueeze(0).half().to(self.device) * 1e2 bottom_region_index_tensor = torch.FloatTensor(bottom_region_index).unsqueeze(0).half().to(self.device) * 1e2 return combine_label_or, top_region_index_tensor, bottom_region_index_tensor, spacing_tensor def sample_one_mask(self, anatomy_size): """ Generate a single synthetic mask. Args: anatomy_size (list): Anatomy size specifications. Returns: torch.Tensor: The generated synthetic mask. """ # generate one synthetic mask synthetic_mask = ldm_conditional_sample_one_mask( self.mask_generation_autoencoder, self.mask_generation_diffusion_unet, self.mask_generation_noise_scheduler, self.mask_generation_scale_factor, anatomy_size, self.device, self.mask_generation_latent_shape, label_dict_remap_json=self.label_dict_remap_json, num_inference_steps=self.mask_generation_num_inference_steps, autoencoder_sliding_window_infer_size=self.autoencoder_sliding_window_infer_size, autoencoder_sliding_window_infer_overlap=self.autoencoder_sliding_window_infer_overlap, ) return synthetic_mask def ensure_output_size_and_spacing(self, labels, check_contains_target_labels=True): """ Ensure the output mask has the correct size and spacing. Args: labels (torch.Tensor): Input label tensor. check_contains_target_labels (bool): Whether to check if the resampled mask contains target labels. Returns: torch.Tensor: Resampled label tensor. Raises: ValueError: If the resampled mask doesn't contain required class labels. """ current_spacing = [labels.affine[0, 0], labels.affine[1, 1], labels.affine[2, 2]] current_shape = list(labels.squeeze().shape) need_resample = False # check spacing for i, j in zip(current_spacing, self.spacing): if i != j: need_resample = True # check output size for i, j in zip(current_shape, self.output_size): if i != j: need_resample = True # resample to target size and spacing if need_resample: logging.info("Resampling mask to target shape and spacing") logging.info(f"Resize Spacing: {current_spacing} -> {self.spacing}") logging.info(f"Output size: {current_shape} -> {self.output_size}") spacing = monai.transforms.Spacing(pixdim=tuple(self.spacing), mode="nearest") pad_crop = monai.transforms.ResizeWithPadOrCrop(spatial_size=tuple(self.output_size)) labels = pad_crop(spacing(labels.squeeze(0))).unsqueeze(0).to(labels.dtype) contained_labels = torch.unique(labels) if check_contains_target_labels: # check if the resampled mask still contains those target labels for anatomy_label in self.anatomy_list: if anatomy_label not in contained_labels: raise ValueError( f"Resampled mask does not contain required class labels {anatomy_label}. Please tune spacing and output size." ) return labels def read_mask_information(self, mask_file): """ Read mask information from a file. Args: mask_file (str): Path to the mask file. Returns: tuple: A tuple containing the mask tensor and associated information. """ val_data = self.val_transforms(mask_file) for key in ["pseudo_label", "spacing", "top_region_index", "bottom_region_index"]: if isinstance(val_data[key], torch.Tensor): val_data[key] = val_data[key].unsqueeze(0).to(self.device) else: val_data[key] = None return ( val_data["pseudo_label"], val_data["top_region_index"], val_data["bottom_region_index"], val_data["spacing"], ) def find_closest_masks(self, num_img): """ Find the closest matching masks from the database. Args: num_img (int): Number of images to generate. Returns: list: List of closest matching mask candidates. Raises: ValueError: If suitable candidates cannot be found. """ # first check the database based on anatomy list candidates = find_masks( self.body_region, self.anatomy_list, self.spacing, self.output_size, False, self.all_mask_files_json, self.data_root, ) if len(candidates) < num_img: raise ValueError(f"candidate masks are less than {num_img}).") # loop through the database and find closest combinations new_candidates = [] for c in candidates: diff = 0 include_c = True for axis in range(3): if abs(c["dim"][axis]) < self.output_size[axis] - 64: # we cannot upsample the mask too much include_c = False break # check diff in FOV, major metric diff += abs( (abs(c["dim"][axis] * c["spacing"][axis]) - self.output_size[axis] * self.spacing[axis]) / 10 ) # check diff in dim diff += abs((abs(c["dim"][axis]) - self.output_size[axis]) / 100) # check diff in spacing diff += abs(abs(c["spacing"][axis]) - self.spacing[axis]) if include_c: new_candidates.append((c, diff)) # choose top-2*num_img candidates (at least 5) num_candidates = max(self.max_try_time * num_img, 5) new_candidates = sorted(new_candidates, key=lambda x: x[1]) final_candidates = [] # check top-2*num_img candidates and update spacing after resampling for c, _ in new_candidates: c = self.resample_mask_check_organ_list(c) if c is not None: final_candidates.append(c) if len(final_candidates) >= num_candidates: break if len(final_candidates) == 0: raise ValueError("Cannot find body region with given organ list.") return final_candidates def resample_mask_check_organ_list(self, mask): """ Resample mask and check if the resampled mask contains the required organ list. Args: mask (dict): input mask. Returns: dict: resampled mask. If None, means the resampled mask does not contain the required organ list Raises: ValueError: If suitable candidates cannot be found. """ image_loader = monai.transforms.LoadImage(image_only=True, ensure_channel_first=True) label = image_loader(mask["pseudo_label"]) try: label = self.ensure_output_size_and_spacing(label.unsqueeze(0)) except ValueError as e: if "Resampled mask does not contain required class labels" in str(e): return None else: raise e # get region_index after resample top_region_index, bottom_region_index = get_body_region_index_from_mask(label) mask["top_region_index"] = top_region_index mask["bottom_region_index"] = bottom_region_index mask["spacing"] = self.spacing mask["dim"] = self.output_size return mask def quality_check(self, image_data, label_data): """ Perform a quality check on the generated image. Args: image_data (np.ndarray): The generated image. label_data (np.ndarray): The corresponding whole body mask. Returns: bool: True if the image passes the quality check, False otherwise. """ outlier_results = is_outlier(self.median_statistics, image_data, label_data, self.label_int_dict) for label, result in outlier_results.items(): if result.get("is_outlier", False): logging.info( f"Generated image quality check for label '{label}' failed: median value {result['median_value']} is outside the acceptable range ({result['low_thresh']} - {result['high_thresh']})." ) return False return True