# Load all the libraries from pathlib import Path import os import dicom2nifti import warnings import requests import zipfile from dotenv import load_dotenv import nibabel as nib import numpy as np import json # Load secret keys from .env file load_dotenv() # Specify the model endopoint within MLIs base_url = os.getenv('BASE_URL') if base_url is not None: print(f"Using model at: {base_url}") else: raise ValueError("BASE_URL was not specified, check your .env file") # Specity the MLIs authorization token mlis_token = os.getenv('MLIS_TOKEN') if mlis_token is not None: print(f"MLIS token loaded successfully!") else: raise ValueError("MLIS_TOKEN was not specified, check your .env file") ## These are the only paths that need to be modified! # If dataset is in Dicom format - Specify the folder where Dicom data reside: DICOM_DATA_PATH = Path('dicom_data') # Specify the folder where all the outputs will be saved ALL_OUTPUTS_PATH = Path('outputs_flare') # Specify the URL where data reside so that Vista3D model can fetch it: # Please refer to https://docs.nvidia.com/nim/medical/vista3d/latest/advanced-usage.html#environment-variables for additional details. # Hint: if data are in 'https://fs2.ingress.pcai0109.dc15.hpecolo.net/califra/dicom_data' # data_path_root must be 'https://fs2.ingress.pcai0109.dc15.hpecolo.net/califra/' data_path_root = 'https://vessels-app.ai-application.hou-pcai.hpecic.net/vista-3d/' # Set the following to True if you want to convert the data Dicom to Nifti CONVERT_DICOM_TO_NIFTI = False # Assess whether the FLARE dataset has been organized correctly if not CONVERT_DICOM_TO_NIFTI: ## Check data were correctly organized assert ( ALL_OUTPUTS_PATH / "nifti_data" / "testing" / "FLARETs_0050.nii.gz" ).is_file(), ( "Ensure the FLARE dataset is organized as expected. Look instructions above - Ignore this message otherwise" ) # If you converted your data from Dicom to Nifti using the code above, then your Nifti data will be in nifti_data folder. nifti_destination_path = Path(ALL_OUTPUTS_PATH) / "nifti_data" if CONVERT_DICOM_TO_NIFTI: # Convert Dicom to Nifti for compatibility with the Vista-3D Nvidia NIM model for dicom_directory in os.listdir(DICOM_DATA_PATH): input_directory = Path(DICOM_DATA_PATH) / f"{dicom_directory}" output_directory = Path(nifti_destination_path) / dicom_directory # If the patient's exam has already been processed, skip it. if output_directory.exists(): warnings.warn(f"{output_directory} already exists, skipping...") continue else: os.makedirs(output_directory, exist_ok=False) print("-"*10) print(f"Saving Nifti files into directory {output_directory}") dicom2nifti.convert_directory(input_directory, output_directory, compression=True, reorient=True) def unzip_file(zip_filepath, output_dir, output_filename=None): """Unzip the first file in a zip archive to a target folder, optionally renaming it.""" with zipfile.ZipFile(zip_filepath, 'r') as zip_ref: info_list = zip_ref.infolist() if not info_list: raise ValueError("Zip file is empty") first_entry = info_list[0] original_name = first_entry.filename # Ensure output directory exists os.makedirs(output_dir, exist_ok=True) # Extract the file to the output directory extracted_path = zip_ref.extract(original_name, path=output_dir) # Optionally rename the extracted file if output_filename: new_path = os.path.join(output_dir, output_filename) os.rename(extracted_path, new_path) return new_path return extracted_path # Location where the model responses and segmentation outputs will land response_dir = Path(ALL_OUTPUTS_PATH) / "responses" os.makedirs(response_dir, exist_ok=True) segmentation_dir = Path(ALL_OUTPUTS_PATH) / "segmentations" os.makedirs(segmentation_dir, exist_ok=True) # Specify the location of the data, they need to be available and accessible online for the model to be able to process them. # Here we use the location within the hosted-trial environment headers = {'Authorization': f'Bearer {mlis_token}'} # Vista-3D model requires passing the data in a dictionary, which has a key `image` that points to a URL to the Nifti file. # e.g. data['image'] = "https://fs2.ingress.pcai0109.dc15.hpecolo.net/califra/outputs/nifti_data/testing/FLARETs_0050_0000.nii.gz" # Please refer to https://docs.nvidia.com/nim/medical/vista3d/latest/advanced-usage.html#environment-variables for additional details. # Note that Vista-3D model will attempt segmenting every Nifti file irrespective of the acquisition orientation (axial, coronal, sagittal). # Zero-shot segmentation quality should be state of the art on all the planes. for path in Path(nifti_destination_path).rglob('*nii.gz'): data = {} # User TODO: # Ensure 'path' does not require custom logic. # This is may be needed if ALL_OUTPUTS_PATH is not a relative path. p = data_path_root + str(path) data['image'] = str(p) print("-" * 10) # Submit a post request to segment the ct scan. response = requests.post(f'{base_url}/v1/vista3d/inference', json=data, headers=headers) if response.status_code == 200: # Success!! file_name = path.name # Create the folders required to have a clean output os.makedirs(Path(response_dir) / path.parent.name, exist_ok=True) os.makedirs(Path(segmentation_dir) / path.parent.name, exist_ok=True) output_segmentation_name = Path(segmentation_dir) / path.parent.name / file_name # Zip the model response, unzip it and rename the files accordingly output_zip_name = Path(response_dir) / path.parent.name / file_name.replace('.nii.gz', '.zip') with open(output_zip_name, 'wb') as f: f.write(response.content) _ = unzip_file( zip_filepath=str(output_zip_name), output_dir = "/".join(str(output_segmentation_name).split("/")[:-1]), output_filename=str(output_segmentation_name).split("/")[-1], ) print(f'{str(path)} {response}') else: print("skipping: ", data['image'], '\n', response.status_code) print(f"Reason: {response.reason}") print(f"Body: {response.text}") # Download from /v1/vista3d/info with open('vista3dmodelinfo.json', 'rb') as f: segmentation_dict = json.load(f)['labels'] segmentation_dict = {v: k for k, v in segmentation_dict.items()} print("Mapping tissue to label:") for tissue, label in segmentation_dict.items(): print(f"{tissue}: {label}") patient_dict = {} for patient_paths in (Path(ALL_OUTPUTS_PATH) / 'nifti_data').rglob('*nii.gz'): patient = patient_paths.parent.name if patient not in patient_dict: patient_dict[patient] = [] if "sa" not in str(patient_paths.resolve()) and "co" not in str(patient_paths.resolve()): patient_dict[patient].append(str(patient_paths.resolve())) # Extract the mapping between segmentation label and segmentation numerical value, using the segmentation dict provided by the Vista-3D model which we previously loaded in memory (i.e. 'vista3dmodelinfo.json') # In this study we are interested on 3 vessels, but could be expanded to any additional vessel/ tissue, provided the segmentation is available. Please refer to the 'vista3dmodelinfo.json' file, or the official Vista-3D model documentation. label_rois = ['aorta', 'left iliac artery', 'right iliac artery'] vessel_labels = {k: int(segmentation_dict[k]) for k in label_rois} for idx, patient_id in enumerate(patient_dict.keys()): patient = patient_id print("-" * 10) print(">> processing patient: ", patient) for series, scan_path in enumerate(patient_dict[patient]): try: scan_path_lower = scan_path.lower() volume_name = Path(patient_dict[patient][series]).name.replace(".nii.gz", "") print(f"processing {volume_name} for patient {patient}") # Load the CT to memory ct_data = nib.load(scan_path) ct_scan = ct_data.get_fdata() affine = ct_data.affine # Load the segmentation to memory segmentation_path = patient_dict[patient][series].replace('nifti_data', 'segmentations') segmentation_data = nib.load(segmentation_path) segmentation_scan = segmentation_data.get_fdata() # Create a new segmentation map with only the vessels of interest in it. new_segmentation_scan = np.zeros_like(segmentation_scan) for n,label in enumerate(vessel_labels.values(), start=1): new_segmentation_scan += (segmentation_scan == label)*n # Create a new dict with the vessel labels new_vessel_labels = { k_key: id_k for id_k, (k_key, k_value) in enumerate(vessel_labels.items(), start=1) } print("After processing segmentation, the new labels mapping is:\n") for tissue, label in new_vessel_labels.items(): print(f"{tissue}: {label}") # ---------------------------------------------------------------------------- # Build the 3D coordinate grid (patient coordinates in mm) # Build coordinate grid scan_size = ct_scan.shape #(4,4) # Meshgrid 3 integer volumes holding voxel indices (i,j,k) (512, 512, 146) each grid_i, grid_j, grid_k = np.meshgrid( np.arange(scan_size[0]), np.arange(scan_size[1]), np.arange(scan_size[2]), indexing='ij' ) # Since affine has shape 4x4, we need to stack a new column of 1s to make it the grid homogenous indices = np.stack([grid_i, grid_j, grid_k, np.ones_like(grid_i)], axis=-1) # Create a tensor, where each voxel's value is the coordinate of that voxel in the x, y, z coordinate space # This uses the affine matrix provided by the Nifti files. coords = indices @ affine.T # Undo the homogenous shape by slicing away the row that we had for math convenience coords_xyz = coords[...,:3] # ---------------------------------------------------------------------------- # Now that the coordinates are available in the x, y, z coordinate we can further process the results by # extracting only the points that belong to vessels' region of interest. mask = new_segmentation_scan > 0 # convert to binary points = coords_xyz[mask] # shape (N,3) # Build colors based on segmentation label colors = np.zeros((points.shape[0], 3)) flat_labels = new_segmentation_scan[mask] colors[flat_labels == 1] = [1, 0, 0] # red for aorta colors[flat_labels == 2] = [0, 1, 0] # green for left iliac colors[flat_labels == 3] = [0, 0, 1] # blue for right iliac # Save pointclouds: pointcloud_dir = Path(ALL_OUTPUTS_PATH) / "pointclouds" os.makedirs(pointcloud_dir, exist_ok=True) patient_pointcloud_npz = Path(pointcloud_dir) / f"{patient_id}" / f"{volume_name}.npz" os.makedirs(patient_pointcloud_npz.parent, exist_ok=True) np.savez_compressed(patient_pointcloud_npz, points=points, colors=colors) print(f"Saved pointcloud with {points.shape[0]} points in \n{patient_pointcloud_npz}") # Save post processed segmentation: processed_segmentation_dir = Path(ALL_OUTPUTS_PATH) / "processed_segmentations" / patient_id vessel_segmentation_npz = Path(processed_segmentation_dir) / f"{volume_name}.npz" os.makedirs(processed_segmentation_dir, exist_ok=True) np.savez_compressed(vessel_segmentation_npz, new_seg_scan=new_segmentation_scan) print(f"{'-'*10}\nProcessed segmentation mask saved in \n{vessel_segmentation_npz}") except Exception as e: print(e) continue