def condense_and_split_components(components, output_shape, malis_neighborhood):
'''
:param components: numpy array with component values
:param output_shape: tuple with spatial dimensions of components
:param malis_neighborhood: array definition of malis neighborhood
:return: numpy array of same shape as components, with new component values
'''
original_shape = components.shape
components_for_malis = components.reshape(output_shape)
affinities = malis.seg_to_affgraph(components_for_malis, malis_neighborhood)
recomputed_components, _ = malis.connected_components_affgraph(affinities.astype(np.int32), malis_neighborhood)
recomputed_components = recomputed_components.reshape(original_shape)
return recomputed_components
def condense_and_split_components(components, output_shape,
malis_neighborhood):
'''
:param components: numpy array with component values
:param output_shape: tuple with spatial dimensions of components
:param malis_neighborhood: array definition of malis neighborhood
:return: numpy array of same shape as components, with new component values
'''
original_shape = components.shape
components_for_malis = components.reshape(output_shape)
affinities = malis.seg_to_affgraph(components_for_malis,
malis_neighborhood)
recomputed_components, _ = malis.connected_components_affgraph(
affinities.astype(np.int32), malis_neighborhood)
recomputed_components = recomputed_components.reshape(original_shape)
return recomputed_components
def ignore_disconnected_components(cells):
global ignore_label
if ignore_label is None:
ignore_label = int(cells.max() + 1)
print("Relabelling connected components...")
simple_neighborhood = malis.mknhood3d()
affinities = malis.seg_to_affgraph(cells, simple_neighborhood)
relabelled, _ = malis.connected_components_affgraph(
affinities, simple_neighborhood)
print("Creating overlay...")
overlay = np.array([cells.flatten(), relabelled.flatten()])
print("Finding unique pairs...")
matches = np.unique(overlay, axis=1)
print("Finding disconnected labels...")
orig_to_new = {}
disconnected = set()
for orig_id, new_id in zip(matches[0], matches[1]):
if orig_id == 0 or new_id == 0:
continue
if orig_id not in orig_to_new:
orig_to_new[orig_id] = [new_id]
else:
orig_to_new[orig_id].append(new_id)
disconnected.add(orig_id)
print("Masking %d disconnected labels..." % len(disconnected))
ignore_mask = replace(cells, np.array([l for l in disconnected]),
np.array([ignore_label], dtype=np.uint64))
ignore_mask = (ignore_mask == ignore_label).astype(np.uint8)
print("done.")
return ignore_mask
#hdf5_raw_file = 'zebrafish_friedrich/raw.hdf5'
#hdf5_gt_file = 'zebrafish_friedrich/labels_2.hdf5'
# hdf5_raw = h5py.File(hdf5_raw_file, 'r')
h5seg = h5py.File(hdf5_gt_file, 'r')
# hdf5_aff = h5py.File(hdf5_aff_file, 'r')
seg = np.asarray(h5seg['main']).astype(np.int32)
print("[" + str(datetime.datetime.now()) + "]" + "Making affinity graph...")
aff = m.seg_to_affgraph(seg, nhood)
print("[" + str(datetime.datetime.now()) + "]" + "Affinity shape:" +
str(aff.shape))
print("[" + str(datetime.datetime.now()) + "]" +
"Computing connected components...")
cc, ccSizes = m.connected_components_affgraph(aff, nhood)
print("[" + str(datetime.datetime.now()) + "]" +
"Making affinity graph again...")
aff2 = m.seg_to_affgraph(cc, nhood)
print("[" + str(datetime.datetime.now()) + "]" +
"Computing connected components...")
cc2, ccSizes2 = m.connected_components_affgraph(aff2, nhood)
print("[" + str(datetime.datetime.now()) + "]" + "Comparing 'seg' and 'cc':")
# ri,fscore,prec,rec = m.rand_index(seg,cc)
# print("\tRand index: %f, fscore: %f, prec: %f, rec: %f" % (ri,fscore,prec,rec)
V_rand, V_rand_split, V_rand_merge = m.compute_V_rand_N2(seg, cc)
print("[" + str(datetime.datetime.now()) + "]" +
"\tV_rand: %f, V_rand_split: %f, V_rand_merge: %f" %
(V_rand, V_rand_split, V_rand_merge))
def train(solver, test_net, data_arrays, train_data_arrays, options):
global data_slices, label_slices, data_offsets
caffe.select_device(options.train_device, False)
net = solver.net
test_eval = None
if (options.test_net != None):
test_eval = TestNetEvaluator(test_net, net, train_data_arrays, options)
input_dims, output_dims, input_padding = get_spatial_io_dims(net)
fmaps_in, fmaps_out = get_fmap_io_dims(net)
dims = len(output_dims)
losses = []
shapes = []
# Raw data slice input (n = 1, f = 1, spatial dims)
shapes += [[1, fmaps_in] + input_dims]
# Label data slice input (n = 1, f = #edges, spatial dims)
shapes += [[1, fmaps_out] + output_dims]
if (options.loss_function == 'malis'):
# Connected components input (n = 1, f = 1, spatial dims)
shapes += [[1, 1] + output_dims]
if (options.loss_function == 'euclid'):
# Error scale input (n = 1, f = #edges, spatial dims)
shapes += [[1, fmaps_out] + output_dims]
# Nhood specifications (n = #edges, f = 3)
if (('nhood' in data_arrays[0]) and (options.loss_function == 'malis')):
shapes += [[1, 1] + list(np.shape(data_arrays[0]['nhood']))]
net_io = NetInputWrapper(net, shapes)
data_sizes = [fmaps_in] + [
output_dims[di] + input_padding[di] for di in range(0, dims)
]
label_sizes = [fmaps_out] + output_dims
# Begin the generation of the training set
training_set = TrainingSetGenerator(data_arrays, options, data_sizes,
label_sizes, input_padding)
training_set.generate_training()
# Loop from current iteration to last iteration
for i in range(solver.iter, solver.max_iter):
if (options.test_net != None and i % options.test_interval == 0):
test_eval.evaluate(i)
# First pick the dataset to train with
dataset = randint(0, len(data_arrays) - 1)
offsets = []
for j in range(0, dims):
offsets.append(
randint(
0, data_arrays[dataset]['data'].shape[1 + j] -
(output_dims[j] + input_padding[j])))
# These are the raw data elements
data_slice_old = slice_data(data_arrays[dataset]['data'],
[0] + offsets, data_sizes)
data_slice = data_slices.get()
# print "Compare sizes of data_slices: {0} and {1}".format(data_slice_old.shape, data_slice.shape)
label_slice = None
components_slice = None
if (options.training_method == 'affinity'):
if ('label' in data_arrays[dataset]):
label_slice_old = slice_data(
data_arrays[dataset]['label'], [0] + [
offsets[di] +
int(math.ceil(input_padding[di] / float(2)))
for di in range(0, dims)
], label_sizes)
label_slice = label_slices.get()
# print "Compare sizes of label_slices: {0} and {1}".format(label_slice_old.shape, label_slice.shape)
# TODO: Not sure about what to do for components_slice
if ('components' in data_arrays[dataset]):
data_offset = data_offsets.get()
components_slice = slice_data(
data_arrays[dataset]['components'][0, :], [
data_offset[di] +
int(math.ceil(input_padding[di] / float(2)))
for di in range(0, dims)
], output_dims)
if (label_slice is None or options.recompute_affinity):
label_slice = malis.seg_to_affgraph(
components_slice,
data_arrays[dataset]['nhood']).astype(float32)
if (components_slice is None or options.recompute_affinity):
components_slice, ccSizes = malis.connected_components_affgraph(
label_slice.astype(int32), data_arrays[dataset]['nhood'])
else:
label_slice_old = slice_data(data_arrays[dataset]['label'], [0] + [
offsets[di] + int(math.ceil(input_padding[di] / float(2)))
for di in range(0, dims)
], [fmaps_out] + output_dims)
label_slice = label_slices.get()
if options.loss_function == 'malis':
# Also recomputing the corresponding labels (connected components)
net_io.setInputs([
data_slice, label_slice, components_slice,
data_arrays[0]['nhood']
])
if options.loss_function == 'euclid':
if (options.scale_error == True):
frac_pos = np.clip(label_slice.mean(), 0.05, 0.95)
w_pos = 1.0 / (2.0 * frac_pos)
w_neg = 1.0 / (2.0 * (1.0 - frac_pos))
else:
w_pos = 1
w_neg = 1
net_io.setInputs([
data_slice, label_slice,
error_scale(label_slice, w_neg, w_pos)
])
if options.loss_function == 'softmax':
# These are the affinity edge values
net_io.setInputs([data_slice, label_slice])
# Single step
loss = solver.step(1)
# sanity_check_net_blobs(net)
while gc.collect():
pass
if options.loss_function == 'euclid' or options.loss_function == 'euclid_aniso':
print("[Iter %i] Loss: %f, frac_pos=%f, w_pos=%f" %
(i, loss, frac_pos, w_pos))
else:
print("[Iter %i] Loss: %f" % (i, loss))
# TODO: Store losses to file
losses += [loss]
def train(solver, test_net, data_arrays, train_data_arrays, options):
global data_slices, label_slices, data_offsets
caffe.select_device(options.train_device, False)
net = solver.net
test_eval = None
if (options.test_net != None):
test_eval = TestNetEvaluator(test_net, net, train_data_arrays, options)
input_dims, output_dims, input_padding = get_spatial_io_dims(net)
fmaps_in, fmaps_out = get_fmap_io_dims(net)
dims = len(output_dims)
losses = []
shapes = []
# Raw data slice input (n = 1, f = 1, spatial dims)
shapes += [[1,fmaps_in] + input_dims]
# Label data slice input (n = 1, f = #edges, spatial dims)
shapes += [[1,fmaps_out] + output_dims]
if (options.loss_function == 'malis'):
# Connected components input (n = 1, f = 1, spatial dims)
shapes += [[1,1] + output_dims]
if (options.loss_function == 'euclid'):
# Error scale input (n = 1, f = #edges, spatial dims)
shapes += [[1,fmaps_out] + output_dims]
# Nhood specifications (n = #edges, f = 3)
if (('nhood' in data_arrays[0]) and (options.loss_function == 'malis')):
shapes += [[1,1] + list(np.shape(data_arrays[0]['nhood']))]
net_io = NetInputWrapper(net, shapes)
data_sizes = [fmaps_in]+[output_dims[di] + input_padding[di] for di in range(0, dims)]
label_sizes = [fmaps_out] + output_dims
# Begin the generation of the training set
training_set = TrainingSetGenerator(data_arrays, options, data_sizes, label_sizes, input_padding)
training_set.generate_training()
# Loop from current iteration to last iteration
for i in range(solver.iter, solver.max_iter):
if (options.test_net != None and i % options.test_interval == 0):
test_eval.evaluate(i)
# First pick the dataset to train with
dataset = randint(0, len(data_arrays) - 1)
offsets = []
for j in range(0, dims):
offsets.append(randint(0, data_arrays[dataset]['data'].shape[1+j] - (output_dims[j] + input_padding[j])))
# These are the raw data elements
data_slice_old = slice_data(data_arrays[dataset]['data'], [0]+offsets, data_sizes)
data_slice = data_slices.get()
# print "Compare sizes of data_slices: {0} and {1}".format(data_slice_old.shape, data_slice.shape)
label_slice = None
components_slice = None
if (options.training_method == 'affinity'):
if ('label' in data_arrays[dataset]):
label_slice_old = slice_data(data_arrays[dataset]['label'], [0] + [offsets[di] + int(math.ceil(input_padding[di] / float(2))) for di in range(0, dims)], label_sizes)
label_slice = label_slices.get()
# print "Compare sizes of label_slices: {0} and {1}".format(label_slice_old.shape, label_slice.shape)
# TODO: Not sure about what to do for components_slice
if ('components' in data_arrays[dataset]):
data_offset = data_offsets.get()
components_slice = slice_data(data_arrays[dataset]['components'][0,:], [data_offset[di] + int(math.ceil(input_padding[di] / float(2))) for di in range(0, dims)], output_dims)
if (label_slice is None or options.recompute_affinity):
label_slice = malis.seg_to_affgraph(components_slice, data_arrays[dataset]['nhood']).astype(float32)
if (components_slice is None or options.recompute_affinity):
components_slice,ccSizes = malis.connected_components_affgraph(label_slice.astype(int32), data_arrays[dataset]['nhood'])
else:
label_slice_old = slice_data(data_arrays[dataset]['label'], [0] + [offsets[di] + int(math.ceil(input_padding[di] / float(2))) for di in range(0, dims)], [fmaps_out] + output_dims)
label_slice = label_slices.get()
if options.loss_function == 'malis':
# Also recomputing the corresponding labels (connected components)
net_io.setInputs([data_slice, label_slice, components_slice, data_arrays[0]['nhood']])
if options.loss_function == 'euclid':
if(options.scale_error == True):
frac_pos = np.clip(label_slice.mean(),0.05,0.95)
w_pos = 1.0/(2.0*frac_pos)
w_neg = 1.0/(2.0*(1.0-frac_pos))
else:
w_pos = 1
w_neg = 1
net_io.setInputs([data_slice, label_slice, error_scale(label_slice,w_neg,w_pos)])
if options.loss_function == 'softmax':
# These are the affinity edge values
net_io.setInputs([data_slice, label_slice])
# Single step
loss = solver.step(1)
# sanity_check_net_blobs(net)
while gc.collect():
pass
if options.loss_function == 'euclid' or options.loss_function == 'euclid_aniso':
print("[Iter %i] Loss: %f, frac_pos=%f, w_pos=%f" % (i,loss,frac_pos,w_pos))
else:
print("[Iter %i] Loss: %f" % (i,loss))
# TODO: Store losses to file
losses += [loss]
def np_aff_to_seg(aff, nhood=malis.mknhood3d(1), threshold=np.array([0.5])):
aff = np.transpose(aff, [3, 0, 1, 2]) # zyx3 to 3zyx
ret = malis.connected_components_affgraph(
(aff > threshold[0]).astype(np.int32), nhood)[0].astype(np.int32)
ret = skimage.measure.label(ret).astype(np.float32)
return ret
def train(solver, test_net, data_arrays, train_data_arrays, options):
caffe.select_device(options.train_device, False)
net = solver.net
test_eval = None
if (options.test_net != None):
test_eval = TestNetEvaluator(test_net, net, train_data_arrays, options)
input_dims, output_dims, input_padding = get_spatial_io_dims(net)
fmaps_in, fmaps_out = get_fmap_io_dims(net)
dims = len(output_dims)
losses = []
shapes = []
# Raw data slice input (n = 1, f = 1, spatial dims)
shapes += [[1,fmaps_in] + input_dims]
# Label data slice input (n = 1, f = #edges, spatial dims)
shapes += [[1,fmaps_out] + output_dims]
if (options.loss_function == 'malis'):
# Connected components input (n = 1, f = 1, spatial dims)
shapes += [[1,1] + output_dims]
if (options.loss_function == 'euclid'):
# Error scale input (n = 1, f = #edges, spatial dims)
shapes += [[1,fmaps_out] + output_dims]
# Nhood specifications (n = #edges, f = 3)
if (('nhood' in data_arrays[0]) and (options.loss_function == 'malis')):
shapes += [[1,1] + list(np.shape(data_arrays[0]['nhood']))]
net_io = NetInputWrapper(net, shapes)
make_dataset_offset = MakeDatasetOffset(input_dims, output_dims)
if data_io.data_loader_should_be_used_with(data_arrays):
using_data_loader = True
# and initialize queue!
loader_size = 20
n_workers = 10
make_dataset_offset = MakeDatasetOffset(dims, output_dims, input_padding)
loader_kwargs = dict(
size=loader_size,
datasets=data_arrays,
input_shape=tuple(input_dims),
output_shape=tuple(output_dims),
n_workers=n_workers,
dataset_offset_func=make_dataset_offset
)
print("creating queue with kwargs {}".format(loader_kwargs))
training_data_loader = data_io.DataLoader(**loader_kwargs)
# start populating the queue
for i in range(loader_size):
if DEBUG:
print("Pre-populating data loader's dataset #{i}/{size}"
.format(i=i, size=training_data_loader.size))
shared_dataset_index, async_result = \
training_data_loader.start_refreshing_shared_dataset(i)
else:
using_data_loader = False
# Loop from current iteration to last iteration
for i in range(solver.iter, solver.max_iter):
start = time.time()
if (options.test_net != None and i % options.test_interval == 1):
test_eval.evaluate(i)
if USE_ONE_THREAD:
# after testing finishes, switch back to the training device
caffe.select_device(options.train_device, False)
if not using_data_loader:
dataset_index, offsets = make_dataset_offset(data_arrays)
dataset = data_arrays[dataset_index]
# These are the raw data elements
data_slice = data_io.util.get_zero_padded_slice_from_array_by_offset(
array=dataset['data'],
origin=[0] + offsets,
shape=[fmaps_in] + input_dims)
label_slice = slice_data(dataset['label'], [0] + [offsets[di] + int(math.ceil(input_padding[di] / float(2))) for di in range(0, dims)], [fmaps_out] + output_dims)
if 'transform' in dataset:
# transform the input
# assumes that the original input pixel values are scaled between (0,1)
if DEBUG:
print("data_slice stats, pre-transform: min", data_slice.min(), "mean", data_slice.mean(),
"max", data_slice.max())
lo, hi = dataset['transform']['scale']
data_slice = 0.5 + (data_slice - 0.5) * np.random.uniform(low=lo, high=hi)
lo, hi = dataset['transform']['shift']
data_slice = data_slice + np.random.uniform(low=lo, high=hi)
else:
dataset, index_of_shared_dataset = training_data_loader.get_dataset()
data_slice = dataset['data']
assert data_slice.shape == (fmaps_in,) + tuple(input_dims)
label_slice = dataset['label']
assert label_slice.shape == (fmaps_out,) + tuple(output_dims)
if DEBUG:
print("Training with next dataset in data loader, which has offset", dataset['offset'])
mask_slice = None
if 'mask' in dataset:
mask_slice = dataset['mask']
if DEBUG:
print("data_slice stats: min", data_slice.min(), "mean", data_slice.mean(), "max", data_slice.max())
if options.loss_function == 'malis':
components_slice, ccSizes = malis.connected_components_affgraph(label_slice.astype(int32), dataset['nhood'])
# Also recomputing the corresponding labels (connected components)
net_io.setInputs([data_slice, label_slice, components_slice, data_arrays[0]['nhood']])
elif options.loss_function == 'euclid':
label_slice_mean = label_slice.mean()
if 'mask' in dataset:
label_slice = label_slice * mask_slice
label_slice_mean = label_slice.mean() / mask_slice.mean()
w_pos = 1.0
w_neg = 1.0
if options.scale_error:
frac_pos = np.clip(label_slice_mean, 0.05, 0.95)
w_pos = w_pos / (2.0 * frac_pos)
w_neg = w_neg / (2.0 * (1.0 - frac_pos))
error_scale_slice = scale_errors(label_slice, w_neg, w_pos)
net_io.setInputs([data_slice, label_slice, error_scale_slice])
elif options.loss_function == 'softmax':
# These are the affinity edge values
net_io.setInputs([data_slice, label_slice])
loss = solver.step(1) # Single step
if using_data_loader:
training_data_loader.start_refreshing_shared_dataset(index_of_shared_dataset)
while gc.collect():
pass
time_of_iteration = time.time() - start
if options.loss_function == 'euclid' or options.loss_function == 'euclid_aniso':
print("[Iter %i] Time: %05.2fs Loss: %f, frac_pos=%f, w_pos=%f" % (i, time_of_iteration, loss, frac_pos, w_pos))
else:
print("[Iter %i] Time: %05.2fs Loss: %f" % (i, time_of_iteration, loss))
losses += [loss]
if hasattr(options, 'loss_snapshot') and ((i % options.loss_snapshot) == 0):
io.savemat('loss.mat',{'loss':losses})
if using_data_loader:
training_data_loader.destroy()
#hdf5_raw_file = 'zebrafish_friedrich/raw.hdf5' #hdf5_gt_file = 'zebrafish_friedrich/labels_2.hdf5' # hdf5_raw = h5py.File(hdf5_raw_file, 'r') h5seg = h5py.File(hdf5_gt_file, 'r') # hdf5_aff = h5py.File(hdf5_aff_file, 'r') seg = np.asarray(h5seg['main']).astype(np.int32) print "[" +str(datetime.datetime.now())+"]" + "Making affinity graph..." aff = m.seg_to_affgraph(seg,nhood) print "[" +str(datetime.datetime.now())+"]" + "Affinity shape:" + str(aff.shape) print "[" +str(datetime.datetime.now())+"]" + "Computing connected components..." cc,ccSizes = m.connected_components_affgraph(aff,nhood) print "[" +str(datetime.datetime.now())+"]" + "Making affinity graph again..." aff2 = m.seg_to_affgraph(cc,nhood) print "[" +str(datetime.datetime.now())+"]" + "Computing connected components..." cc2,ccSizes2 = m.connected_components_affgraph(aff2,nhood) print "[" +str(datetime.datetime.now())+"]" + "Comparing 'seg' and 'cc':" # ri,fscore,prec,rec = m.rand_index(seg,cc) # print "\tRand index: %f, fscore: %f, prec: %f, rec: %f" % (ri,fscore,prec,rec) V_rand,V_rand_split,V_rand_merge = m.compute_V_rand_N2(seg,cc) print "[" +str(datetime.datetime.now())+"]" + "\tV_rand: %f, V_rand_split: %f, V_rand_merge: %f" % (V_rand,V_rand_split,V_rand_merge) print "[" +str(datetime.datetime.now())+"]" + "Comparing 'cc' and 'cc2':" # ri,fscore,prec,rec = m.rand_index(cc,cc2) # print "\tRand index: %f, fscore: %f, prec: %f, rec: %f" % (ri,fscore,prec,rec) V_rand,V_rand_split,V_rand_merge = m.compute_V_rand_N2(cc,cc2)
def get_outputs(original_dataset, output_slice):
output_shape = tuple(
[slice_.stop - slice_.start for slice_ in output_slice])
n_spatial_dimensions = len(output_slice)
components_shape = (1, ) + output_shape
mask_shape = (1, ) + output_shape
affinities_shape = (n_spatial_dimensions, ) + output_shape
component_slices = [
slice(0, l) for l in original_dataset['components'].shape
]
component_slices[-n_spatial_dimensions:] = output_slice
logger.debug("component_slices: {}".format(component_slices))
components_array = get_zero_padded_array_slice(
original_dataset['components'], component_slices)
source_class = type(original_dataset['components'])
components_are_from_dvid = source_class in dvid_classes
exclude_strings = original_dataset.get('body_names_to_exclude', [])
if exclude_strings and components_are_from_dvid:
dvid_uuid = original_dataset['components'].uuid
components_to_keep = get_good_components(dvid_uuid, exclude_strings)
logger.debug("components before: {}".format(
list(np.unique(components_array))))
components_array = replace_array_except_whitelist(
components_array, 0, components_to_keep)
logger.debug("components after: {}".format(
list(np.unique(components_array))))
minimum_component_size = original_dataset.get('minimum_component_size', 0)
if minimum_component_size > 0:
components_array = replace_infrequent_values(components_array,
minimum_component_size, 0)
component_erosion_steps = original_dataset.get('component_erosion_steps',
0)
if component_erosion_steps > 0:
components_array = erode_value_blobs(components_array,
steps=component_erosion_steps,
values_to_ignore=(0, ))
components_for_malis = components_array.reshape(output_shape)
affinities_from_components = malis.seg_to_affgraph(
components_for_malis, original_dataset['nhood'])
components_array, _ = malis.connected_components_affgraph(
affinities_from_components, original_dataset['nhood'])
components_array = shift_up_component_values(components_array)
components_array = components_array.reshape(components_shape)
if 'label' in original_dataset:
label_shape = original_dataset['label'].shape
label_slices = [slice(0, l) for l in label_shape]
label_slices[-n_spatial_dimensions:] = output_slice
affinities_array = get_zero_padded_array_slice(
original_dataset['label'], label_slices)
else:
# compute affinities from components
logger.debug(
"Computing affinity labels from components because 'label' wasn't provided in data source."
)
affinities_array = affinities_from_components
assert affinities_array.shape == affinities_shape, \
"affinities_array.shape is {actual} but should be {desired}".format(
actual=str(affinities_array.shape), desired=str(affinities_shape))
if 'mask' in original_dataset:
mask_array = get_zero_padded_array_slice(original_dataset['mask'],
output_slice)
else:
if components_are_from_dvid:
# infer mask values: 1 if component is nonzero, 0 otherwise
mask_array = np.not_equal(components_array, 0)
logger.debug(
"No mask provided. Setting to 1 where components != 0.")
else:
# assume no masking
mask_array = np.ones_like(components_array, dtype=np.uint8)
logger.debug("No mask provided. Setting to 1 where outputs exist.")
mask_dilation_steps = original_dataset.get('mask_dilation_steps', 0)
if mask_dilation_steps > 0:
mask_array = ndimage.binary_dilation(mask_array,
iterations=mask_dilation_steps)
mask_array = mask_array.astype(np.uint8)
mask_array = mask_array.reshape(mask_shape)
return components_array, affinities_array, mask_array
def train(solver,
test_net,
data_arrays,
train_data_arrays,
options,
random_blacks=False):
caffe.select_device(options.train_device, False)
net = solver.net
#pdb.set_trace()
clwt = None
test_eval = None
if options.scale_error == 2:
clwt = ClassWeight(data_arrays, solver.iter)
test_eval = TestNetEvaluator(test_net, net, data_arrays, options)
test_eval2 = None
if (options.test_net != None):
test_eval2 = TestNetEvaluator(test_net, net, train_data_arrays,
options)
input_dims, output_dims, input_padding = get_spatial_io_dims(net)
fmaps_in, fmaps_out = get_fmap_io_dims(net)
dims = len(output_dims)
losses = []
shapes = []
# Raw data slice input (n = 1, f = 1, spatial dims)
shapes += [[1, fmaps_in] + input_dims]
# Label data slice input (n = 1, f = #edges, spatial dims)
shapes += [[1, fmaps_out] + output_dims]
if (options.loss_function == 'malis'):
# Connected components input (n = 1, f = 1, spatial dims)
shapes += [[1, 1] + output_dims]
if (options.loss_function == 'euclid'):
# Error scale input (n = 1, f = #edges, spatial dims)
shapes += [[1, fmaps_out] + output_dims]
# Nhood specifications (n = #edges, f = 3)
if (('nhood' in data_arrays[0]) and (options.loss_function == 'malis')):
shapes += [[1, 1] + list(np.shape(data_arrays[0]['nhood']))]
net_io = NetInputWrapper(net, shapes)
weight_vec = []
if (options.loss_function == 'softmax'
or options.loss_function == 'euclid') and options.scale_error == 1:
#pdb.set_trace()
weight_vec = class_balance_distribution(data_arrays[0]['label'])
#weight_vec[2] = weight_vec[1]*4.0 #for 3 class, inversed during weighting
#pdb.set_trace()
# Loop from current iteration to last iteration
for i in range(solver.iter, solver.max_iter):
if (options.test_net != None and i % options.test_interval == 0
and i > 1):
#pdb.set_trace()
test_eval2.evaluate(i)
if options.scale_error == 2:
test_eval.evaluate(i)
clwt.recompute_weight(test_eval.pred_arrays_samesize, i)
# First pick the dataset to train with
dataset = randint(0, len(data_arrays) - 1)
if dims == 3:
offsets = []
for j in range(0, dims):
offsets.append(
randint(
0, data_arrays[dataset]['data'].shape[1 + j] -
(output_dims[j] + input_padding[j])))
# These are the raw data elements
#pdb.set_trace()
data_slice = slice_data(
data_arrays[dataset]['data'], [0] + offsets, [fmaps_in] +
[output_dims[di] + input_padding[di] for di in range(0, dims)])
label_slice = slice_data(data_arrays[dataset]['label'], [0] + [
offsets[di] + int(math.ceil(input_padding[di] / float(2)))
for di in range(0, dims)
], [fmaps_out] + output_dims)
if options.scale_error == 2 and clwt != None:
weight_slice = slice_data(clwt.class_weights[dataset], [0] + [
offsets[di] + int(math.ceil(input_padding[di] / float(2)))
for di in range(0, dims)
], [fmaps_out] + output_dims)
elif dims == 2:
offsets = []
offsets.append(
randint(0, data_arrays[dataset]['data'].shape[1] - 1))
for j in range(0, dims):
offsets.append(
randint(
0, data_arrays[dataset]['data'].shape[2 + j] -
(output_dims[j] + input_padding[j])))
# These are the raw data elements
#pdb.set_trace()
data_slice = slice_data(
data_arrays[dataset]['data'], [0] + offsets, [fmaps_in, 1] +
[output_dims[di] + input_padding[di] for di in range(0, dims)])
label_slice = slice_data(
data_arrays[dataset]['label'], [0, offsets[0]] + [
offsets[di + 1] +
int(math.ceil(input_padding[di] / float(2)))
for di in range(0, dims)
], [fmaps_out, 1] + output_dims)
data_slice = np.squeeze(data_slice)
label_slice = np.squeeze(label_slice)
#offsets=np.zeros(dims);
if (data_slice.shape[0] < 1) or (label_slice.shape[0] < 2):
pp = 1
pdb.set_trace()
#pdb.set_trace()
#if(np.unique(label_slice).shape[0]<2):
# continue;
# transform the input
# this code assumes that the original input pixel values are scaled between (0,1)
if 'transform' in data_arrays[dataset]:
# print('Pre:',(data_slice.min(),data_slice.mean(),data_slice.max()))
data_slice_mean = data_slice.mean()
lo, hi = data_arrays[dataset]['transform']['scale']
data_slice = data_slice_mean + (
data_slice - data_slice_mean) * np.random.uniform(low=lo,
high=hi)
lo, hi = data_arrays[dataset]['transform']['shift']
data_slice = data_slice + np.random.uniform(low=lo, high=hi)
# print('Post:',(data_slice.min(),data_slice.mean(),data_slice.max()))
data_slice = np.clip(data_slice, 0.0, 0.95)
#pdb.set_trace()
if random_blacks and np.random.random() < 0.25:
data_slice[randint(0, data_slice.shape[0] - 1), ...] = 0.0
if options.loss_function == 'malis':
components_slice, ccSizes = malis.connected_components_affgraph(
label_slice.astype(int32), data_arrays[dataset]['nhood'])
# Also recomputing the corresponding labels (connected components)
net_io.setInputs([
data_slice, label_slice, components_slice,
data_arrays[0]['nhood']
])
if options.loss_function == 'euclid':
###if(options.scale_error == True):
###frac_pos = np.clip(label_slice.mean(),0.05,0.95) #for binary labels
###w_pos = 1.0/(2.0*frac_pos)
###w_neg = 1.0/(2.0*(1.0-frac_pos))
###else:
###w_pos = 1
###w_neg = 1
###net_io.setInputs([data_slice, label_slice, error_scale(label_slice,w_neg,w_pos)])
if (options.scale_error == 3):
frac_pos = np.clip(label_slice.mean(), 0.01,
0.99) #for binary labels
w_pos = 1.0 / (2.0 * frac_pos)
w_neg = 1.0 / (2.0 * (1.0 - frac_pos))
net_io.setInputs([
data_slice, label_slice,
error_scale(label_slice, w_neg, w_pos)
])
elif (options.scale_error == 1):
frac_pos = weight_vec[0]
w_pos = 1. / frac_pos
label_weights = error_scale_overall(label_slice, weight_vec)
net_io.setInputs([data_slice, label_slice, label_weights])
elif options.scale_error == 2:
net_io.setInputs([data_slice, label_slice, weight_slice])
elif options.scale_error == 0:
net_io.setInputs([data_slice, label_slice])
if options.loss_function == 'softmax':
net_io.setInputs([data_slice, label_slice])
# Single step
loss = solver.step(1)
# sanity_check_net_blobs(net)
while gc.collect():
pass
if (options.loss_function == 'euclid' or options.loss_function
== 'euclid_aniso') and options.scale_error == 1:
print("[Iter %i] Loss: %f, frac_pos=%f, w_pos=%f" %
(i, loss, frac_pos, w_pos))
else:
print("[Iter %i] Loss: %f" % (i, loss))
# TODO: Store losses to file
losses += [loss]
if hasattr(options, 'loss_snapshot') and ((i % options.loss_snapshot)
== 0):
io.savemat('loss.mat', {'loss': losses})
def get_outputs(original_dataset, output_slice):
output_shape = tuple([slice_.stop - slice_.start for slice_ in output_slice])
n_spatial_dimensions = len(output_slice)
components_shape = (1,) + output_shape
mask_shape = (1,) + output_shape
affinities_shape = (n_spatial_dimensions,) + output_shape
component_slices = [slice(0, l) for l in original_dataset['components'].shape]
component_slices[-n_spatial_dimensions:] = output_slice
logger.debug("component_slices: {}".format(component_slices))
components_array = get_zero_padded_array_slice(original_dataset['components'], component_slices)
source_class = type(original_dataset['components'])
components_are_from_dvid = source_class in dvid_classes
exclude_strings = original_dataset.get('body_names_to_exclude', [])
if exclude_strings and components_are_from_dvid:
dvid_uuid = original_dataset['components'].uuid
components_to_keep = get_good_components(dvid_uuid, exclude_strings)
logger.debug("components before: {}".format(list(np.unique(components_array))))
components_array = replace_array_except_whitelist(components_array, 0, components_to_keep)
logger.debug("components after: {}".format(list(np.unique(components_array))))
minimum_component_size = original_dataset.get('minimum_component_size', 0)
if minimum_component_size > 0:
components_array = replace_infrequent_values(components_array, minimum_component_size, 0)
component_erosion_steps = original_dataset.get('component_erosion_steps', 0)
if component_erosion_steps > 0:
components_array = erode_value_blobs(
components_array,
steps=component_erosion_steps,
values_to_ignore=(0,))
components_for_malis = components_array.reshape(output_shape)
affinities_from_components = malis.seg_to_affgraph(
components_for_malis,
original_dataset['nhood'])
components_array, _ = malis.connected_components_affgraph(
affinities_from_components,
original_dataset['nhood'])
components_array = shift_up_component_values(components_array)
components_array = components_array.reshape(components_shape)
if 'label' in original_dataset:
label_shape = original_dataset['label'].shape
label_slices = [slice(0, l) for l in label_shape]
label_slices[-n_spatial_dimensions:] = output_slice
affinities_array = get_zero_padded_array_slice(original_dataset['label'], label_slices)
else:
# compute affinities from components
logger.debug("Computing affinity labels from components because 'label' wasn't provided in data source.")
affinities_array = affinities_from_components
assert affinities_array.shape == affinities_shape, \
"affinities_array.shape is {actual} but should be {desired}".format(
actual=str(affinities_array.shape), desired=str(affinities_shape))
if 'mask' in original_dataset:
mask_array = get_zero_padded_array_slice(original_dataset['mask'], output_slice)
else:
if components_are_from_dvid:
# infer mask values: 1 if component is nonzero, 0 otherwise
mask_array = np.not_equal(components_array, 0)
logger.debug("No mask provided. Setting to 1 where components != 0.")
else:
# assume no masking
mask_array = np.ones_like(components_array, dtype=np.uint8)
logger.debug("No mask provided. Setting to 1 where outputs exist.")
mask_dilation_steps = original_dataset.get('mask_dilation_steps', 0)
if mask_dilation_steps > 0:
mask_array = ndimage.binary_dilation(mask_array, iterations=mask_dilation_steps)
mask_array = mask_array.astype(np.uint8)
mask_array = mask_array.reshape(mask_shape)
return components_array, affinities_array, mask_array
def train(solver, test_net, data_arrays, train_data_arrays, options):
if DEBUG:
data_io.logger.setLevel(logging.DEBUG)
else:
data_io.logger.setLevel(logging.INFO)
caffe.select_device(options.train_device, False)
net = solver.net
test_eval = None
if (options.test_net != None):
test_eval = TestNetEvaluator(test_net, net, train_data_arrays, options)
input_dims, output_dims, input_padding = get_spatial_io_dims(net)
fmaps_in, fmaps_out = get_fmap_io_dims(net)
dims = len(output_dims)
losses = []
shapes = []
# Raw data slice input (n = 1, f = 1, spatial dims)
shapes += [[1,fmaps_in] + input_dims]
# Label data slice input (n = 1, f = #edges, spatial dims)
shapes += [[1,fmaps_out] + output_dims]
if (options.loss_function == 'malis'):
# Connected components input. 2 channels, one for each phase of computation
shapes += [[1, 2] + output_dims]
if (options.loss_function == 'euclid'):
# Error scale input (n = 1, f = #edges, spatial dims)
shapes += [[1,fmaps_out] + output_dims]
# Nhood specifications (n = #edges, f = 3)
if (('nhood' in data_arrays[0]) and (options.loss_function == 'malis')):
shapes += [[1,1] + list(np.shape(data_arrays[0]['nhood']))]
net_io = NetInputWrapper(net, shapes)
if DEBUG:
for key in net.blobs.keys():
print(key, net.blobs[key].data.shape)
make_dataset_offset = MakeDatasetOffset(input_dims, output_dims)
if data_io.data_loader_should_be_used_with(data_arrays):
using_data_loader = True
# and initialize queue!
loader_size = 20
n_workers = 10
loader_kwargs = dict(
size=loader_size,
datasets=data_arrays,
input_shape=tuple(input_dims),
output_shape=tuple(output_dims),
n_workers=n_workers,
dataset_offset_func=make_dataset_offset
)
print("creating queue with kwargs {}".format(loader_kwargs))
training_data_loader = data_io.DataLoader(**loader_kwargs)
# start populating the queue
for i in range(loader_size):
if DEBUG:
print("Pre-populating data loader's dataset #{i}/{size}"
.format(i=i, size=training_data_loader.size))
shared_dataset_index, async_result = \
training_data_loader.start_refreshing_shared_dataset(i, wait=True)
else:
using_data_loader = False
# Loop from current iteration to last iteration
for i in range(solver.iter, solver.max_iter):
start = time.time()
if (options.test_net != None and i % options.test_interval == 1):
if not SAVE_IMAGES:
test_eval.evaluate(i)
if USE_ONE_THREAD:
# after testing finishes, switch back to the training device
caffe.select_device(options.train_device, False)
if not using_data_loader:
dataset_index, offsets = make_dataset_offset(data_arrays)
dataset = data_arrays[dataset_index]
# These are the raw data elements
data_slice = data_io.util.get_zero_padded_slice_from_array_by_offset(
array=dataset['data'],
origin=[0] + offsets,
shape=[fmaps_in] + input_dims)
label_slice = slice_data(dataset['label'], [0] + [offsets[di] + int(math.ceil(input_padding[di] / float(2))) for di in range(0, dims)], [fmaps_out] + output_dims)
components_slice, ccSizes = malis.connected_components_affgraph(label_slice.astype(np.int32), dataset['nhood'])
components_shape = (1,) + tuple(output_dims)
components_slice = components_slice.reshape(components_shape)
mask_slice = np.ones_like(components_slice, dtype=np.uint8)
mask_mean = 1
if 'transform' in dataset:
# transform the input
# assumes that the original input pixel values are scaled between (0,1)
if DEBUG:
print("data_slice stats, pre-transform: min", data_slice.min(), "mean", data_slice.mean(),
"max", data_slice.max())
lo, hi = dataset['transform']['scale']
data_slice = 0.5 + (data_slice - 0.5) * np.random.uniform(low=lo, high=hi)
lo, hi = dataset['transform']['shift']
data_slice = data_slice + np.random.uniform(low=lo, high=hi)
else:
dataset, index_of_shared_dataset = training_data_loader.get_dataset()
data_slice = dataset['data']
label_slice = dataset['label']
if DEBUG:
print("Training with next dataset in data loader, which has offset", dataset['offset'])
mask_slice = dataset['mask']
mask_mean = np.mean(mask_slice)
components_slice = dataset['components']
assert data_slice.shape == (fmaps_in,) + tuple(input_dims)
assert label_slice.shape == (fmaps_out,) + tuple(output_dims)
assert mask_slice.shape == (1,) + tuple(output_dims)
assert components_slice.shape == (1,) + tuple(output_dims)
if DEBUG:
print("data_slice stats: min", data_slice.min(), "mean", data_slice.mean(), "max", data_slice.max())
print("mask_mean: ", mask_mean)
if options.loss_function == 'malis':
try:
use_simple_malis_components =\
mask_mean == 1 or not options.malis_split_component_phases
except AttributeError:
use_simple_malis_components = mask_mean == 1
if use_simple_malis_components:
components_negative_slice = components_slice
else:
'''
assumes that...
* mask_slice is 1 at voxels containing good components, with a small dilation
* components_slice does not contain component values equal to 1. (Original values were incremented.)
'''
assert 1 not in components_slice, "components_slice can't contain a component value of 1. " \
"That's used for masked voxels in 'negative' MALIS training."
mask_inverse = np.ones_like(mask_slice) - mask_slice
# assert mask_inverse.shape == components_slice.shape
mask_inverse = mask_inverse.astype(components_slice.dtype)
# assert mask_inverse.dtype == components_slice.dtype
components_negative_slice = components_slice + mask_inverse
components_positive_slice = components_slice
components_malis_slice = np.concatenate(
(components_negative_slice, components_positive_slice),
axis=0)
net_io.setInputs([data_slice, label_slice, components_malis_slice, data_arrays[0]['nhood']])
elif options.loss_function == 'euclid':
if mask_mean < 1:
label_slice = label_slice * mask_slice
label_slice_mean = label_slice.mean() / mask_mean
else:
label_slice_mean = label_slice.mean()
w_pos = 1.0
w_neg = 1.0
if options.scale_error:
frac_pos = np.clip(label_slice_mean, 0.05, 0.95)
w_pos = w_pos / (2.0 * frac_pos)
w_neg = w_neg / (2.0 * (1.0 - frac_pos))
error_scale_slice = scale_errors(label_slice, w_neg, w_pos)
if mask_mean < 1:
error_scale_slice *= mask_slice
net_io.setInputs([data_slice, label_slice, error_scale_slice])
elif options.loss_function == 'softmax':
# These are the affinity edge values
net_io.setInputs([data_slice, label_slice])
loss = solver.step(1) # Single step
try:
save_image = SAVE_IMAGES or i % options.save_image_snapshot_period == 0
except AttributeError:
save_image = SAVE_IMAGES
if save_image:
dataset_to_show = dict()
net_prediction_shape = (1, fmaps_out,) + tuple(output_dims)
prediction = np.zeros(net_prediction_shape, np.float32)
for blob_key in reversed(solver.net.blobs.keys()):
try:
blob_shape = solver.net.blobs[blob_key].data.shape
except:
blob_shape = None
if blob_shape == net_prediction_shape:
prediction = solver.net.blobs[blob_key].data.copy()
break # stop checking blobs
dataset_to_show['pred'] = prediction.reshape((fmaps_out,) + tuple(output_dims))
try:
import zwatershed
(s, V) = zwatershed.zwatershed_and_metrics(
components_slice.reshape(output_dims).astype(np.uint32),
dataset_to_show['pred'],
[50000],
[50000])
components_prediction = s[0]
except:
components_prediction = np.zeros(shape=output_dims, dtype=np.int32)
dataset_to_show['predseg'] = components_prediction
dataset_to_show['data'] = data_slice.reshape(input_dims)
if components_slice is None:
components_slice, _ = malis.connected_components_affgraph(label_slice.astype(np.int32), dataset['nhood'])
dataset_to_show['components'] = components_slice.reshape(output_dims)
dataset_to_show['label'] = label_slice
dataset_to_show['mask'] = mask_slice.reshape(output_dims)
try:
dataset_to_show['components_negative'] = components_negative_slice.reshape(output_dims)
dataset_to_show['components_positive'] = components_positive_slice.reshape(output_dims)
except UnboundLocalError:
# variables weren't declared, because malis isn't being computed
pass
try:
dataset_to_show['error_scale_slice'] = error_scale_slice
except UnboundLocalError:
pass
assert dataset_to_show['data'].shape == tuple(input_dims), dataset_to_show['data'].shape
assert dataset_to_show['label'].shape == (3,) + tuple(output_dims), dataset_to_show['label'].shape
assert dataset_to_show['components'].shape == tuple(output_dims), dataset_to_show['components'].shape
assert dataset_to_show['pred'].shape == (3,) + tuple(output_dims), dataset_to_show['pred'].shape
assert dataset_to_show['predseg'].shape == tuple(output_dims), dataset_to_show['predseg'].shape
f = visualization.showme(dataset_to_show, int(input_dims[0] / 2))
if not os.path.exists('snapshots'):
os.mkdir('snapshots')
f.savefig('snapshots/%08d.png' % i)
plt.close()
while gc.collect():
pass
time_of_iteration = time.time() - start
if 'mask' in dataset:
mask_fraction_str = "%07.5f" % np.mean(dataset['mask'])
else:
mask_fraction_str = "no mask"
if options.loss_function == 'euclid' or options.loss_function == 'euclid_aniso':
print("[Iter %06i]" % i,
"Time: %05.2fs" % time_of_iteration,
"Loss: %08.6f" % loss,
"Mask:", mask_fraction_str,
"frac_pos=%08.6f" % frac_pos,
"w_pos=%08.6f" % w_pos,
)
else:
print("[Iter %06i]" % i,
"Time: %05.2fs" % time_of_iteration,
"Loss: %08.6f" % loss,
"Mask:", mask_fraction_str,
)
losses += [loss]
if hasattr(options, 'loss_snapshot') and ((i % options.loss_snapshot) == 0):
io.savemat('loss.mat',{'loss':losses})
if using_data_loader:
training_data_loader.start_refreshing_shared_dataset(index_of_shared_dataset)
if using_data_loader:
training_data_loader.destroy()
def train(solver, test_net, data_arrays, train_data_arrays, options):
caffe.select_device(options.train_device, False)
net = solver.net
test_eval = None
if (options.test_net != None):
test_eval = TestNetEvaluator(test_net, net, train_data_arrays, options)
input_dims, output_dims, input_padding = get_spatial_io_dims(net)
fmaps_in, fmaps_out = get_fmap_io_dims(net)
dims = len(output_dims)
losses = []
shapes = []
# Raw data slice input (n = 1, f = 1, spatial dims)
shapes += [[1,fmaps_in] + input_dims]
# Label data slice input (n = 1, f = #edges, spatial dims)
shapes += [[1,fmaps_out] + output_dims]
if (options.loss_function == 'malis'):
# Connected components input (n = 1, f = 1, spatial dims)
shapes += [[1,1] + output_dims]
if (options.loss_function == 'euclid'):
# Error scale input (n = 1, f = #edges, spatial dims)
shapes += [[1,fmaps_out] + output_dims]
# Nhood specifications (n = #edges, f = 3)
if (('nhood' in data_arrays[0]) and (options.loss_function == 'malis')):
shapes += [[1,1] + list(np.shape(data_arrays[0]['nhood']))]
net_io = NetInputWrapper(net, shapes)
# Loop from current iteration to last iteration
for i in range(solver.iter, solver.max_iter):
if (options.test_net != None and i % options.test_interval == 0):
test_eval.evaluate(i)
# First pick the dataset to train with
dataset = randint(0, len(data_arrays) - 1)
offsets = []
for j in range(0, dims):
offsets.append(randint(0, data_arrays[dataset]['data'].shape[1+j] - (output_dims[j] + input_padding[j])))
# These are the raw data elements
data_slice = slice_data(data_arrays[dataset]['data'], [0]+offsets, [fmaps_in]+[output_dims[di] + input_padding[di] for di in range(0, dims)])
label_slice = slice_data(data_arrays[dataset]['label'], [0] + [offsets[di] + int(math.ceil(input_padding[di] / float(2))) for di in range(0, dims)], [fmaps_out] + output_dims)
# transform the input
# this code assumes that the original input pixel values are scaled between (0,1)
if 'transform' in data_arrays[dataset]:
# print('Pre:',(data_slice.min(),data_slice.mean(),data_slice.max()))
lo, hi = data_arrays[dataset]['transform']['scale']
data_slice = 0.5 + (data_slice-0.5)*np.random.uniform(low=lo,high=hi)
lo, hi = data_arrays[dataset]['transform']['shift']
data_slice = data_slice + np.random.uniform(low=lo,high=hi)
# print('Post:',(data_slice.min(),data_slice.mean(),data_slice.max()))
if options.loss_function == 'malis':
components_slice,ccSizes = malis.connected_components_affgraph(label_slice.astype(int32), data_arrays[dataset]['nhood'])
# Also recomputing the corresponding labels (connected components)
net_io.setInputs([data_slice, label_slice, components_slice, data_arrays[0]['nhood']])
if options.loss_function == 'euclid':
if(options.scale_error == True):
frac_pos = np.clip(label_slice.mean(),0.05,0.95)
w_pos = 1.0/(2.0*frac_pos)
w_neg = 1.0/(2.0*(1.0-frac_pos))
else:
w_pos = 1
w_neg = 1
net_io.setInputs([data_slice, label_slice, error_scale(label_slice,w_neg,w_pos)])
if options.loss_function == 'softmax':
# These are the affinity edge values
net_io.setInputs([data_slice, label_slice])
# Single step
loss = solver.step(1)
# sanity_check_net_blobs(net)
while gc.collect():
pass
if options.loss_function == 'euclid' or options.loss_function == 'euclid_aniso':
print("[Iter %i] Loss: %f, frac_pos=%f, w_pos=%f" % (i,loss,frac_pos,w_pos))
else:
print("[Iter %i] Loss: %f" % (i,loss))
# TODO: Store losses to file
losses += [loss]
if hasattr(options, 'loss_snapshot') and ((i % options.loss_snapshot) == 0):
io.savemat('loss.mat',{'loss':losses})
def train(solver, test_net, data_arrays, train_data_arrays, options):
caffe.select_device(options.train_device, False)
print('====> in training....')
net = solver.net
net.debug_info = True
#pdb.set_trace()
'''
clwt=None
test_eval = None
if options.scale_error == 2:
clwt = ClassWeight(data_arrays, solver.iter)
test_eval = TestNetEvaluator(test_net, net, data_arrays, options)
test_eval2 = None
if (options.test_net != None):
test_eval2 = TestNetEvaluator(test_net, net, train_data_arrays, options)
'''
input_dims, output_dims, input_padding = get_spatial_io_dims(net)
fmaps_in, fmaps_out = get_fmap_io_dims(net)
print('input_dims:', input_dims)
print('output_dims:', output_dims)
print('input_padding:', input_padding)
print('fmaps_out:', fmaps_out)
dims = len(output_dims)
losses = []
shapes = []
# Raw data slice input (n = 1, f = 1, spatial dims)
shapes += [[1, fmaps_in] + input_dims]
# Label data slice input (n = 1, f = #edges, spatial dims)
shapes += [[1, fmaps_out] + output_dims]
if (options.loss_function == 'malis'):
# Connected components input (n = 1, f = 1, spatial dims)
shapes += [[1, 1] + output_dims]
if (options.loss_function == 'euclid'):
# Error scale input (n = 1, f = #edges, spatial dims)
shapes += [[1, fmaps_out] + output_dims]
# Nhood specifications (n = #edges, f = 3)
if (('nhood' in data_arrays[0]) and (options.loss_function == 'malis')):
shapes += [[1, 1] + list(np.shape(data_arrays[0]['nhood']))]
net_io = NetInputWrapper(net, shapes)
weight_vec = []
if (options.loss_function == 'softmax'
or options.loss_function == 'euclid') and options.scale_error == 1:
#pdb.set_trace()
weight_vec = class_balance_distribution(data_arrays[0]['label'])
#weight_vec[2] = weight_vec[1]*4.0 #for 3 class, inversed during weighting
#pdb.set_trace()
'''
dims3d = dims + 1
output_dims3d = output_dims + [output_dims[-1]]
input_padding3d = input_padding + [input_padding[-1]]
output_dims3d = output_dims + [output_dims[-1]]
'''
n_slices = output_dims[-1]
i_slice = 0
# Loop from current iteration to last iteration
for i in range(solver.iter, solver.max_iter):
'''
if (options.test_net != None and i % options.test_interval == 0 and i>1):
#pdb.set_trace()
test_eval2.evaluate(i)
if options.scale_error == 2:
test_eval.evaluate(i)
clwt.recompute_weight(test_eval.pred_arrays_samesize, i)
'''
# First pick the dataset to train with
dataset = randint(0, len(data_arrays) - 1)
#print('dataset shape:', data_arrays[dataset]['data'].shape)
if i_slice == 0 or i_slice == n_slices:
i_slice = 0
offsets = []
for j in range(0, dims):
offsets.append(
randint(
0, data_arrays[dataset]['data'].shape[1 + j] -
(output_dims[j] + input_padding[j])))
data_slice = slice_data(
data_arrays[dataset]['data'], [0] + offsets, [fmaps_in] +
[output_dims[di] + input_padding[di] for di in range(0, dims)])
label_slice = slice_data(data_arrays[dataset]['label'], [0] + [
offsets[di] + int(math.ceil(input_padding[di] / float(2)))
for di in range(0, dims)
], [fmaps_out] + output_dims)
#print(data_slice.shape)
#print(label_slice.shape)
#data_slice = np.squeeze(data_slice)
label_slice = np.squeeze(label_slice)
#print(label_slice)
#offsets=np.zeros(dims);
if (data_slice.shape[0] < 1) or (label_slice.shape[0] < 2):
pp = 1
# transform the input
# this code assumes that the original input pixel values are scaled between (0,1)
if 'transform' in data_arrays[dataset]:
# print('Pre:',(data_slice.min(),data_slice.mean(),data_slice.max()))
data_slice_mean = data_slice.mean()
lo, hi = data_arrays[dataset]['transform']['scale']
data_slice = data_slice_mean + (
data_slice - data_slice_mean) * np.random.uniform(low=lo,
high=hi)
lo, hi = data_arrays[dataset]['transform']['shift']
data_slice = data_slice + np.random.uniform(low=lo, high=hi)
# print('Post:',(data_slice.min(),data_slice.mean(),data_slice.max()))
data_slice = np.clip(data_slice, 0.0, 0.95)
if options.loss_function == 'malis':
components_slice, ccSizes = malis.connected_components_affgraph(
label_slice.astype(int32), data_arrays[dataset]['nhood'])
# Also recomputing the corresponding labels (connected components)
net_io.setInputs([
data_slice, label_slice, components_slice,
data_arrays[0]['nhood']
])
if options.loss_function == 'euclid':
###net_io.setInputs([data_slice, label_slice, error_scale(label_slice,w_neg,w_pos)])
if (options.scale_error == 3):
frac_pos = np.clip(label_slice.mean(), 0.01,
0.99) #for binary labels
w_pos = 1.0 / (2.0 * frac_pos)
w_neg = 1.0 / (2.0 * (1.0 - frac_pos))
net_io.setInputs([
data_slice, label_slice,
error_scale(label_slice, w_neg, w_pos)
])
elif (options.scale_error == 1):
frac_pos = weight_vec[0]
w_pos = 1. / frac_pos
label_weights = error_scale_overall(label_slice, weight_vec)
net_io.setInputs([data_slice, label_slice, label_weights])
elif options.scale_error == 2:
net_io.setInputs([data_slice, label_slice, weight_slice])
elif options.scale_error == 0:
net_io.setInputs([data_slice, label_slice])
if options.loss_function == 'softmax':
net_io.setInputs([data_slice, label_slice])
#pdb.set_trace()
#print('training slice#: ', i_slice)
# Single step
n_slices = output_dims[-1]
#loss = solver.stepForward(1)
#loss = solver.stepParallel( n_slices )
loss = solver.stepForward(n_slices)
solver.stepBackward(n_slices)
#solver.stepBackward(1)
i_slice = n_slices
# do backward when all slices have been processed.
#if i_slice == n_slices:
# solver.stepBackward(1)
#loss = solver.step(1) #n_slices)
#for i in range(n_slices):
# loss = solver.stepForward(1)
#solver.stepBackward()
# sanity_check_net_blobs(net)
if (options.loss_function == 'euclid' or options.loss_function
== 'euclid_aniso') and options.scale_error == 1:
print("[Iter %i] Loss: %f, frac_pos=%f, w_pos=%f" %
(i, loss, frac_pos, w_pos))
else:
print("[Iter %i] Loss: %f" % (i, loss))
sys.stdout.flush()
# TODO: Store losses to file
losses += [loss]
if hasattr(options, 'loss_snapshot') and ((i % options.loss_snapshot)
== 0):
io.savemat('loss.mat', {'loss': losses})
#pdb.set_trace()
while gc.collect():
pass
def np_func(aff, nhood, threshold):
aff = np.transpose(aff, [3, 0, 1, 2]) # zyx3 to 3zyx
ret = malis.connected_components_affgraph(
(aff > threshold[0]).astype(np.int32), nhood)[0].astype(np.int32)
ret = skimage.measure.label(ret).astype(np.int32)
return ret
def train(solver, data_arrays, label_arrays, mode='malis'):
losses = []
net = solver.net
if mode == 'malis':
nhood = malis.mknhood3d()
if mode == 'euclid':
nhood = malis.mknhood3d()
if mode == 'malis_aniso':
nhood = malis.mknhood3d_aniso()
if mode == 'euclid_aniso':
nhood = malis.mknhood3d_aniso()
data_slice_cont = np.zeros((1,1,132,132,132), dtype=float32)
label_slice_cont = np.zeros((1,1,44,44,44), dtype=float32)
aff_slice_cont = np.zeros((1,3,44,44,44), dtype=float32)
nhood_cont = np.zeros((1,1,3,3), dtype=float32)
error_scale_cont = np.zeros((1,1,44,44,44), dtype=float32)
dummy_slice = np.ascontiguousarray([0]).astype(float32)
# Loop from current iteration to last iteration
for i in range(solver.iter, solver.max_iter):
# First pick the dataset to train with
dataset = randint(0, len(data_arrays) - 1)
data_array = data_arrays[dataset]
label_array = label_arrays[dataset]
# affinity_array = affinity_arrays[dataset]
offsets = []
for j in range(0, dims):
offsets.append(randint(0, data_array.shape[j] - (config.output_dims[j] + config.input_padding[j])))
# These are the raw data elements
data_slice = slice_data(data_array, offsets, [config.output_dims[di] + config.input_padding[di] for di in range(0, dims)])
# These are the labels (connected components)
label_slice = slice_data(label_array, [offsets[di] + int(math.ceil(config.input_padding[di] / float(2))) for di in range(0, dims)], config.output_dims)
# These are the affinity edge values
# Also recomputing the corresponding labels (connected components)
aff_slice = malis.seg_to_affgraph(label_slice,nhood)
label_slice,ccSizes = malis.connected_components_affgraph(aff_slice,nhood)
print (data_slice[None, None, :]).shape
print (label_slice[None, None, :]).shape
print (aff_slice[None, :]).shape
print (nhood).shape
if mode == 'malis':
np.copyto(data_slice_cont, np.ascontiguousarray(data_slice[None, None, :]).astype(float32))
np.copyto(label_slice_cont, np.ascontiguousarray(label_slice[None, None, :]).astype(float32))
np.copyto(aff_slice_cont, np.ascontiguousarray(aff_slice[None, :]).astype(float32))
np.copyto(nhood_cont, np.ascontiguousarray(nhood[None, None, :]).astype(float32))
net.set_input_arrays(0, data_slice_cont, dummy_slice)
net.set_input_arrays(1, label_slice_cont, dummy_slice)
net.set_input_arrays(2, aff_slice_cont, dummy_slice)
net.set_input_arrays(3, nhood_cont, dummy_slice)
# We pass the raw and affinity array only
if mode == 'euclid':
net.set_input_arrays(0, np.ascontiguousarray(data_slice[None, None, :]).astype(float32), np.ascontiguousarray(dummy_slice).astype(float32))
net.set_input_arrays(1, np.ascontiguousarray(aff_slice[None, :]).astype(float32), np.ascontiguousarray(dummy_slice).astype(float32))
net.set_input_arrays(2, np.ascontiguousarray(error_scale(aff_slice[None, :],1.0,0.045)).astype(float32), np.ascontiguousarray(dummy_slice).astype(float32))
if mode == 'softmax':
net.set_input_arrays(0, np.ascontiguousarray(data_slice[None, None, :]).astype(float32), np.ascontiguousarray(dummy_slice).astype(float32))
net.set_input_arrays(1, np.ascontiguousarray(label_slice[None, None, :]).astype(float32), np.ascontiguousarray(dummy_slice).astype(float32))
# Single step
loss = solver.step(1)
# Memory clean up and report
print("Memory usage (before GC): %d MiB" % ((resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) / (1024)))
while gc.collect():
pass
print("Memory usage (after GC): %d MiB" % ((resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) / (1024)))
# m = volume_slicer.VolumeSlicer(data=np.squeeze((net.blobs['Convolution18'].data[0])[0,:,:]))
# m.configure_traits()
print("Loss: %s" % loss)
losses += [loss]
def train(solver, data_arrays, label_arrays, mode='malis'):
losses = []
net = solver.net
if mode == 'malis':
nhood = malis.mknhood3d()
if mode == 'euclid':
nhood = malis.mknhood3d()
if mode == 'malis_aniso':
nhood = malis.mknhood3d_aniso()
if mode == 'euclid_aniso':
nhood = malis.mknhood3d_aniso()
data_slice_cont = np.zeros((1, 1, 132, 132, 132), dtype=float32)
label_slice_cont = np.zeros((1, 1, 44, 44, 44), dtype=float32)
aff_slice_cont = np.zeros((1, 3, 44, 44, 44), dtype=float32)
nhood_cont = np.zeros((1, 1, 3, 3), dtype=float32)
error_scale_cont = np.zeros((1, 1, 44, 44, 44), dtype=float32)
dummy_slice = np.ascontiguousarray([0]).astype(float32)
# Loop from current iteration to last iteration
for i in range(solver.iter, solver.max_iter):
# First pick the dataset to train with
dataset = randint(0, len(data_arrays) - 1)
data_array = data_arrays[dataset]
label_array = label_arrays[dataset]
# affinity_array = affinity_arrays[dataset]
offsets = []
for j in range(0, dims):
offsets.append(
randint(
0, data_array.shape[j] -
(config.output_dims[j] + config.input_padding[j])))
# These are the raw data elements
data_slice = slice_data(data_array, offsets, [
config.output_dims[di] + config.input_padding[di]
for di in range(0, dims)
])
# These are the labels (connected components)
label_slice = slice_data(label_array, [
offsets[di] + int(math.ceil(config.input_padding[di] / float(2)))
for di in range(0, dims)
], config.output_dims)
# These are the affinity edge values
# Also recomputing the corresponding labels (connected components)
aff_slice = malis.seg_to_affgraph(label_slice, nhood)
label_slice, ccSizes = malis.connected_components_affgraph(
aff_slice, nhood)
print(data_slice[None, None, :]).shape
print(label_slice[None, None, :]).shape
print(aff_slice[None, :]).shape
print(nhood).shape
if mode == 'malis':
np.copyto(
data_slice_cont,
np.ascontiguousarray(data_slice[None,
None, :]).astype(float32))
np.copyto(
label_slice_cont,
np.ascontiguousarray(label_slice[None,
None, :]).astype(float32))
np.copyto(aff_slice_cont,
np.ascontiguousarray(aff_slice[None, :]).astype(float32))
np.copyto(
nhood_cont,
np.ascontiguousarray(nhood[None, None, :]).astype(float32))
net.set_input_arrays(0, data_slice_cont, dummy_slice)
net.set_input_arrays(1, label_slice_cont, dummy_slice)
net.set_input_arrays(2, aff_slice_cont, dummy_slice)
net.set_input_arrays(3, nhood_cont, dummy_slice)
# We pass the raw and affinity array only
if mode == 'euclid':
net.set_input_arrays(
0,
np.ascontiguousarray(data_slice[None,
None, :]).astype(float32),
np.ascontiguousarray(dummy_slice).astype(float32))
net.set_input_arrays(
1,
np.ascontiguousarray(aff_slice[None, :]).astype(float32),
np.ascontiguousarray(dummy_slice).astype(float32))
net.set_input_arrays(
2,
np.ascontiguousarray(
error_scale(aff_slice[None, :], 1.0,
0.045)).astype(float32),
np.ascontiguousarray(dummy_slice).astype(float32))
if mode == 'softmax':
net.set_input_arrays(
0,
np.ascontiguousarray(data_slice[None,
None, :]).astype(float32),
np.ascontiguousarray(dummy_slice).astype(float32))
net.set_input_arrays(
1,
np.ascontiguousarray(label_slice[None,
None, :]).astype(float32),
np.ascontiguousarray(dummy_slice).astype(float32))
# Single step
loss = solver.step(1)
# Memory clean up and report
print("Memory usage (before GC): %d MiB" %
((resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) / (1024)))
while gc.collect():
pass
print("Memory usage (after GC): %d MiB" %
((resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) / (1024)))
# m = volume_slicer.VolumeSlicer(data=np.squeeze((net.blobs['Convolution18'].data[0])[0,:,:]))
# m.configure_traits()
print("Loss: %s" % loss)
losses += [loss]
