ho11laqe's picture
init
ecf08bc
# 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.
import shutil
from multiprocessing 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 disk
from skimage.morphology import erosion
from skimage.transform import resize
from nnunet.paths import nnUNet_raw_data
def load_bmp_convert_to_nifti_borders_2d(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)[None])
img_itk.SetSpacing(list(spacing)[::-1] + [999])
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_2d(l, spacing, border_thickness)
l[l > 0] = 1
l[borders == 1] = 2
l_itk = sitk.GetImageFromArray(l.astype(np.uint8)[None])
l_itk.SetSpacing(list(spacing)[::-1] + [999])
sitk.WriteImage(l_itk, anno_out)
def generate_disk(spacing, radius, dtype=int):
radius_in_voxels = np.round(radius / np.array(spacing)).astype(int)
n = 2 * radius_in_voxels + 1
disk_iso = disk(max(n) * 2, dtype=np.float64)
disk_resampled = resize(disk_iso, n, 1, 'constant', 0, clip=True, anti_aliasing=False, preserve_range=True)
disk_resampled[disk_resampled > 0.5] = 1
disk_resampled[disk_resampled <= 0.5] = 0
return disk_resampled.astype(dtype)
def generate_border_as_suggested_by_twollmann_2d(label_img: np.ndarray, spacing,
border_thickness: float = 2) -> np.ndarray:
border = np.zeros_like(label_img)
selem = generate_disk(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 prepare_task(base, task_id, task_name, spacing, border_thickness: float = 15):
p = Pool(16)
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_2d,
((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_2d,
((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 convert_to_instance_seg(arr: np.ndarray, spacing: tuple = (0.125, 0.125), small_center_threshold: int = 30,
isolated_border_as_separate_instance_threshold=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))
spacing = np.array(spacing)
while np.sum(remaining_border) > 0:
strel_size = [0, 0]
maximum_dilation = max(dilated_mm)
for i in range(2):
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 = disk(1)
if strel_size[0] == 0: ball_here = ball_here[1:2]
if strel_size[1] == 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(2)]
# 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_convert_to_instance_save(file_in: str, file_out: str, spacing):
img = sitk.ReadImage(file_in)
img_npy = sitk.GetArrayFromImage(img)
out = convert_to_instance_seg(img_npy[0], spacing)[None]
out_itk = sitk.GetImageFromArray(out.astype(np.int16))
out_itk.CopyInformation(img)
sitk.WriteImage(out_itk, file_out)
def convert_folder_to_instanceseg(folder_in: str, folder_out: str, spacing, processes: int = 12):
input_files = subfiles(folder_in, suffix=".nii.gz", join=False)
maybe_mkdir_p(folder_out)
output_files = [join(folder_out, i) for i in input_files]
input_files = [join(folder_in, i) for i in input_files]
p = Pool(processes)
r = []
for i, o in zip(input_files, output_files):
r.append(
p.starmap_async(
load_convert_to_instance_save,
((i, o, spacing),)
)
)
_ = [i.get() for i in r]
p.close()
p.join()
def convert_to_tiff(nifti_image: str, output_name: str):
npy = sitk.GetArrayFromImage(sitk.ReadImage(nifti_image))
imsave(output_name, npy[0].astype(np.uint16), compress=6)
if __name__ == "__main__":
base = "/home/fabian/Downloads/Fluo-N2DH-SIM+"
task_name = 'Fluo-N2DH-SIM'
spacing = (0.125, 0.125)
task_id = 999
border_thickness = 0.7
prepare_task(base, task_id, task_name, spacing, border_thickness)
task_id = 89
additional_time_steps = 4
task_name = 'Fluo-N2DH-SIM_thickborder_time'
full_taskname = 'Task%03.0d_' % task_id + task_name
output_raw = join(nnUNet_raw_data, full_taskname)
shutil.rmtree(output_raw)
shutil.copytree(join(nnUNet_raw_data, 'Task999_Fluo-N2DH-SIM_thickborder'), output_raw)
shutil.rmtree(join(nnUNet_raw_data, 'Task999_Fluo-N2DH-SIM_thickborder'))
# now add additional time information
for fld in ['imagesTr', 'imagesTs']:
curr = join(output_raw, fld)
for seq in ['01', '02']:
images = subfiles(curr, prefix=seq, join=False)
for i in images:
current_timestep = int(i.split('_')[1][1:])
renamed = join(curr, i.replace("_0000", "_%04.0d" % additional_time_steps))
shutil.move(join(curr, i), renamed)
for previous_timestep in range(-additional_time_steps, 0):
# previous time steps will already have been processed and renamed!
expected_filename = join(curr, seq + "_t%03.0d" % (
current_timestep + previous_timestep) + "_%04.0d" % additional_time_steps + ".nii.gz")
if not isfile(expected_filename):
# create empty image
img = sitk.ReadImage(renamed)
empty = sitk.GetImageFromArray(np.zeros_like(sitk.GetArrayFromImage(img)))
empty.CopyInformation(img)
sitk.WriteImage(empty, join(curr, i.replace("_0000", "_%04.0d" % (
additional_time_steps + previous_timestep))))
else:
shutil.copy(expected_filename, join(curr, i.replace("_0000", "_%04.0d" % (
additional_time_steps + previous_timestep))))
dataset = load_json(join(output_raw, 'dataset.json'))
dataset['modality'] = {
'0': 't_minus 4',
'1': 't_minus 3',
'2': 't_minus 2',
'3': 't_minus 1',
'4': 'frame of interest',
}
save_json(dataset, join(output_raw, 'dataset.json'))
# we do not need custom splits since we train on all training cases
# test set predictions are converted to instance seg with convert_folder_to_instanceseg
# test set predictions are converted to tiff with convert_to_tiff