def readRadiopaedia(all_imgs, all_lesion_ctrs, all_lung_ctrs):
data_folder = "/data/UM/COVID/Kaggle_1_Radiopedia"
file_names = listdir(join(data_folder, "ct_scans"))
file_names.sort()
for id, c_file_name in enumerate(file_names):
c_img_file = join(data_folder, "ct_scans", c_file_name)
if c_file_name.find("radiopaedia") != -1:
c_file_name = c_file_name.replace("org_covid-19-pneumonia-", "")
c_file_name = c_file_name.replace("-dcm", "")
c_mask_file = join(data_folder, "infection_mask", c_file_name)
c_lung_file = join(data_folder, "lung_mask", c_file_name)
else:
c_file_name = c_file_name.replace("org_", "")
c_mask_file = join(data_folder, "infection_mask", c_file_name)
c_lung_file = join(data_folder, "lung_mask", c_file_name)
itk_img = sitk.ReadImage(c_img_file)
itk_mask = sitk.ReadImage(c_mask_file)
itk_lung = sitk.ReadImage(c_lung_file)
np_mask = sitk.GetArrayFromImage(itk_mask)
np_mask[np_mask >= .01] = 1 # Not sure why but it has decimals
all_imgs.append(itk_img)
all_lesion_ctrs.append(copyItkImage(itk_img, np_mask))
all_lung_ctrs.append(itk_lung)
return all_imgs, all_lesion_ctrs, all_lung_ctrs
print("Done!")
def correct_adc_itk(img_itk,
low_threshold=1,
max_value=4500,
fix_value='mean'):
"""
This function receives and outputs an itk image.
:param img_itk:
:param low_threshold:
:param max_value:
:param fix_value:
:return:
"""
img_np = sitk.GetArrayFromImage(img_itk)
img_np_norm = correct_adc_np(img_np, low_threshold, max_value, fix_value)
return copyItkImage(img_itk, img_np_norm)
def readMedSeg2(all_imgs, all_lesion_ctrs, all_lung_ctrs):
data_folder = "/data/UM/COVID/Medicalsegmentation_V2"
for c_id in range(1, 10):
c_mask_file = join(data_folder, "rp_msk", F"{c_id}.nii")
c_img_file = join(data_folder, "rp_im", F"{c_id}.nii")
c_lung_file = join(data_folder, "rp_lung_msk", F"{c_id}.nii")
itk_img = sitk.ReadImage(c_img_file)
itk_mask = sitk.ReadImage(c_mask_file)
itk_lung = sitk.ReadImage(c_lung_file)
np_mask = sitk.GetArrayFromImage(itk_mask)
# np_mask[np_mask == 3] = 0
np_mask[np_mask >= 1] = 1
all_imgs.append(itk_img)
all_lesion_ctrs.append(copyItkImage(itk_img, np_mask))
all_lung_ctrs.append(itk_lung)
print("Done!")
return all_imgs, all_lesion_ctrs, all_lung_ctrs
# -------------- Test example -----------------
# input_image = sitk.ReadImage("/data/UM/COVID/Kaggle_Mosmed/studies/study_0001.nii")
# segmentation = mask.apply(input_image) # default model is U-net(R231)
# viz_obj = MedicalImageVisualizer(disp_images=False, output_folder="/home/olmozavala/Dropbox/MyProjects/UM/Covid19_Prognosis/COVID_DL_Segmentation/LungSegmentationSoftware/OUTPUT")
# viz_obj.plot_img_and_ctrs_np(sitk.GetArrayFromImage(input_image), [segmentation],
# title="Test Lung seg")
# Visualizing Mosmed data
data_folder = "/data/UM/COVID/Kaggle_Mosmed"
output_folder = "/data/UM/COVID/Kaggle_Mosmed/lung_masks"
viz_obj = MedicalImageVisualizer(output_folder=join(data_folder,
"output_imgs"),
disp_images=False)
mask_files = os.listdir(join(data_folder, "masks"))
mask_files.sort()
for c_mask_file_name in mask_files:
c_case = int(c_mask_file_name.split("_")[1])
c_img_file_name = F"study_0{c_case:03d}.nii"
c_img_file = join(data_folder, "studies", c_img_file_name)
print(F"--- Working with {c_img_file} ---")
itk_img = sitk.ReadImage(c_img_file)
np_lung_mask = mask.apply(itk_img) # default model is U-net(R231)
itk_lung_mask = copyItkImage(itk_img, np_lung_mask)
viz_obj.plot_img_and_ctrs_itk(itk_img, [itk_lung_mask],
title=c_img_file_name,
draw_only_ctrs=True,
file_name_prefix=c_img_file_name)
save_image(itk_lung_mask, output_folder, c_mask_file_name)
def make_3d_segmentation(config):
"""
:param config:
:return:
"""
# *********** Reads the parameters ***********
cases = config[ClassificationParams.cases]
save_segmented_ctrs = config[ClassificationParams.save_segmented_ctrs]
input_folder = config[ClassificationParams.input_folder]
input_img_names = config[ClassificationParams.input_img_file_names]
output_folder = config[ClassificationParams.output_folder]
output_imgs_folder = config[ClassificationParams.output_imgs_folder]
output_file_name = config[ClassificationParams.output_file_name]
model_weights_file = config[ClassificationParams.model_weights_file]
compute_metrics = config[ClassificationParams.compute_metrics]
compute_original_resolution = config[
ClassificationParams.compute_original_resolution]
save_imgs = config[ClassificationParams.save_imgs]
if save_imgs:
save_imgs_planes = config[ClassificationParams.save_img_planes]
save_imgs_slices = config[ClassificationParams.save_img_slices]
# Builds the visualization object
viz_obj = MedicalImageVisualizer(
disp_images=config[ClassificationParams.show_imgs],
output_folder=output_imgs_folder)
if compute_metrics:
output_ctr_file_names = config[
ClassificationParams.output_ctr_file_names]
else:
output_ctr_file_names = []
# *********** Chooses the proper model ***********
print('Reading model ....')
model = select_3d_model(config)
# *********** Reads the weights***********
print('Reading weights ....')
model.load_weights(model_weights_file)
examples = select_cases_from_folder(input_folder, cases)
create_folder(output_imgs_folder)
# *********** Makes a dataframe to contain the DSC information **********
metrics_params = config[ClassificationParams.metrics]
metrics_dict = {met.name: met.value for met in metrics_params}
# Check if the output fiels already exist, in thtat case read the df from it.
if os.path.exists(join(output_imgs_folder, output_file_name)):
data = pd.read_csv(join(output_imgs_folder, output_file_name),
index_col=0)
else:
data_columns = list(metrics_dict.values())
if compute_original_resolution:
# In this case we add all the desired metrics, but append 'original at the beginnig'
data_columns = {
*data_columns,
*[F'{ORIGINAL_TXT}_{col}' for col in data_columns]
}
data = DataFrame(index=examples, columns=data_columns)
# *********** Iterates over each case *********
segmentation_type = config[ClassificationParams.segmentation_type]
for id_folder, current_folder in enumerate(examples):
print(F'******* Computing folder {current_folder} ************')
t0 = time.time()
try:
# -------------------- Reading data -------------
print('\t Reading data....')
# All these names are predefined, for any other 3d segmentation we will need to create a different configuration
imgs_itk, ctrs_itk, size_roi, start_roi, _ = read_preproc_imgs_and_ctrs_itk(
input_folder,
folders_to_read=[current_folder],
img_names=input_img_names,
ctr_names=output_ctr_file_names)
imgs_np = [sitk.GetArrayFromImage(c_img) for c_img in imgs_itk[0]
] # The 0 is because we read a single fold
ctrs_np = [sitk.GetArrayFromImage(c_img) for c_img in ctrs_itk[0]]
# If we want to visualize the input images
# viz_obj.plot_imgs_and_ctrs_itk(imgs_itk[0], ctrs_itk=ctrs_itk[0])
# ------------------- Making prediction -----------
print('\t Making prediction....')
input_array = format_for_nn_classification(imgs_np)
output_nn_all = model.predict(input_array, verbose=1)
output_nn_np = output_nn_all[0, :, :, :, 0]
# For visualizing the output of the network
# viz_obj.plot_img_and_ctrs_np(img_np=output_nn_np)
# ------------------- Postprocessing -----------
print('\t Postprocessing prediction....')
threshold = .5
output_nn_itk = copyItkImage(imgs_itk[0][0], output_nn_np)
print(F'\t\t Threshold NN output to {threshold} ....')
output_nn_itk = binaryThresholdImage(output_nn_itk, threshold)
if segmentation_type == SegmentationTypes.PROSTATE or segmentation_type == SegmentationTypes.PZ:
print(
F'\t\t Restricting to largest connected component only ....'
)
output_nn_itk = getLargestConnectedComponents(output_nn_itk)
output_nn_np = sitk.GetArrayViewFromImage(output_nn_itk)
if compute_original_resolution:
print('\t Recovering original resolution...')
print('\t\t Reading original resolution images....')
img_names = [
config[
ClassificationParams.resampled_resolution_image_name],
config[ClassificationParams.original_resolution_image_name]
]
ctr_name = config[
ClassificationParams.original_resolution_ctr_name]
imgs_itk_original_temp, ctrs_itk_original_temp, _, _, _ = read_preproc_imgs_and_ctrs_itk(
input_folder,
folders_to_read=[current_folder],
img_names=img_names,
ctr_names=[ctr_name])
gt_ctr_original_itk = ctrs_itk_original_temp[0][
0] # Retrieves the gt ctr at the original resolution
img_original_resampled_itk = imgs_itk_original_temp[0][0]
img_original_itk = imgs_itk_original_temp[0][1]
print('\t\t Resampling to original....')
output_nn_original_itk = recover_original_resolution(
roi_np=output_nn_np,
resampled_itk=img_original_resampled_itk,
original_itk=img_original_itk,
start_positions=start_roi[0],
size_roi=size_roi[0])
output_nn_original_itk = binaryThresholdImage(
output_nn_original_itk, threshold)
if segmentation_type == SegmentationTypes.PROSTATE or segmentation_type == SegmentationTypes.PZ:
print(
F'\t\t\t Restricting to largest connected component only ....'
)
output_nn_original_itk = getLargestConnectedComponents(
output_nn_original_itk)
output_nn_original_np = sitk.GetArrayViewFromImage(
output_nn_original_itk)
if save_segmented_ctrs:
print('\t Saving Prediction...')
create_folder(join(output_folder, current_folder))
# TODO at some point we will need to see if we can output more than one ctr
sitk.WriteImage(
output_nn_itk,
join(output_folder, current_folder,
output_ctr_file_names[0]))
if compute_original_resolution:
sitk.WriteImage(
output_nn_original_itk,
join(output_folder, current_folder,
F'{ORIGINAL_TXT}_{output_ctr_file_names[0]}'))
if compute_metrics:
# Compute metrics
print('\t Computing metrics....')
for c_metric in metrics_params: # Here we can add more metrics
if c_metric == ClassificationMetrics.DSC_3D:
metric_value = numpy_dice(output_nn_np, ctrs_np[0])
data.loc[current_folder][c_metric.value] = metric_value
print(F'\t\t ----- DSC: {metric_value:.3f} -----')
if compute_original_resolution:
metric_value = numpy_dice(
output_nn_original_np,
sitk.GetArrayViewFromImage(
gt_ctr_original_itk))
data.loc[current_folder][
F'{ORIGINAL_TXT}_{c_metric.value}'] = metric_value
print(F'\t\t ----- DSC: {metric_value:.3f} -----')
# Saving the results every 10 steps
if id_folder % 10 == 0:
save_metrics_images(data,
metric_names=list(
metrics_dict.values()),
viz_obj=viz_obj)
data.to_csv(join(output_folder, output_file_name))
if save_imgs:
print('\t Plotting images....')
plot_intermediate_results(current_folder,
data_columns,
imgs_itk=imgs_itk[0],
gt_ctr_itk=ctrs_itk[0][0],
nn_ctr_itk=output_nn_itk,
data=data,
viz_obj=viz_obj,
slices=save_imgs_slices,
compute_metrics=compute_metrics)
if compute_original_resolution:
plot_intermediate_results(
current_folder,
data_columns,
imgs_itk=[img_original_itk],
gt_ctr_itk=gt_ctr_original_itk,
nn_ctr_itk=output_nn_original_itk,
data=data,
viz_obj=viz_obj,
slices=save_imgs_slices,
compute_metrics=compute_metrics,
prefix_name=ORIGINAL_TXT)
except Exception as e:
print(
"---------------------------- Failed {} error: {} ----------------"
.format(current_folder, e))
print(F'\t Done! Elapsed time {time.time()-t0:0.2f} seg')
if compute_metrics:
save_metrics_images(data,
metric_names=list(metrics_dict.values()),
viz_obj=viz_obj)
data.to_csv(join(output_folder, output_file_name))
c_img_crop.SetSpacing(c_imgs_res[0].GetSpacing())
# Copying metadata to lesion contour and setting its value to 1
c_ctr_crop_np = cropped_ctrs_np[0]
# Making final correction of decimal values
c_ctr_crop_np[c_ctr_crop_np > .001] = 1
tmp_ctr_lesion_np = c_ctr_crop_np.copy(
) # Just a temporal variable used to save the lungs separated from the lesion in 2D
tmp_lung_ctr_crop_np = cropped_ctrs_np[1]
tmp_lung_ctr_crop_np[
tmp_lung_ctr_crop_np >
.001] = 1 # Setiting the two lungs to a value of 1
c_ctr_crop_np += tmp_lung_ctr_crop_np
c_ctr_crop = copyItkImage(c_img_crop, c_ctr_crop_np)
# np_mask = tmp_lung_ctr_crop_np
# midlay = int(np_mask.shape[0]/2)
# plt.imshow(np_mask[midlay, :, :])
# plt.show()
print("Done!")
print("------------ Saving images 3D --------------")
# save_image(c_img_crop, join(output_folder_3d, F"Case-{id:04d}"), "img.nrrd")
# save_image(c_ctr_crop, join(output_folder_3d, F"Case-{id:04d}"), "ctr.nrrd")
print("------------ Saving images 2D --------------")
tot_slices = c_img_crop_np.shape[0]
for c_slice in range(tot_slices):
# Only saving slices where there is a contour
if np.amax(c_ctr_crop_np[c_slice, :, :]) > 1: