Commit bbc3c0fb authored by Wenqi Li's avatar Wenqi Li

Merge branch 'unet-demo' into 'dev'

Unet demo

See merge request CMIC/NiftyNet!272
parents 0039a418 bd2229e7
Pipeline #12143 passed with stages
in 84 minutes and 12 seconds
......@@ -17,6 +17,7 @@ testing_data*
*.csv
*.swp
data/brain_parcellation/OASIS
data/u-net
dist
build
*.egg-info
......@@ -27,3 +28,6 @@ doc/source/niftynet*
doc/source/modules.rst
.vscode/
.ipynb_checkpoints
demos/unet/run_U373.sh
demos/unet/run_hela.sh
[cells]
path_to_search = ./data/u-net/DIC-C2DH-HeLa/niftynet_data
filename_contains = img_
spatial_window_size = (572, 572, 1)
interp_order = 3
loader = skimage
[label]
path_to_search = ./data/u-net/DIC-C2DH-HeLa/niftynet_data
filename_contains = bin_seg_
spatial_window_size = (388, 388, 1)
interp_order = 0
loader = skimage
[xent_weights]
path_to_search = ./data/u-net/DIC-C2DH-HeLa/niftynet_data
filename_contains = weight_
spatial_window_size = (388, 388, 1)
interp_order = 3
loader = skimage
[SYSTEM]
cuda_devices = ""
num_threads = 6
num_gpus = 1
[NETWORK]
name = unet_2d
activation_function = relu
batch_size = 4
# volume level preprocessing
volume_padding_size = (92, 92, 0)
volume_padding_mode = symmetric
whitening = True
normalise_foreground_only=False
queue_length = 20
window_sampling = uniform
[TRAINING]
sample_per_volume = 2
random_flipping_axes=-1
lr = 0.0003
loss_type = CrossEntropy
starting_iter = 0
save_every_n = 200
max_iter = 2000
max_checkpoints = 10
do_elastic_deformation = False
deformation_sigma = 50
num_ctrl_points = 6
proportion_to_deform=0.9
validation_every_n = 10
validation_max_iter = 1
[INFERENCE]
border = (92, 92, 0)
inference_iter = -1
save_seg_dir = ./output
output_interp_order = 0
spatial_window_size = (572,572,1)
############################ custom configuration sections
[SEGMENTATION]
image = cells
label = label
output_prob = False
num_classes = 2
label_normalisation = False
weight = xent_weights
This diff is collapsed.
[cells]
path_to_search = ./data/u-net/PhC-C2DH-U373/niftynet_data
filename_contains = img_
filename_not_contains =
spatial_window_size = (572, 572, 1)
interp_order = 3
loader = skimage
[label]
path_to_search = ./data/u-net/PhC-C2DH-U373/niftynet_data
filename_contains = bin_seg_
filename_not_contains =
spatial_window_size = (388, 388, 1)
interp_order = 0
loader = skimage
[xent_weights]
path_to_search = ./data/u-net/PhC-C2DH-U373/niftynet_data
filename_contains = weight_
filename_not_contains =
spatial_window_size = (388, 388, 1)
interp_order = 3
loader = skimage
[SYSTEM]
cuda_devices = ""
num_threads = 6
num_gpus = 1
[NETWORK]
name = unet_2d
activation_function = relu
batch_size = 4
# volume level preprocessing
volume_padding_size = (92, 92, 0)
volume_padding_mode = symmetric
normalisation = False
whitening = True
normalise_foreground_only=False
queue_length = 20
window_sampling = uniform
[TRAINING]
sample_per_volume = 2
random_flipping_axes=0,1
lr = 0.0003
loss_type = CrossEntropy
starting_iter = 0
save_every_n = 500
max_iter = 10000
max_checkpoints = 15
do_elastic_deformation = True
deformation_sigma = 50
num_ctrl_points = 6
proportion_to_deform=0.9
validation_every_n = 10
validation_max_iter = 1
[INFERENCE]
border = (92, 92, 0)
inference_iter = -1
save_seg_dir = ./output
output_interp_order = 0
spatial_window_size = (572,572,1)
############################ custom configuration sections
[SEGMENTATION]
image = cells
label = label
output_prob = False
num_classes = 2
label_normalisation = False
weight = xent_weights
import os
from shutil import copy
import argparse
def get_user_args():
parser = argparse.ArgumentParser()
parser.add_argument('--file_dir',
default=os.path.join('data', 'u-net'),
help='The directory containing the cell tracking data.',
)
parser.add_argument('--experiment_names',
default=['DIC-C2DH-HeLa', 'PhC-C2DH-U373'],
help='The names of the cell tracking experiments.',
type=list
)
return parser.parse_args()
def main():
args = get_user_args()
# for each specified folder
for experiment_name in args.experiment_names:
out_dir = os.path.join(args.file_dir, experiment_name, 'niftynet_data')
if not os.path.isdir(out_dir):
os.makedirs(out_dir)
for root, _, files in os.walk(os.path.join(args.file_dir, experiment_name)):
for name in [f for f in files if 'track' not in f]:
if 'niftynet_data' not in root: # don't look at the ones that are done already
cell_id = root.split(os.sep)[root.split('/').index(experiment_name) + 1][:2]
out_name = name.replace('t0', 'img_0').replace('t1', 'img_1').replace('man_seg', 'seg_')
out_name = ''.join([out_name.split('.')[0] + '_', cell_id, '.tif'])
out_path = os.path.join(out_dir, out_name)
copy(os.path.join(root, name), out_path)
if __name__ == "__main__":
main()
import itertools
class ExperimentProtocol(object):
def __init__(self, base_command, conditions, model_dir_prefix):
"""
:param base_command: the base command used for NiftyNet
:param conditions: dictionary of experiment conditions,
{"name" : {"--command_line_command": "values (iterable)"} }
e.g. {"elastic": {"--do_elastic_deformation": ["True", "False"]}}
:param model_dir_prefix: prefix for the experimental directory
"""
self.base_command = base_command
self.conditions = conditions
self.model_dir_prefix = model_dir_prefix
self.commands = []
def add_condition(self, new_condition):
for cond in new_condition:
self.conditions[cond] = new_condition[cond]
def to_file(self, file_name):
with open(file_name, 'w') as f:
for line in self.commands:
f.write(line + "\n")
def generate_commands(self):
self.commands = []
combinations = list(dict_product(self.conditions))
for i, combo in enumerate(combinations):
str_command = [self.base_command]
for condition in combo:
str_command += [str(condition), combo[condition]]
str_command += ["--model_dir",
self.model_dir_prefix + '_' + str(i).zfill(len(combinations) // 10 + 1)]
self.commands += [" ".join(str_command)]
def __str__(self):
return "\n".join(self.commands)
def dict_product(dicts):
"""
from https://stackoverflow.com/questions/5228158/cartesian-product-of-a-dictionary-of-lists
"""
return (dict(zip(dicts, x)) for x in itertools.product(*dicts.values()))
def main():
base_command = "python net_segment.py train -c ./demos/unet/U373.ini"
model_dir_prefix = "./models/U373"
# note: these paths are relative to the model directory
d_split_files = ["../../demos/unet/u373_d_split_%i.csv" % i for i in [1, 2]]
conditions = {"--do_elastic_deformation": ["True", "False"],
"--random_flipping_axes": ["'0,1'", "-1"],
"--dataset_split_file": d_split_files}
u373_experiments = ExperimentProtocol(base_command, conditions, model_dir_prefix=model_dir_prefix)
u373_experiments.generate_commands()
u373_experiments.to_file("./run_U373.sh")
base_command = "python net_segment.py train -c ./demos/unet/HeLa.ini"
model_dir_prefix = "./models/HeLa"
# note: these paths are relative to the model directory
d_split_files = ["../../demos/unet/hela_d_split_%i.csv" % i for i in [1, 2]]
conditions = {"--do_elastic_deformation": ["True", "False"],
"--random_flipping_axes": ["'0,1'", "-1"],
"--dataset_split_file": d_split_files}
hela_experiments = ExperimentProtocol(base_command, conditions, model_dir_prefix=model_dir_prefix)
hela_experiments.generate_commands()
hela_experiments.to_file("./run_hela.sh")
if __name__ == "__main__":
main()
import os
import numpy as np
import matplotlib.pyplot as plt # for visualising and debugging
from scipy.ndimage.morphology import distance_transform_edt
from skimage.io import imsave, imread
from skimage.segmentation import find_boundaries
from demos.unet.file_sorter import get_user_args
W_0, SIGMA = 10, 5
def construct_weights_and_mask(img):
seg_boundaries = find_boundaries(img, mode='inner')
bin_img = img > 0
# take segmentations, ignore boundaries
binary_with_borders = np.bitwise_xor(bin_img, seg_boundaries)
foreground_weight = 1 - binary_with_borders.sum() / binary_with_borders.size
background_weight = 1 - foreground_weight
# build euclidean distances maps for each cell:
cell_ids = [x for x in np.unique(img) if x > 0]
distances = np.zeros((img.shape[0], img.shape[1], len(cell_ids)))
for i, cell_id in enumerate(cell_ids):
distances[..., i] = distance_transform_edt(img != cell_id)
# we need to look at the two smallest distances
distances.sort(axis=-1)
weight_map = W_0 * np.exp(-(1 / (2 * SIGMA ** 2)) * ((distances[..., 0] + distances[..., 1]) ** 2))
weight_map[binary_with_borders] = foreground_weight
weight_map[~binary_with_borders] += background_weight
return weight_map, binary_with_borders
def main():
args = get_user_args()
for experiment_name in args.experiment_names:
file_dir = os.path.join(args.file_dir, experiment_name, 'niftynet_data')
for f_name in [f for f in os.listdir(file_dir) if f.startswith('seg') and f.endswith('.tif')]:
img = imread(os.path.join(file_dir, f_name))
weight_map, binary_seg = construct_weights_and_mask(img)
imsave(os.path.join(file_dir, f_name).replace('seg', 'weight'), weight_map)
imsave(os.path.join(file_dir, f_name).replace('seg', 'bin_seg'), binary_seg.astype(np.uint8))
if __name__ == "__main__":
main()
import os
import nibabel as nib
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from skimage.io import imread
import re
def dice_score(gt, img):
"""Calculate the dice score for evaluation purposes"""
gt, img = [x > 0 for x in (gt, img)]
num = 2 * np.sum(gt & img)
den = gt.sum() + img.sum()
return num / den
def results(ground_truths, est_dirs):
"""Collates the dice scores from various experiments"""
result = {e: [] for e in est_dirs}
result['ids'] = []
for f in ground_truths:
r = re.search('.*bin_seg_(.*_\d+)', f)
if r:
gt = imread(f)
subj_id = r.group(1)
result['ids'].append(subj_id)
for exp_name in est_dirs:
est_path = os.path.join(est_dirs[exp_name], subj_id + '_niftynet_out.nii.gz')
est = nib.load(est_path).get_data().squeeze()
result[exp_name].append(dice_score(gt, est))
df = pd.DataFrame(result)
return df
def results_long(df, csv_name):
"""Labels the results as from train or validation datasets"""
d_split = pd.read_csv(csv_name)
d_split.columns = ('ids', 'fold')
merged_df = pd.merge(df, d_split)
df_long = pd.melt(merged_df, id_vars=['ids', 'fold'])
return df_long
def add_experiment_info_to_datasets(df, est_dirs):
"""adds the experimental information from the training settings"""
experiment_numbers, flipping, dataset_splits, deforming = [], [], [], []
for est_dir_key in est_dirs:
# getting the dataset_split file from the settings_train txt file:
train_settings = ' '.join([l.strip() for l in open(est_dirs[est_dir_key] + '../settings_train.txt', 'r')])
experiment_numbers.append(est_dir_key)
r = re.search('dataset_split_file:\s.*(\d).csv', train_settings)
dataset_splits.append(r.group(1))
r = re.search('flipping_axes:\s\((.*?)\)', train_settings)
flip = 'False' if '-1' in r.group(1) else 'True'
flipping.append(flip)
r = re.search('elastic_deformation:\s(\w+)', train_settings)
deforming.append(r.group(1))
data_dict = {'variable': experiment_numbers,
'flip': flipping,
'deform': deforming,
'train_split': dataset_splits,
'augmentations': ['_'.join(['flip', x[0], 'def', y[0]]) for x, y in zip(flipping, deforming)]
}
conditions_df = pd.DataFrame(data_dict)
combined_df = pd.merge(df, conditions_df)
return combined_df
def get_and_plot_results(ground_truths, est_dirs, subj_ids):
df = None
for est_dir_key in est_dirs:
# getting the dataset_split file from the settings_train txt file:
train_settings = [l.strip() for l in open(est_dirs[est_dir_key] + '../settings_train.txt', 'r')]
dataset_split_file = [x.split(':')[1].strip() for x in train_settings if 'dataset_split' in x][0]
new_df = results(ground_truths, {est_dir_key: est_dirs[est_dir_key]})
new_df_long = results_long(new_df, dataset_split_file)
f, axes = plt.subplots(2, 1, figsize=(9, 5))
f.suptitle("Experiment %s" % est_dir_key)
show_model_outputs(ground_truths, new_df_long, {est_dir_key: est_dirs[est_dir_key]}, subj_ids, axes)
if df is None:
df = new_df_long
else:
df = pd.concat([df, new_df_long])
combined_df = add_experiment_info_to_datasets(df, est_dirs)
return combined_df
def show_model_outputs(ground_truths, df, est_dirs, subj_ids, axes):
"""Plots the results for visualisation"""
for est_dir in est_dirs.values():
for i, sid in enumerate(subj_ids):
a = imread([f for f in ground_truths if sid in f][0])
b = nib.load(est_dir + '/' + sid + '_niftynet_out.nii.gz').get_data().squeeze()
axes[i].imshow(np.hstack([a, b, a - b]), cmap='gray')
axes[i].set_axis_off()
train_or_val = df[df['ids'] == sid]['fold'].values[0]
axes[i].set_title('{} Fold: Ground truth, Estimate and Difference. Dice Score = {:.2f}'.format(
train_or_val, dice_score(a, b)))
......@@ -122,7 +122,8 @@ class RegressionApplication(BaseApplication):
if self.net_param.volume_padding_size:
volume_padding_layer.append(PadLayer(
image_name=SUPPORTED_INPUT,
border=self.net_param.volume_padding_size))
border=self.net_param.volume_padding_size,
mode=self.net_param.volume_padding_mode))
for reader in self.readers:
reader.add_preprocessing_layers(volume_padding_layer +
normalisation_layers +
......
......@@ -171,7 +171,9 @@ class SegmentationApplication(BaseApplication):
if self.net_param.volume_padding_size:
volume_padding_layer.append(PadLayer(
image_name=SUPPORTED_INPUT,
border=self.net_param.volume_padding_size))
border=self.net_param.volume_padding_size,
mode=self.net_param.volume_padding_mode
))
# only add augmentation to first reader (not validation reader)
self.readers[0].add_preprocessing_layers(
......
......@@ -17,7 +17,7 @@ class PadLayer(Layer, Invertible):
therefore assumes the input has at least three spatial dims.
"""
def __init__(self, image_name, border, name='pad'):
def __init__(self, image_name, border, name='pad', mode='minimum'):
super(PadLayer, self).__init__(name=name)
try:
spatial_border = tuple(map(lambda x: (x,), border))
......@@ -26,11 +26,12 @@ class PadLayer(Layer, Invertible):
raise
self.border = spatial_border
self.image_name = set(image_name)
self.mode = mode
def layer_op(self, input_image, mask=None):
if not isinstance(input_image, dict):
full_border = match_ndim(self.border, input_image.ndim)
input_image = np.pad(input_image, full_border, mode='minimum')
input_image = np.pad(input_image, full_border, mode=self.mode)
return input_image, mask
for name, image in input_image.items():
......@@ -40,7 +41,7 @@ class PadLayer(Layer, Invertible):
name, self.image_name)
continue
full_border = match_ndim(self.border, image.ndim)
input_image[name] = np.pad(image, full_border, mode='minimum')
input_image[name] = np.pad(image, full_border, mode=self.mode)
return input_image, mask
def inverse_op(self, input_image, mask=None):
......
......@@ -270,6 +270,16 @@ def add_network_args(parser):
type=spatialnumarray,
default=(0, 0, 0))
parser.add_argument(
"--volume_padding_mode",
metavar='',
help="Set which type of numpy padding to do, see "
"https://docs.scipy.org/doc/numpy-1.14.0/"
"reference/generated/numpy.pad.html "
"for details",
type=str,
default='minimum')
parser.add_argument(
"--window_sampling",
metavar='TYPE_STR',
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment