diff --git a/README.md b/README.md index d73b165..90799be 100644 --- a/README.md +++ b/README.md @@ -83,10 +83,11 @@ bin/C2C spine -i ### Usage ```bash -bin/C2C aortic_calcium -i +bin/C2C aortic_calcium -i -o --threshold ``` -- input_path should contain a DICOM series or subfolders that contain DICOM series, or a nifti file. -- Aortic calcifications are detected through an adaptive threshold in a dilated aorta mask +- The input path should contain a DICOM series or subfolders that contain DICOM series, or a nifti file. +- The threshold can be controlled with --threshold and be either an integer HU threshold, "adataptive" or "agatson". + - If "agatson" is used, agatson score is calculated and the a threshold of 130 HU is used - Aortic calcifications are divided into abdominal and thoracic at the T12 level - Segmentation masks for aortic calcium, the dilated aorta mask and the T12 seperation plane are saved in ./segmentation_masks/ - Metrics on an aggregated and individual level for the calcifications are written to .csv files in ./metrics/ diff --git a/comp2comp/aortic_calcium/aortic_calcium.py b/comp2comp/aortic_calcium/aortic_calcium.py index 71f20e1..1a5a83c 100644 --- a/comp2comp/aortic_calcium/aortic_calcium.py +++ b/comp2comp/aortic_calcium/aortic_calcium.py @@ -9,10 +9,11 @@ import time from pathlib import Path from typing import Union -import matplotlib.pyplot as plt +import matplotlib.pyplot as plt import numpy as np from scipy import ndimage + # from totalsegmentator.libs import ( # download_pretrained_weights, # nostdout, @@ -20,7 +21,6 @@ # ) from totalsegmentatorv2.python_api import totalsegmentator - from comp2comp.inference_class_base import InferenceClass @@ -69,18 +69,18 @@ def aorta_seg( print("Segmenting aorta...") st = time.time() os.environ["SCRATCH"] = self.model_dir - + seg = totalsegmentator( - input = os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz"), - output = os.path.join(self.output_dir_segmentations, "segmentation.nii"), - task_ids = [293], - ml = True, - nr_thr_resamp = 1, - nr_thr_saving = 6, - fast = False, - nora_tag = "None", - preview = False, - task = "total", + input=os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz"), + output=os.path.join(self.output_dir_segmentations, "segmentation.nii"), + task_ids=[293], + ml=True, + nr_thr_resamp=1, + nr_thr_saving=6, + fast=False, + nora_tag="None", + preview=False, + task="total", # roi_subset = [ # "vertebrae_T12", # "vertebrae_L1", @@ -89,24 +89,24 @@ def aorta_seg( # "vertebrae_L4", # "vertebrae_L5", # ], - roi_subset = None, - statistics = False, - radiomics = False, - crop_path = None, - body_seg = False, - force_split = False, - output_type = "nifti", - quiet = False, - verbose = False, - test = 0, - skip_saving = True, - device = "gpu", - license_number = None, - statistics_exclude_masks_at_border = True, - no_derived_masks = False, - v1_order = False, + roi_subset=None, + statistics=False, + radiomics=False, + crop_path=None, + body_seg=False, + force_split=False, + output_type="nifti", + quiet=False, + verbose=False, + test=0, + skip_saving=True, + device="gpu", + license_number=None, + statistics_exclude_masks_at_border=True, + no_derived_masks=False, + v1_order=False, ) - + end = time.time() # Log total time for spine segmentation @@ -126,8 +126,12 @@ def __call__(self, inference_pipeline): self.output_dir = inference_pipeline.output_dir self.output_dir_images_organs = os.path.join(self.output_dir, "images/") inference_pipeline.output_dir_images_organs = self.output_dir_images_organs - self.output_dir_segmentation_masks = os.path.join(self.output_dir, "segmentation_masks/") - inference_pipeline.output_dir_segmentation_masks = self.output_dir_segmentation_masks + self.output_dir_segmentation_masks = os.path.join( + self.output_dir, "segmentation_masks/" + ) + inference_pipeline.output_dir_segmentation_masks = ( + self.output_dir_segmentation_masks + ) if not os.path.exists(self.output_dir_images_organs): os.makedirs(self.output_dir_images_organs) @@ -135,60 +139,64 @@ def __call__(self, inference_pipeline): os.makedirs(self.output_dir_segmentation_masks) if not os.path.exists(os.path.join(self.output_dir, "metrics/")): os.makedirs(os.path.join(self.output_dir, "metrics/")) - + ct = inference_pipeline.medical_volume.get_fdata() aorta_mask = inference_pipeline.segmentation.get_fdata().astype(np.int8) == 52 spine_mask = inference_pipeline.spine_segmentation.get_fdata() > 0 - - # Determine the target number of pixels + + # Determine the target number of pixels pix_size = np.array(inference_pipeline.medical_volume.header.get_zooms()) # target: 1 mm - target_aorta_dil = round(1/pix_size[0]) + target_aorta_dil = round(1 / pix_size[0]) # target: 3 mm - target_exclude_dil = round(3/pix_size[0]) + target_exclude_dil = round(3 / pix_size[0]) # target: 7 mm - target_aorta_erode = round(7/pix_size[0]) - - # Run calcification detection pipeline + target_aorta_erode = round(7 / pix_size[0]) + + # Run calcification detection pipeline calcification_results = self.detectCalcifications( - ct, - aorta_mask, - exclude_mask=spine_mask, - remove_size=3, + ct, + aorta_mask, + exclude_mask=spine_mask, + remove_size=3, return_dilated_mask=True, threshold=inference_pipeline.args.threshold, - dilation_iteration = target_aorta_dil, - dilation_iteration_exclude = target_exclude_dil, - aorta_erode_iteration = target_aorta_erode, + dilation_iteration=target_aorta_dil, + dilation_iteration_exclude=target_exclude_dil, + aorta_erode_iteration=target_aorta_erode, ) - - inference_pipeline.calc_mask = calcification_results['calc_mask'] - inference_pipeline.calcium_threshold = calcification_results['threshold'] - + + inference_pipeline.calc_mask = calcification_results["calc_mask"] + inference_pipeline.calcium_threshold = calcification_results["threshold"] + # save masks inference_pipeline.saveArrToNifti( inference_pipeline.calc_mask, - os.path.join(inference_pipeline.output_dir_segmentation_masks, - "calcium_segmentations.nii.gz") - ) + os.path.join( + inference_pipeline.output_dir_segmentation_masks, + "calcium_segmentations.nii.gz", + ), + ) inference_pipeline.saveArrToNifti( - calcification_results['dilated_mask'], - os.path.join(inference_pipeline.output_dir_segmentation_masks, - "dilated_aorta_mask.nii.gz") - ) + calcification_results["dilated_mask"], + os.path.join( + inference_pipeline.output_dir_segmentation_masks, + "dilated_aorta_mask.nii.gz", + ), + ) inference_pipeline.saveArrToNifti( aorta_mask, - os.path.join(inference_pipeline.output_dir_segmentation_masks, - "aorta_mask.nii.gz") - ) + os.path.join( + inference_pipeline.output_dir_segmentation_masks, "aorta_mask.nii.gz" + ), + ) inference_pipeline.saveArrToNifti( ct, - os.path.join(inference_pipeline.output_dir_segmentation_masks, - "ct.nii.gz") - ) + os.path.join(inference_pipeline.output_dir_segmentation_masks, "ct.nii.gz"), + ) return {} - + def detectCalcifications( self, ct, @@ -207,9 +215,9 @@ def detectCalcifications( exclude_center_aorta=True, return_eroded_aorta=False, aorta_erode_iteration=6, - threshold = 'adaptive', - agatson_failsafe = 100, - generate_plots = True, + threshold="adaptive", + agatson_failsafe=100, + generate_plots=True, ): """ Function that takes in a CT image and aorta segmentation (and optionally volumes to use @@ -251,7 +259,7 @@ def detectCalcifications( aorta_erode_iteration (int, optional): Number of iterations for the strcturing element. Defaults to 6. threshold: (str, int): - Can either be 'adaptive', 'agatson', or int. Choosing 'agatson' + Can either be 'adaptive', 'agatson', or int. Choosing 'agatson' Will mean a threshold of 130 HU. agatson_failsafe: (int): A fail-safe raising an error if the mean HU of the aorta is too high @@ -261,10 +269,10 @@ def detectCalcifications( results: array of only the mask is returned, or dict if other volumes are also returned. """ - - ''' + + """ Remove the ascending aorta if present - ''' + """ # remove parts that are not the abdominal aorta labelled_aorta, num_classes = ndimage.label(aorta_mask) if num_classes > 1: @@ -279,10 +287,9 @@ def detectCalcifications( biggest_idx = np.argmax(aorta_vols) + 1 aorta_mask[labelled_aorta != biggest_idx] = 0 - - ''' + """ Erode the center aorta to get statistics from the blood pool - ''' + """ t0 = time.time() struct = ndimage.generate_binary_structure(3, 1) @@ -294,80 +301,95 @@ def detectCalcifications( num_iteration=aorta_erode_iteration, operation="erode", ) - - eroded_ct_points = ct[aorta_eroded==1] + + eroded_ct_points = ct[aorta_eroded == 1] eroded_ct_points_mean = eroded_ct_points.mean() eroded_ct_points_std = eroded_ct_points.std() - + if generate_plots: # save the statistics of the eroded aorta for reference - with open(os.path.join(self.output_dir, 'metrics/eroded_aorta_statistics.csv'), 'w') as f: - f.write('metric,value\n') - f.write('mean,{:.1f}\n'.format(eroded_ct_points_mean)) - f.write('std,{:.1f}\n'.format(eroded_ct_points_std)) - + with open( + os.path.join(self.output_dir, "metrics/eroded_aorta_statistics.csv"), + "w", + ) as f: + f.write("metric,value\n") + f.write("mean,{:.1f}\n".format(eroded_ct_points_mean)) + f.write("std,{:.1f}\n".format(eroded_ct_points_std)) + # save a histogram: fig, axx = plt.subplots(1) axx.hist(eroded_ct_points, bins=100) - axx.set_ylabel('Counts') - axx.set_xlabel('HU') - axx.set_title('Histogram of eroded aorta') + axx.set_ylabel("Counts") + axx.set_xlabel("HU") + axx.set_title("Histogram of eroded aorta") axx.grid() plt.tight_layout() - fig.savefig(os.path.join(self.output_dir, 'images/histogram_eroded_aorta.png')) - + fig.savefig( + os.path.join(self.output_dir, "images/histogram_eroded_aorta.png") + ) + # Perform the fail-safe check if the method is agatson - if threshold == 'agatson' and eroded_ct_points_mean > agatson_failsafe: - raise ValueError('The mean HU in the center aorta is {:.0f}, and the Agatson method will provide unreliable results (fail-safe threshold is {})'.format( - eroded_ct_points_mean, agatson_failsafe - )) - + if threshold == "agatson" and eroded_ct_points_mean > agatson_failsafe: + raise ValueError( + "The mean HU in the center aorta is {:.0f}, and the Agatson method will provide unreliable results (fail-safe threshold is {})".format( + eroded_ct_points_mean, agatson_failsafe + ) + ) + # calc_mask = calc_mask * (aorta_eroded == 0) if show_time: print("exclude center aorta time: {:.2f} sec".format(time.time() - t0)) - - ''' + + """ Choose threshold - ''' - - if threshold == 'adaptive': + """ + + if threshold == "adaptive": # calc_thres = eroded_ct_points.max() # Get aortic CT point to set adaptive threshold aorta_ct_points = ct[aorta_mask == 1] - + # equal to one standard deviation to the left of the curve quant = 0.158 quantile_median_dist = np.median(aorta_ct_points) - np.quantile( aorta_ct_points, q=quant ) calc_thres = np.median(aorta_ct_points) + quantile_median_dist * num_std - - elif threshold == 'agatson': + + elif threshold == "agatson": calc_thres = 130 - + counter = self.slicedSizeCount(aorta_eroded, ct, remove_size, calc_thres) # if num_features >= 10: # raise ValueError('Too many pixels above 130 in blood pool, found: {}'.format(num_features)) if verbose: - print('{} calc over threshold of {}'.format(counter, remove_size)) - + print("{} calc over threshold of {}".format(counter, remove_size)) + if generate_plots: # save the statistics of the eroded aorta for reference - with open(os.path.join(self.output_dir, 'metrics/eroded_aorta_statistics.csv'), 'a') as f: - f.write('num calcification blood pool,{}\n'.format(counter)) + with open( + os.path.join( + self.output_dir, "metrics/eroded_aorta_statistics.csv" + ), + "a", + ) as f: + f.write("num calcification blood pool,{}\n".format(counter)) else: try: calc_thres = int(threshold) except: - raise ValueError('Error in threshold value for aortic calcium segmentaiton. \ - Should be \'adaptive\', \'agatson\' or int, but got: ' + str(threshold)) + raise ValueError( + "Error in threshold value for aortic calcium segmentaiton. \ + Should be 'adaptive', 'agatson' or int, but got: " + + str(threshold) + ) - ''' + """ Dilate aorta before using threshold to segment calcifications - ''' + """ t0 = time.time() if dilation is not None: struct = ndimage.generate_binary_structure(*dilation) @@ -382,7 +404,7 @@ def detectCalcifications( if show_time: print("dilation mask time: {:.2f}".format(time.time() - t0)) - + t0 = time.time() # make threshold calc_mask = np.logical_and(aorta_dilated == 1, ct >= calc_thres) @@ -412,7 +434,7 @@ def detectCalcifications( print("exclude dilation time: {:.2f}".format(time.time() - t0)) t0 = time.time() - calc_mask = calc_mask * (exclude_mask == 0) + calc_mask = calc_mask * (exclude_mask == 0) if show_time: print("exclude time: {:.2f}".format(time.time() - t0)) @@ -421,22 +443,21 @@ def detectCalcifications( print("Excluding calcifications under {} pixels".format(remove_size)) t0 = time.time() - + if calc_mask.sum() != 0: # perform the exclusion on a slice for speed - arr_slices = self.getSmallestArraySlice(calc_mask, margin = 1) + arr_slices = self.getSmallestArraySlice(calc_mask, margin=1) labels, num_features = ndimage.label(calc_mask[arr_slices]) - + counter = 0 for n in range(1, num_features + 1): idx_tmp = labels == n if idx_tmp.sum() <= remove_size: labels[idx_tmp] = 0 counter += 1 - + calc_mask[arr_slices] = labels > 0 - - + if show_time: print("Size exclusion time: {:.1f} sec".format(time.time() - t0)) @@ -456,30 +477,33 @@ def detectCalcifications( return results - def slicedDilationOrErosion(self, input_mask, struct, num_iteration, operation): """ Perform the dilation on the smallest slice that will fit the segmentation """ - + if num_iteration < 1: return input_mask - + margin = 2 if num_iteration is None else num_iteration + 1 x_idx = np.where(input_mask.sum(axis=(1, 2)))[0] if len(x_idx) > 0: - x_start, x_end = max(x_idx[0] - margin, 0), min(x_idx[-1] + margin, input_mask.shape[0]) - + x_start, x_end = max(x_idx[0] - margin, 0), min( + x_idx[-1] + margin, input_mask.shape[0] + ) + y_idx = np.where(input_mask.sum(axis=(0, 2)))[0] if len(y_idx) > 0: - y_start, y_end = max(y_idx[0] - margin, 0), min(y_idx[-1] + margin, input_mask.shape[1]) - - # Don't dilate the aorta at the bifurcation + y_start, y_end = max(y_idx[0] - margin, 0), min( + y_idx[-1] + margin, input_mask.shape[1] + ) + + # Don't dilate the aorta at the bifurcation z_idx = np.where(input_mask.sum(axis=(0, 1)))[0] z_start, z_end = z_idx[0], z_idx[-1] - + if operation == "dilate": mask_slice = ndimage.binary_dilation( input_mask[x_start:x_end, y_start:y_end, :], structure=struct @@ -493,54 +517,56 @@ def slicedDilationOrErosion(self, input_mask, struct, num_iteration, operation): output_mask = input_mask.copy() # insert dilated mask, but restrain to undilated z_start - output_mask[x_start:x_end, y_start:y_end, z_start:] = mask_slice[:,:,z_start:] + output_mask[x_start:x_end, y_start:y_end, z_start:] = mask_slice[:, :, z_start:] return output_mask def slicedSizeCount(self, aorta_eroded, ct, remove_size, calc_thres): - ''' + """ Counts the number of calcifications over the size threshold in the eroded aorta on the smallest slice that fits the aorta. - ''' + """ eroded_calc_mask = np.logical_and(aorta_eroded == 1, ct >= calc_thres) - + if eroded_calc_mask.sum() != 0: # Perfom the counts on a slice of the aorta for speed - arr_slices = self.getSmallestArraySlice(eroded_calc_mask, margin = 1) + arr_slices = self.getSmallestArraySlice(eroded_calc_mask, margin=1) labels, num_features = ndimage.label(eroded_calc_mask[arr_slices]) counter = 0 for n in range(1, num_features + 1): idx_tmp = labels == n if idx_tmp.sum() > remove_size: counter += 1 - + return counter else: return 0 - - def getSmallestArraySlice(self, input_mask, margin = 0): - ''' + def getSmallestArraySlice(self, input_mask, margin=0): + """ Generated the smallest slice that will fit the mask plus the given margin - and return a touple of slice objects - ''' - + and return a touple of slice objects + """ + x_idx = np.where(input_mask.sum(axis=(1, 2)))[0] if len(x_idx) > 0: - x_start, x_end = max(x_idx[0] - margin, 0), min(x_idx[-1] + margin, input_mask.shape[0]) - + x_start, x_end = max(x_idx[0] - margin, 0), min( + x_idx[-1] + margin, input_mask.shape[0] + ) + y_idx = np.where(input_mask.sum(axis=(0, 2)))[0] if len(y_idx) > 0: - y_start, y_end = max(y_idx[0] - margin, 0), min(y_idx[-1] + margin, input_mask.shape[1]) - + y_start, y_end = max(y_idx[0] - margin, 0), min( + y_idx[-1] + margin, input_mask.shape[1] + ) + z_idx = np.where(input_mask.sum(axis=(0, 1)))[0] if len(z_idx) > 0: - z_start, z_end = max(z_idx[0] - margin, 0), min(z_idx[-1] + margin, input_mask.shape[2]) - - - return (slice(x_start, x_end), slice(y_start, y_end), slice(z_start, z_end)) - + z_start, z_end = max(z_idx[0] - margin, 0), min( + z_idx[-1] + margin, input_mask.shape[2] + ) + return (slice(x_start, x_end), slice(y_start, y_end), slice(z_start, z_end)) class AorticCalciumMetrics(InferenceClass): @@ -551,8 +577,8 @@ def __init__(self): def __call__(self, inference_pipeline): calc_mask = inference_pipeline.calc_mask - spine_mask = inference_pipeline.spine_segmentation.get_fdata().astype(np.int8) - ''' 26: "vertebrae_S1", + spine_mask = inference_pipeline.spine_segmentation.get_fdata().astype(np.int8) + """ 26: "vertebrae_S1", 27: "vertebrae_L5", 28: "vertebrae_L4", 29: "vertebrae_L3", @@ -576,54 +602,54 @@ def __call__(self, inference_pipeline): 47: "vertebrae_C4", 48: "vertebrae_C3", 49: "vertebrae_C2", - 50: "vertebrae_C1",''' - - t12_level = np.where((spine_mask == 32).sum(axis=(0,1)))[0] - l1_level = np.where((spine_mask == 31).sum(axis=(0,1)))[0] - - + 50: "vertebrae_C1",""" + + t12_level = np.where((spine_mask == 32).sum(axis=(0, 1)))[0] + l1_level = np.where((spine_mask == 31).sum(axis=(0, 1)))[0] + if len(t12_level) != 0 and len(l1_level) != 0: sep_plane = round(np.mean([t12_level[0], l1_level[-1]])) elif len(t12_level) == 0 and len(l1_level) != 0: - print('WARNNG: could not locate T12, using L1 only..') + print("WARNNG: could not locate T12, using L1 only..") sep_plane = l1_level[-1] elif len(t12_level) != 0 and len(l1_level) == 0: - print('WARNNG: could not locate L1, using T12 only..') + print("WARNNG: could not locate L1, using T12 only..") sep_plane = t12_level[0] - else: - raise ValueError('Could not locate either T12 or L1, aborting..') - + else: + raise ValueError("Could not locate either T12 or L1, aborting..") + planes = np.zeros_like(spine_mask, dtype=np.int8) - planes[:,:,sep_plane] = 1 + planes[:, :, sep_plane] = 1 planes[spine_mask == 32] = 2 planes[spine_mask == 31] = 3 - + inference_pipeline.saveArrToNifti( planes, - os.path.join(inference_pipeline.output_dir_segmentation_masks, - "t12_plane.nii.gz") - ) - + os.path.join( + inference_pipeline.output_dir_segmentation_masks, "t12_plane.nii.gz" + ), + ) + inference_pipeline.pix_dims = inference_pipeline.medical_volume.header[ "pixdim" ][1:4] # divided with 10 to get in cm inference_pipeline.vol_per_pixel = np.prod(inference_pipeline.pix_dims / 10) - + all_regions = {} - region_names = ['Abdominal', 'Thoracic'] - + region_names = ["Abdominal", "Thoracic"] + ct_full = inference_pipeline.medical_volume.get_fdata() - for i in range(len(region_names)): + for i in range(len(region_names)): # count statistics for individual calcifications if i == 0: - calc_mask_region = calc_mask[:,:,:sep_plane] - ct = ct_full[:,:,:sep_plane] + calc_mask_region = calc_mask[:, :, :sep_plane] + ct = ct_full[:, :, :sep_plane] elif i == 1: - calc_mask_region = calc_mask[:,:,sep_plane:] - ct = ct_full[:,:,sep_plane:] - + calc_mask_region = calc_mask[:, :, sep_plane:] + ct = ct_full[:, :, sep_plane:] + labelled_calc, num_lesions = ndimage.label(calc_mask_region) metrics = { @@ -632,7 +658,7 @@ def __call__(self, inference_pipeline): "median_hu": [], "max_hu": [], } - + if num_lesions == 0: metrics["volume"].append(0) metrics["mean_hu"].append(0) @@ -641,39 +667,42 @@ def __call__(self, inference_pipeline): else: for j in range(1, num_lesions + 1): tmp_mask = labelled_calc == j - + tmp_ct_vals = ct[tmp_mask] - + metrics["volume"].append( len(tmp_ct_vals) * inference_pipeline.vol_per_pixel ) metrics["mean_hu"].append(np.mean(tmp_ct_vals)) metrics["median_hu"].append(np.median(tmp_ct_vals)) metrics["max_hu"].append(np.max(tmp_ct_vals)) - + # Volume of calcificaitons calc_vol = np.sum(metrics["volume"]) metrics["volume_total"] = calc_vol - + metrics["num_calc"] = num_lesions - if inference_pipeline.args.threshold == 'agatson': + if inference_pipeline.args.threshold == "agatson": if num_lesions == 0: metrics["agatson_score"] = 0 else: - metrics["agatson_score"] = self.CalculateAgatsonScore(calc_mask_region, ct, inference_pipeline.pix_dims) - + metrics["agatson_score"] = self.CalculateAgatsonScore( + calc_mask_region, ct, inference_pipeline.pix_dims + ) + all_regions[region_names[i]] = metrics - + inference_pipeline.metrics = all_regions return {} - + def CalculateAgatsonScore(self, calc_mask_region, ct, pix_dims): - ''' + """ Original Agatson papers says need to be >= 1mm^2, other papers use at least 3 face-linked pixels. - ''' + """ + def get_hu_factor(max_hu): # if max_hu >< if max_hu < 200: @@ -685,28 +714,30 @@ def get_hu_factor(max_hu): elif max_hu >= 400: factor = 4 else: - raise ValueError('Could determine factor, got: ' + str(max_hu)) - + raise ValueError("Could determine factor, got: " + str(max_hu)) + return factor - + # dims are in mm here area_per_pixel = pix_dims[0] * pix_dims[1] agatson = 0 - - for i in range(calc_mask_region.shape[2]): - tmp_slice = calc_mask_region[:,:,i] - tmp_ct_slice = ct[:,:,i] - + + for i in range(calc_mask_region.shape[2]): + tmp_slice = calc_mask_region[:, :, i] + tmp_ct_slice = ct[:, :, i] + labelled_calc, num_lesions = ndimage.label(tmp_slice) - + for j in range(1, num_lesions + 1): tmp_mask = labelled_calc == j - + tmp_area = tmp_mask.sum() * area_per_pixel # exclude if less than 1 mm^2 if tmp_area <= 1: continue else: - agatson += tmp_area * get_hu_factor(int(tmp_ct_slice[tmp_mask].max())) - - return agatson \ No newline at end of file + agatson += tmp_area * get_hu_factor( + int(tmp_ct_slice[tmp_mask].max()) + ) + + return agatson diff --git a/comp2comp/aortic_calcium/aortic_calcium_visualization.py b/comp2comp/aortic_calcium/aortic_calcium_visualization.py index 1ed8d52..0f5a532 100644 --- a/comp2comp/aortic_calcium/aortic_calcium_visualization.py +++ b/comp2comp/aortic_calcium/aortic_calcium_visualization.py @@ -25,30 +25,29 @@ def __init__(self): super().__init__() def __call__(self, inference_pipeline): - + all_metrics = inference_pipeline.metrics inference_pipeline.csv_output_dir = os.path.join( inference_pipeline.output_dir, "metrics" ) os.makedirs(inference_pipeline.csv_output_dir, exist_ok=True) - - + # Write metrics to CSV file with open( os.path.join(inference_pipeline.csv_output_dir, "aortic_calcification.csv"), "w", ) as f: f.write("Volume (cm^3),Mean HU,Median HU,Max HU\n") - + with open( - os.path.join(inference_pipeline.csv_output_dir, "aortic_calcification.csv"), - "a", - ) as f: + os.path.join(inference_pipeline.csv_output_dir, "aortic_calcification.csv"), + "a", + ) as f: for region, metrics in all_metrics.items(): - f.write(region + ',,,\n') - + f.write(region + ",,,\n") + for vol, mean, median, max in zip( metrics["volume"], metrics["mean_hu"], @@ -56,8 +55,7 @@ def __call__(self, inference_pipeline): metrics["max_hu"], ): f.write("{},{:.1f},{:.1f},{:.1f}\n".format(vol, mean, median, max)) - - + # Write total results with open( os.path.join( @@ -66,60 +64,66 @@ def __call__(self, inference_pipeline): "w", ) as f: for region, metrics in all_metrics.items(): - f.write(region + ',\n') - + f.write(region + ",\n") + f.write("Total number,{}\n".format(metrics["num_calc"])) f.write("Total volume (cm^3),{:.3f}\n".format(metrics["volume_total"])) - f.write("Threshold (HU),{:.1f}\n".format(inference_pipeline.calcium_threshold)) - - f.write("{},{:.1f}+/-{:.1f}\n".format( + f.write( + "Threshold (HU),{:.1f}\n".format( + inference_pipeline.calcium_threshold + ) + ) + + f.write( + "{},{:.1f}+/-{:.1f}\n".format( "Mean HU", np.mean(metrics["mean_hu"]), np.std(metrics["mean_hu"]), ) ) - f.write("{},{:.1f}+/-{:.1f}\n".format( + f.write( + "{},{:.1f}+/-{:.1f}\n".format( "Median HU", np.mean(metrics["median_hu"]), np.std(metrics["median_hu"]), ) ) - f.write("{},{:.1f}+/-{:.1f}\n".format( + f.write( + "{},{:.1f}+/-{:.1f}\n".format( "Max HU", np.mean(metrics["max_hu"]), np.std(metrics["max_hu"]), ) ) - f.write("{},{:.3f}+/-{:.3f}\n".format( + f.write( + "{},{:.3f}+/-{:.3f}\n".format( "Mean volume (cm³):", np.mean(metrics["volume"]), np.std(metrics["volume"]), ) ) - f.write("{},{:.3f}\n".format( - "Median volume (cm³)", np.median(metrics["volume"]) + f.write( + "{},{:.3f}\n".format( + "Median volume (cm³)", np.median(metrics["volume"]) ) ) - f.write("{},{:.3f}\n".format( - "Max volume (cm³)", np.max(metrics["volume"]) - ) + f.write( + "{},{:.3f}\n".format("Max volume (cm³)", np.max(metrics["volume"])) ) - f.write("{},{:.3f}\n".format( - "Min volume (cm³):", np.min(metrics["volume"]) - ) + f.write( + "{},{:.3f}\n".format("Min volume (cm³):", np.min(metrics["volume"])) ) - - if inference_pipeline.args.threshold == 'agatson': + + if inference_pipeline.args.threshold == "agatson": f.write("Agatson score,{:.1f}\n".format(metrics["agatson_score"])) - distance = 25 print("\n") print("Statistics on aortic calcifications:") - + for region, metrics in all_metrics.items(): - print(region + ':') - + print(region + ":") + if metrics["num_calc"] == 0: print("No aortic calcifications were found.\n") else: @@ -178,18 +182,18 @@ def __call__(self, inference_pipeline): ) print( "{:<{}}{:.3f}".format( - "Threshold (HU):", distance, inference_pipeline.calcium_threshold + "Threshold (HU):", + distance, + inference_pipeline.calcium_threshold, ) ) - if inference_pipeline.args.threshold == 'agatson': + if inference_pipeline.args.threshold == "agatson": print( "{:<{}}{:.1f}".format( "Agatson score:", distance, metrics["agatson_score"] ) ) - - + print("\n") return {} - diff --git a/comp2comp/inference_pipeline.py b/comp2comp/inference_pipeline.py index 0676569..0aca5da 100644 --- a/comp2comp/inference_pipeline.py +++ b/comp2comp/inference_pipeline.py @@ -6,15 +6,18 @@ import os from typing import Dict, List +import nibabel as nib + from comp2comp.inference_class_base import InferenceClass from comp2comp.io.io import DicomLoader, NiftiSaver -import nibabel as nib class InferencePipeline(InferenceClass): """Inference pipeline.""" - def __init__(self, inference_classes: List = None, config: Dict = None, args: Dict = None): + def __init__( + self, inference_classes: List = None, config: Dict = None, args: Dict = None + ): self.config = config self.args = args # assign values from config to attributes @@ -81,11 +84,8 @@ def __call__(self, inference_pipeline=None, **kwargs): return output - def saveArrToNifti( - self, - arr, - path): - ''' + def saveArrToNifti(self, arr, path): + """ Saves an array to nifti using the CT as reference Args: @@ -95,11 +95,11 @@ def saveArrToNifti( Returns: None. - ''' - img = nib.Nifti1Image(arr, - self.medical_volume.affine, - self.medical_volume.header) - nib.save(img, path) + """ + img = nib.Nifti1Image( + arr, self.medical_volume.affine, self.medical_volume.header + ) + nib.save(img, path) if __name__ == "__main__": diff --git a/comp2comp/io/io.py b/comp2comp/io/io.py index 2c7eff9..63491b2 100644 --- a/comp2comp/io/io.py +++ b/comp2comp/io/io.py @@ -1,6 +1,7 @@ """ @author: louisblankemeier """ + import os import shutil from pathlib import Path @@ -8,9 +9,9 @@ # import dicom2nifti import dosma as dm +import nibabel as nib import pydicom import SimpleITK as sitk -import nibabel as nib from comp2comp.inference_class_base import InferenceClass @@ -114,8 +115,10 @@ def __call__(self, inference_pipeline): os.path.join(segmentations_output_dir, "converted_dcm.nii.gz"), ) - inference_pipeline.medical_volume = nib.load(os.path.join(segmentations_output_dir, "converted_dcm.nii.gz")) - + inference_pipeline.medical_volume = nib.load( + os.path.join(segmentations_output_dir, "converted_dcm.nii.gz") + ) + return {} @@ -130,7 +133,9 @@ def series_selector(dicom_path, pipeline_name=None): if ds.ImageOrientationPatient != [1, 0, 0, 0, 1, 0]: raise ValueError("Image orientation is not axial") else: - print(f"Skipping primary, original, and orientation image type check for the {pipeline_name} pipeline.") + print( + f"Skipping primary, original, and orientation image type check for the {pipeline_name} pipeline." + ) # if any("gsi" in s.lower() for s in image_type_list): # raise ValueError("GSI image type") return ds diff --git a/comp2comp/io/io_utils.py b/comp2comp/io/io_utils.py index e1f9190..261d782 100644 --- a/comp2comp/io/io_utils.py +++ b/comp2comp/io/io_utils.py @@ -1,8 +1,10 @@ """ @author: louisblankemeier """ + import csv import os + import nibabel as nib import pydicom @@ -48,18 +50,30 @@ def get_dicom_or_nifti_paths_and_num(path): list: List of paths. """ dicom_nifti_paths = [] - + if path.endswith(".nii") or path.endswith(".nii.gz"): - dicom_nifti_paths.append( (path, getNumSlicesNifti(path)) ) - elif path.endswith('.txt'): + dicom_nifti_paths.append((path, getNumSlicesNifti(path))) + elif path.endswith(".txt"): dicom_nifti_paths = [] - with open(path, 'r') as f: + with open(path, "r") as f: for dicom_folder_path in f: - if dicom_folder_path.endswith(".nii") or path.dicom_folder_path(".nii.gz"): - dicom_nifti_paths.append( (dicom_folder_path.strip(), getNumSlicesNifti(dicom_folder_path.strip()))) + if dicom_folder_path.endswith(".nii") or path.dicom_folder_path( + ".nii.gz" + ): + dicom_nifti_paths.append( + ( + dicom_folder_path.strip(), + getNumSlicesNifti(dicom_folder_path.strip()), + ) + ) else: - dicom_nifti_paths.append( (dicom_folder_path.strip(), len(os.listdir(dicom_folder_path.strip())))) - else: + dicom_nifti_paths.append( + ( + dicom_folder_path.strip(), + len(os.listdir(dicom_folder_path.strip())), + ) + ) + else: for root, dirs, files in os.walk(path): if len(files) > 0: # if all(file.endswith(".dcm") or file.endswith(".dicom") for file in files): @@ -80,9 +94,9 @@ def write_dicom_metadata_to_csv(ds, csv_filename): continue value = str(element.value) csvwriter.writerow([tag, keyword, value]) - + + def getNumSlicesNifti(path): img = nib.load(path) img = nib.as_closest_canonical(img) return img.shape[2] - diff --git a/comp2comp/liver_spleen_pancreas/liver_spleen_pancreas.py b/comp2comp/liver_spleen_pancreas/liver_spleen_pancreas.py index 7265b39..5ec2932 100644 --- a/comp2comp/liver_spleen_pancreas/liver_spleen_pancreas.py +++ b/comp2comp/liver_spleen_pancreas/liver_spleen_pancreas.py @@ -3,12 +3,8 @@ from time import time from typing import Union -from totalsegmentator.libs import ( - download_pretrained_weights, - nostdout, - setup_nnunet, -) from totalsegmentatorv2.python_api import totalsegmentator + from comp2comp.inference_class_base import InferenceClass @@ -25,17 +21,17 @@ def __call__(self, inference_pipeline): self.output_dir_segmentations = os.path.join(self.output_dir, "segmentations/") if not os.path.exists(self.output_dir_segmentations): os.makedirs(self.output_dir_segmentations) - + self.model_dir = inference_pipeline.model_dir - + seg = self.organ_seg( os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz"), self.output_dir_segmentations + "organs.nii.gz", inference_pipeline.model_dir, ) - + inference_pipeline.segmentation = seg - + return {} def organ_seg( @@ -53,18 +49,18 @@ def organ_seg( os.environ["SCRATCH"] = self.model_dir seg = totalsegmentator( - input = input_path, - output = output_path, + input=input_path, + output=output_path, # input = os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz"), # output = os.path.join(self.output_dir_segmentations, "segmentation.nii"), - task_ids = [291], - ml = True, - nr_thr_resamp = 1, - nr_thr_saving = 6, - fast = False, - nora_tag = "None", - preview = False, - task = "total", + task_ids=[291], + ml=True, + nr_thr_resamp=1, + nr_thr_saving=6, + fast=False, + nora_tag="None", + preview=False, + task="total", # roi_subset = [ # "vertebrae_T12", # "vertebrae_L1", @@ -73,22 +69,22 @@ def organ_seg( # "vertebrae_L4", # "vertebrae_L5", # ], - roi_subset = None, - statistics = False, - radiomics = False, - crop_path = None, - body_seg = False, - force_split = False, - output_type = "nifti", - quiet = False, - verbose = False, - test = 0, - skip_saving = True, - device = "gpu", - license_number = None, - statistics_exclude_masks_at_border = True, - no_derived_masks = False, - v1_order = False, + roi_subset=None, + statistics=False, + radiomics=False, + crop_path=None, + body_seg=False, + force_split=False, + output_type="nifti", + quiet=False, + verbose=False, + test=0, + skip_saving=True, + device="gpu", + license_number=None, + statistics_exclude_masks_at_border=True, + no_derived_masks=False, + v1_order=False, ) # Setup nnunet @@ -100,7 +96,7 @@ def organ_seg( # setup_nnunet() # download_pretrained_weights(task_id[0]) - + # from totalsegmentator.nnunet import nnUNet_predict_image # with nostdout(): diff --git a/comp2comp/liver_spleen_pancreas/liver_spleen_pancreas_visualization.py b/comp2comp/liver_spleen_pancreas/liver_spleen_pancreas_visualization.py index 8e6dbf8..ac0f2bd 100644 --- a/comp2comp/liver_spleen_pancreas/liver_spleen_pancreas_visualization.py +++ b/comp2comp/liver_spleen_pancreas/liver_spleen_pancreas_visualization.py @@ -3,6 +3,7 @@ import os +import nibabel as nib import numpy as np from comp2comp.inference_class_base import InferenceClass @@ -10,8 +11,7 @@ generate_liver_spleen_pancreas_report, generate_slice_images, ) -import nibabel as nib -import numpy as np + class LiverSpleenPancreasVisualizer(InferenceClass): def __init__(self): @@ -30,33 +30,35 @@ def __init__(self): def __call__(self, inference_pipeline): self.output_dir = inference_pipeline.output_dir self.output_dir_images_organs = os.path.join(self.output_dir, "images/") - inference_pipeline.output_dir_images_organs = ( - self.output_dir_images_organs - ) + inference_pipeline.output_dir_images_organs = self.output_dir_images_organs if not os.path.exists(self.output_dir_images_organs): os.makedirs(self.output_dir_images_organs) - + # make folder for volumes self.output_dir_volumes = os.path.join(self.output_dir, "volumes/") if not os.path.exists(self.output_dir_volumes): os.makedirs(self.output_dir_volumes) - - # save the volume to disk in nifti format - nib.save(inference_pipeline.medical_volume, os.path.join(self.output_dir_volumes, 'ct.nii.gz')) - - segmentation_subset = np.zeros(inference_pipeline.medical_volume.shape, dtype=np.int8) + + # save the volume to disk in nifti format + nib.save( + inference_pipeline.medical_volume, + os.path.join(self.output_dir_volumes, "ct.nii.gz"), + ) + + segmentation_subset = np.zeros( + inference_pipeline.medical_volume.shape, dtype=np.int8 + ) tmp_seg = inference_pipeline.segmentation.get_fdata().astype(np.int8) - + for i, c in enumerate(self.class_nums, start=1): segmentation_subset[tmp_seg == c] = i - + inference_pipeline.saveArrToNifti( segmentation_subset, - os.path.join(self.output_dir_volumes, - "liver_spleen_pancreas_mask.nii.gz") - ) - + os.path.join(self.output_dir_volumes, "liver_spleen_pancreas_mask.nii.gz"), + ) + inference_pipeline.medical_volume_arr = np.flip( inference_pipeline.medical_volume.get_fdata(), axis=1 ) @@ -83,7 +85,7 @@ def __call__(self, inference_pipeline): ) inference_pipeline.organ_metrics = self.organ_metrics - + generate_liver_spleen_pancreas_report( self.output_dir_images_organs, self.organ_names ) diff --git a/comp2comp/liver_spleen_pancreas/visualization_utils.py b/comp2comp/liver_spleen_pancreas/visualization_utils.py index db93117..050313a 100644 --- a/comp2comp/liver_spleen_pancreas/visualization_utils.py +++ b/comp2comp/liver_spleen_pancreas/visualization_utils.py @@ -31,11 +31,12 @@ def extract_axial_mid_slice(ct, mask, crop=True): ct_slice_z = np.transpose(ct[:, :, max_extent_idx], axes=(1, 0)) mask_slice_z = np.transpose(mask[:, :, max_extent_idx], axes=(1, 0)) - ct_slice_z = np.flip(ct_slice_z, axis=(0,1)) - mask_slice_z = np.flip(mask_slice_z, axis=(0,1)) + ct_slice_z = np.flip(ct_slice_z, axis=(0, 1)) + mask_slice_z = np.flip(mask_slice_z, axis=(0, 1)) return ct_slice_z, mask_slice_z, max_extent_idx + # def extract_axial_mid_slice(ct, mask, crop=True): # # find the slice with max surface area of the organ # slice_idx = np.argmax(mask.sum(axis=(0, 1))) @@ -187,25 +188,25 @@ def slicedDilationOrErosion(input_mask, num_iteration, operation): Perform the dilation on the smallest slice that will fit the segmentation """ - + # if empty, don't do dilation if input_mask.sum() == 0: return input_mask - + margin = 2 if num_iteration is None else num_iteration + 1 - + x_size, y_size, z_size = input_mask.shape # find the minimum volume enclosing the organ x_idx = np.where(input_mask.sum(axis=(1, 2)))[0] x_start, x_end = max(0, x_idx[0] - margin), min(x_idx[-1] + margin, x_size) - + y_idx = np.where(input_mask.sum(axis=(0, 2)))[0] y_start, y_end = max(0, y_idx[0] - margin), min(y_idx[-1] + margin, y_size) - + z_idx = np.where(input_mask.sum(axis=(0, 1)))[0] z_start, z_end = max(0, z_idx[0] - margin), min(z_idx[-1] + margin, z_size) - + struct = scipy.ndimage.generate_binary_structure(3, 1) struct = scipy.ndimage.iterate_structure(struct, num_iteration) @@ -269,20 +270,29 @@ def generate_slice_images( colors = [1, 3, 4] # create the txt files for the slices idx - with open(os.path.join(root, 'slice_idx.csv'), 'w') as f: - f.write('organ,mean_HU,axial_idx\n') - + with open(os.path.join(root, "slice_idx.csv"), "w") as f: + f.write("organ,mean_HU,axial_idx\n") + for i, c_num in enumerate(class_nums): organ_name = class_map_part_organs[c_num] axial_path = os.path.join(root, organ_name.lower() + "_axial.png") coronal_path = os.path.join(root, organ_name.lower() + "_coronal.png") - ct_slice_z, liver_slice_z, slice_idx_z = extract_axial_mid_slice(ct, all_masks == c_num) - with open(os.path.join(root, 'slice_idx.csv'), 'a') as f: + ct_slice_z, liver_slice_z, slice_idx_z = extract_axial_mid_slice( + ct, all_masks == c_num + ) + with open(os.path.join(root, "slice_idx.csv"), "a") as f: mean_hu = ct_slice_z[liver_slice_z == 1].mean() - f.write(organ_name.lower()+ ',' + '{:.1f}'.format(mean_hu) + ',' + str(slice_idx_z) + '\n') - + f.write( + organ_name.lower() + + "," + + "{:.1f}".format(mean_hu) + + "," + + str(slice_idx_z) + + "\n" + ) + results = extract_organ_metrics( ct, all_masks, class_num=c_num, vol_per_pixel=vol_per_pixel ) @@ -320,7 +330,7 @@ def generate_slice_images( def generate_liver_spleen_pancreas_report(root, organ_names): - + axial_imgs = [ Image.open(os.path.join(root, organ + "_axial.png")) for organ in organ_names ] @@ -357,121 +367,121 @@ def generate_liver_spleen_pancreas_report(root, organ_names): # from https://github.com/wasserth/TotalSegmentator/blob/master/totalsegmentator/map_to_binary.py class_map_part_organs = { - 1: "spleen", - 2: "kidney_right", - 3: "kidney_left", - 4: "gallbladder", - 5: "liver", - 6: "stomach", - 7: "pancreas", - 8: "adrenal_gland_right", - 9: "adrenal_gland_left", - 10: "lung_upper_lobe_left", - 11: "lung_lower_lobe_left", - 12: "lung_upper_lobe_right", - 13: "lung_middle_lobe_right", - 14: "lung_lower_lobe_right", - 15: "esophagus", - 16: "trachea", - 17: "thyroid_gland", - 18: "small_bowel", - 19: "duodenum", - 20: "colon", - 21: "urinary_bladder", - 22: "prostate", - 23: "kidney_cyst_left", - 24: "kidney_cyst_right", - 25: "sacrum", - 26: "vertebrae_S1", - 27: "vertebrae_L5", - 28: "vertebrae_L4", - 29: "vertebrae_L3", - 30: "vertebrae_L2", - 31: "vertebrae_L1", - 32: "vertebrae_T12", - 33: "vertebrae_T11", - 34: "vertebrae_T10", - 35: "vertebrae_T9", - 36: "vertebrae_T8", - 37: "vertebrae_T7", - 38: "vertebrae_T6", - 39: "vertebrae_T5", - 40: "vertebrae_T4", - 41: "vertebrae_T3", - 42: "vertebrae_T2", - 43: "vertebrae_T1", - 44: "vertebrae_C7", - 45: "vertebrae_C6", - 46: "vertebrae_C5", - 47: "vertebrae_C4", - 48: "vertebrae_C3", - 49: "vertebrae_C2", - 50: "vertebrae_C1", - 51: "heart", - 52: "aorta", - 53: "pulmonary_vein", - 54: "brachiocephalic_trunk", - 55: "subclavian_artery_right", - 56: "subclavian_artery_left", - 57: "common_carotid_artery_right", - 58: "common_carotid_artery_left", - 59: "brachiocephalic_vein_left", - 60: "brachiocephalic_vein_right", - 61: "atrial_appendage_left", - 62: "superior_vena_cava", - 63: "inferior_vena_cava", - 64: "portal_vein_and_splenic_vein", - 65: "iliac_artery_left", - 66: "iliac_artery_right", - 67: "iliac_vena_left", - 68: "iliac_vena_right", - 69: "humerus_left", - 70: "humerus_right", - 71: "scapula_left", - 72: "scapula_right", - 73: "clavicula_left", - 74: "clavicula_right", - 75: "femur_left", - 76: "femur_right", - 77: "hip_left", - 78: "hip_right", - 79: "spinal_cord", - 80: "gluteus_maximus_left", - 81: "gluteus_maximus_right", - 82: "gluteus_medius_left", - 83: "gluteus_medius_right", - 84: "gluteus_minimus_left", - 85: "gluteus_minimus_right", - 86: "autochthon_left", - 87: "autochthon_right", - 88: "iliopsoas_left", - 89: "iliopsoas_right", - 90: "brain", - 91: "skull", - 92: "rib_left_1", - 93: "rib_left_2", - 94: "rib_left_3", - 95: "rib_left_4", - 96: "rib_left_5", - 97: "rib_left_6", - 98: "rib_left_7", - 99: "rib_left_8", - 100: "rib_left_9", - 101: "rib_left_10", - 102: "rib_left_11", - 103: "rib_left_12", - 104: "rib_right_1", - 105: "rib_right_2", - 106: "rib_right_3", - 107: "rib_right_4", - 108: "rib_right_5", - 109: "rib_right_6", - 110: "rib_right_7", - 111: "rib_right_8", - 112: "rib_right_9", - 113: "rib_right_10", - 114: "rib_right_11", - 115: "rib_right_12", - 116: "sternum", - 117: "costal_cartilages" - } + 1: "spleen", + 2: "kidney_right", + 3: "kidney_left", + 4: "gallbladder", + 5: "liver", + 6: "stomach", + 7: "pancreas", + 8: "adrenal_gland_right", + 9: "adrenal_gland_left", + 10: "lung_upper_lobe_left", + 11: "lung_lower_lobe_left", + 12: "lung_upper_lobe_right", + 13: "lung_middle_lobe_right", + 14: "lung_lower_lobe_right", + 15: "esophagus", + 16: "trachea", + 17: "thyroid_gland", + 18: "small_bowel", + 19: "duodenum", + 20: "colon", + 21: "urinary_bladder", + 22: "prostate", + 23: "kidney_cyst_left", + 24: "kidney_cyst_right", + 25: "sacrum", + 26: "vertebrae_S1", + 27: "vertebrae_L5", + 28: "vertebrae_L4", + 29: "vertebrae_L3", + 30: "vertebrae_L2", + 31: "vertebrae_L1", + 32: "vertebrae_T12", + 33: "vertebrae_T11", + 34: "vertebrae_T10", + 35: "vertebrae_T9", + 36: "vertebrae_T8", + 37: "vertebrae_T7", + 38: "vertebrae_T6", + 39: "vertebrae_T5", + 40: "vertebrae_T4", + 41: "vertebrae_T3", + 42: "vertebrae_T2", + 43: "vertebrae_T1", + 44: "vertebrae_C7", + 45: "vertebrae_C6", + 46: "vertebrae_C5", + 47: "vertebrae_C4", + 48: "vertebrae_C3", + 49: "vertebrae_C2", + 50: "vertebrae_C1", + 51: "heart", + 52: "aorta", + 53: "pulmonary_vein", + 54: "brachiocephalic_trunk", + 55: "subclavian_artery_right", + 56: "subclavian_artery_left", + 57: "common_carotid_artery_right", + 58: "common_carotid_artery_left", + 59: "brachiocephalic_vein_left", + 60: "brachiocephalic_vein_right", + 61: "atrial_appendage_left", + 62: "superior_vena_cava", + 63: "inferior_vena_cava", + 64: "portal_vein_and_splenic_vein", + 65: "iliac_artery_left", + 66: "iliac_artery_right", + 67: "iliac_vena_left", + 68: "iliac_vena_right", + 69: "humerus_left", + 70: "humerus_right", + 71: "scapula_left", + 72: "scapula_right", + 73: "clavicula_left", + 74: "clavicula_right", + 75: "femur_left", + 76: "femur_right", + 77: "hip_left", + 78: "hip_right", + 79: "spinal_cord", + 80: "gluteus_maximus_left", + 81: "gluteus_maximus_right", + 82: "gluteus_medius_left", + 83: "gluteus_medius_right", + 84: "gluteus_minimus_left", + 85: "gluteus_minimus_right", + 86: "autochthon_left", + 87: "autochthon_right", + 88: "iliopsoas_left", + 89: "iliopsoas_right", + 90: "brain", + 91: "skull", + 92: "rib_left_1", + 93: "rib_left_2", + 94: "rib_left_3", + 95: "rib_left_4", + 96: "rib_left_5", + 97: "rib_left_6", + 98: "rib_left_7", + 99: "rib_left_8", + 100: "rib_left_9", + 101: "rib_left_10", + 102: "rib_left_11", + 103: "rib_left_12", + 104: "rib_right_1", + 105: "rib_right_2", + 106: "rib_right_3", + 107: "rib_right_4", + 108: "rib_right_5", + 109: "rib_right_6", + 110: "rib_right_7", + 111: "rib_right_8", + 112: "rib_right_9", + 113: "rib_right_10", + 114: "rib_right_11", + 115: "rib_right_12", + 116: "sternum", + 117: "costal_cartilages", +} diff --git a/comp2comp/muscle_adipose_tissue/muscle_adipose_tissue.py b/comp2comp/muscle_adipose_tissue/muscle_adipose_tissue.py index bfeeb50..6ec41cb 100644 --- a/comp2comp/muscle_adipose_tissue/muscle_adipose_tissue.py +++ b/comp2comp/muscle_adipose_tissue/muscle_adipose_tissue.py @@ -81,7 +81,9 @@ def __call__(self, inference_pipeline): if self.model_name == "stanford_v0.0.2": self.download_muscle_adipose_tissue_model(inference_pipeline.model_dir) nifti_path = os.path.join( - inference_pipeline.output_dir, "segmentations", "converted_dcm_multilevel.nii.gz" + inference_pipeline.output_dir, + "segmentations", + "converted_dcm_multilevel.nii.gz", ) output_path = os.path.join( inference_pipeline.output_dir, diff --git a/comp2comp/spine/spine.py b/comp2comp/spine/spine.py index 8928522..9e42801 100644 --- a/comp2comp/spine/spine.py +++ b/comp2comp/spine/spine.py @@ -15,14 +15,12 @@ import pandas as pd import wget from PIL import Image - from totalsegmentatorv2.python_api import totalsegmentator from comp2comp.inference_class_base import InferenceClass from comp2comp.models.models import Models from comp2comp.spine import spine_utils from comp2comp.visualization.dicom import to_dicom -from comp2comp.io import io_utils # from totalsegmentator.libs import ( # download_pretrained_weights, @@ -31,8 +29,6 @@ # ) - - class SpineSegmentation(InferenceClass): """Spine segmentation.""" @@ -314,10 +310,12 @@ def __call__(self, inference_pipeline): # Compute ROIs inference_pipeline.spine_model_type = self.spine_model_type - (spine_hus, rois, segmentation_hus, centroids_3d, spine_masks) = spine_utils.compute_rois( - inference_pipeline.segmentation, - inference_pipeline.medical_volume, - self.spine_model_type, + (spine_hus, rois, segmentation_hus, centroids_3d, spine_masks) = ( + spine_utils.compute_rois( + inference_pipeline.segmentation, + inference_pipeline.medical_volume, + self.spine_model_type, + ) ) inference_pipeline.spine_hus = spine_hus diff --git a/comp2comp/spine/spine_utils.py b/comp2comp/spine/spine_utils.py index 680341d..4979d89 100644 --- a/comp2comp/spine/spine_utils.py +++ b/comp2comp/spine/spine_utils.py @@ -5,14 +5,14 @@ import logging import math import os +import time from typing import Dict, List import cv2 import nibabel as nib import numpy as np -from scipy.ndimage import zoom import scipy -import time +from scipy.ndimage import zoom from comp2comp.spine import spine_visualization @@ -229,7 +229,7 @@ def roi_from_mask(img, centroid: np.ndarray, seg: np.ndarray, slice: np.ndarray) lower_z_idx = updated_z_center - ((length_k * 1.5) // 2) upper_z_idx = updated_z_center + ((length_k * 1.5) // 2) for idx in range(int(lower_z_idx), int(upper_z_idx) + 1): - # take multiple to increase robustness + # take multiple to increase robustness posterior_anterior_lines = [ slice[:, idx], slice[:, idx + 1], diff --git a/comp2comp/utils/process.py b/comp2comp/utils/process.py index 740a177..7feef82 100644 --- a/comp2comp/utils/process.py +++ b/comp2comp/utils/process.py @@ -9,19 +9,22 @@ import traceback from datetime import datetime from pathlib import Path + from comp2comp.io import io_utils + def find_common_root(paths): - paths_with_sep = [path if path.endswith('/') else path + '/' for path in paths] + paths_with_sep = [path if path.endswith("/") else path + "/" for path in paths] # Find common prefix, ensuring it ends with a directory separator common_root = os.path.commonprefix(paths_with_sep) common_root - if not common_root.endswith('/'): + if not common_root.endswith("/"): # Find the last separator to correctly identify the common root directory - common_root = common_root[:common_root.rfind('/')+1] - - return common_root + common_root = common_root[: common_root.rfind("/") + 1] + + return common_root + def process_2d(args, pipeline_builder): output_dir = Path( @@ -60,15 +63,14 @@ def process_3d(args, pipeline_builder): output_path = os.path.join(output_path, date_time) path_and_num = io_utils.get_dicom_or_nifti_paths_and_num(args.input_path) - + # in case input is a .txt file we need to find the common root of the files - if args.input_path.endswith('.txt'): + if args.input_path.endswith(".txt"): all_paths = [p[0] for p in path_and_num] common_root = find_common_root(all_paths) - - + for path, num in path_and_num: - + try: st = time.time() @@ -103,13 +105,11 @@ def process_3d(args, pipeline_builder): ) else: - if args.input_path.endswith('.txt'): + if args.input_path.endswith(".txt"): output_dir = Path( os.path.join( output_path, - os.path.relpath( - os.path.normpath(path), common_root - ), + os.path.relpath(os.path.normpath(path), common_root), ) ) else: @@ -118,16 +118,17 @@ def process_3d(args, pipeline_builder): output_path, Path(os.path.basename(os.path.normpath(args.input_path))), os.path.relpath( - os.path.normpath(path), os.path.normpath(args.input_path) + os.path.normpath(path), + os.path.normpath(args.input_path), ), ) ) if not os.path.exists(output_dir): output_dir.mkdir(parents=True) - + pipeline = pipeline_builder(path, args) - + pipeline(output_dir=output_dir, model_dir=model_dir) if not args.save_segmentations: @@ -137,10 +138,9 @@ def process_3d(args, pipeline_builder): shutil.rmtree(segmentations_dir) print(f"Finished processing {path} in {time.time() - st:.1f} seconds\n") - print('Output was saved to:') + print("Output was saved to:") print(output_dir) - except Exception: print(f"ERROR PROCESSING {path}\n") traceback.print_exc()