def make_data_provider(provider_string):
data_providers = []
# data_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data/realigned'
# file_pattern = '*merged*fixed-offset-fixed-mask.h5'
pattern = provider_string.split(':')[0]
paths = {**DEFAULT_PATHS}
paths.update(**{entry.split('=')[0].lower() : entry.split('=')[1] for entry in provider_string.split(':')[1:]})
for data in glob.glob(pattern):
h5_source = Hdf5Source(
data,
datasets={
RAW_KEY: paths['raw'],
GT_LABELS_KEY: paths['labels'],
GT_MASK_KEY: paths['mask']
},
array_specs={
GT_MASK_KEY: ArraySpec(interpolatable=False)
}
)
data_providers.append(h5_source)
return tuple(data_providers)
def test_pipeline3(self):
array_key = ArrayKey("TEST_ARRAY")
points_key = GraphKey("TEST_POINTS")
voxel_size = Coordinate((1, 1))
spec = ArraySpec(voxel_size=voxel_size, interpolatable=True)
hdf5_source = Hdf5Source(self.fake_data_file, {array_key: "testdata"},
array_specs={array_key: spec})
csv_source = CsvPointsSource(
self.fake_points_file,
points_key,
GraphSpec(roi=Roi(shape=Coordinate((100, 100)), offset=(0, 0))),
)
request = BatchRequest()
shape = Coordinate((60, 60))
request.add(array_key, shape, voxel_size=Coordinate((1, 1)))
request.add(points_key, shape)
shift_node = ShiftAugment(prob_slip=0.2,
prob_shift=0.2,
sigma=5,
shift_axis=0)
pipeline = ((hdf5_source, csv_source) + MergeProvider() +
RandomLocation(ensure_nonempty=points_key) + shift_node)
with build(pipeline) as b:
request = b.request_batch(request)
# print(request[points_key])
target_vals = [
self.fake_data[point[0]][point[1]] for point in self.fake_points
]
result_data = request[array_key].data
result_points = list(request[points_key].nodes)
result_vals = [
result_data[int(point.location[0])][int(point.location[1])]
for point in result_points
]
for result_val in result_vals:
self.assertTrue(
result_val in target_vals,
msg=
"result value {} at points {} not in target values {} at points {}"
.format(
result_val,
list(result_points),
target_vals,
self.fake_points,
),
)
def test_prepare1(self):
key = ArrayKey("TEST_ARRAY")
spec = ArraySpec(voxel_size=Coordinate((1, 1)), interpolatable=True)
hdf5_source = Hdf5Source(self.fake_data_file, {key: "testdata"},
array_specs={key: spec})
request = BatchRequest()
shape = Coordinate((3, 3))
request.add(key, shape, voxel_size=Coordinate((1, 1)))
shift_node = ShiftAugment(sigma=1, shift_axis=0)
with build((hdf5_source + shift_node)):
shift_node.prepare(request)
self.assertTrue(shift_node.ndim == 2)
self.assertTrue(shift_node.shift_sigmas == tuple([0.0, 1.0]))
def test_pipeline2(self):
key = ArrayKey("TEST_ARRAY")
spec = ArraySpec(voxel_size=Coordinate((3, 1)), interpolatable=True)
hdf5_source = Hdf5Source(self.fake_data_file, {key: "testdata"},
array_specs={key: spec})
request = BatchRequest()
shape = Coordinate((3, 3))
request.add(key, shape, voxel_size=Coordinate((3, 1)))
shift_node = ShiftAugment(prob_slip=0.2,
prob_shift=0.2,
sigma=1,
shift_axis=0)
with build((hdf5_source + shift_node)) as b:
b.request_batch(request)
def make_data_provider(provider_string, required_paths, default_paths,
lower_key_identifier_to_key):
data_providers = []
# data_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data/realigned'
# file_pattern = '*merged*fixed-offset-fixed-mask.h5'
pattern = provider_string.split(':')[0]
paths = {r: default_paths[r] for r in required_paths if r in default_paths}
paths.update(
**{
entry.split('=')[0].lower(): entry.split('=')[1]
for entry in provider_string.split(':')[1:]
})
for r in required_paths:
assert r in paths, "Required path {} not provided in {}".format(
r, paths)
for k, _ in paths.items():
assert k in lower_key_identifier_to_key, "Identifier %s not in array key map %s" % (
k, lower_key_identifier_to_key)
datasets = {lower_key_identifier_to_key[k]: v for k, v in paths.items()}
specs = {
MASK_KEY: ArraySpec(interpolatable=False),
GLIA_MASK_KEY: ArraySpec(interpolatable=False)
}
for data in glob.glob(pattern):
if data.endswith('h5') or data.endswith('hdf'):
source = Hdf5Source(data, datasets=datasets, array_specs=specs)
else:
source = Z5Source(data,
datasets=datasets,
array_specs=specs,
revert=False)
data_providers.append(source)
return tuple(data_providers)
def predict():
iteration = 18000
prototxt = '/groups/turaga/home/grisaitisw/src/syntist/gunpowder/tasks/training/net_test.prototxt'
weights = os.path.join(
'/groups/turaga/home/grisaitisw/experiments/20170524/090000',
'net_iter_%d.caffemodel' % iteration)
input_size = Coordinate((236, 236, 236))
output_size = Coordinate((148, 148, 148))
pipeline = (Hdf5Source(
'/groups/turaga/home/grisaitisw/src/gunpowder/examples/fib25/tstvol-520-1-h5_y0_x0_xy0_angle000.0.h5',
raw_dataset='volumes/raw',
resolution=(8, 8, 8),
) + Normalize() + Padding() + Predict(prototxt, weights, use_gpu=0) +
Snapshot(every=1,
output_dir=os.path.join('chunks', '%d' % iteration),
output_filename='chunk_{id}.hdf',
with_timestamp=True) + PrintProfilingStats() + Chunk(
BatchSpec(
input_size,
output_size,
resolution=(8, 8, 8),
)) +
Snapshot(every=1,
output_dir=os.path.join('processed',
'%d' % iteration),
output_filename='tstvol-520-1-h5_2_{id}.hdf'))
# request a "batch" of the size of the whole dataset
with gunpowder.build(pipeline) as p:
shape = p.get_spec().roi.get_shape()
p.request_batch(
BatchSpec(
shape,
shape - (input_size - output_size),
resolution=(8, 8, 8),
))
def visualize_augmentations_paintera(args=None):
data_providers = []
data_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data'
# data_dir = os.path.expanduser('~/Dropbox/cremi-upsampled/')
file_pattern = 'sample_A_padded_20160501-2-additional-sections-fixed-offset.h5'
file_pattern = 'sample_B_padded_20160501-2-additional-sections-fixed-offset.h5'
file_pattern = 'sample_C_padded_20160501-2-additional-sections-fixed-offset.h5'
defect_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data/defects'
artifact_source = (
Hdf5Source(
os.path.join(defect_dir, 'sample_ABC_padded_20160501.defects.hdf'),
datasets={
RAW: 'defect_sections/raw',
ALPHA_MASK: 'defect_sections/mask',
},
array_specs={
RAW: ArraySpec(voxel_size=tuple(d * 9 for d in (40, 4, 4))),
ALPHA_MASK: ArraySpec(voxel_size=tuple(d * 9
for d in (40, 4, 4))),
}) + RandomLocation(min_masked=0.05, mask=ALPHA_MASK) +
Normalize(RAW) +
IntensityAugment(RAW, 0.9, 1.1, -0.1, 0.1, z_section_wise=True) +
ElasticAugment(voxel_size=(360, 36, 36),
control_point_spacing=(4, 40, 40),
control_point_displacement_sigma=(0, 2 * 36, 2 * 36),
rotation_interval=(0, np.pi / 2.0),
subsample=8) + SimpleAugment(transpose_only=[1, 2]))
for data in glob.glob(os.path.join(data_dir, file_pattern)):
h5_source = Hdf5Source(data,
datasets={
RAW:
'volumes/raw',
GT_LABELS:
'volumes/labels/neuron_ids-downsampled',
})
data_providers.append(h5_source)
input_resolution = (360.0, 36.0, 36.0)
output_resolution = Coordinate((120.0, 108.0, 108.0))
offset = (13640, 10932, 10932)
output_shape = Coordinate((60.0, 100.0, 100.0)) * output_resolution
output_offset = (13320 + 3600, 32796 + 36 + 10800, 32796 + 36 + 10800)
overhang = Coordinate((360.0, 108.0, 108.0)) * 16
input_shape = output_shape + overhang * 2
input_offset = Coordinate(output_offset) - overhang
output_roi = Roi(offset=output_offset, shape=output_shape)
input_roi = Roi(offset=input_offset, shape=input_shape)
augmentations = (
Snapshot(dataset_names={
GT_LABELS: 'volumes/gt',
RAW: 'volumes/raw'
},
output_dir='.',
output_filename='snapshot-before.h5',
attributes_callback=Snapshot.default_attributes_callback()),
ElasticAugment(voxel_size=(360.0, 36.0, 36.0),
control_point_spacing=(4, 40, 40),
control_point_displacement_sigma=(0, 5 * 2 * 36,
5 * 2 * 36),
rotation_interval=(0 * np.pi / 8, 0 * 2 * np.pi),
subsample=8,
augmentation_probability=1.0,
seed=None),
Misalign(z_resolution=360,
prob_slip=0.2,
prob_shift=0.5,
max_misalign=(3600, 0),
seed=100,
ignore_keys_for_slip=(GT_LABELS, )),
Snapshot(dataset_names={
GT_LABELS: 'volumes/gt',
RAW: 'volumes/raw'
},
output_dir='.',
output_filename='snapshot-after.h5',
attributes_callback=Snapshot.default_attributes_callback()))
keys = (RAW, GT_LABELS)[:]
batch, snapshot = run_augmentations(
data_providers=data_providers,
roi=lambda key: output_roi.copy()
if key == GT_LABELS else input_roi.copy(),
augmentations=augmentations,
keys=keys,
voxel_size=lambda key: input_resolution
if key == RAW else (output_resolution if key == GT_LABELS else None))
args = get_parser().parse_args() if args is None else args
jnius_config.add_options('-Xmx{}'.format(args.max_heap_size))
import payntera.jfx
from jnius import autoclass
payntera.jfx.init_platform()
PainteraBaseView = autoclass(
'org.janelia.saalfeldlab.paintera.PainteraBaseView')
viewer = PainteraBaseView.defaultView()
pbv = viewer.baseView
scene, stage = payntera.jfx.start_stage(viewer.paneWithStatus.getPane())
screen_scale_setter = lambda: pbv.orthogonalViews().setScreenScales(
[0.3, 0.1, 0.03])
# payntera.jfx.invoke_on_jfx_application_thread(screen_scale_setter)
snapshot_states = add_to_viewer(
snapshot, keys=keys, name=lambda key: '%s-snapshot' % key.identifier)
states = add_to_viewer(batch, keys=keys)
viewer.keyTracker.installInto(scene)
scene.addEventFilter(
autoclass('javafx.scene.input.MouseEvent').ANY, viewer.mouseTracker)
while stage.isShowing():
time.sleep(0.1)
def train_until(
data_providers,
affinity_neighborhood,
meta_graph_filename,
stop,
input_shape,
output_shape,
loss,
optimizer,
tensor_affinities,
tensor_affinities_mask,
tensor_glia,
tensor_glia_mask,
summary,
save_checkpoint_every,
pre_cache_size,
pre_cache_num_workers,
snapshot_every,
balance_labels,
renumber_connected_components,
network_inputs,
ignore_labels_for_slip,
grow_boundaries,
mask_out_labels,
snapshot_dir):
ignore_keys_for_slip = (LABELS_KEY, GT_MASK_KEY, GT_GLIA_KEY, GLIA_MASK_KEY, UNLABELED_KEY) if ignore_labels_for_slip else ()
defect_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data/defects'
if tf.train.latest_checkpoint('.'):
trained_until = int(tf.train.latest_checkpoint('.').split('_')[-1])
print('Resuming training from', trained_until)
else:
trained_until = 0
print('Starting fresh training')
input_voxel_size = Coordinate((120, 12, 12)) * 3
output_voxel_size = Coordinate((40, 36, 36)) * 3
input_size = Coordinate(input_shape) * input_voxel_size
output_size = Coordinate(output_shape) * output_voxel_size
num_affinities = sum(len(nh) for nh in affinity_neighborhood)
gt_affinities_size = Coordinate((num_affinities,) + tuple(s for s in output_size))
print("gt affinities size", gt_affinities_size)
# TODO why is GT_AFFINITIES three-dimensional? compare to
# TODO https://github.com/funkey/gunpowder/blob/master/examples/cremi/train.py#L35
# TODO Use glia scale somehow, probably not possible with tensorflow 1.3 because it does not know uint64...
# specifiy which Arrays should be requested for each batch
request = BatchRequest()
request.add(RAW_KEY, input_size, voxel_size=input_voxel_size)
request.add(LABELS_KEY, output_size, voxel_size=output_voxel_size)
request.add(GT_AFFINITIES_KEY, output_size, voxel_size=output_voxel_size)
request.add(AFFINITIES_MASK_KEY, output_size, voxel_size=output_voxel_size)
request.add(GT_MASK_KEY, output_size, voxel_size=output_voxel_size)
request.add(GLIA_MASK_KEY, output_size, voxel_size=output_voxel_size)
request.add(GLIA_KEY, output_size, voxel_size=output_voxel_size)
request.add(GT_GLIA_KEY, output_size, voxel_size=output_voxel_size)
request.add(UNLABELED_KEY, output_size, voxel_size=output_voxel_size)
if balance_labels:
request.add(AFFINITIES_SCALE_KEY, output_size, voxel_size=output_voxel_size)
# always balance glia labels!
request.add(GLIA_SCALE_KEY, output_size, voxel_size=output_voxel_size)
network_inputs[tensor_affinities_mask] = AFFINITIES_SCALE_KEY if balance_labels else AFFINITIES_MASK_KEY
network_inputs[tensor_glia_mask] = GLIA_SCALE_KEY#GLIA_SCALE_KEY if balance_labels else GLIA_MASK_KEY
# create a tuple of data sources, one for each HDF file
data_sources = tuple(
provider +
Normalize(RAW_KEY) + # ensures RAW is in float in [0, 1]
# zero-pad provided RAW and GT_MASK to be able to draw batches close to
# the boundary of the available data
# size more or less irrelevant as followed by Reject Node
Pad(RAW_KEY, None) +
Pad(GT_MASK_KEY, None) +
Pad(GLIA_MASK_KEY, None) +
Pad(LABELS_KEY, size=NETWORK_OUTPUT_SHAPE / 2, value=np.uint64(-3)) +
Pad(GT_GLIA_KEY, size=NETWORK_OUTPUT_SHAPE / 2) +
# Pad(LABELS_KEY, None) +
# Pad(GT_GLIA_KEY, None) +
RandomLocation() + # chose a random location inside the provided arrays
Reject(mask=GT_MASK_KEY, min_masked=0.5) +
Reject(mask=GLIA_MASK_KEY, min_masked=0.5) +
MapNumpyArray(lambda array: np.require(array, dtype=np.int64), GT_GLIA_KEY) # this is necessary because gunpowder 1.3 only understands int64, not uint64
for provider in data_providers)
# TODO figure out what this is for
snapshot_request = BatchRequest({
LOSS_GRADIENT_KEY : request[GT_AFFINITIES_KEY],
AFFINITIES_KEY : request[GT_AFFINITIES_KEY],
})
# no need to do anything here. random sections will be replaced with sections from this source (only raw)
artifact_source = (
Hdf5Source(
os.path.join(defect_dir, 'sample_ABC_padded_20160501.defects.hdf'),
datasets={
RAW_KEY : 'defect_sections/raw',
DEFECT_MASK_KEY : 'defect_sections/mask',
},
array_specs={
RAW_KEY : ArraySpec(voxel_size=input_voxel_size),
DEFECT_MASK_KEY : ArraySpec(voxel_size=input_voxel_size),
}
) +
RandomLocation(min_masked=0.05, mask=DEFECT_MASK_KEY) +
Normalize(RAW_KEY) +
IntensityAugment(RAW_KEY, 0.9, 1.1, -0.1, 0.1, z_section_wise=True) +
ElasticAugment(
voxel_size=(360, 36, 36),
control_point_spacing=(4, 40, 40),
control_point_displacement_sigma=(0, 2 * 36, 2 * 36),
rotation_interval=(0, math.pi / 2.0),
subsample=8
) +
SimpleAugment(transpose_only=[1,2])
)
train_pipeline = data_sources
train_pipeline += RandomProvider()
train_pipeline += ElasticAugment(
voxel_size=(360, 36, 36),
control_point_spacing=(4, 40, 40),
control_point_displacement_sigma=(0, 2 * 36, 2 * 36),
rotation_interval=(0, math.pi / 2.0),
augmentation_probability=0.5,
subsample=8
)
# train_pipeline += Log.log_numpy_array_stats_after_process(GT_MASK_KEY, 'min', 'max', 'dtype', logging_prefix='%s: before misalign: ' % GT_MASK_KEY)
train_pipeline += Misalign(z_resolution=360, prob_slip=0.05, prob_shift=0.05, max_misalign=(360,) * 2, ignore_keys_for_slip=ignore_keys_for_slip)
# train_pipeline += Log.log_numpy_array_stats_after_process(GT_MASK_KEY, 'min', 'max', 'dtype', logging_prefix='%s: after misalign: ' % GT_MASK_KEY)
train_pipeline += SimpleAugment(transpose_only=[1,2])
train_pipeline += IntensityAugment(RAW_KEY, 0.9, 1.1, -0.1, 0.1, z_section_wise=True)
train_pipeline += DefectAugment(RAW_KEY,
prob_missing=0.03,
prob_low_contrast=0.01,
prob_artifact=0.03,
artifact_source=artifact_source,
artifacts=RAW_KEY,
artifacts_mask=DEFECT_MASK_KEY,
contrast_scale=0.5)
train_pipeline += IntensityScaleShift(RAW_KEY, 2, -1)
train_pipeline += ZeroOutConstSections(RAW_KEY)
if grow_boundaries > 0:
train_pipeline += GrowBoundary(LABELS_KEY, GT_MASK_KEY, steps=grow_boundaries, only_xy=True)
_logger.info("Renumbering connected components? %s", renumber_connected_components)
if renumber_connected_components:
train_pipeline += RenumberConnectedComponents(labels=LABELS_KEY)
train_pipeline += NewKeyFromNumpyArray(lambda array: 1 - array, GT_GLIA_KEY, UNLABELED_KEY)
if len(mask_out_labels) > 0:
train_pipeline += MaskOutLabels(label_key=LABELS_KEY, mask_key=GT_MASK_KEY, ids_to_be_masked=mask_out_labels)
# labels_mask: anything that connects into labels_mask will be zeroed out
# unlabelled: anyhing that points into unlabeled will have zero affinity;
# affinities within unlabelled will be masked out
train_pipeline += AddAffinities(
affinity_neighborhood=affinity_neighborhood,
labels=LABELS_KEY,
labels_mask=GT_MASK_KEY,
affinities=GT_AFFINITIES_KEY,
affinities_mask=AFFINITIES_MASK_KEY,
unlabelled=UNLABELED_KEY
)
snapshot_datasets = {
RAW_KEY: 'volumes/raw',
LABELS_KEY: 'volumes/labels/neuron_ids',
GT_AFFINITIES_KEY: 'volumes/affinities/gt',
GT_GLIA_KEY: 'volumes/labels/glia_gt',
UNLABELED_KEY: 'volumes/labels/unlabeled',
AFFINITIES_KEY: 'volumes/affinities/prediction',
LOSS_GRADIENT_KEY: 'volumes/loss_gradient',
AFFINITIES_MASK_KEY: 'masks/affinities',
GLIA_KEY: 'volumes/labels/glia_pred',
GT_MASK_KEY: 'masks/gt',
GLIA_MASK_KEY: 'masks/glia'}
if balance_labels:
train_pipeline += BalanceLabels(labels=GT_AFFINITIES_KEY, scales=AFFINITIES_SCALE_KEY, mask=AFFINITIES_MASK_KEY)
snapshot_datasets[AFFINITIES_SCALE_KEY] = 'masks/affinity-scale'
train_pipeline += BalanceLabels(labels=GT_GLIA_KEY, scales=GLIA_SCALE_KEY, mask=GLIA_MASK_KEY)
snapshot_datasets[GLIA_SCALE_KEY] = 'masks/glia-scale'
if (pre_cache_size > 0 and pre_cache_num_workers > 0):
train_pipeline += PreCache(cache_size=pre_cache_size, num_workers=pre_cache_num_workers)
train_pipeline += Train(
summary=summary,
graph=meta_graph_filename,
save_every=save_checkpoint_every,
optimizer=optimizer,
loss=loss,
inputs=network_inputs,
log_dir='log',
outputs={tensor_affinities: AFFINITIES_KEY, tensor_glia: GLIA_KEY},
gradients={tensor_affinities: LOSS_GRADIENT_KEY},
array_specs={
AFFINITIES_KEY : ArraySpec(voxel_size=output_voxel_size),
LOSS_GRADIENT_KEY : ArraySpec(voxel_size=output_voxel_size),
AFFINITIES_MASK_KEY : ArraySpec(voxel_size=output_voxel_size),
GT_MASK_KEY : ArraySpec(voxel_size=output_voxel_size),
AFFINITIES_SCALE_KEY : ArraySpec(voxel_size=output_voxel_size),
GLIA_MASK_KEY : ArraySpec(voxel_size=output_voxel_size),
GLIA_SCALE_KEY : ArraySpec(voxel_size=output_voxel_size),
GLIA_KEY : ArraySpec(voxel_size=output_voxel_size)
}
)
train_pipeline += Snapshot(
snapshot_datasets,
every=snapshot_every,
output_filename='batch_{iteration}.hdf',
output_dir=snapshot_dir,
additional_request=snapshot_request,
attributes_callback=Snapshot.default_attributes_callback())
train_pipeline += PrintProfilingStats(every=50)
print("Starting training...")
with build(train_pipeline) as b:
for i in range(trained_until, stop):
b.request_batch(request)
print("Training finished")
def train_until(
data_providers,
affinity_neighborhood,
meta_graph_filename,
stop,
input_shape,
output_shape,
loss,
optimizer,
tensor_affinities,
tensor_affinities_nn,
tensor_affinities_mask,
summary,
save_checkpoint_every,
pre_cache_size,
pre_cache_num_workers,
snapshot_every,
balance_labels,
renumber_connected_components,
network_inputs,
ignore_labels_for_slip,
grow_boundaries):
ignore_keys_for_slip = (GT_LABELS_KEY, GT_MASK_KEY) if ignore_labels_for_slip else ()
defect_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data/defects'
if tf.train.latest_checkpoint('.'):
trained_until = int(tf.train.latest_checkpoint('.').split('_')[-1])
print('Resuming training from', trained_until)
else:
trained_until = 0
print('Starting fresh training')
input_voxel_size = Coordinate((120, 12, 12)) * 3
output_voxel_size = Coordinate((40, 36, 36)) * 3
input_size = Coordinate(input_shape) * input_voxel_size
output_size = Coordinate(output_shape) * output_voxel_size
output_size_nn = Coordinate(s - 2 for s in output_shape) * output_voxel_size
num_affinities = sum(len(nh) for nh in affinity_neighborhood)
gt_affinities_size = Coordinate((num_affinities,) + tuple(s for s in output_size))
print("gt affinities size", gt_affinities_size)
# TODO why is GT_AFFINITIES three-dimensional? compare to
# TODO https://github.com/funkey/gunpowder/blob/master/examples/cremi/train.py#L35
# specifiy which Arrays should be requested for each batch
request = BatchRequest()
request.add(RAW_KEY, input_size, voxel_size=input_voxel_size)
request.add(GT_LABELS_KEY, output_size, voxel_size=output_voxel_size)
request.add(GT_AFFINITIES_KEY, output_size, voxel_size=output_voxel_size)
request.add(AFFINITIES_MASK_KEY, output_size, voxel_size=output_voxel_size)
request.add(GT_MASK_KEY, output_size, voxel_size=output_voxel_size)
request.add(AFFINITIES_NN_KEY, output_size_nn, voxel_size=output_voxel_size)
if balance_labels:
request.add(AFFINITIES_SCALE_KEY, output_size, voxel_size=output_voxel_size)
network_inputs[tensor_affinities_mask] = AFFINITIES_SCALE_KEY if balance_labels else AFFINITIES_MASK_KEY
# create a tuple of data sources, one for each HDF file
data_sources = tuple(
provider +
Normalize(RAW_KEY) + # ensures RAW is in float in [0, 1]
# zero-pad provided RAW and GT_MASK to be able to draw batches close to
# the boundary of the available data
# size more or less irrelevant as followed by Reject Node
Pad(RAW_KEY, None) +
Pad(GT_MASK_KEY, None) +
RandomLocation() + # chose a random location inside the provided arrays
Reject(GT_MASK_KEY) + # reject batches wich do contain less than 50% labelled data
Reject(GT_LABELS_KEY, min_masked=0.0, reject_probability=0.95)
for provider in data_providers)
# TODO figure out what this is for
snapshot_request = BatchRequest({
LOSS_GRADIENT_KEY : request[GT_AFFINITIES_KEY],
AFFINITIES_KEY : request[GT_AFFINITIES_KEY],
AFFINITIES_NN_KEY : request[AFFINITIES_NN_KEY]
})
# no need to do anything here. random sections will be replaced with sections from this source (only raw)
artifact_source = (
Hdf5Source(
os.path.join(defect_dir, 'sample_ABC_padded_20160501.defects.hdf'),
datasets={
RAW_KEY : 'defect_sections/raw',
ALPHA_MASK_KEY : 'defect_sections/mask',
},
array_specs={
RAW_KEY : ArraySpec(voxel_size=input_voxel_size),
ALPHA_MASK_KEY : ArraySpec(voxel_size=input_voxel_size),
}
) +
RandomLocation(min_masked=0.05, mask=ALPHA_MASK_KEY) +
Normalize(RAW_KEY) +
IntensityAugment(RAW_KEY, 0.9, 1.1, -0.1, 0.1, z_section_wise=True) +
ElasticAugment(
voxel_size=(360, 36, 36),
control_point_spacing=(4, 40, 40),
control_point_displacement_sigma=(0, 2 * 36, 2 * 36),
rotation_interval=(0, math.pi / 2.0),
subsample=8
) +
SimpleAugment(transpose_only=[1,2])
)
train_pipeline = data_sources
train_pipeline += RandomProvider()
train_pipeline += ElasticAugment(
voxel_size=(360, 36, 36),
control_point_spacing=(4, 40, 40),
control_point_displacement_sigma=(0, 2 * 36, 2 * 36),
rotation_interval=(0, math.pi / 2.0),
augmentation_probability=0.5,
subsample=8
)
train_pipeline += Misalign(z_resolution=360, prob_slip=0.05, prob_shift=0.05, max_misalign=(360,) * 2, ignore_keys_for_slip=ignore_keys_for_slip)
train_pipeline += SimpleAugment(transpose_only=[1,2])
train_pipeline += IntensityAugment(RAW_KEY, 0.9, 1.1, -0.1, 0.1, z_section_wise=True)
train_pipeline += DefectAugment(RAW_KEY,
prob_missing=0.03,
prob_low_contrast=0.01,
prob_artifact=0.03,
artifact_source=artifact_source,
artifacts=RAW_KEY,
artifacts_mask=ALPHA_MASK_KEY,
contrast_scale=0.5)
train_pipeline += IntensityScaleShift(RAW_KEY, 2, -1)
train_pipeline += ZeroOutConstSections(RAW_KEY)
if grow_boundaries > 0:
train_pipeline += GrowBoundary(GT_LABELS_KEY, GT_MASK_KEY, steps=grow_boundaries, only_xy=True)
if renumber_connected_components:
train_pipeline += RenumberConnectedComponents(labels=GT_LABELS_KEY)
train_pipeline += AddAffinities(
affinity_neighborhood=affinity_neighborhood,
labels=GT_LABELS_KEY,
labels_mask=GT_MASK_KEY,
affinities=GT_AFFINITIES_KEY,
affinities_mask=AFFINITIES_MASK_KEY
)
if balance_labels:
train_pipeline += BalanceLabels(labels=GT_AFFINITIES_KEY, scales=AFFINITIES_SCALE_KEY, mask=AFFINITIES_MASK_KEY)
train_pipeline += PreCache(cache_size=pre_cache_size, num_workers=pre_cache_num_workers)
train_pipeline += Train(
summary=summary,
graph=meta_graph_filename,
save_every=save_checkpoint_every,
optimizer=optimizer,
loss=loss,
inputs=network_inputs,
log_dir='log',
outputs={tensor_affinities: AFFINITIES_KEY, tensor_affinities_nn: AFFINITIES_NN_KEY},
gradients={tensor_affinities: LOSS_GRADIENT_KEY},
array_specs={
AFFINITIES_KEY : ArraySpec(voxel_size=output_voxel_size),
LOSS_GRADIENT_KEY : ArraySpec(voxel_size=output_voxel_size),
AFFINITIES_MASK_KEY : ArraySpec(voxel_size=output_voxel_size),
GT_MASK_KEY : ArraySpec(voxel_size=output_voxel_size),
AFFINITIES_SCALE_KEY : ArraySpec(voxel_size=output_voxel_size),
AFFINITIES_NN_KEY : ArraySpec(voxel_size=output_voxel_size)
}
)
train_pipeline += Snapshot(
dataset_names={
RAW_KEY : 'volumes/raw',
GT_LABELS_KEY : 'volumes/labels/neuron_ids',
GT_AFFINITIES_KEY : 'volumes/affinities/gt',
AFFINITIES_KEY : 'volumes/affinities/prediction',
LOSS_GRADIENT_KEY : 'volumes/loss_gradient',
AFFINITIES_MASK_KEY : 'masks/affinities',
AFFINITIES_NN_KEY : 'volumes/affinities/prediction-nn'
},
every=snapshot_every,
output_filename='batch_{iteration}.hdf',
output_dir='snapshots/',
additional_request=snapshot_request,
attributes_callback=Snapshot.default_attributes_callback())
train_pipeline += PrintProfilingStats(every=50)
print("Starting training...")
with build(train_pipeline) as b:
for i in range(trained_until, stop):
b.request_batch(request)
print("Training finished")
def visualize_original_elastic_augment(args=None):
data_providers = []
data_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data'
file_pattern = 'sample_A_padded_20160501-2-additional-sections-fixed-offset.h5'
file_pattern = 'sample_B_padded_20160501-2-additional-sections-fixed-offset.h5'
file_pattern = 'sample_C_padded_20160501-2-additional-sections-fixed-offset.h5'
for data in glob.glob(os.path.join(data_dir, file_pattern)):
h5_source = Hdf5Source(data, datasets={
RAW: 'volumes/raw',
})
data_providers.append(h5_source)
input_resolution = (360, 36, 36)
output_resolution = Coordinate((120, 108, 108))
offset = (13640, 10932, 10932)
output_shape = Coordinate((12, 100, 100)) * output_resolution
output_offset = (13320 + 3600, 32796 + 36 + 10800, 32796 + 36 + 10800)
overhang = Coordinate((360, 108, 108)) * 16
input_shape = output_shape + overhang * 2
input_offset = Coordinate(output_offset) - overhang
output_roi = Roi(offset=output_offset, shape=output_shape)
input_roi = Roi(offset=input_offset, shape=input_shape)
augmentations = (
# SimpleAugment(transpose_only=[1,2]),
# ElasticAugmentNonMatchingVoxelSize(control_point_spacing=(1, 1, 1), jitter_sigma=(0.0, 3.0, 3.0), rotation_interval=(0, np.pi/2.0)),
ElasticAugment(
control_point_spacing=(4, 40, 40),
#jitter_sigma=(0, 1 * 2 * 36, 0 * 36),
jitter_sigma=(0, 2, 2),
rotation_interval=(2 * np.pi / 8, 0 * 2 * np.pi),
subsample=1), )
batch, snapshot = fuse.util.run_augmentations(
data_providers=data_providers,
roi=lambda key: input_roi.copy(),
augmentations=augmentations,
keys=(RAW, ),
voxel_size=lambda key: input_resolution)
args = get_parser().parse_args() if args is None else args
jnius_config.add_options('-Xmx{}'.format(args.max_heap_size))
import payntera.jfx
import imglyb
from jnius import autoclass
payntera.jfx.init_platform()
PainteraBaseView = autoclass(
'org.janelia.saalfeldlab.paintera.PainteraBaseView')
viewer = PainteraBaseView.defaultView()
pbv = viewer.baseView
scene, stage = payntera.jfx.start_stage(viewer.paneWithStatus.getPane())
payntera.jfx.invoke_on_jfx_application_thread(
lambda: pbv.orthogonalViews().setScreenScales([0.3, 0.1, 0.03]))
keys_to_show = (RAW, )
snapshot_states = add_to_viewer(
snapshot,
keys=keys_to_show,
name=lambda key: '%s-snapshot' % key.identifier)
states = add_to_viewer(batch, keys=keys_to_show)
viewer.keyTracker.installInto(scene)
scene.addEventFilter(
autoclass('javafx.scene.input.MouseEvent').ANY, viewer.mouseTracker)
while stage.isShowing():
time.sleep(0.1)
data_providers = []
data_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data'
# data_dir = os.path.expanduser('~/Dropbox/cremi-upsampled/')
file_pattern = 'sample_A_padded_20160501-2-additional-sections-fixed-offset.h5'
file_pattern = 'sample_B_padded_20160501-2-additional-sections-fixed-offset.h5'
file_pattern = 'sample_C_padded_20160501-2-additional-sections-fixed-offset.h5'
defect_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data/defects'
artifact_source = (
Hdf5Source(
os.path.join(defect_dir, 'sample_ABC_padded_20160501.defects.hdf'),
datasets={
RAW: 'defect_sections/raw',
ALPHA_MASK: 'defect_sections/mask',
},
array_specs={
RAW: ArraySpec(voxel_size=tuple(d * 9 for d in (40, 4, 4))),
ALPHA_MASK: ArraySpec(voxel_size=tuple(d * 9 for d in (40, 4, 4))),
}) + RandomLocation(min_masked=0.05, mask=ALPHA_MASK) +
Normalize(RAW) +
IntensityAugment(RAW, 0.9, 1.1, -0.1, 0.1, z_section_wise=True) +
ElasticAugment(voxel_size=(360, 36, 36),
control_point_spacing=(4, 40, 40),
control_point_displacement_sigma=(0, 2 * 36, 2 * 36),
rotation_interval=(0, np.pi / 2.0),
subsample=8) + SimpleAugment(transpose_only=[1, 2]))
for data in glob.glob(os.path.join(data_dir, file_pattern)):
h5_source = Hdf5Source(data,
datasets={
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--max-heap-size", default="16g")
args = parser.parse_args()
data_providers = []
data_dir = '/groups/saalfeld/home/hanslovskyp/experiments/quasi-isotropic/data'
file_pattern = 'sample_A_padded_20160501-2-additional-sections-fixed-offset.h5'
file_pattern = 'sample_B_padded_20160501-2-additional-sections-fixed-offset.h5'
file_pattern = 'sample_C_padded_20160501-2-additional-sections-fixed-offset.h5'
for data in glob.glob(os.path.join(data_dir, file_pattern)):
h5_source = Hdf5Source(data, datasets={
RAW: 'volumes/raw',
})
data_providers.append(h5_source)
input_resolution = (360, 36, 36)
output_resolution = Coordinate((120, 108, 108))
offset = (13640, 10932, 10932)
output_shape = Coordinate((12, 100, 100)) * output_resolution
output_offset = (13320 + 3600, 32796 + 36 + 10800, 32796 + 36 + 10800)
overhang = Coordinate((360, 108, 108)) * 16
input_shape = output_shape + overhang * 2
input_offset = Coordinate(output_offset) - overhang
