def generate_data(num_batches):
neighborhood = [[-1, 0, 0], [0, -1, 0], [0, 0, -1]]
neighborhood_opp = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
# define array-keys
labels_key = ArrayKey('LABELS')
raw_affs_key = ArrayKey('RAW_AFFINITIES')
raw_joined_affs_key = ArrayKey('RAW_JOINED_AFFINITIES')
raw_key = ArrayKey('RAW')
affs_key = ArrayKey('AFFINITIES')
affs_opp_key = ArrayKey('AFFINITIES_OPP')
joined_affs_key = ArrayKey('JOINED_AFFINITIES')
joined_affs_opp_key = ArrayKey('JOINED_AFFINITIES_OPP')
merged_labels_key = ArrayKey('MERGED_LABELS')
gt_affs_key = ArrayKey('GT_AFFINITIES')
gt_affs_in_key = ArrayKey('GT_AFFINITIES_IN')
gt_affs_mask_key = ArrayKey('GT_AFFINITIES_MASK')
voxel_size = Coordinate((1, 1, 1))
input_shape = Coordinate((132, 132, 132)) * voxel_size
input_affs_shape = Coordinate([i + 1
for i in (132, 132, 132)]) * voxel_size
output_shape = Coordinate((44, 44, 44)) * voxel_size
output_affs_shape = Coordinate([i + 1 for i in (44, 44, 44)]) * voxel_size
print("input_shape: ", input_shape)
print("output_shape: ", output_shape)
request = BatchRequest()
request.add(labels_key, output_shape)
request.add(raw_key, input_shape)
request.add(raw_affs_key, input_shape)
request.add(raw_joined_affs_key, input_shape)
request.add(affs_key, input_affs_shape)
request.add(affs_opp_key, input_affs_shape)
request.add(joined_affs_key, input_affs_shape)
request.add(joined_affs_opp_key, input_affs_shape)
request.add(merged_labels_key, output_shape)
request.add(gt_affs_key, output_shape)
request.add(gt_affs_in_key, input_shape)
request.add(gt_affs_mask_key, output_shape)
# offset = Coordinate((input_size[i]-output_size[i])/2 for i in range(len(input_size)))
# crop_roi = Roi(offset, output_size)
# print("crop_roi: ", crop_roi)
# print ("input_affinities_key: ", input_affinities_key)
pipeline = ()
# print ("iteration: ", iteration)
pipeline += ToyNeuronSegmentationGenerator(array_key=labels_key,
n_objects=15,
points_per_skeleton=8,
smoothness=3,
noise_strength=1,
interpolation="linear")
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=raw_affs_key)
pipeline += AddJoinedAffinities(input_affinities=raw_affs_key,
joined_affinities=raw_joined_affs_key)
pipeline += AddRealism(joined_affinities=raw_joined_affs_key,
raw=raw_key,
sp=0.25,
sigma=1,
contrast=0.7)
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=affs_key)
pipeline += AddAffinities(affinity_neighborhood=neighborhood_opp,
labels=labels_key,
affinities=affs_opp_key)
pipeline += AddJoinedAffinities(input_affinities=affs_key,
joined_affinities=joined_affs_key)
pipeline += AddJoinedAffinities(input_affinities=affs_opp_key,
joined_affinities=joined_affs_opp_key)
pipeline += MergeLabels(labels=labels_key,
joined_affinities=joined_affs_key,
joined_affinities_opp=joined_affs_opp_key,
merged_labels=merged_labels_key,
every=2)
# # pipeline += GrowBoundary(merged_labels_key, steps=1, only_xy=True)
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=merged_labels_key,
affinities=gt_affs_in_key)
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=merged_labels_key,
affinities=gt_affs_key,
affinities_mask=gt_affs_mask_key)
pipeline += PreCache(cache_size=4, num_workers=2)
pipeline += Snapshot(dataset_names={
labels_key: 'volumes/labels',
merged_labels_key: 'volumes/merged_labels',
raw_key: 'volumes/raw',
raw_affs_key: 'volumes/raw_affs',
gt_affs_key: 'volumes/gt_affs',
gt_affs_in_key: 'volumes/gt_affs_in'
},
output_filename='test_edges.hdf',
every=1,
dataset_dtypes={
merged_labels_key: np.uint64,
labels_key: np.uint64,
raw_key: np.float32
})
pipeline += PrintProfilingStats(every=100)
hashes = []
with build(pipeline) as p:
for i in range(num_batches):
print("iteration: ", i)
req = p.request_batch(request)
def predict(checkpoint, iterations):
print("iterations:", iterations)
labels_key = ArrayKey('GT_LABELS')
joined_affinities_key = ArrayKey('GT_JOINED_AFFINITIES')
raw_affinities_key = ArrayKey('RAW_AFFINITIES_KEY')
raw_key = ArrayKey('RAW')
pred_affinities_key = ArrayKey('PREDICTED_AFFS')
debug_key = ArrayKey("DEBUG")
voxel_size = Coordinate((1, 1, 1))
input_shape = Coordinate(config['input_shape']) * voxel_size
output_shape = Coordinate(config['output_shape']) * voxel_size
debug_shape = Coordinate((1, 1, 1)) * voxel_size
print("input_size: ", input_shape)
print("output_size: ", output_shape)
request = BatchRequest()
# request.add(labels_key, input_shape) # TODO: why does adding this request cause a duplication of generations?
request.add(joined_affinities_key, input_shape)
request.add(raw_affinities_key, input_shape)
request.add(raw_key, input_shape)
request.add(pred_affinities_key, output_shape)
request.add(debug_key, debug_shape)
pipeline = ()
pipeline += ToyNeuronSegmentationGenerator(array_key=labels_key,
n_objects=7,
points_per_skeleton=8,
smoothness=3,
noise_strength=1,
interpolation="random")
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=raw_affinities_key)
pipeline += AddJoinedAffinities(input_affinities=raw_affinities_key,
joined_affinities=joined_affinities_key)
pipeline += AddRealism(joined_affinities=joined_affinities_key,
raw=raw_key,
sp=0.65,
sigma=1,
contrast=0.7)
# Pad(raw_key, size=None) +
# Crop(raw_key, read_roi) +
# Normalize(raw_key) +
pipeline += IntensityScaleShift(raw_key, 2, -1)
predict = Predict(
checkpoint=os.path.join(setup_dir,
'train_net_checkpoint_%d' % checkpoint),
inputs={config['raw']: raw_key},
outputs={
config['pred_affs']: pred_affinities_key,
config['debug']: debug_key,
},
# graph=os.path.join(setup_dir, 'predict_net.meta')
)
pipeline += predict
pipeline += IntensityScaleShift(array=raw_key, scale=0.5, shift=0.5)
pipeline += Snapshot(dataset_names={
labels_key: 'volumes/labels/labels',
raw_affinities_key: 'volumes/raw_affs',
raw_key: 'volumes/raw',
pred_affinities_key: 'volumes/pred_affs'
},
output_filename='unet/prediction.hdf',
every=1,
dataset_dtypes={
raw_key: np.float32,
pred_affinities_key: np.float32,
labels_key: np.uint64
})
pipeline += PrintProfilingStats(every=20)
print("Starting prediction...")
with build(pipeline) as p:
for i in range(iterations):
print("iteration: ", i)
req = p.request_batch(request)
debug = req[debug_key].data
# pred_affs = req[pred_affinities_key].data
print("debug", debug)
# print("pred_affs", pred_affs)
# with predict.session.as_default():
# d = predict.graph.get_tensor_by_name('debug:0')
# print(d.eval())
print("Prediction finished")
def train(iterations):
# tf.reset_default_graph()
if tf.train.latest_checkpoint('.'):
trained_until = int(tf.train.latest_checkpoint('.').split('_')[-1])
else:
trained_until = 0
if trained_until >= iterations:
return
if trained_until < phase_switch and iterations > phase_switch:
train(phase_switch)
phase = 'euclid' if iterations <= phase_switch else 'malis'
print("Training in phase %s until %i" % (phase, iterations))
# define array-keys
labels_key = ArrayKey('LABELS')
raw_affs_key = ArrayKey('RAW_AFFINITIES')
raw_joined_affs_key = ArrayKey('RAW_JOINED_AFFINITIES')
raw_key = ArrayKey('RAW')
affs_key = ArrayKey('AFFINITIES')
affs_opp_key = ArrayKey('AFFINITIES_OPP')
joined_affs_key = ArrayKey('JOINED_AFFINITIES')
joined_affs_opp_key = ArrayKey('JOINED_AFFINITIES_OPP')
merged_labels_key = ArrayKey('MERGED_LABELS')
gt_affs_key = ArrayKey('GT_AFFINITIES')
gt_affs_in_key = ArrayKey('GT_AFFINITIES_IN')
gt_affs_mask_key = ArrayKey('GT_AFFINITIES_MASK')
gt_affs_scale_key = ArrayKey('GT_AFFINITIES_SCALE')
pred_affs_key = ArrayKey('PRED_AFFS')
pred_affs_gradient_key = ArrayKey('PRED_AFFS_GRADIENT')
voxel_size = Coordinate((1, 1, 1))
input_shape = Coordinate(config['input_shape']) * voxel_size
input_affs_shape = Coordinate([i + 1 for i in config['input_shape']
]) * voxel_size
output_shape = Coordinate(config['output_shape']) * voxel_size
output_affs_shape = Coordinate([i + 1 for i in config['output_shape']
]) * voxel_size
print("input_shape: ", input_shape)
print("input_affs_shape: ", input_affs_shape)
print("output_shape: ", output_shape)
print("output_affs_shape: ", output_affs_shape)
request = BatchRequest()
request.add(labels_key, output_shape)
request.add(raw_key, input_shape)
request.add(raw_affs_key, input_shape)
request.add(raw_joined_affs_key, input_shape)
request.add(affs_key, input_affs_shape)
request.add(affs_opp_key, input_affs_shape)
request.add(joined_affs_key, input_affs_shape)
request.add(joined_affs_opp_key, input_affs_shape)
request.add(merged_labels_key, output_shape)
request.add(gt_affs_key, output_shape)
request.add(gt_affs_in_key, input_shape)
request.add(gt_affs_mask_key, output_shape)
request.add(gt_affs_scale_key, output_shape)
request.add(pred_affs_key, output_shape)
# offset = Coordinate((input_shape[i]-output_shape[i])/2 for i in range(len(input_shape)))
# crop_roi = Roi(offset, output_shape)
# print("crop_roi: ", crop_roi)
pipeline = ()
# print ("iteration: ", iteration)
pipeline += ToyNeuronSegmentationGenerator(array_key=labels_key,
n_objects=50,
points_per_skeleton=8,
smoothness=3,
noise_strength=1,
interpolation="linear")
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=raw_affs_key)
pipeline += AddJoinedAffinities(input_affinities=raw_affs_key,
joined_affinities=raw_joined_affs_key)
pipeline += AddRealism(joined_affinities=raw_joined_affs_key,
raw=raw_key,
sp=0.65,
sigma=1,
contrast=0.7)
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=affs_key)
pipeline += AddAffinities(affinity_neighborhood=neighborhood_opp,
labels=labels_key,
affinities=affs_opp_key)
pipeline += AddJoinedAffinities(input_affinities=affs_key,
joined_affinities=joined_affs_key)
pipeline += AddJoinedAffinities(input_affinities=affs_opp_key,
joined_affinities=joined_affs_opp_key)
pipeline += MergeLabels(labels=labels_key,
joined_affinities=joined_affs_key,
joined_affinities_opp=joined_affs_opp_key,
merged_labels=merged_labels_key,
every=2)
# pipeline += GrowBoundary(merged_labels_key, steps=1, only_xy=True)
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=merged_labels_key,
affinities=gt_affs_in_key)
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=merged_labels_key,
affinities=gt_affs_key,
affinities_mask=gt_affs_mask_key)
# if phase == 'euclid':
pipeline += BalanceLabels(labels=gt_affs_key, scales=gt_affs_scale_key)
pipeline += DefectAugment(intensities=raw_key,
prob_missing=0.03,
prob_low_contrast=0.01,
contrast_scale=0.5,
axis=0)
pipeline += IntensityScaleShift(raw_key, 2, -1)
if phase == 'malis':
# pipeline += Crop(
# key=merged_labels_key,
# roi=crop_roi)
pipeline += RenumberConnectedComponents(labels=merged_labels_key)
pipeline += PreCache(cache_size=32, num_workers=8)
train_inputs = {
config['raw']: raw_key,
config['gt_affs_in']: gt_affs_in_key,
config['gt_affs_out']: gt_affs_key,
config['pred_affs_loss_weights']: gt_affs_scale_key
}
if phase == 'euclid':
train_loss = config['loss']
train_optimizer = config['optimizer']
train_summary = config['summary']
# train_inputs[config['pred_affs_loss_weights']] = input_affinities_scale_key
else:
train_loss = None
train_optimizer = add_malis_loss
train_inputs['gt_seg:0'] = merged_labels_key # XXX question
train_inputs['gt_affs_mask:0'] = gt_affs_mask_key
train_summary = 'Merge/MergeSummary:0'
pipeline += Train(graph=setup_dir + 'train_net',
optimizer=train_optimizer,
loss=train_loss,
inputs=train_inputs,
outputs={config['pred_affs']: pred_affs_key},
gradients={config['pred_affs']: pred_affs_gradient_key},
summary=train_summary,
log_dir='log/prob_unet/' + setup_name,
save_every=2000)
pipeline += IntensityScaleShift(array=raw_key, scale=0.5, shift=0.5)
pipeline += Snapshot(dataset_names={
merged_labels_key: 'volumes/labels',
raw_key: 'volumes/raw',
pred_affs_key: 'volumes/pred_affs',
gt_affs_key: 'volumes/gt_affs'
},
output_filename='prob_unet/' + setup_name +
'/batch_{iteration}.hdf',
every=4000,
dataset_dtypes={
merged_labels_key: np.uint64,
raw_key: np.float32
})
pipeline += PrintProfilingStats(every=20)
print("Starting training...")
with build(pipeline) as p:
for i in range(iterations - trained_until):
req = p.request_batch(request)
print("Training finished")
def train(iterations):
tf.reset_default_graph()
if tf.train.latest_checkpoint('.'):
trained_until = int(tf.train.latest_checkpoint('.').split('_')[-1])
else:
trained_until = 0
if trained_until >= iterations:
return
# define array-keys
labels_key = ArrayKey('GT_LABELS')
gt_affs_in_key = ArrayKey('GT_AFFINITIES_IN')
gt_affs_out_key = ArrayKey('GT_AFFINITIES_OUT')
joined_affinities_key = ArrayKey('GT_JOINED_AFFINITIES')
raw_key = ArrayKey('RAW')
input_affinities_scale_key = ArrayKey('GT_AFFINITIES_SCALE')
pred_affinities_key = ArrayKey('PREDICTED_AFFS')
pred_affinities_gradient_key = ArrayKey('AFFS_GRADIENT')
gt_affs_mask = ArrayKey('GT_AFFINITIES_MASK')
voxel_size = Coordinate((1, 1, 1))
input_shape = Coordinate(config['input_shape']) * voxel_size
output_shape = Coordinate(config['output_shape']) * voxel_size
print ("input_shape: ", input_shape)
print ("output_shape: ", output_shape)
# define requests
request = BatchRequest()
request.add(labels_key, output_shape)
request.add(gt_affs_in_key, input_shape)
request.add(joined_affinities_key, input_shape)
request.add(raw_key, input_shape)
request.add(gt_affs_out_key, output_shape)
request.add(input_affinities_scale_key, output_shape)
request.add(pred_affinities_key, output_shape)
request.add(gt_affs_mask, output_shape)
offset = Coordinate((input_shape[i]-output_shape[i])/2 for i in range(len(input_shape)))
crop_roi = Roi(offset, output_shape)
# print("crop_roi: ", crop_roi)
pipeline = (
ToyNeuronSegmentationGenerator(
array_key=labels_key,
n_objects=50,
points_per_skeleton=8,
smoothness=3,
noise_strength=1,
interpolation="linear") +
# ElasticAugment(
# control_point_spacing=[4,40,40],
# jitter_sigma=[0,2,2],
# rotation_interval=[0,math.pi/2.0],
# prob_slip=0.05,
# prob_shift=0.05,
# max_misalign=10,
# subsample=8) +
# SimpleAugment(transpose_only=[1, 2]) +
# IntensityAugment(labels, 0.9, 1.1, -0.1, 0.1, z_section_wise=True) +
AddAffinities(
affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=gt_affs_in_key) +
GrowBoundary(labels_key, steps=1, only_xy=True) +
AddAffinities(
affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=gt_affs_out_key,
affinities_mask=gt_affs_mask) +
AddJoinedAffinities(
input_affinities=gt_affs_in_key,
joined_affinities=joined_affinities_key) +
AddRealism(
joined_affinities = joined_affinities_key,
raw = raw_key,
sp=0.25,
sigma=1,
contrast=0.7) +
BalanceLabels(
labels=gt_affs_out_key,
scales=input_affinities_scale_key) +
DefectAugment(
intensities=raw_key,
prob_missing=0.03,
prob_low_contrast=0.01,
contrast_scale=0.5,
axis=0) +
IntensityScaleShift(raw_key, 2,-1) +
Crop(
key=labels_key,
roi=crop_roi) +
RenumberConnectedComponents(labels=labels_key) +
PreCache(
cache_size=32,
num_workers=8) +
Train(
graph='train/prob_unet/setup_1/train_net',
optimizer=add_loss,
loss=None,
inputs={
"gt_seg:0": labels_key,
"gt_affs_mask:0": gt_affs_mask,
config['raw']: raw_key,
config['gt_affs_in']: gt_affs_in_key,
config['gt_affs_out']: gt_affs_out_key,
config['pred_affs_loss_weights']: input_affinities_scale_key
},
outputs={
config['pred_affs']: pred_affinities_key
},
gradients={
config['pred_affs']: pred_affinities_gradient_key
},
summary="Merge/MergeSummary:0",
log_dir='log/prob_unet/setup_2',
save_every=500) +
IntensityScaleShift(
array=raw_key,
scale=0.5,
shift=0.5) +
Snapshot(
dataset_names={
labels_key: 'volumes/labels',
raw_key: 'volumes/raw',
pred_affinities_key: 'volumes/pred_affs',
gt_affs_out_key: 'volumes/output_affs'
},
output_filename='prob_unet/setup_2/batch_{iteration}.hdf',
every=1000,
dataset_dtypes={
labels_key: np.uint64,
raw_key: np.float32
}) +
PrintProfilingStats(every=20)
)
print("Starting training...")
with build(pipeline) as p:
for i in range(iterations - trained_until):
req = p.request_batch(request)
# print ("labels: ", req[labels_key].data.shape)
# print ("affinities_in: ", req[gt_affs_in_key].data.shape)
# print ("affinities_out: ", req[gt_affs_out_key].data.shape)
# print ("affinities_joined: ", req[joined_affinities_key].data.shape)
# print ("raw: ", req[raw_key].data.shape)
# print ("affinities_in_scale: ", req[input_affinities_scale_key].data.shape)
print("Training finished")
def generate_full_samples(num_batches):
neighborhood = [[-1, 0, 0], [0, -1, 0], [0, 0, -1]]
neighborhood_opp = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
# define array-keys
labels_key = ArrayKey('LABELS')
gt_affs_key = ArrayKey('RAW_AFFINITIES')
gt_joined_affs_key = ArrayKey('RAW_JOINED_AFFINITIES')
raw_key = ArrayKey('RAW')
merged_labels_key = []
merged_affs_key = []
affs_neg_key = ArrayKey('AFFINITIES')
affs_pos_key = ArrayKey('AFFINITIES_OPP')
joined_affs_neg_key = ArrayKey('JOINED_AFFINITIES')
joined_affs_pos_key = ArrayKey('JOINED_AFFINITIES_OPP')
num_merges = 3
for i in range(num_merges):
merged_labels_key.append(ArrayKey('MERGED_LABELS_%i' % (i + 1)))
merged_affs_key.append(ArrayKey('MERGED_AFFINITIES_IN_%i' % (i + 1)))
voxel_size = Coordinate((1, 1, 1))
input_shape = Coordinate((132, 132, 132)) * voxel_size
input_affs_shape = Coordinate([i + 1
for i in (132, 132, 132)]) * voxel_size
output_shape = Coordinate((44, 44, 44)) * voxel_size
output_affs_shape = Coordinate([i + 1 for i in (44, 44, 44)]) * voxel_size
print("input_shape: ", input_shape)
print("output_shape: ", output_shape)
request = BatchRequest()
request.add(labels_key, input_shape)
# request.add(gt_affs_key, input_shape)
# request.add(gt_joined_affs_key, input_shape)
# request.add(raw_key, input_shape)
request.add(affs_neg_key, input_shape)
request.add(affs_pos_key, input_shape)
request.add(joined_affs_neg_key, input_shape)
request.add(joined_affs_pos_key, input_shape)
for i in range(num_merges):
request.add(merged_labels_key[i], input_shape)
# request.add(merged_affs_key[i], input_shape)
pipeline = ()
pipeline += ToyNeuronSegmentationGenerator(array_key=labels_key,
n_objects=50,
points_per_skeleton=8,
smoothness=3,
noise_strength=1,
interpolation="linear")
# pipeline += AddAffinities(
# affinity_neighborhood=neighborhood,
# labels=labels_key,
# affinities=gt_affs_key)
# pipeline += AddJoinedAffinities(
# input_affinities=gt_affs_key,
# joined_affinities=gt_joined_affs_key)
# pipeline += AddRealism(
# joined_affinities = gt_joined_affs_key,
# raw = raw_key,
# sp=0.25,
# sigma=1,
# contrast=0.7)
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=affs_neg_key)
pipeline += AddAffinities(affinity_neighborhood=neighborhood_opp,
labels=labels_key,
affinities=affs_pos_key)
pipeline += AddJoinedAffinities(input_affinities=affs_neg_key,
joined_affinities=joined_affs_neg_key)
pipeline += AddJoinedAffinities(input_affinities=affs_pos_key,
joined_affinities=joined_affs_pos_key)
for i in range(num_merges):
pipeline += MergeLabels(labels=labels_key,
joined_affinities=joined_affs_neg_key,
joined_affinities_opp=joined_affs_pos_key,
merged_labels=merged_labels_key[i],
cropped_roi=None,
every=1)
# pipeline += AddAffinities(
# affinity_neighborhood=neighborhood,
# labels=merged_labels_key[i],
# affinities=merged_affs_key[i])
pipeline += PreCache(cache_size=32, num_workers=8)
dataset_names = {
labels_key: 'volumes/labels',
# affs_neg_key: 'volumes/gt_affs'
}
dataset_dtypes = {
labels_key: np.uint16,
# affs_neg_key: np.float32
}
for i in range(num_merges):
dataset_names[merged_labels_key[i]] = 'volumes/merged_labels_%i' % (i +
1)
# dataset_names[merged_affs_key[i]] = 'volumes/merged_affs_%i'%(i+1)
dataset_dtypes[merged_labels_key[i]] = np.uint16
# dataset_dtypes[merged_affs_key[i]] = np.float32
pipeline += Snapshot(
dataset_names=dataset_names,
output_filename='gt_1_merge_3/batch_{id}.hdf',
# output_filename='test_affs.hdf',
every=1,
dataset_dtypes=dataset_dtypes)
pipeline += PrintProfilingStats(every=10)
hashes = []
with build(pipeline) as p:
for i in range(num_batches):
req = p.request_batch(request)
label_hash = np.sum(req[labels_key].data)
def predict(checkpoint, iterations):
print ("checkpoint: ", checkpoint)
labels_key = ArrayKey('GT_LABELS')
joined_affinities_key = ArrayKey('GT_JOINED_AFFINITIES')
raw_affinities_key = ArrayKey('RAW_AFFINITIES_KEY')
raw_key = ArrayKey('RAW')
pred_affinities_key = ArrayKey('PREDICTED_AFFS')
voxel_size = Coordinate((1, 1, 1))
input_shape = Coordinate(config['input_shape']) * voxel_size
output_shape = Coordinate(config['output_shape']) * voxel_size
print ("input_size: ", input_shape)
print ("output_size: ", output_shape)
request = BatchRequest()
# request.add(labels_key, input_shape)
request.add(joined_affinities_key, input_shape)
request.add(raw_affinities_key, input_shape)
request.add(raw_key, input_shape)
request.add(pred_affinities_key, output_shape)
pipeline = (
ToyNeuronSegmentationGenerator(
array_key=labels_key,
n_objects=50,
points_per_skeleton=8,
smoothness=3,
noise_strength=1,
interpolation="random",
seed=0) +
AddAffinities(
affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=raw_affinities_key) +
AddJoinedAffinities(
input_affinities=raw_affinities_key,
joined_affinities=joined_affinities_key) +
AddRealism(
joined_affinities = joined_affinities_key,
raw = raw_key,
sp=0.65,
sigma=1,
contrast=0.7) +
# Pad(raw_key, size=None) +
# Crop(raw_key, read_roi) +
# Normalize(raw_key) +
IntensityScaleShift(raw_key, 2,-1) +
Predict(
checkpoint = os.path.join(setup_dir, 'train_net_checkpoint_%d' % checkpoint),
inputs={
config['raw']: raw_key
},
outputs={
config['pred_affs']: pred_affinities_key
},
graph=os.path.join(setup_dir, 'predict_net.meta')
) +
IntensityScaleShift(
array=raw_key,
scale=0.5,
shift=0.5) +
Snapshot(
dataset_names={
labels_key: 'volumes/labels',
raw_affinities_key: 'volumes/raw_affs',
raw_key: 'volumes/raw',
pred_affinities_key: 'volumes/pred_affs',
},
output_filename='prob_unet/prediction_{id}.hdf',
every=1,
dataset_dtypes={
labels_key: np.uint16,
raw_key: np.float32,
pred_affinities_key: np.float32,
sample_z_key: np.float32
})
# PrintProfilingStats(every=1)
)
print("Starting prediction...")
with build(pipeline) as p:
for i in range(iterations):
p.request_batch(request)
print("Prediction finished")
def train(iterations):
# tf.reset_default_graph()
if tf.train.latest_checkpoint('.'):
trained_until = int(tf.train.latest_checkpoint('.').split('_')[-1])
else:
trained_until = 0
if trained_until >= iterations:
return
if trained_until < phase_switch and iterations > phase_switch:
train(phase_switch)
phase = 'euclid' if iterations <= phase_switch else 'malis'
print("Training in phase %s until %i" % (phase, iterations))
# define array-keys
labels_key = ArrayKey('LABELS')
raw_affs_key = ArrayKey('RAW_AFFINITIES')
raw_joined_affs_key = ArrayKey('RAW_JOINED_AFFINITIES')
raw_key = ArrayKey('RAW')
affs_neg_key = ArrayKey('AFFINITIES')
affs_pos_key = ArrayKey('AFFINITIES_OPP')
joined_affs_neg_key = ArrayKey('JOINED_AFFINITIES')
joined_affs_pos_key = ArrayKey('JOINED_AFFINITIES_OPP')
gt_affs_out_key = ArrayKey('GT_AFFINITIES')
gt_affs_in_key = ArrayKey('GT_AFFINITIES_IN')
gt_affs_mask_key = ArrayKey('GT_AFFINITIES_MASK')
gt_affs_scale_key = ArrayKey('GT_AFFINITIES_SCALE')
pred_affs_key = ArrayKey('PRED_AFFS')
pred_affs_gradient_key = ArrayKey('PRED_AFFS_GRADIENT')
sample_z_key = ArrayKey("SAMPLE_Z")
broadcast_key = ArrayKey("BROADCAST")
pred_logits_key = ArrayKey("PRED_LOGITS")
sample_out_key = ArrayKey("SAMPLE_OUT")
debug_key = ArrayKey("DEBUG")
voxel_size = Coordinate((1, 1, 1))
input_shape = Coordinate(config['input_shape']) * voxel_size
input_affs_shape = Coordinate([i + 1 for i in config['input_shape']
]) * voxel_size
output_shape = Coordinate(config['output_shape']) * voxel_size
output_affs_shape = Coordinate([i + 1 for i in config['output_shape']
]) * voxel_size
sample_shape = Coordinate((1, 1, config['latent_dims'])) * voxel_size
debug_shape = Coordinate((1, 1, 5)) * voxel_size
print("input_shape: ", input_shape)
print("input_affs_shape: ", input_affs_shape)
print("output_shape: ", output_shape)
print("output_affs_shape: ", output_affs_shape)
request = BatchRequest()
request.add(labels_key, input_shape)
request.add(raw_affs_key, input_shape)
request.add(raw_joined_affs_key, input_shape)
request.add(raw_key, input_shape)
request.add(gt_affs_out_key, output_shape)
request.add(gt_affs_in_key, input_shape)
request.add(gt_affs_mask_key, output_shape)
request.add(gt_affs_scale_key, output_shape)
request.add(pred_affs_key, output_shape)
request.add(pred_affs_gradient_key, output_shape)
request.add(broadcast_key, output_shape)
request.add(sample_z_key, sample_shape)
request.add(pred_logits_key, output_shape)
request.add(sample_out_key, sample_shape)
request.add(debug_key, debug_shape)
pipeline = ()
# print ("iteration: ", iteration)
pipeline += ToyNeuronSegmentationGenerator(array_key=labels_key,
n_objects=50,
points_per_skeleton=8,
smoothness=3,
noise_strength=1,
interpolation="linear")
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=raw_affs_key)
pipeline += AddJoinedAffinities(input_affinities=raw_affs_key,
joined_affinities=raw_joined_affs_key)
pipeline += AddRealism(joined_affinities=raw_joined_affs_key,
raw=raw_key,
sp=0.25,
sigma=1,
contrast=0.7)
pipeline += RenumberConnectedComponents(labels=labels_key)
pipeline += GrowBoundary(labels_key, steps=1, only_xy=True)
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=gt_affs_in_key)
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=gt_affs_out_key,
affinities_mask=gt_affs_mask_key)
pipeline += BalanceLabels(labels=gt_affs_out_key, scales=gt_affs_scale_key)
pipeline += DefectAugment(intensities=raw_key,
prob_missing=0.03,
prob_low_contrast=0.01,
contrast_scale=0.5,
axis=0)
pipeline += IntensityScaleShift(raw_key, 2, -1)
pipeline += PreCache(cache_size=8, num_workers=4)
train_inputs = {
config['raw']: raw_key,
config['gt_affs_in']: gt_affs_in_key,
config['gt_affs_out']: gt_affs_out_key,
config['pred_affs_loss_weights']: gt_affs_scale_key
}
if phase == 'euclid':
train_loss = config['loss']
train_optimizer = config['optimizer']
train_summary = config['summary']
else:
train_loss = None
train_optimizer = add_malis_loss
train_inputs['gt_seg:0'] = labels_key
train_inputs['gt_affs_mask:0'] = gt_affs_mask_key
train_summary = 'Merge/MergeSummary:0'
pipeline += Train(graph=setup_dir + 'train_net',
optimizer=train_optimizer,
loss=train_loss,
inputs=train_inputs,
outputs={
config['pred_affs']: pred_affs_key,
config['broadcast']: broadcast_key,
config['sample_z']: sample_z_key,
config['pred_logits']: pred_logits_key,
config['sample_out']: sample_out_key,
config['debug']: debug_key
},
gradients={config['pred_affs']: pred_affs_gradient_key},
summary=train_summary,
log_dir='log/prob_unet/' + setup_name,
save_every=2000)
pipeline += IntensityScaleShift(array=raw_key, scale=0.5, shift=0.5)
pipeline += Snapshot(dataset_names={
labels_key: 'volumes/labels',
raw_affs_key: 'volumes/raw_affs',
raw_key: 'volumes/raw',
gt_affs_in_key: 'volumes/gt_affs_in',
gt_affs_out_key: 'volumes/gt_affs_out',
pred_affs_key: 'volumes/pred_affs',
pred_logits_key: 'volumes/pred_logits'
},
output_filename='prob_unet/' + setup_name +
'/batch_{iteration}.hdf',
every=2000,
dataset_dtypes={
labels_key: np.uint16,
raw_affs_key: np.float32,
raw_key: np.float32,
gt_affs_in_key: np.float32,
gt_affs_out_key: np.float32,
pred_affs_key: np.float32,
pred_logits_key: np.float32
})
pipeline += PrintProfilingStats(every=20)
print("Starting training...")
with build(pipeline) as p:
for i in range(iterations - trained_until):
req = p.request_batch(request)
# sample_z = req[sample_z_key].data
# broadcast_sample = req[broadcast_key].data
# sample_out = req[sample_out_key].data
# debug = req[debug_key].data
# print("debug", debug)
# print("sample_z: ", sample_z)
# print("sample_out:", sample_out)
# print("Z - 0")
# print(np.unique(broadcast_sample[0, 0, :, :, :]))
# print("Z - 1")
# print(np.unique(broadcast_sample[0, 1, :, :, :]))
# print("Z - 2")
# print(np.unique(broadcast_sample[0, 2, :, :, :]))
# print("Z - 3")
# print(np.unique(broadcast_sample[0, 3, :, :, :]))
# print("Z - 4")
# print(np.unique(broadcast_sample[0, 4, :, :, :]))
# print("Z - 5")
# print(np.unique(broadcast_sample[0, 5, :, :, :]))
print("Training finished")
def train(iterations):
tf.reset_default_graph()
if tf.train.latest_checkpoint('.'):
trained_until = int(tf.train.latest_checkpoint('.').split('_')[-1])
else:
trained_until = 0
if trained_until >= iterations:
return
# define array-keys
labels_key = ArrayKey('GT_LABELS')
gt_affs_in_key = ArrayKey('GT_AFFINITIES_IN')
gt_affs_out_key = ArrayKey('GT_AFFINITIES_OUT')
joined_affinities_key = ArrayKey('GT_JOINED_AFFINITIES')
raw_key = ArrayKey('RAW')
input_affinities_scale_key = ArrayKey('GT_AFFINITIES_SCALE')
pred_affinities_key = ArrayKey('PREDICTED_AFFS')
pred_affinities_gradient_key = ArrayKey('AFFS_GRADIENT')
gt_affs_mask = ArrayKey('GT_AFFINITIES_MASK')
debug_key = ArrayKey("DEBUG")
voxel_size = Coordinate((1, 1, 1))
input_shape = Coordinate(config['input_shape']) * voxel_size
output_shape = Coordinate(config['output_shape']) * voxel_size
debug_shape = Coordinate((1, 1, 5)) * voxel_size
print("input_shape: ", input_shape)
print("output_shape: ", output_shape)
# define requests
request = BatchRequest()
request.add(labels_key, output_shape)
request.add(gt_affs_in_key, input_shape)
request.add(joined_affinities_key, input_shape)
request.add(raw_key, input_shape)
request.add(gt_affs_out_key, output_shape)
request.add(input_affinities_scale_key, output_shape)
request.add(pred_affinities_key, output_shape)
request.add(gt_affs_mask, output_shape)
request.add(debug_key, debug_shape)
offset = Coordinate((input_shape[i] - output_shape[i]) / 2
for i in range(len(input_shape)))
crop_roi = Roi(offset, output_shape)
# print("crop_roi: ", crop_roi)
pipeline = ()
pipeline += ToyNeuronSegmentationGenerator(array_key=labels_key,
n_objects=50,
points_per_skeleton=8,
smoothness=3,
noise_strength=1,
interpolation="linear")
# ElasticAugment(
# control_point_spacing=[4,40,40],
# jitter_sigma=[0,2,2],
# rotation_interval=[0,math.pi/2.0],
# prob_slip=0.05,
# prob_shift=0.05,
# max_misalign=10,
# subsample=8) +
# SimpleAugment(transpose_only=[1, 2]) +
# IntensityAugment(labels, 0.9, 1.1, -0.1, 0.1, z_section_wise=True) +
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=gt_affs_in_key)
pipeline += GrowBoundary(labels_key, steps=1, only_xy=True)
pipeline += AddAffinities(affinity_neighborhood=neighborhood,
labels=labels_key,
affinities=gt_affs_out_key,
affinities_mask=gt_affs_mask)
pipeline += AddJoinedAffinities(input_affinities=gt_affs_in_key,
joined_affinities=joined_affinities_key)
pipeline += AddRealism(joined_affinities=joined_affinities_key,
raw=raw_key,
sp=0.25,
sigma=1,
contrast=0.7)
pipeline += BalanceLabels(labels=gt_affs_out_key,
scales=input_affinities_scale_key)
pipeline += DefectAugment(intensities=raw_key,
prob_missing=0.03,
prob_low_contrast=0.01,
contrast_scale=0.5,
axis=0)
pipeline += IntensityScaleShift(raw_key, 2, -1)
pipeline += PreCache(cache_size=32, num_workers=8)
pipeline += Crop(key=labels_key, roi=crop_roi)
pipeline += RenumberConnectedComponents(labels=labels_key)
train = Train(
graph='train/unet/train_net',
# optimizer=config['optimizer'],
optimizer=add_malis_loss,
# loss=config['loss'],
loss=None,
inputs={
config['raw']: raw_key,
"gt_seg:0": labels_key,
"gt_affs_mask:0": gt_affs_mask,
config['gt_affs']: gt_affs_out_key,
config['pred_affs_loss_weights']: input_affinities_scale_key,
},
outputs={
config['pred_affs']: pred_affinities_key,
config['debug']: debug_key,
},
gradients={config['pred_affs']: pred_affinities_gradient_key},
summary="malis_loss:0",
log_dir='log/unet',
save_every=1)
pipeline += train
pipeline += IntensityScaleShift(array=raw_key, scale=0.5, shift=0.5)
# Snapshot(
# dataset_names={
# labels_key: 'volumes/labels',
# raw_key: 'volumes/raw',
# gt_affs_out_key: 'volumes/gt_affs',
# pred_affinities_key: 'volumes/pred_affs'
# },
# output_filename='unet/train/batch_{iteration}.hdf',
# every=100,
# dataset_dtypes={
# raw_key: np.float32,
# labels_key: np.uint64
# }) +
pipeline += PrintProfilingStats(every=8)
print("Starting training... COOL BEANS")
with build(pipeline) as p:
for i in range(iterations - trained_until):
req = p.request_batch(request)
pred_affs = req[pred_affinities_key].data
debug = req[debug_key].data
print("debug", debug)
print('pred_affs', pred_affs)
print("name of pred_adds: ", req[pred_affinities_key])
# print("train session: ", train.session)
# print ("all vars: ", [n.name for n in tf.get_default_graph().as_graph_def().node])
# graph_def = tf.graph_util.convert_variables_to_constants(train.session, tf.get_default_graph().as_graph_def(), ["pred_affs:0".split(':')[0]])
# print ("labels: ", req[labels_key].data.shape)
# print ("affinities_out: ", req[gt_affs_out_key].data.shape)
# print ("affinities_joined: ", req[joined_affinities_key].data.shape)
# print ("raw: ", req[raw_key].data.shape)
# print ("affinities_in_scale: ", req[input_affinities_scale_key].data.shape)
print("Training finished")