Spaces:

ho11laqe
/

nnUNet_calvingfront_detection

Sleeping

App Files Files Community

nnUNet_calvingfront_detection / nnunet /dataset_conversion /Task076_Fluo_N3DH_SIM.py

ho11laqe

init

ecf08bc about 1 year ago

raw history blame

No virus

12.3 kB

	# Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
	#
	# 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.


	from multiprocessing import Pool
	from multiprocessing.dummy import Pool

	import SimpleITK as sitk
	import numpy as np
	from batchgenerators.utilities.file_and_folder_operations import *
	from skimage.io import imread
	from skimage.io import imsave
	from skimage.morphology import ball
	from skimage.morphology import erosion
	from skimage.transform import resize

	from nnunet.paths import nnUNet_raw_data
	from nnunet.paths import preprocessing_output_dir


	def load_bmp_convert_to_nifti_borders(img_file, lab_file, img_out_base, anno_out, spacing, border_thickness=0.7):
	img = imread(img_file)
	img_itk = sitk.GetImageFromArray(img.astype(np.float32))
	img_itk.SetSpacing(np.array(spacing)[::-1])
	sitk.WriteImage(img_itk, join(img_out_base + "_0000.nii.gz"))

	if lab_file is not None:
	l = imread(lab_file)
	borders = generate_border_as_suggested_by_twollmann(l, spacing, border_thickness)
	l[l > 0] = 1
	l[borders == 1] = 2
	l_itk = sitk.GetImageFromArray(l.astype(np.uint8))
	l_itk.SetSpacing(np.array(spacing)[::-1])
	sitk.WriteImage(l_itk, anno_out)


	def generate_ball(spacing, radius, dtype=int):
	radius_in_voxels = np.round(radius / np.array(spacing)).astype(int)
	n = 2 * radius_in_voxels + 1
	ball_iso = ball(max(n) * 2, dtype=np.float64)
	ball_resampled = resize(ball_iso, n, 1, 'constant', 0, clip=True, anti_aliasing=False, preserve_range=True)
	ball_resampled[ball_resampled > 0.5] = 1
	ball_resampled[ball_resampled <= 0.5] = 0
	return ball_resampled.astype(dtype)


	def generate_border_as_suggested_by_twollmann(label_img: np.ndarray, spacing, border_thickness: float = 2) -> np.ndarray:
	border = np.zeros_like(label_img)
	selem = generate_ball(spacing, border_thickness)
	for l in np.unique(label_img):
	if l == 0: continue
	mask = (label_img == l).astype(int)
	eroded = erosion(mask, selem)
	border[(eroded == 0) & (mask != 0)] = 1
	return border


	def find_differences(labelstr1, labelstr2):
	for n in subfiles(labelstr1, suffix='.nii.gz', join=False):
	a = sitk.GetArrayFromImage(sitk.ReadImage(join(labelstr1, n)))
	b = sitk.GetArrayFromImage(sitk.ReadImage(join(labelstr2, n)))
	print(n, np.sum(a != b))


	def prepare_task(base, task_id, task_name, spacing, border_thickness: float = 15, processes: int = 16):
	p = Pool(processes)

	foldername = "Task%03.0d_%s" % (task_id, task_name)

	out_base = join(nnUNet_raw_data, foldername)
	imagestr = join(out_base, "imagesTr")
	imagests = join(out_base, "imagesTs")
	labelstr = join(out_base, "labelsTr")
	maybe_mkdir_p(imagestr)
	maybe_mkdir_p(imagests)
	maybe_mkdir_p(labelstr)

	train_patient_names = []
	test_patient_names = []
	res = []

	for train_sequence in [i for i in subfolders(base + "_train", join=False) if not i.endswith("_GT")]:
	train_cases = subfiles(join(base + '_train', train_sequence), suffix=".tif", join=False)
	for t in train_cases:
	casename = train_sequence + "_" + t[:-4]
	img_file = join(base + '_train', train_sequence, t)
	lab_file = join(base + '_train', train_sequence + "_GT", "SEG", "man_seg" + t[1:])
	if not isfile(lab_file):
	continue
	img_out_base = join(imagestr, casename)
	anno_out = join(labelstr, casename + ".nii.gz")
	res.append(
	p.starmap_async(load_bmp_convert_to_nifti_borders, ((img_file, lab_file, img_out_base, anno_out, spacing, border_thickness),)))
	train_patient_names.append(casename)

	for test_sequence in [i for i in subfolders(base + "_test", join=False) if not i.endswith("_GT")]:
	test_cases = subfiles(join(base + '_test', test_sequence), suffix=".tif", join=False)
	for t in test_cases:
	casename = test_sequence + "_" + t[:-4]
	img_file = join(base + '_test', test_sequence, t)
	lab_file = None
	img_out_base = join(imagests, casename)
	anno_out = None
	res.append(
	p.starmap_async(load_bmp_convert_to_nifti_borders, ((img_file, lab_file, img_out_base, anno_out, spacing, border_thickness),)))
	test_patient_names.append(casename)

	_ = [i.get() for i in res]

	json_dict = {}
	json_dict['name'] = task_name
	json_dict['description'] = ""
	json_dict['tensorImageSize'] = "4D"
	json_dict['reference'] = ""
	json_dict['licence'] = ""
	json_dict['release'] = "0.0"
	json_dict['modality'] = {
	"0": "BF",
	}
	json_dict['labels'] = {
	"0": "background",
	"1": "cell",
	"2": "border",
	}

	json_dict['numTraining'] = len(train_patient_names)
	json_dict['numTest'] = len(test_patient_names)
	json_dict['training'] = [{'image': "./imagesTr/%s.nii.gz" % i, "label": "./labelsTr/%s.nii.gz" % i} for i in
	train_patient_names]
	json_dict['test'] = ["./imagesTs/%s.nii.gz" % i for i in test_patient_names]

	save_json(json_dict, os.path.join(out_base, "dataset.json"))
	p.close()
	p.join()


	def plot_images(folder, output_folder):
	maybe_mkdir_p(output_folder)
	import matplotlib.pyplot as plt
	for i in subfiles(folder, suffix='.nii.gz', join=False):
	img = sitk.GetArrayFromImage(sitk.ReadImage(join(folder, i)))
	center_slice = img[img.shape[0]//2]
	plt.imsave(join(output_folder, i[:-7] + '.png'), center_slice)


	def convert_to_tiff(nifti_image: str, output_name: str):
	npy = sitk.GetArrayFromImage(sitk.ReadImage(nifti_image))
	imsave(output_name, npy.astype(np.uint16), compress=6)


	def convert_to_instance_seg(arr: np.ndarray, spacing: tuple = (0.2, 0.125, 0.125)):
	from skimage.morphology import label, dilation
	# 1 is core, 2 is border
	objects = label((arr == 1).astype(int))
	final = np.copy(objects)
	remaining_border = arr == 2
	current = np.copy(objects)
	dilated_mm = np.array((0, 0, 0))
	spacing = np.array(spacing)

	while np.sum(remaining_border) > 0:
	strel_size = [0, 0, 0]
	maximum_dilation = max(dilated_mm)
	for i in range(3):
	if spacing[i] == min(spacing):
	strel_size[i] = 1
	continue
	if dilated_mm[i] + spacing[i] / 2 < maximum_dilation:
	strel_size[i] = 1
	ball_here = ball(1)

	if strel_size[0] == 0: ball_here = ball_here[1:2]
	if strel_size[1] == 0: ball_here = ball_here[:, 1:2]
	if strel_size[2] == 0: ball_here = ball_here[:, :, 1:2]

	#print(1)
	dilated = dilation(current, ball_here)
	diff = (current == 0) & (dilated != current)
	final[diff & remaining_border] = dilated[diff & remaining_border]
	remaining_border[diff] = 0
	current = dilated
	dilated_mm = [dilated_mm[i] + spacing[i] if strel_size[i] == 1 else dilated_mm[i] for i in range(3)]
	return final.astype(np.uint32)


	def convert_to_instance_seg2(arr: np.ndarray, spacing: tuple = (0.2, 0.125, 0.125), small_center_threshold=30,
	isolated_border_as_separate_instance_threshold: int = 15):
	from skimage.morphology import label, dilation
	# we first identify centers that are too small and set them to be border. This should remove false positive instances
	objects = label((arr == 1).astype(int))
	for o in np.unique(objects):
	if o > 0 and np.sum(objects == o) <= small_center_threshold:
	arr[objects == o] = 2

	# 1 is core, 2 is border
	objects = label((arr == 1).astype(int))
	final = np.copy(objects)
	remaining_border = arr == 2
	current = np.copy(objects)
	dilated_mm = np.array((0, 0, 0))
	spacing = np.array(spacing)

	while np.sum(remaining_border) > 0:
	strel_size = [0, 0, 0]
	maximum_dilation = max(dilated_mm)
	for i in range(3):
	if spacing[i] == min(spacing):
	strel_size[i] = 1
	continue
	if dilated_mm[i] + spacing[i] / 2 < maximum_dilation:
	strel_size[i] = 1
	ball_here = ball(1)

	if strel_size[0] == 0: ball_here = ball_here[1:2]
	if strel_size[1] == 0: ball_here = ball_here[:, 1:2]
	if strel_size[2] == 0: ball_here = ball_here[:, :, 1:2]

	#print(1)
	dilated = dilation(current, ball_here)
	diff = (current == 0) & (dilated != current)
	final[diff & remaining_border] = dilated[diff & remaining_border]
	remaining_border[diff] = 0
	current = dilated
	dilated_mm = [dilated_mm[i] + spacing[i] if strel_size[i] == 1 else dilated_mm[i] for i in range(3)]

	# what can happen is that a cell is so small that the network only predicted border and no core. This cell will be
	# fused with the nearest other instance, which we don't want. Therefore we identify isolated border predictions and
	# give them a separate instance id
	# we identify isolated border predictions by checking each foreground object in arr and see whether this object
	# also contains label 1
	max_label = np.max(final)

	foreground_objects = label((arr != 0).astype(int))
	for i in np.unique(foreground_objects):
	if i > 0 and (1 not in np.unique(arr[foreground_objects==i])):
	size_of_object = np.sum(foreground_objects==i)
	if size_of_object >= isolated_border_as_separate_instance_threshold:
	final[foreground_objects == i] = max_label + 1
	max_label += 1
	#print('yeah boi')

	return final.astype(np.uint32)


	def load_instanceseg_save(in_file: str, out_file:str, better: bool):
	itk_img = sitk.ReadImage(in_file)
	if not better:
	instanceseg = convert_to_instance_seg(sitk.GetArrayFromImage(itk_img))
	else:
	instanceseg = convert_to_instance_seg2(sitk.GetArrayFromImage(itk_img))
	itk_out = sitk.GetImageFromArray(instanceseg)
	itk_out.CopyInformation(itk_img)
	sitk.WriteImage(itk_out, out_file)


	def convert_all_to_instance(input_folder: str, output_folder: str, processes: int = 24, better: bool = False):
	maybe_mkdir_p(output_folder)
	p = Pool(processes)
	files = subfiles(input_folder, suffix='.nii.gz', join=False)
	output_files = [join(output_folder, i) for i in files]
	input_files = [join(input_folder, i) for i in files]
	better = [better] * len(files)
	r = p.starmap_async(load_instanceseg_save, zip(input_files, output_files, better))
	_ = r.get()
	p.close()
	p.join()


	if __name__ == "__main__":
	base = "/home/fabian/data/Fluo-N3DH-SIM"
	task_id = 76
	task_name = 'Fluo_N3DH_SIM'
	spacing = (0.2, 0.125, 0.125)
	border_thickness = 0.5

	prepare_task(base, task_id, task_name, spacing, border_thickness, 12)

	# we need custom splits
	task_name = "Task076_Fluo_N3DH_SIM"
	labelsTr = join(nnUNet_raw_data, task_name, "labelsTr")
	cases = subfiles(labelsTr, suffix='.nii.gz', join=False)
	splits = []
	splits.append(
	{'train': [i[:-7] for i in cases if i.startswith('01_')],
	'val': [i[:-7] for i in cases if i.startswith('02_')]}
	)
	splits.append(
	{'train': [i[:-7] for i in cases if i.startswith('02_')],
	'val': [i[:-7] for i in cases if i.startswith('01_')]}
	)

	maybe_mkdir_p(join(preprocessing_output_dir, task_name))

	save_pickle(splits, join(preprocessing_output_dir, task_name, "splits_final.pkl"))

	# test set was converted to instance seg with convert_all_to_instance with better=True

	# convert to tiff with convert_to_tiff