def add_augmentation_pipeline(
pipeline,
raw,
simple=None,
elastic=None,
blur=None,
noise=None):
'''Add an augmentation pipeline to an existing pipeline.
All optional arguments are kwargs for the corresponding augmentation node.
If not given, those augmentations are not added.
'''
if simple is not None:
pipeline = pipeline + gp.SimpleAugment(**simple)
if elastic is not None:
pipeline = pipeline + gp.ElasticAugment(**elastic)
if blur is not None:
pipeline = pipeline + Blur(raw, **blur)
if noise is not None:
pipeline = pipeline + gp.NoiseAugment(raw, **noise)
return pipeline
def _augmentation_pipeline(self, raw, source):
if 'elastic' in self.params and self.params['elastic']:
source = source + gp.ElasticAugment(
**self.params["elastic_params"])
if 'blur' in self.params and self.params['blur']:
source = source + Blur(raw, **self.params["blur_params"])
if 'simple' in self.params and self.params['simple']:
source = source + gp.SimpleAugment(**self.params["simple_params"])
if 'noise' in self.params and self.params['noise']:
source = source + gp.NoiseAugment(raw, **
self.params['noise_params'])
return source
def train(iterations):
##################
# DECLARE ARRAYS #
##################
# raw intensities
raw = gp.ArrayKey('RAW')
# objects labelled with unique IDs
gt_labels = gp.ArrayKey('LABELS')
# array of per-voxel affinities to direct neighbors
gt_affs = gp.ArrayKey('AFFINITIES')
# weights to use to balance the loss
loss_weights = gp.ArrayKey('LOSS_WEIGHTS')
# the predicted affinities
pred_affs = gp.ArrayKey('PRED_AFFS')
# the gredient of the loss wrt to the predicted affinities
pred_affs_gradients = gp.ArrayKey('PRED_AFFS_GRADIENTS')
####################
# DECLARE REQUESTS #
####################
with open('train_net_config.json', 'r') as f:
net_config = json.load(f)
# get the input and output size in world units (nm, in this case)
voxel_size = gp.Coordinate((8, 8, 8))
input_size = gp.Coordinate(net_config['input_shape']) * voxel_size
output_size = gp.Coordinate(net_config['output_shape']) * voxel_size
# formulate the request for what a batch should (at least) contain
request = gp.BatchRequest()
request.add(raw, input_size)
request.add(gt_affs, output_size)
request.add(loss_weights, output_size)
# when we make a snapshot for inspection (see below), we also want to
# request the predicted affinities and gradients of the loss wrt the
# affinities
snapshot_request = gp.BatchRequest()
snapshot_request[pred_affs] = request[gt_affs]
snapshot_request[pred_affs_gradients] = request[gt_affs]
##############################
# ASSEMBLE TRAINING PIPELINE #
##############################
pipeline = (
# a tuple of sources, one for each sample (A, B, and C) provided by the
# CREMI challenge
tuple(
# read batches from the HDF5 file
gp.Hdf5Source(os.path.join(data_dir, 'fib.hdf'),
datasets={
raw: 'volumes/raw',
gt_labels: 'volumes/labels/neuron_ids'
}) +
# convert raw to float in [0, 1]
gp.Normalize(raw) +
# chose a random location for each requested batch
gp.RandomLocation()) +
# chose a random source (i.e., sample) from the above
gp.RandomProvider() +
# elastically deform the batch
gp.ElasticAugment([8, 8, 8], [0, 2, 2], [0, math.pi / 2.0],
prob_slip=0.05,
prob_shift=0.05,
max_misalign=25) +
# apply transpose and mirror augmentations
gp.SimpleAugment(transpose_only=[1, 2]) +
# scale and shift the intensity of the raw array
gp.IntensityAugment(raw,
scale_min=0.9,
scale_max=1.1,
shift_min=-0.1,
shift_max=0.1,
z_section_wise=True) +
# grow a boundary between labels
gp.GrowBoundary(gt_labels, steps=3, only_xy=True) +
# convert labels into affinities between voxels
gp.AddAffinities([[-1, 0, 0], [0, -1, 0], [0, 0, -1]], gt_labels,
gt_affs) +
# create a weight array that balances positive and negative samples in
# the affinity array
gp.BalanceLabels(gt_affs, loss_weights) +
# pre-cache batches from the point upstream
gp.PreCache(cache_size=10, num_workers=5) +
# perform one training iteration for each passing batch (here we use
# the tensor names earlier stored in train_net.config)
gp.tensorflow.Train(
'train_net',
net_config['optimizer'],
net_config['loss'],
inputs={
net_config['raw']: raw,
net_config['gt_affs']: gt_affs,
net_config['loss_weights']: loss_weights
},
outputs={net_config['pred_affs']: pred_affs},
gradients={net_config['pred_affs']: pred_affs_gradients},
save_every=10000) +
# save the passing batch as an HDF5 file for inspection
gp.Snapshot(
{
raw: '/volumes/raw',
gt_labels: '/volumes/labels/neuron_ids',
gt_affs: '/volumes/labels/affs',
pred_affs: '/volumes/pred_affs',
pred_affs_gradients: '/volumes/pred_affs_gradients'
},
output_dir='snapshots',
output_filename='batch_{iteration}.hdf',
every=1000,
additional_request=snapshot_request,
compression_type='gzip') +
# show a summary of time spend in each node every 10 iterations
gp.PrintProfilingStats(every=1000))
#########
# TRAIN #
#########
print("Training for", iterations, "iterations")
with gp.build(pipeline):
for i in range(iterations):
pipeline.request_batch(request)
print("Finished")
scale=voxel_size,
transpose=(2, 1, 0),
transform_file=str((filename / "transform.txt").absolute()),
ignore_human_nodes=False,
),
) + gp.MergeProvider() + gp.RandomLocation(
ensure_nonempty=swcs, ensure_centered=True) + RasterizeSkeleton(
points=swcs,
array=labels,
array_spec=gp.ArraySpec(
interpolatable=False, voxel_size=voxel_size, dtype=np.uint32),
) + GrowLabels(labels, radius=20)
# augment
+ gp.ElasticAugment(
[40, 10, 10], [0.25, 1, 1], [0, math.pi / 2.0],
subsample=4) + gp.SimpleAugment(
mirror_only=[1, 2], transpose_only=[1, 2]) +
gp.Normalize(raw) +
gp.IntensityAugment(raw, 0.9, 1.1, -0.001, 0.001)
for filename in path_to_data.iterdir()
if "2018-07-02" in filename.name)
pipeline = (data_sources + gp.RandomProvider() + GetNeuronPair(
swcs,
raw,
labels,
(swc_base, swc_add),
(raw_base, raw_add),
(labels_base, labels_add),
seperate_by=SEPERATE_DISTANCE,
shift_attempts=50,
request_attempts=10,
def train_until(**kwargs):
if tf.train.latest_checkpoint(kwargs['output_folder']):
trained_until = int(
tf.train.latest_checkpoint(kwargs['output_folder']).split('_')[-1])
else:
trained_until = 0
if trained_until >= kwargs['max_iteration']:
return
anchor = gp.ArrayKey('ANCHOR')
raw = gp.ArrayKey('RAW')
raw_cropped = gp.ArrayKey('RAW_CROPPED')
gt_threeclass = gp.ArrayKey('GT_THREECLASS')
loss_weights_threeclass = gp.ArrayKey('LOSS_WEIGHTS_THREECLASS')
pred_threeclass = gp.ArrayKey('PRED_THREECLASS')
pred_threeclass_gradients = gp.ArrayKey('PRED_THREECLASS_GRADIENTS')
with open(
os.path.join(kwargs['output_folder'],
kwargs['name'] + '_config.json'), 'r') as f:
net_config = json.load(f)
with open(
os.path.join(kwargs['output_folder'],
kwargs['name'] + '_names.json'), 'r') as f:
net_names = json.load(f)
voxel_size = gp.Coordinate(kwargs['voxel_size'])
input_shape_world = gp.Coordinate(net_config['input_shape']) * voxel_size
output_shape_world = gp.Coordinate(net_config['output_shape']) * voxel_size
# formulate the request for what a batch should (at least) contain
request = gp.BatchRequest()
request.add(raw, input_shape_world)
request.add(raw_cropped, output_shape_world)
request.add(gt_threeclass, output_shape_world)
request.add(anchor, output_shape_world)
request.add(loss_weights_threeclass, output_shape_world)
# when we make a snapshot for inspection (see below), we also want to
# request the predicted affinities and gradients of the loss wrt the
# affinities
snapshot_request = gp.BatchRequest()
snapshot_request.add(raw_cropped, output_shape_world)
snapshot_request.add(gt_threeclass, output_shape_world)
snapshot_request.add(pred_threeclass, output_shape_world)
# snapshot_request.add(pred_threeclass_gradients, output_shape_world)
if kwargs['input_format'] != "hdf" and kwargs['input_format'] != "zarr":
raise NotImplementedError("train node for {} not implemented".format(
kwargs['input_format']))
fls = []
shapes = []
for f in kwargs['data_files']:
fls.append(os.path.splitext(f)[0])
if kwargs['input_format'] == "hdf":
vol = h5py.File(f, 'r')['volumes/raw']
elif kwargs['input_format'] == "zarr":
vol = zarr.open(f, 'r')['volumes/raw']
print(f, vol.shape, vol.dtype)
shapes.append(vol.shape)
if vol.dtype != np.float32:
print("please convert to float32")
ln = len(fls)
print("first 5 files: ", fls[0:4])
# padR = 46
# padGT = 32
if kwargs['input_format'] == "hdf":
sourceNode = gp.Hdf5Source
elif kwargs['input_format'] == "zarr":
sourceNode = gp.ZarrSource
augmentation = kwargs['augmentation']
pipeline = (
tuple(
# read batches from the HDF5 file
sourceNode(
fls[t] + "." + kwargs['input_format'],
datasets={
raw: 'volumes/raw',
gt_threeclass: 'volumes/gt_threeclass',
anchor: 'volumes/gt_threeclass',
},
array_specs={
raw: gp.ArraySpec(interpolatable=True),
gt_threeclass: gp.ArraySpec(interpolatable=False),
anchor: gp.ArraySpec(interpolatable=False)
}
)
+ gp.MergeProvider()
+ gp.Pad(raw, None)
+ gp.Pad(gt_threeclass, None)
+ gp.Pad(anchor, gp.Coordinate((2,2,2)))
# chose a random location for each requested batch
+ gp.RandomLocation()
for t in range(ln)
) +
# chose a random source (i.e., sample) from the above
gp.RandomProvider() +
# elastically deform the batch
(gp.ElasticAugment(
augmentation['elastic']['control_point_spacing'],
augmentation['elastic']['jitter_sigma'],
[augmentation['elastic']['rotation_min']*np.pi/180.0,
augmentation['elastic']['rotation_max']*np.pi/180.0],
subsample=augmentation['elastic'].get('subsample', 1)) \
if augmentation.get('elastic') is not None else NoOp()) +
# apply transpose and mirror augmentations
gp.SimpleAugment(mirror_only=augmentation['simple'].get("mirror"),
transpose_only=augmentation['simple'].get("transpose")) +
# # scale and shift the intensity of the raw array
gp.IntensityAugment(
raw,
scale_min=augmentation['intensity']['scale'][0],
scale_max=augmentation['intensity']['scale'][1],
shift_min=augmentation['intensity']['shift'][0],
shift_max=augmentation['intensity']['shift'][1],
z_section_wise=False) +
# grow a boundary between labels
# TODO: check
# gp.GrowBoundary(
# gt_threeclass,
# steps=1,
# only_xy=False) +
gp.BalanceLabels(
gt_threeclass,
loss_weights_threeclass,
num_classes=3) +
# pre-cache batches from the point upstream
gp.PreCache(
cache_size=kwargs['cache_size'],
num_workers=kwargs['num_workers']) +
# perform one training iteration for each passing batch (here we use
# the tensor names earlier stored in train_net.config)
gp.tensorflow.Train(
os.path.join(kwargs['output_folder'], kwargs['name']),
optimizer=net_names['optimizer'],
summary=net_names['summaries'],
log_dir=kwargs['output_folder'],
loss=net_names['loss'],
inputs={
net_names['raw']: raw,
net_names['anchor']: anchor,
net_names['gt_threeclass']: gt_threeclass,
net_names['loss_weights_threeclass']: loss_weights_threeclass
},
outputs={
net_names['pred_threeclass']: pred_threeclass,
net_names['raw_cropped']: raw_cropped,
},
gradients={
net_names['pred_threeclass']: pred_threeclass_gradients,
},
save_every=kwargs['checkpoints']) +
# save the passing batch as an HDF5 file for inspection
gp.Snapshot(
{
raw: '/volumes/raw',
raw_cropped: 'volumes/raw_cropped',
gt_threeclass: '/volumes/gt_threeclass',
pred_threeclass: '/volumes/pred_threeclass',
},
output_dir=os.path.join(kwargs['output_folder'], 'snapshots'),
output_filename='batch_{iteration}.hdf',
every=kwargs['snapshots'],
additional_request=snapshot_request,
compression_type='gzip') +
# show a summary of time spend in each node every 10 iterations
gp.PrintProfilingStats(every=kwargs['profiling'])
)
#########
# TRAIN #
#########
print("Starting training...")
with gp.build(pipeline):
print(pipeline)
for i in range(trained_until, kwargs['max_iteration']):
# print("request", request)
start = time.time()
pipeline.request_batch(request)
time_of_iteration = time.time() - start
logger.info("Batch: iteration=%d, time=%f", i, time_of_iteration)
# exit()
print("Training finished")
def train(until):
model = SpineUNet()
loss = torch.nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
input_size = (8, 96, 96)
raw = gp.ArrayKey('RAW')
labels = gp.ArrayKey('LABELS')
affs = gp.ArrayKey('AFFS')
affs_predicted = gp.ArrayKey('AFFS_PREDICTED')
pipeline = (
(
gp.ZarrSource(
'data/20200201.zarr',
{
raw: 'train/sample1/raw',
labels: 'train/sample1/labels'
}),
gp.ZarrSource(
'data/20200201.zarr',
{
raw: 'train/sample2/raw',
labels: 'train/sample2/labels'
}),
gp.ZarrSource(
'data/20200201.zarr',
{
raw: 'train/sample3/raw',
labels: 'train/sample3/labels'
})
) +
gp.RandomProvider() +
gp.Normalize(raw) +
gp.RandomLocation() +
gp.SimpleAugment(transpose_only=(1, 2)) +
gp.ElasticAugment((2, 10, 10), (0.0, 0.5, 0.5), [0, math.pi]) +
gp.AddAffinities(
[(1, 0, 0), (0, 1, 0), (0, 0, 1)],
labels,
affs) +
gp.Normalize(affs, factor=1.0) +
#gp.PreCache(num_workers=1) +
# raw: (d, h, w)
# affs: (3, d, h, w)
gp.Stack(1) +
# raw: (1, d, h, w)
# affs: (1, 3, d, h, w)
AddChannelDim(raw) +
# raw: (1, 1, d, h, w)
# affs: (1, 3, d, h, w)
gp_torch.Train(
model,
loss,
optimizer,
inputs={'x': raw},
outputs={0: affs_predicted},
loss_inputs={0: affs_predicted, 1: affs},
save_every=10000) +
RemoveChannelDim(raw) +
RemoveChannelDim(raw) +
RemoveChannelDim(affs) +
RemoveChannelDim(affs_predicted) +
# raw: (d, h, w)
# affs: (3, d, h, w)
# affs_predicted: (3, d, h, w)
gp.Snapshot(
{
raw: 'raw',
labels: 'labels',
affs: 'affs',
affs_predicted: 'affs_predicted'
},
every=500,
output_filename='iteration_{iteration}.hdf')
)
request = gp.BatchRequest()
request.add(raw, input_size)
request.add(labels, input_size)
request.add(affs, input_size)
request.add(affs_predicted, input_size)
with gp.build(pipeline):
for i in range(until):
pipeline.request_batch(request)
def train_until(**kwargs):
if tf.train.latest_checkpoint(kwargs['output_folder']):
trained_until = int(
tf.train.latest_checkpoint(kwargs['output_folder']).split('_')[-1])
else:
trained_until = 0
if trained_until >= kwargs['max_iteration']:
return
anchor = gp.ArrayKey('ANCHOR')
raw = gp.ArrayKey('RAW')
raw_cropped = gp.ArrayKey('RAW_CROPPED')
points = gp.PointsKey('POINTS')
gt_cp = gp.ArrayKey('GT_CP')
pred_cp = gp.ArrayKey('PRED_CP')
pred_cp_gradients = gp.ArrayKey('PRED_CP_GRADIENTS')
with open(
os.path.join(kwargs['output_folder'],
kwargs['name'] + '_config.json'), 'r') as f:
net_config = json.load(f)
with open(
os.path.join(kwargs['output_folder'],
kwargs['name'] + '_names.json'), 'r') as f:
net_names = json.load(f)
voxel_size = gp.Coordinate(kwargs['voxel_size'])
input_shape_world = gp.Coordinate(net_config['input_shape']) * voxel_size
output_shape_world = gp.Coordinate(net_config['output_shape']) * voxel_size
# formulate the request for what a batch should (at least) contain
request = gp.BatchRequest()
request.add(raw, input_shape_world)
request.add(raw_cropped, output_shape_world)
request.add(gt_cp, output_shape_world)
request.add(anchor, output_shape_world)
# when we make a snapshot for inspection (see below), we also want to
# request the predicted affinities and gradients of the loss wrt the
# affinities
snapshot_request = gp.BatchRequest()
snapshot_request.add(raw_cropped, output_shape_world)
snapshot_request.add(gt_cp, output_shape_world)
snapshot_request.add(pred_cp, output_shape_world)
# snapshot_request.add(pred_cp_gradients, output_shape_world)
if kwargs['input_format'] != "hdf" and kwargs['input_format'] != "zarr":
raise NotImplementedError("train node for %s not implemented yet",
kwargs['input_format'])
fls = []
shapes = []
mn = []
mx = []
for f in kwargs['data_files']:
fls.append(os.path.splitext(f)[0])
if kwargs['input_format'] == "hdf":
vol = h5py.File(f, 'r')['volumes/raw']
elif kwargs['input_format'] == "zarr":
vol = zarr.open(f, 'r')['volumes/raw']
print(f, vol.shape, vol.dtype)
shapes.append(vol.shape)
mn.append(np.min(vol))
mx.append(np.max(vol))
if vol.dtype != np.float32:
print("please convert to float32")
ln = len(fls)
print("first 5 files: ", fls[0:4])
if kwargs['input_format'] == "hdf":
sourceNode = gp.Hdf5Source
elif kwargs['input_format'] == "zarr":
sourceNode = gp.ZarrSource
augmentation = kwargs['augmentation']
sources = tuple(
(sourceNode(fls[t] + "." + kwargs['input_format'],
datasets={
raw: 'volumes/raw',
anchor: 'volumes/gt_fgbg',
},
array_specs={
raw: gp.ArraySpec(interpolatable=True),
anchor: gp.ArraySpec(interpolatable=False)
}),
gp.CsvIDPointsSource(fls[t] + ".csv",
points,
points_spec=gp.PointsSpec(
roi=gp.Roi(gp.Coordinate((
0, 0, 0)), gp.Coordinate(shapes[t]))))) +
gp.MergeProvider()
# + Clip(raw, mn=mn[t], mx=mx[t])
# + NormalizeMinMax(raw, mn=mn[t], mx=mx[t])
+ gp.Pad(raw, None) + gp.Pad(points, None)
# chose a random location for each requested batch
+ gp.RandomLocation() for t in range(ln))
pipeline = (
sources +
# chose a random source (i.e., sample) from the above
gp.RandomProvider() +
# elastically deform the batch
(gp.ElasticAugment(
augmentation['elastic']['control_point_spacing'],
augmentation['elastic']['jitter_sigma'],
[augmentation['elastic']['rotation_min']*np.pi/180.0,
augmentation['elastic']['rotation_max']*np.pi/180.0],
subsample=augmentation['elastic'].get('subsample', 1)) \
if augmentation.get('elastic') is not None else NoOp()) +
# apply transpose and mirror augmentations
gp.SimpleAugment(mirror_only=augmentation['simple'].get("mirror"),
transpose_only=augmentation['simple'].get("transpose")) +
# (gp.SimpleAugment(
# mirror_only=augmentation['simple'].get("mirror"),
# transpose_only=augmentation['simple'].get("transpose")) \
# if augmentation.get('simple') is not None and \
# augmentation.get('simple') != {} else NoOp()) +
# # scale and shift the intensity of the raw array
(gp.IntensityAugment(
raw,
scale_min=augmentation['intensity']['scale'][0],
scale_max=augmentation['intensity']['scale'][1],
shift_min=augmentation['intensity']['shift'][0],
shift_max=augmentation['intensity']['shift'][1],
z_section_wise=False) \
if augmentation.get('intensity') is not None and \
augmentation.get('intensity') != {} else NoOp()) +
gp.RasterizePoints(
points,
gt_cp,
array_spec=gp.ArraySpec(voxel_size=voxel_size),
settings=gp.RasterizationSettings(
radius=(2, 2, 2),
mode='peak')) +
# pre-cache batches from the point upstream
gp.PreCache(
cache_size=kwargs['cache_size'],
num_workers=kwargs['num_workers']) +
# perform one training iteration for each passing batch (here we use
# the tensor names earlier stored in train_net.config)
gp.tensorflow.Train(
os.path.join(kwargs['output_folder'], kwargs['name']),
optimizer=net_names['optimizer'],
summary=net_names['summaries'],
log_dir=kwargs['output_folder'],
loss=net_names['loss'],
inputs={
net_names['raw']: raw,
net_names['gt_cp']: gt_cp,
net_names['anchor']: anchor,
},
outputs={
net_names['pred_cp']: pred_cp,
net_names['raw_cropped']: raw_cropped,
},
gradients={
# net_names['pred_cp']: pred_cp_gradients,
},
save_every=kwargs['checkpoints']) +
# save the passing batch as an HDF5 file for inspection
gp.Snapshot(
{
raw: '/volumes/raw',
raw_cropped: 'volumes/raw_cropped',
gt_cp: '/volumes/gt_cp',
pred_cp: '/volumes/pred_cp',
# pred_cp_gradients: '/volumes/pred_cp_gradients',
},
output_dir=os.path.join(kwargs['output_folder'], 'snapshots'),
output_filename='batch_{iteration}.hdf',
every=kwargs['snapshots'],
additional_request=snapshot_request,
compression_type='gzip') +
# show a summary of time spend in each node every 10 iterations
gp.PrintProfilingStats(every=kwargs['profiling'])
)
#########
# TRAIN #
#########
print("Starting training...")
with gp.build(pipeline):
print(pipeline)
for i in range(trained_until, kwargs['max_iteration']):
# print("request", request)
start = time.time()
pipeline.request_batch(request)
time_of_iteration = time.time() - start
logger.info("Batch: iteration=%d, time=%f", i, time_of_iteration)
# exit()
print("Training finished")
def train_until(max_iteration):
# get the latest checkpoint
if tf.train.latest_checkpoint('.'):
trained_until = int(tf.train.latest_checkpoint('.').split('_')[-1])
else:
trained_until = 0
if trained_until >= max_iteration:
return
# array keys for fused volume
raw = gp.ArrayKey('RAW')
labels = gp.ArrayKey('LABELS')
labels_fg = gp.ArrayKey('LABELS_FG')
# array keys for base volume
raw_base = gp.ArrayKey('RAW_BASE')
labels_base = gp.ArrayKey('LABELS_BASE')
swc_base = gp.PointsKey('SWC_BASE')
swc_center_base = gp.PointsKey('SWC_CENTER_BASE')
# array keys for add volume
raw_add = gp.ArrayKey('RAW_ADD')
labels_add = gp.ArrayKey('LABELS_ADD')
swc_add = gp.PointsKey('SWC_ADD')
swc_center_add = gp.PointsKey('SWC_CENTER_ADD')
# output data
fg = gp.ArrayKey('FG')
gradient_fg = gp.ArrayKey('GRADIENT_FG')
loss_weights = gp.ArrayKey('LOSS_WEIGHTS')
voxel_size = gp.Coordinate((3, 3, 3))
input_size = gp.Coordinate(net_config['input_shape']) * voxel_size
output_size = gp.Coordinate(net_config['output_shape']) * voxel_size
# add request
request = gp.BatchRequest()
request.add(raw, input_size)
request.add(labels, output_size)
request.add(labels_fg, output_size)
request.add(loss_weights, output_size)
request.add(swc_center_base, output_size)
request.add(swc_base, input_size)
request.add(swc_center_add, output_size)
request.add(swc_add, input_size)
# add snapshot request
snapshot_request = gp.BatchRequest()
snapshot_request.add(fg, output_size)
snapshot_request.add(labels_fg, output_size)
snapshot_request.add(gradient_fg, output_size)
snapshot_request.add(raw_base, input_size)
snapshot_request.add(raw_add, input_size)
snapshot_request.add(labels_base, input_size)
snapshot_request.add(labels_add, input_size)
# data source for "base" volume
data_sources_base = tuple()
data_sources_base += tuple(
(gp.Hdf5Source(file,
datasets={
raw_base: '/volume',
},
array_specs={
raw_base:
gp.ArraySpec(interpolatable=True,
voxel_size=voxel_size,
dtype=np.uint16),
},
channels_first=False),
SwcSource(filename=file,
dataset='/reconstruction',
points=(swc_center_base, swc_base),
scale=voxel_size)) + gp.MergeProvider() +
gp.RandomLocation(ensure_nonempty=swc_center_base) + RasterizeSkeleton(
points=swc_base,
array=labels_base,
array_spec=gp.ArraySpec(
interpolatable=False, voxel_size=voxel_size, dtype=np.uint32),
iteration=10) for file in files)
data_sources_base += gp.RandomProvider()
# data source for "add" volume
data_sources_add = tuple()
data_sources_add += tuple(
(gp.Hdf5Source(file,
datasets={
raw_add: '/volume',
},
array_specs={
raw_add:
gp.ArraySpec(interpolatable=True,
voxel_size=voxel_size,
dtype=np.uint16),
},
channels_first=False),
SwcSource(filename=file,
dataset='/reconstruction',
points=(swc_center_add, swc_add),
scale=voxel_size)) + gp.MergeProvider() +
gp.RandomLocation(ensure_nonempty=swc_center_add) + RasterizeSkeleton(
points=swc_add,
array=labels_add,
array_spec=gp.ArraySpec(
interpolatable=False, voxel_size=voxel_size, dtype=np.uint32),
iteration=1) for file in files)
data_sources_add += gp.RandomProvider()
data_sources = tuple([data_sources_base, data_sources_add
]) + gp.MergeProvider()
pipeline = (
data_sources + FusionAugment(raw_base,
raw_add,
labels_base,
labels_add,
raw,
labels,
blend_mode='labels_mask',
blend_smoothness=10,
num_blended_objects=0) +
# augment
gp.ElasticAugment([10, 10, 10], [1, 1, 1], [0, math.pi / 2.0],
subsample=8) +
gp.SimpleAugment(mirror_only=[2], transpose_only=[]) +
gp.Normalize(raw) + gp.IntensityAugment(raw, 0.9, 1.1, -0.001, 0.001) +
BinarizeGt(labels, labels_fg) +
gp.BalanceLabels(labels_fg, loss_weights) +
# train
gp.PreCache(cache_size=40, num_workers=10) +
gp.tensorflow.Train('./train_net',
optimizer=net_names['optimizer'],
loss=net_names['loss'],
inputs={
net_names['raw']: raw,
net_names['labels_fg']: labels_fg,
net_names['loss_weights']: loss_weights,
},
outputs={
net_names['fg']: fg,
},
gradients={
net_names['fg']: gradient_fg,
},
save_every=100) +
# visualize
gp.Snapshot(output_filename='snapshot_{iteration}.hdf',
dataset_names={
raw: 'volumes/raw',
raw_base: 'volumes/raw_base',
raw_add: 'volumes/raw_add',
labels: 'volumes/labels',
labels_base: 'volumes/labels_base',
labels_add: 'volumes/labels_add',
fg: 'volumes/fg',
labels_fg: 'volumes/labels_fg',
gradient_fg: 'volumes/gradient_fg',
},
additional_request=snapshot_request,
every=10) + gp.PrintProfilingStats(every=100))
with gp.build(pipeline):
print("Starting training...")
for i in range(max_iteration - trained_until):
pipeline.request_batch(request)
def random_point_pairs_pipeline(model,
loss,
optimizer,
dataset,
augmentation_parameters,
point_density,
out_dir,
normalize_factor=None,
checkpoint_interval=5000,
snapshot_interval=5000):
raw_0 = gp.ArrayKey('RAW_0')
points_0 = gp.GraphKey('POINTS_0')
locations_0 = gp.ArrayKey('LOCATIONS_0')
emb_0 = gp.ArrayKey('EMBEDDING_0')
raw_1 = gp.ArrayKey('RAW_1')
points_1 = gp.GraphKey('POINTS_1')
locations_1 = gp.ArrayKey('LOCATIONS_1')
emb_1 = gp.ArrayKey('EMBEDDING_1')
# TODO parse this key from somewhere
key = 'train/raw/0'
data = daisy.open_ds(dataset.filename, key)
source_roi = gp.Roi(data.roi.get_offset(), data.roi.get_shape())
voxel_size = gp.Coordinate(data.voxel_size)
emb_voxel_size = voxel_size
# Get in and out shape
in_shape = gp.Coordinate(model.in_shape)
out_shape = gp.Coordinate(model.out_shape)
logger.info(f"source roi: {source_roi}")
logger.info(f"in_shape: {in_shape}")
logger.info(f"out_shape: {out_shape}")
logger.info(f"voxel_size: {voxel_size}")
request = gp.BatchRequest()
request.add(raw_0, in_shape)
request.add(raw_1, in_shape)
request.add(points_0, out_shape)
request.add(points_1, out_shape)
request[locations_0] = gp.ArraySpec(nonspatial=True)
request[locations_1] = gp.ArraySpec(nonspatial=True)
snapshot_request = gp.BatchRequest()
snapshot_request[emb_0] = gp.ArraySpec(roi=request[points_0].roi)
snapshot_request[emb_1] = gp.ArraySpec(roi=request[points_1].roi)
# Let's hardcode this for now
# TODO read actual number from zarr file keys
n_samples = 447
batch_size = 1
dim = 2
padding = (100, 100)
sources = []
for i in range(n_samples):
ds_key = f'train/raw/{i}'
image_sources = tuple(
gp.ZarrSource(
dataset.filename, {raw: ds_key},
{raw: gp.ArraySpec(interpolatable=True, voxel_size=(1, 1))}) +
gp.Pad(raw, None) for raw in [raw_0, raw_1])
random_point_generator = RandomPointGenerator(density=point_density,
repetitions=2)
point_sources = tuple(
(RandomPointSource(points_0,
dim,
random_point_generator=random_point_generator),
RandomPointSource(points_1,
dim,
random_point_generator=random_point_generator)))
# TODO: get augmentation parameters from some config file!
points_and_image_sources = tuple(
(img_source, point_source) + gp.MergeProvider() + \
gp.SimpleAugment() + \
gp.ElasticAugment(
spatial_dims=2,
control_point_spacing=(10, 10),
jitter_sigma=(0.0, 0.0),
rotation_interval=(0, math.pi/2)) + \
gp.IntensityAugment(r,
scale_min=0.8,
scale_max=1.2,
shift_min=-0.2,
shift_max=0.2,
clip=False) + \
gp.NoiseAugment(r, var=0.01, clip=False)
for r, img_source, point_source
in zip([raw_0, raw_1], image_sources, point_sources))
sample_source = points_and_image_sources + gp.MergeProvider()
data = daisy.open_ds(dataset.filename, ds_key)
source_roi = gp.Roi(data.roi.get_offset(), data.roi.get_shape())
sample_source += gp.Crop(raw_0, source_roi)
sample_source += gp.Crop(raw_1, source_roi)
sample_source += gp.Pad(raw_0, padding)
sample_source += gp.Pad(raw_1, padding)
sample_source += gp.RandomLocation()
sources.append(sample_source)
sources = tuple(sources)
pipeline = sources + gp.RandomProvider()
pipeline += gp.Unsqueeze([raw_0, raw_1])
pipeline += PrepareBatch(raw_0, raw_1, points_0, points_1, locations_0,
locations_1)
# How does prepare batch relate to Stack?????
pipeline += RejectArray(ensure_nonempty=locations_1)
pipeline += RejectArray(ensure_nonempty=locations_0)
# batch content
# raw_0: (1, h, w)
# raw_1: (1, h, w)
# locations_0: (n, 2)
# locations_1: (n, 2)
pipeline += gp.Stack(batch_size)
# batch content
# raw_0: (b, 1, h, w)
# raw_1: (b, 1, h, w)
# locations_0: (b, n, 2)
# locations_1: (b, n, 2)
pipeline += gp.PreCache(num_workers=10)
pipeline += gp.torch.Train(
model,
loss,
optimizer,
inputs={
'raw_0': raw_0,
'raw_1': raw_1
},
loss_inputs={
'emb_0': emb_0,
'emb_1': emb_1,
'locations_0': locations_0,
'locations_1': locations_1
},
outputs={
2: emb_0,
3: emb_1
},
array_specs={
emb_0: gp.ArraySpec(voxel_size=emb_voxel_size),
emb_1: gp.ArraySpec(voxel_size=emb_voxel_size)
},
checkpoint_basename=os.path.join(out_dir, 'model'),
save_every=checkpoint_interval)
pipeline += gp.Snapshot(
{
raw_0: 'raw_0',
raw_1: 'raw_1',
emb_0: 'emb_0',
emb_1: 'emb_1',
# locations_0 : 'locations_0',
# locations_1 : 'locations_1',
},
every=snapshot_interval,
additional_request=snapshot_request)
return pipeline, request
def train_until(max_iteration):
# get the latest checkpoint
if tf.train.latest_checkpoint("."):
trained_until = int(tf.train.latest_checkpoint(".").split("_")[-1])
else:
trained_until = 0
if trained_until >= max_iteration:
return
# array keys for data sources
raw = gp.ArrayKey("RAW")
swcs = gp.PointsKey("SWCS")
labels = gp.ArrayKey("LABELS")
# array keys for base volume
raw_base = gp.ArrayKey("RAW_BASE")
labels_base = gp.ArrayKey("LABELS_BASE")
swc_base = gp.PointsKey("SWC_BASE")
# array keys for add volume
raw_add = gp.ArrayKey("RAW_ADD")
labels_add = gp.ArrayKey("LABELS_ADD")
swc_add = gp.PointsKey("SWC_ADD")
# array keys for fused volume
raw_fused = gp.ArrayKey("RAW_FUSED")
labels_fused = gp.ArrayKey("LABELS_FUSED")
swc_fused = gp.PointsKey("SWC_FUSED")
# output data
fg = gp.ArrayKey("FG")
labels_fg = gp.ArrayKey("LABELS_FG")
labels_fg_bin = gp.ArrayKey("LABELS_FG_BIN")
gradient_fg = gp.ArrayKey("GRADIENT_FG")
loss_weights = gp.ArrayKey("LOSS_WEIGHTS")
voxel_size = gp.Coordinate((10, 3, 3))
input_size = gp.Coordinate(net_config["input_shape"]) * voxel_size
output_size = gp.Coordinate(net_config["output_shape"]) * voxel_size
# add request
request = gp.BatchRequest()
request.add(raw_fused, input_size)
request.add(labels_fused, input_size)
request.add(swc_fused, input_size)
request.add(labels_fg, output_size)
request.add(labels_fg_bin, output_size)
request.add(loss_weights, output_size)
# add snapshot request
# request.add(fg, output_size)
# request.add(labels_fg, output_size)
request.add(gradient_fg, output_size)
request.add(raw_base, input_size)
request.add(raw_add, input_size)
request.add(labels_base, input_size)
request.add(labels_add, input_size)
request.add(swc_base, input_size)
request.add(swc_add, input_size)
data_sources = tuple(
(
gp.N5Source(
filename=str(
(
filename
/ "consensus-neurons-with-machine-centerpoints-labelled-as-swcs-carved.n5"
).absolute()
),
datasets={raw: "volume"},
array_specs={
raw: gp.ArraySpec(
interpolatable=True, voxel_size=voxel_size, dtype=np.uint16
)
},
),
MouselightSwcFileSource(
filename=str(
(
filename
/ "consensus-neurons-with-machine-centerpoints-labelled-as-swcs"
).absolute()
),
points=(swcs,),
scale=voxel_size,
transpose=(2, 1, 0),
transform_file=str((filename / "transform.txt").absolute()),
ignore_human_nodes=True
),
)
+ gp.MergeProvider()
+ gp.RandomLocation(
ensure_nonempty=swcs, ensure_centered=True
)
+ RasterizeSkeleton(
points=swcs,
array=labels,
array_spec=gp.ArraySpec(
interpolatable=False, voxel_size=voxel_size, dtype=np.uint32
),
)
+ GrowLabels(labels, radius=10)
# augment
+ gp.ElasticAugment(
[40, 10, 10],
[0.25, 1, 1],
[0, math.pi / 2.0],
subsample=4,
)
+ gp.SimpleAugment(mirror_only=[1, 2], transpose_only=[1, 2])
+ gp.Normalize(raw)
+ gp.IntensityAugment(raw, 0.9, 1.1, -0.001, 0.001)
for filename in Path(sample_dir).iterdir()
if "2018-08-01" in filename.name # or "2018-07-02" in filename.name
)
pipeline = (
data_sources
+ gp.RandomProvider()
+ GetNeuronPair(
swcs,
raw,
labels,
(swc_base, swc_add),
(raw_base, raw_add),
(labels_base, labels_add),
seperate_by=150,
shift_attempts=50,
request_attempts=10,
)
+ FusionAugment(
raw_base,
raw_add,
labels_base,
labels_add,
swc_base,
swc_add,
raw_fused,
labels_fused,
swc_fused,
blend_mode="labels_mask",
blend_smoothness=10,
num_blended_objects=0,
)
+ Crop(labels_fused, labels_fg)
+ BinarizeGt(labels_fg, labels_fg_bin)
+ gp.BalanceLabels(labels_fg_bin, loss_weights)
# train
+ gp.PreCache(cache_size=40, num_workers=10)
+ gp.tensorflow.Train(
"./train_net",
optimizer=net_names["optimizer"],
loss=net_names["loss"],
inputs={
net_names["raw"]: raw_fused,
net_names["labels_fg"]: labels_fg_bin,
net_names["loss_weights"]: loss_weights,
},
outputs={net_names["fg"]: fg},
gradients={net_names["fg"]: gradient_fg},
save_every=100000,
)
+ gp.Snapshot(
output_filename="snapshot_{iteration}.hdf",
dataset_names={
raw_fused: "volumes/raw_fused",
raw_base: "volumes/raw_base",
raw_add: "volumes/raw_add",
labels_fused: "volumes/labels_fused",
labels_base: "volumes/labels_base",
labels_add: "volumes/labels_add",
labels_fg_bin: "volumes/labels_fg_bin",
fg: "volumes/pred_fg",
gradient_fg: "volumes/gradient_fg",
},
every=100,
)
+ gp.PrintProfilingStats(every=10)
)
with gp.build(pipeline):
logging.info("Starting training...")
for i in range(max_iteration - trained_until):
logging.info("requesting batch {}".format(i))
batch = pipeline.request_batch(request)
"""
def train_until(max_iteration):
in_channels = 1
num_fmaps = 12
fmap_inc_factors = 6
downsample_factors = [(1, 3, 3), (1, 3, 3), (3, 3, 3)]
unet = UNet(in_channels,
num_fmaps,
fmap_inc_factors,
downsample_factors,
constant_upsample=True)
model = Convolve(unet, 12, 1)
loss = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-6)
# start of gunpowder part:
raw = gp.ArrayKey('RAW')
points = gp.GraphKey('POINTS')
groundtruth = gp.ArrayKey('RASTER')
prediction = gp.ArrayKey('PRED_POINT')
grad = gp.ArrayKey('GRADIENT')
voxel_size = gp.Coordinate((40, 4, 4))
input_shape = (96, 430, 430)
output_shape = (60, 162, 162)
input_size = gp.Coordinate(input_shape) * voxel_size
output_size = gp.Coordinate(output_shape) * voxel_size
request = gp.BatchRequest()
request.add(raw, input_size)
request.add(points, output_size)
request.add(groundtruth, output_size)
request.add(prediction, output_size)
request.add(grad, output_size)
pos_sources = tuple(
gp.ZarrSource(filename, {raw: 'volumes/raw'},
{raw: gp.ArraySpec(interpolatable=True)}) +
AddCenterPoint(points, raw) + gp.Pad(raw, None) +
gp.RandomLocation(ensure_nonempty=points)
for filename in pos_samples) + gp.RandomProvider()
neg_sources = tuple(
gp.ZarrSource(filename, {raw: 'volumes/raw'},
{raw: gp.ArraySpec(interpolatable=True)}) +
AddNoPoint(points, raw) + gp.RandomLocation()
for filename in neg_samples) + gp.RandomProvider()
data_sources = (pos_sources, neg_sources)
data_sources += gp.RandomProvider(probabilities=[0.9, 0.1])
data_sources += gp.Normalize(raw)
train_pipeline = data_sources
train_pipeline += gp.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)
train_pipeline += gp.SimpleAugment(transpose_only=[1, 2])
train_pipeline += gp.IntensityAugment(raw, 0.9, 1.1, -0.1, 0.1, \
z_section_wise=True)
train_pipeline += gp.RasterizePoints(
points,
groundtruth,
array_spec=gp.ArraySpec(voxel_size=voxel_size),
settings=gp.RasterizationSettings(radius=(100, 100, 100), mode='peak'))
train_pipeline += gp.PreCache(cache_size=40, num_workers=10)
train_pipeline += Reshape(raw, (1, 1) + input_shape)
train_pipeline += Reshape(groundtruth, (1, 1) + output_shape)
train_pipeline += gp_torch.Train(model=model,
loss=loss,
optimizer=optimizer,
inputs={'x': raw},
outputs={0: prediction},
loss_inputs={
0: prediction,
1: groundtruth
},
gradients={0: grad},
save_every=1000,
log_dir='log')
train_pipeline += Reshape(raw, input_shape)
train_pipeline += Reshape(groundtruth, output_shape)
train_pipeline += Reshape(prediction, output_shape)
train_pipeline += Reshape(grad, output_shape)
train_pipeline += gp.Snapshot(
{
raw: 'volumes/raw',
groundtruth: 'volumes/groundtruth',
prediction: 'volumes/prediction',
grad: 'volumes/gradient'
},
every=500,
output_filename='test_{iteration}.hdf')
train_pipeline += gp.PrintProfilingStats(every=10)
with gp.build(train_pipeline):
for i in range(max_iteration):
train_pipeline.request_batch(request)
def train_until(max_iteration):
# get the latest checkpoint
if tf.train.latest_checkpoint("."):
trained_until = int(tf.train.latest_checkpoint(".").split("_")[-1])
else:
trained_until = 0
if trained_until >= max_iteration:
return
# array keys for fused volume
raw = gp.ArrayKey("RAW")
labels = gp.ArrayKey("LABELS")
labels_fg = gp.ArrayKey("LABELS_FG")
# array keys for base volume
raw_base = gp.ArrayKey("RAW_BASE")
labels_base = gp.ArrayKey("LABELS_BASE")
swc_base = gp.PointsKey("SWC_BASE")
swc_center_base = gp.PointsKey("SWC_CENTER_BASE")
# array keys for add volume
raw_add = gp.ArrayKey("RAW_ADD")
labels_add = gp.ArrayKey("LABELS_ADD")
swc_add = gp.PointsKey("SWC_ADD")
swc_center_add = gp.PointsKey("SWC_CENTER_ADD")
# output data
fg = gp.ArrayKey("FG")
gradient_fg = gp.ArrayKey("GRADIENT_FG")
loss_weights = gp.ArrayKey("LOSS_WEIGHTS")
voxel_size = gp.Coordinate((4, 1, 1))
input_size = gp.Coordinate(net_config["input_shape"]) * voxel_size
output_size = gp.Coordinate(net_config["output_shape"]) * voxel_size
# add request
request = gp.BatchRequest()
request.add(raw, input_size)
request.add(labels, output_size)
request.add(labels_fg, output_size)
request.add(loss_weights, output_size)
request.add(swc_center_base, output_size)
request.add(swc_center_add, output_size)
# add snapshot request
snapshot_request = gp.BatchRequest()
snapshot_request.add(fg, output_size)
snapshot_request.add(labels_fg, output_size)
snapshot_request.add(gradient_fg, output_size)
snapshot_request.add(raw_base, input_size)
snapshot_request.add(raw_add, input_size)
snapshot_request.add(labels_base, input_size)
snapshot_request.add(labels_add, input_size)
# data source for "base" volume
data_sources_base = tuple(
(
gp.Hdf5Source(
filename,
datasets={raw_base: "/volume"},
array_specs={
raw_base:
gp.ArraySpec(interpolatable=True,
voxel_size=voxel_size,
dtype=np.uint16)
},
channels_first=False,
),
SwcSource(
filename=filename,
dataset="/reconstruction",
points=(swc_center_base, swc_base),
scale=voxel_size,
),
) + gp.MergeProvider() +
gp.RandomLocation(ensure_nonempty=swc_center_base) + RasterizeSkeleton(
points=swc_base,
array=labels_base,
array_spec=gp.ArraySpec(
interpolatable=False, voxel_size=voxel_size, dtype=np.uint32),
radius=5.0,
) for filename in files)
# data source for "add" volume
data_sources_add = tuple(
(
gp.Hdf5Source(
file,
datasets={raw_add: "/volume"},
array_specs={
raw_add:
gp.ArraySpec(interpolatable=True,
voxel_size=voxel_size,
dtype=np.uint16)
},
channels_first=False,
),
SwcSource(
filename=file,
dataset="/reconstruction",
points=(swc_center_add, swc_add),
scale=voxel_size,
),
) + gp.MergeProvider() +
gp.RandomLocation(ensure_nonempty=swc_center_add) + RasterizeSkeleton(
points=swc_add,
array=labels_add,
array_spec=gp.ArraySpec(
interpolatable=False, voxel_size=voxel_size, dtype=np.uint32),
radius=5.0,
) for file in files)
data_sources = (
(data_sources_base + gp.RandomProvider()),
(data_sources_add + gp.RandomProvider()),
) + gp.MergeProvider()
pipeline = (
data_sources + FusionAugment(
raw_base,
raw_add,
labels_base,
labels_add,
raw,
labels,
blend_mode="labels_mask",
blend_smoothness=10,
num_blended_objects=0,
) +
# augment
gp.ElasticAugment([40, 10, 10], [0.25, 1, 1], [0, math.pi / 2.0],
subsample=4) +
gp.SimpleAugment(mirror_only=[1, 2], transpose_only=[1, 2]) +
gp.Normalize(raw) + gp.IntensityAugment(raw, 0.9, 1.1, -0.001, 0.001) +
BinarizeGt(labels, labels_fg) +
gp.BalanceLabels(labels_fg, loss_weights) +
# train
gp.PreCache(cache_size=40, num_workers=10) + gp.tensorflow.Train(
"./train_net",
optimizer=net_names["optimizer"],
loss=net_names["loss"],
inputs={
net_names["raw"]: raw,
net_names["labels_fg"]: labels_fg,
net_names["loss_weights"]: loss_weights,
},
outputs={net_names["fg"]: fg},
gradients={net_names["fg"]: gradient_fg},
save_every=100000,
) +
# visualize
gp.Snapshot(
output_filename="snapshot_{iteration}.hdf",
dataset_names={
raw: "volumes/raw",
raw_base: "volumes/raw_base",
raw_add: "volumes/raw_add",
labels: "volumes/labels",
labels_base: "volumes/labels_base",
labels_add: "volumes/labels_add",
fg: "volumes/fg",
labels_fg: "volumes/labels_fg",
gradient_fg: "volumes/gradient_fg",
},
additional_request=snapshot_request,
every=100,
) + gp.PrintProfilingStats(every=100))
with gp.build(pipeline):
print("Starting training...")
for i in range(max_iteration - trained_until):
pipeline.request_batch(request)
def node(self, _raw_key=None, _gt_key=None, _mask_key=None):
return gp.SimpleAugment()
def train_until(max_iteration, name='train_net', output_folder='.', clip_max=2000):
# get the latest checkpoint
if tf.train.latest_checkpoint(output_folder):
trained_until = int(tf.train.latest_checkpoint(output_folder).split('_')[-1])
else:
trained_until = 0
if trained_until >= max_iteration:
return
with open(os.path.join(output_folder, name + '_config.json'), 'r') as f:
net_config = json.load(f)
with open(os.path.join(output_folder, name + '_names.json'), 'r') as f:
net_names = json.load(f)
# array keys
raw = gp.ArrayKey('RAW')
gt_instances = gp.ArrayKey('GT_INSTANCES')
gt_mask = gp.ArrayKey('GT_MASK')
pred_mask = gp.ArrayKey('PRED_MASK')
#loss_weights = gp.ArrayKey('LOSS_WEIGHTS')
loss_gradients = gp.ArrayKey('LOSS_GRADIENTS')
# array keys for base and add volume
raw_base = gp.ArrayKey('RAW_BASE')
gt_instances_base = gp.ArrayKey('GT_INSTANCES_BASE')
gt_mask_base = gp.ArrayKey('GT_MASK_BASE')
raw_add = gp.ArrayKey('RAW_ADD')
gt_instances_add = gp.ArrayKey('GT_INSTANCES_ADD')
gt_mask_add = gp.ArrayKey('GT_MASK_ADD')
voxel_size = gp.Coordinate((1, 1, 1))
input_shape = gp.Coordinate(net_config['input_shape'])
output_shape = gp.Coordinate(net_config['output_shape'])
context = gp.Coordinate(input_shape - output_shape) / 2
request = gp.BatchRequest()
request.add(raw, input_shape)
request.add(gt_instances, output_shape)
request.add(gt_mask, output_shape)
#request.add(loss_weights, output_shape)
request.add(raw_base, input_shape)
request.add(raw_add, input_shape)
request.add(gt_mask_base, output_shape)
request.add(gt_mask_add, output_shape)
snapshot_request = gp.BatchRequest()
snapshot_request.add(raw, input_shape)
#snapshot_request.add(raw_base, input_shape)
#snapshot_request.add(raw_add, input_shape)
snapshot_request.add(gt_mask, output_shape)
#snapshot_request.add(gt_mask_base, output_shape)
#snapshot_request.add(gt_mask_add, output_shape)
snapshot_request.add(pred_mask, output_shape)
snapshot_request.add(loss_gradients, output_shape)
# specify data source
# data source for base volume
data_sources_base = tuple()
for data_file in data_files:
current_path = os.path.join(data_dir, data_file)
with h5py.File(current_path, 'r') as f:
data_sources_base += tuple(
gp.Hdf5Source(
current_path,
datasets={
raw_base: sample + '/raw',
gt_instances_base: sample + '/gt',
gt_mask_base: sample + '/fg',
},
array_specs={
raw_base: gp.ArraySpec(interpolatable=True, dtype=np.uint16, voxel_size=voxel_size),
gt_instances_base: gp.ArraySpec(interpolatable=False, dtype=np.uint16, voxel_size=voxel_size),
gt_mask_base: gp.ArraySpec(interpolatable=False, dtype=np.bool, voxel_size=voxel_size),
}
) +
Convert(gt_mask_base, np.uint8) +
gp.Pad(raw_base, context) +
gp.Pad(gt_instances_base, context) +
gp.Pad(gt_mask_base, context) +
gp.RandomLocation(min_masked=0.005, mask=gt_mask_base)
#gp.Reject(gt_mask_base, min_masked=0.005, reject_probability=1.)
for sample in f)
data_sources_base += gp.RandomProvider()
# data source for add volume
data_sources_add = tuple()
for data_file in data_files:
current_path = os.path.join(data_dir, data_file)
with h5py.File(current_path, 'r') as f:
data_sources_add += tuple(
gp.Hdf5Source(
current_path,
datasets={
raw_add: sample + '/raw',
gt_instances_add: sample + '/gt',
gt_mask_add: sample + '/fg',
},
array_specs={
raw_add: gp.ArraySpec(interpolatable=True, dtype=np.uint16, voxel_size=voxel_size),
gt_instances_add: gp.ArraySpec(interpolatable=False, dtype=np.uint16, voxel_size=voxel_size),
gt_mask_add: gp.ArraySpec(interpolatable=False, dtype=np.bool, voxel_size=voxel_size),
}
) +
Convert(gt_mask_add, np.uint8) +
gp.Pad(raw_add, context) +
gp.Pad(gt_instances_add, context) +
gp.Pad(gt_mask_add, context) +
gp.RandomLocation() +
gp.Reject(gt_mask_add, min_masked=0.005, reject_probability=0.95)
for sample in f)
data_sources_add += gp.RandomProvider()
data_sources = tuple([data_sources_base, data_sources_add]) + gp.MergeProvider()
pipeline = (
data_sources +
nl.FusionAugment(
raw_base, raw_add, gt_instances_base, gt_instances_add, raw, gt_instances,
blend_mode='labels_mask', blend_smoothness=5, num_blended_objects=0
) +
BinarizeLabels(gt_instances, gt_mask) +
nl.Clip(raw, 0, clip_max) +
gp.Normalize(raw, factor=1.0/clip_max) +
gp.ElasticAugment(
control_point_spacing=[20, 20, 20],
jitter_sigma=[1, 1, 1],
rotation_interval=[0, math.pi/2.0],
subsample=4) +
gp.SimpleAugment(mirror_only=[1, 2], transpose_only=[1, 2]) +
gp.IntensityAugment(raw, 0.9, 1.1, -0.1, 0.1) +
gp.IntensityScaleShift(raw, 2, -1) +
#gp.BalanceLabels(gt_mask, loss_weights) +
# train
gp.PreCache(
cache_size=40,
num_workers=10) +
gp.tensorflow.Train(
os.path.join(output_folder, name),
optimizer=net_names['optimizer'],
loss=net_names['loss'],
inputs={
net_names['raw']: raw,
net_names['gt']: gt_mask,
#net_names['loss_weights']: loss_weights,
},
outputs={
net_names['pred']: pred_mask,
},
gradients={
net_names['output']: loss_gradients,
},
save_every=5000) +
# visualize
gp.Snapshot({
raw: 'volumes/raw',
pred_mask: 'volumes/pred_mask',
gt_mask: 'volumes/gt_mask',
#loss_weights: 'volumes/loss_weights',
loss_gradients: 'volumes/loss_gradients',
},
output_filename=os.path.join(output_folder, 'snapshots', 'batch_{iteration}.hdf'),
additional_request=snapshot_request,
every=2500) +
gp.PrintProfilingStats(every=1000)
)
with gp.build(pipeline):
print("Starting training...")
for i in range(max_iteration - trained_until):
pipeline.request_batch(request)
def build_pipeline(parameter, augment=True):
voxel_size = gp.Coordinate(parameter['voxel_size'])
# Array Specifications.
raw = gp.ArrayKey('RAW')
gt_neurons = gp.ArrayKey('GT_NEURONS')
gt_postpre_vectors = gp.ArrayKey('GT_POSTPRE_VECTORS')
gt_post_indicator = gp.ArrayKey('GT_POST_INDICATOR')
post_loss_weight = gp.ArrayKey('POST_LOSS_WEIGHT')
vectors_mask = gp.ArrayKey('VECTORS_MASK')
pred_postpre_vectors = gp.ArrayKey('PRED_POSTPRE_VECTORS')
pred_post_indicator = gp.ArrayKey('PRED_POST_INDICATOR')
grad_syn_indicator = gp.ArrayKey('GRAD_SYN_INDICATOR')
grad_partner_vectors = gp.ArrayKey('GRAD_PARTNER_VECTORS')
# Points specifications
dummypostsyn = gp.PointsKey('DUMMYPOSTSYN')
postsyn = gp.PointsKey('POSTSYN')
presyn = gp.PointsKey('PRESYN')
trg_context = 140 # AddPartnerVectorMap context in nm - pre-post distance
with open('train_net_config.json', 'r') as f:
net_config = json.load(f)
input_size = gp.Coordinate(net_config['input_shape']) * voxel_size
output_size = gp.Coordinate(net_config['output_shape']) * voxel_size
request = gp.BatchRequest()
request.add(raw, input_size)
request.add(gt_neurons, output_size)
request.add(gt_postpre_vectors, output_size)
request.add(gt_post_indicator, output_size)
request.add(post_loss_weight, output_size)
request.add(vectors_mask, output_size)
request.add(dummypostsyn, output_size)
for (key, request_spec) in request.items():
print(key)
print(request_spec.roi)
request_spec.roi.contains(request_spec.roi)
# slkfdms
snapshot_request = gp.BatchRequest({
pred_post_indicator:
request[gt_postpre_vectors],
pred_postpre_vectors:
request[gt_postpre_vectors],
grad_syn_indicator:
request[gt_postpre_vectors],
grad_partner_vectors:
request[gt_postpre_vectors],
vectors_mask:
request[gt_postpre_vectors]
})
postsyn_rastersetting = gp.RasterizationSettings(
radius=parameter['blob_radius'],
mask=gt_neurons,
mode=parameter['blob_mode'])
pipeline = tuple([
create_source(sample, raw, presyn, postsyn, dummypostsyn, parameter,
gt_neurons) for sample in samples
])
pipeline += gp.RandomProvider()
if augment:
pipeline += gp.ElasticAugment([4, 40, 40], [0, 2, 2],
[0, math.pi / 2.0],
prob_slip=0.05,
prob_shift=0.05,
max_misalign=10,
subsample=8)
pipeline += gp.SimpleAugment(transpose_only=[1, 2], mirror_only=[1, 2])
pipeline += gp.IntensityAugment(raw,
0.9,
1.1,
-0.1,
0.1,
z_section_wise=True)
pipeline += gp.IntensityScaleShift(raw, 2, -1)
pipeline += gp.RasterizePoints(
postsyn, gt_post_indicator,
gp.ArraySpec(voxel_size=voxel_size, dtype=np.int32),
postsyn_rastersetting)
spec = gp.ArraySpec(voxel_size=voxel_size)
pipeline += AddPartnerVectorMap(
src_points=postsyn,
trg_points=presyn,
array=gt_postpre_vectors,
radius=parameter['d_blob_radius'],
trg_context=trg_context, # enlarge
array_spec=spec,
mask=gt_neurons,
pointmask=vectors_mask)
pipeline += gp.BalanceLabels(labels=gt_post_indicator,
scales=post_loss_weight,
slab=(-1, -1, -1),
clipmin=parameter['cliprange'][0],
clipmax=parameter['cliprange'][1])
if parameter['d_scale'] != 1:
pipeline += gp.IntensityScaleShift(gt_postpre_vectors,
scale=parameter['d_scale'],
shift=0)
pipeline += gp.PreCache(cache_size=40, num_workers=10)
pipeline += gp.tensorflow.Train(
'./train_net',
optimizer=net_config['optimizer'],
loss=net_config['loss'],
summary=net_config['summary'],
log_dir='./tensorboard/',
save_every=30000, # 10000
log_every=100,
inputs={
net_config['raw']:
raw,
net_config['gt_partner_vectors']:
gt_postpre_vectors,
net_config['gt_syn_indicator']:
gt_post_indicator,
net_config['vectors_mask']:
vectors_mask,
# Loss weights --> mask
net_config['indicator_weight']:
post_loss_weight, # Loss weights
},
outputs={
net_config['pred_partner_vectors']: pred_postpre_vectors,
net_config['pred_syn_indicator']: pred_post_indicator,
},
gradients={
net_config['pred_partner_vectors']: grad_partner_vectors,
net_config['pred_syn_indicator']: grad_syn_indicator,
},
)
# Visualize.
pipeline += gp.IntensityScaleShift(raw, 0.5, 0.5)
pipeline += gp.Snapshot(
{
raw: 'volumes/raw',
gt_neurons: 'volumes/labels/neuron_ids',
gt_post_indicator: 'volumes/gt_post_indicator',
gt_postpre_vectors: 'volumes/gt_postpre_vectors',
pred_postpre_vectors: 'volumes/pred_postpre_vectors',
pred_post_indicator: 'volumes/pred_post_indicator',
post_loss_weight: 'volumes/post_loss_weight',
grad_syn_indicator: 'volumes/post_indicator_gradients',
grad_partner_vectors: 'volumes/partner_vectors_gradients',
vectors_mask: 'volumes/vectors_mask'
},
every=1000,
output_filename='batch_{iteration}.hdf',
compression_type='gzip',
additional_request=snapshot_request)
pipeline += gp.PrintProfilingStats(every=100)
print("Starting training...")
max_iteration = parameter['max_iteration']
with gp.build(pipeline) as b:
for i in range(max_iteration):
b.request_batch(request)
def train_until(**kwargs):
print("cuda visibile devices", os.environ["CUDA_VISIBLE_DEVICES"])
if tf.train.latest_checkpoint(kwargs['output_folder']):
trained_until = int(
tf.train.latest_checkpoint(kwargs['output_folder']).split('_')[-1])
else:
trained_until = 0
if trained_until >= kwargs['max_iteration']:
return
raw = gp.ArrayKey('RAW')
raw_cropped = gp.ArrayKey('RAW_CROPPED')
gt_labels = gp.ArrayKey('GT_LABELS')
gt_instances = gp.ArrayKey('GT_INSTANCES')
gt_affs = gp.ArrayKey('GT_AFFS')
gt_numinst = gp.ArrayKey('GT_NUMINST')
gt_sample_mask = gp.ArrayKey('GT_SAMPLE_MASK')
pred_affs = gp.ArrayKey('PRED_AFFS')
pred_affs_gradients = gp.ArrayKey('PRED_AFFS_GRADIENTS')
pred_numinst = gp.ArrayKey('PRED_NUMINST')
with open(os.path.join(kwargs['output_folder'],
kwargs['name'] + '_config.json'), 'r') as f:
net_config = json.load(f)
with open(os.path.join(kwargs['output_folder'],
kwargs['name'] + '_names.json'), 'r') as f:
net_names = json.load(f)
voxel_size = gp.Coordinate(kwargs['voxel_size'])
input_shape_world = gp.Coordinate(net_config['input_shape'])*voxel_size
output_shape_world = gp.Coordinate(net_config['output_shape'])*voxel_size
context = gp.Coordinate(input_shape_world - output_shape_world) / 2
# formulate the request for what a batch should (at least) contain
request = gp.BatchRequest()
request.add(raw, input_shape_world)
request.add(raw_cropped, output_shape_world)
request.add(gt_labels, output_shape_world)
request.add(gt_instances, output_shape_world)
request.add(gt_sample_mask, output_shape_world)
request.add(gt_affs, output_shape_world)
if kwargs['overlapping_inst']:
request.add(gt_numinst, output_shape_world)
# request.add(loss_weights_affs, output_shape_world)
# when we make a snapshot for inspection (see below), we also want to
# request the predicted affinities and gradients of the loss wrt the
# affinities
snapshot_request = gp.BatchRequest()
snapshot_request.add(raw_cropped, output_shape_world)
snapshot_request.add(pred_affs, output_shape_world)
if kwargs['overlapping_inst']:
snapshot_request.add(pred_numinst, output_shape_world)
# snapshot_request.add(pred_affs_gradients, output_shape_world)
if kwargs['input_format'] != "hdf" and kwargs['input_format'] != "zarr":
raise NotImplementedError("train node for %s not implemented yet",
kwargs['input_format'])
raw_key = kwargs.get('raw_key', 'volumes/raw')
print('raw key: ', raw_key)
fls = []
shapes = []
for f in kwargs['data_files']:
fls.append(os.path.splitext(f)[0])
if kwargs['input_format'] == "hdf":
vol = h5py.File(f, 'r')[raw_key]
elif kwargs['input_format'] == "zarr":
vol = zarr.open(f, 'r')[raw_key]
# print(f, vol.shape, vol.dtype)
shapes.append(vol.shape)
if vol.dtype != np.float32:
print("please convert to float32")
ln = len(fls)
print("first 5 files: ", fls[0:4])
if kwargs['input_format'] == "hdf":
sourceNode = gp.Hdf5Source
elif kwargs['input_format'] == "zarr":
sourceNode = gp.ZarrSource
neighborhood = []
psH = np.array(kwargs['patchshape'])//2
for i in range(-psH[1], psH[1]+1, kwargs['patchstride'][1]):
for j in range(-psH[2], psH[2]+1, kwargs['patchstride'][2]):
neighborhood.append([i,j])
datasets = {
raw: raw_key,
gt_labels: 'volumes/gt_labels',
gt_instances: 'volumes/gt_instances'
}
array_specs = {
raw: gp.ArraySpec(interpolatable=True),
gt_labels: gp.ArraySpec(interpolatable=False),
gt_instances: gp.ArraySpec(interpolatable=False)
}
inputs = {
net_names['raw']: raw,
net_names['gt_affs']: gt_affs,
# net_names['loss_weights_affs']: loss_weights_affs,
}
outputs = {
net_names['pred_affs']: pred_affs,
net_names['raw_cropped']: raw_cropped,
}
snapshot = {
raw: '/volumes/raw',
raw_cropped: 'volumes/raw_cropped',
gt_affs: '/volumes/gt_affs',
pred_affs: '/volumes/pred_affs',
pred_affs_gradients: '/volumes/pred_affs_gradients',
}
if kwargs['overlapping_inst']:
datasets[gt_numinst] = 'volumes/gt_numinst'
array_specs[gt_numinst] = gp.ArraySpec(interpolatable=False)
inputs[net_names['gt_numinst']] = gt_numinst
outputs[net_names['pred_numinst']] = pred_numinst
snapshot[gt_numinst] = '/volumes/gt_numinst'
snapshot[pred_numinst] = '/volumes/pred_numinst'
augmentation = kwargs['augmentation']
sampling = kwargs['sampling']
source_fg = tuple(
sourceNode(
fls[t] + "." + kwargs['input_format'],
datasets=datasets,
array_specs=array_specs
) +
gp.Pad(raw, context) +
# chose a random location for each requested batch
nl.CountOverlap(gt_labels, gt_sample_mask, maxnuminst=1) +
gp.RandomLocation(
min_masked=sampling['min_masked'],
mask=gt_sample_mask
)
for t in range(ln)
)
source_fg += gp.RandomProvider()
source_overlap = tuple(
sourceNode(
fls[t] + "." + kwargs['input_format'],
datasets=datasets,
array_specs=array_specs
) +
gp.Pad(raw, context) +
# chose a random location for each requested batch
nl.MaskCloseDistanceToOverlap(
gt_labels, gt_sample_mask,
sampling['overlap_min_dist'],
sampling['overlap_max_dist']
) +
gp.RandomLocation(
min_masked=sampling['min_masked_overlap'],
mask=gt_sample_mask
)
for t in range(ln)
)
source_overlap += gp.RandomProvider()
pipeline = (
(source_fg, source_overlap) +
# chose a random source (i.e., sample) from the above
gp.RandomProvider(probabilities=[sampling['probability_fg'],
sampling['probability_overlap']]) +
# elastically deform the batch
gp.ElasticAugment(
augmentation['elastic']['control_point_spacing'],
augmentation['elastic']['jitter_sigma'],
[augmentation['elastic']['rotation_min']*np.pi/180.0,
augmentation['elastic']['rotation_max']*np.pi/180.0]) +
# apply transpose and mirror augmentations
gp.SimpleAugment(
mirror_only=augmentation['simple'].get("mirror"),
transpose_only=augmentation['simple'].get("transpose")) +
# # scale and shift the intensity of the raw array
gp.IntensityAugment(
raw,
scale_min=augmentation['intensity']['scale'][0],
scale_max=augmentation['intensity']['scale'][1],
shift_min=augmentation['intensity']['shift'][0],
shift_max=augmentation['intensity']['shift'][1],
z_section_wise=False) +
gp.IntensityScaleShift(raw, 2, -1) +
# convert labels into affinities between voxels
nl.AddAffinities(
neighborhood,
gt_labels,
gt_affs,
multiple_labels=kwargs['overlapping_inst']) +
# pre-cache batches from the point upstream
gp.PreCache(
cache_size=kwargs['cache_size'],
num_workers=kwargs['num_workers']) +
# perform one training iteration for each passing batch (here we use
# the tensor names earlier stored in train_net.config)
gp.tensorflow.Train(
os.path.join(kwargs['output_folder'], kwargs['name']),
optimizer=net_names['optimizer'],
summary=net_names['summaries'],
log_dir=kwargs['output_folder'],
loss=net_names['loss'],
inputs=inputs,
outputs=outputs,
gradients={
net_names['pred_affs']: pred_affs_gradients,
},
save_every=kwargs['checkpoints']) +
# save the passing batch as an HDF5 file for inspection
gp.Snapshot(
snapshot,
output_dir=os.path.join(kwargs['output_folder'], 'snapshots'),
output_filename='batch_{iteration}.hdf',
every=kwargs['snapshots'],
additional_request=snapshot_request,
compression_type='gzip') +
# show a summary of time spend in each node every 10 iterations
gp.PrintProfilingStats(every=kwargs['profiling'])
)
#########
# TRAIN #
#########
print("Starting training...")
with gp.build(pipeline):
print(pipeline)
for i in range(trained_until, kwargs['max_iteration']):
# print("request", request)
start = time.time()
pipeline.request_batch(request)
time_of_iteration = time.time() - start
logger.info(
"Batch: iteration=%d, time=%f",
i, time_of_iteration)
# exit()
print("Training finished")
def train_until(max_iteration, name='train_net', output_folder='.', clip_max=2000):
# get the latest checkpoint
if tf.train.latest_checkpoint(output_folder):
trained_until = int(tf.train.latest_checkpoint(output_folder).split('_')[-1])
else:
trained_until = 0
if trained_until >= max_iteration:
return
with open(os.path.join(output_folder, name + '_config.json'), 'r') as f:
net_config = json.load(f)
with open(os.path.join(output_folder, name + '_names.json'), 'r') as f:
net_names = json.load(f)
# array keys
raw = gp.ArrayKey('RAW')
gt_mask = gp.ArrayKey('GT_MASK')
gt_dt = gp.ArrayKey('GT_DT')
pred_dt = gp.ArrayKey('PRED_DT')
loss_gradient = gp.ArrayKey('LOSS_GRADIENT')
voxel_size = gp.Coordinate((1, 1, 1))
input_shape = gp.Coordinate(net_config['input_shape'])
output_shape = gp.Coordinate(net_config['output_shape'])
context = gp.Coordinate(input_shape - output_shape) / 2
request = gp.BatchRequest()
request.add(raw, input_shape)
request.add(gt_mask, output_shape)
request.add(gt_dt, output_shape)
snapshot_request = gp.BatchRequest()
snapshot_request.add(raw, input_shape)
snapshot_request.add(gt_mask, output_shape)
snapshot_request.add(gt_dt, output_shape)
snapshot_request.add(pred_dt, output_shape)
snapshot_request.add(loss_gradient, output_shape)
# specify data source
data_sources = tuple()
for data_file in data_files:
current_path = os.path.join(data_dir, data_file)
with h5py.File(current_path, 'r') as f:
data_sources += tuple(
gp.Hdf5Source(
current_path,
datasets={
raw: sample + '/raw',
gt_mask: sample + '/fg'
},
array_specs={
raw: gp.ArraySpec(interpolatable=True, dtype=np.uint16, voxel_size=voxel_size),
gt_mask: gp.ArraySpec(interpolatable=False, dtype=np.bool, voxel_size=voxel_size),
}
) +
Convert(gt_mask, np.uint8) +
gp.Pad(raw, context) +
gp.Pad(gt_mask, context) +
gp.RandomLocation()
for sample in f)
pipeline = (
data_sources +
gp.RandomProvider() +
gp.Reject(gt_mask, min_masked=0.005, reject_probability=1.) +
DistanceTransform(gt_mask, gt_dt, 3) +
nl.Clip(raw, 0, clip_max) +
gp.Normalize(raw, factor=1.0/clip_max) +
gp.ElasticAugment(
control_point_spacing=[20, 20, 20],
jitter_sigma=[1, 1, 1],
rotation_interval=[0, math.pi/2.0],
subsample=4) +
gp.SimpleAugment(mirror_only=[1,2], transpose_only=[1,2]) +
gp.IntensityAugment(raw, 0.9, 1.1, -0.1, 0.1) +
gp.IntensityScaleShift(raw, 2,-1) +
# train
gp.PreCache(
cache_size=40,
num_workers=5) +
gp.tensorflow.Train(
os.path.join(output_folder, name),
optimizer=net_names['optimizer'],
loss=net_names['loss'],
inputs={
net_names['raw']: raw,
net_names['gt_dt']: gt_dt,
},
outputs={
net_names['pred_dt']: pred_dt,
},
gradients={
net_names['pred_dt']: loss_gradient,
},
save_every=5000) +
# visualize
gp.Snapshot({
raw: 'volumes/raw',
gt_mask: 'volumes/gt_mask',
gt_dt: 'volumes/gt_dt',
pred_dt: 'volumes/pred_dt',
loss_gradient: 'volumes/gradient',
},
output_filename=os.path.join(output_folder, 'snapshots', 'batch_{iteration}.hdf'),
additional_request=snapshot_request,
every=2000) +
gp.PrintProfilingStats(every=500)
)
with gp.build(pipeline):
print("Starting training...")
for i in range(max_iteration - trained_until):
pipeline.request_batch(request)
def train_until(**kwargs):
print("cuda visibile devices", os.environ["CUDA_VISIBLE_DEVICES"])
if tf.train.latest_checkpoint(kwargs['output_folder']):
trained_until = int(
tf.train.latest_checkpoint(kwargs['output_folder']).split('_')[-1])
else:
trained_until = 0
if trained_until >= kwargs['max_iteration']:
return
anchor = gp.ArrayKey('ANCHOR')
raw = gp.ArrayKey('RAW')
raw_cropped = gp.ArrayKey('RAW_CROPPED')
gt_labels = gp.ArrayKey('GT_LABELS')
gt_affs = gp.ArrayKey('GT_AFFS')
pred_affs = gp.ArrayKey('PRED_AFFS')
pred_affs_gradients = gp.ArrayKey('PRED_AFFS_GRADIENTS')
with open(
os.path.join(kwargs['output_folder'],
kwargs['name'] + '_config.json'), 'r') as f:
net_config = json.load(f)
with open(
os.path.join(kwargs['output_folder'],
kwargs['name'] + '_names.json'), 'r') as f:
net_names = json.load(f)
voxel_size = gp.Coordinate(kwargs['voxel_size'])
input_shape_world = gp.Coordinate(net_config['input_shape']) * voxel_size
output_shape_world = gp.Coordinate(net_config['output_shape']) * voxel_size
# formulate the request for what a batch should (at least) contain
request = gp.BatchRequest()
# when we make a snapshot for inspection (see below), we also want to
# request the predicted affinities and gradients of the loss wrt the
# affinities
snapshot_request = gp.BatchRequest()
snapshot_request.add(raw_cropped, output_shape_world)
snapshot_request.add(pred_affs, output_shape_world)
snapshot_request.add(gt_affs, output_shape_world)
if kwargs['input_format'] != "hdf" and kwargs['input_format'] != "zarr":
raise NotImplementedError("train node for %s not implemented yet",
kwargs['input_format'])
fls = []
for f in kwargs['data_files']:
fls.append(os.path.splitext(f)[0])
ln = len(fls)
print("first 5 files: ", fls[0:4])
if kwargs['input_format'] == "hdf":
sourceNode = gp.Hdf5Source
elif kwargs['input_format'] == "zarr":
sourceNode = gp.ZarrSource
neighborhood = []
psH = np.array(kwargs['patchshape']) // 2
for i in range(-psH[0], psH[0] + 1, kwargs['patchstride'][0]):
for j in range(-psH[1], psH[1] + 1, kwargs['patchstride'][1]):
for k in range(-psH[2], psH[2] + 1, kwargs['patchstride'][2]):
neighborhood.append([i, j, k])
datasets = {
raw: 'volumes/raw',
gt_labels: 'volumes/gt_labels',
anchor: 'volumes/gt_fgbg',
}
input_specs = {
raw:
gp.ArraySpec(roi=gp.Roi((0, ) * len(input_shape_world),
input_shape_world),
interpolatable=True,
dtype=np.float32),
gt_labels:
gp.ArraySpec(roi=gp.Roi((0, ) * len(output_shape_world),
output_shape_world),
interpolatable=False,
dtype=np.uint16),
anchor:
gp.ArraySpec(roi=gp.Roi((0, ) * len(output_shape_world),
output_shape_world),
interpolatable=False,
dtype=np.uint8),
gt_affs:
gp.ArraySpec(roi=gp.Roi((0, ) * len(output_shape_world),
output_shape_world),
interpolatable=False,
dtype=np.uint8)
}
inputs = {
net_names['raw']: raw,
net_names['gt_affs']: gt_affs,
net_names['anchor']: anchor,
}
outputs = {
net_names['pred_affs']: pred_affs,
net_names['raw_cropped']: raw_cropped,
}
snapshot = {
raw_cropped: 'volumes/raw_cropped',
gt_affs: '/volumes/gt_affs',
pred_affs: '/volumes/pred_affs',
}
optimizer_args = None
if kwargs['auto_mixed_precision']:
optimizer_args = (kwargs['optimizer'], {
'args': kwargs['args'],
'kwargs': kwargs['kwargs']
})
augmentation = kwargs['augmentation']
pipeline = (
tuple(
sourceNode(
fls[t] + "." + kwargs['input_format'],
datasets=datasets,
# array_specs=array_specs
)
+ gp.Pad(raw, None)
+ gp.Pad(gt_labels, None)
# chose a random location for each requested batch
+ gp.RandomLocation()
for t in range(ln)
) +
# chose a random source (i.e., sample) from the above
gp.RandomProvider() +
# elastically deform the batch
gp.ElasticAugment(
augmentation['elastic']['control_point_spacing'],
augmentation['elastic']['jitter_sigma'],
[augmentation['elastic']['rotation_min']*np.pi/180.0,
augmentation['elastic']['rotation_max']*np.pi/180.0],
subsample=4) +
# apply transpose and mirror augmentations
gp.SimpleAugment(mirror_only=augmentation['simple'].get("mirror"),
transpose_only=augmentation['simple'].get("transpose")) +
# scale and shift the intensity of the raw array
gp.IntensityAugment(
raw,
scale_min=augmentation['intensity']['scale'][0],
scale_max=augmentation['intensity']['scale'][1],
shift_min=augmentation['intensity']['shift'][0],
shift_max=augmentation['intensity']['shift'][1],
z_section_wise=False) +
# grow a boundary between labels
gp.GrowBoundary(
gt_labels,
steps=1,
only_xy=False) +
# convert labels into affinities between voxels
gp.AddAffinities(
neighborhood,
gt_labels,
gt_affs) +
# create a weight array that balances positive and negative samples in
# the affinity array
# gp.BalanceLabels(
# gt_affs,
# loss_weights_affs) +
# pre-cache batches from the point upstream
gp.PreCache(
cache_size=kwargs['cache_size'],
num_workers=kwargs['num_workers']) +
# pre-fetch batches from the point upstream
(gp.tensorflow.TFData() \
if kwargs.get('use_tf_data') else NoOp()) +
# perform one training iteration for each passing batch (here we use
# the tensor names earlier stored in train_net.config)
gp.tensorflow.Train(
os.path.join(kwargs['output_folder'], kwargs['name']),
optimizer=net_names['optimizer'],
summary=net_names['summaries'],
log_dir=kwargs['output_folder'],
loss=net_names['loss'],
inputs=inputs,
outputs=outputs,
array_specs=input_specs,
gradients={
net_names['pred_affs']: pred_affs_gradients,
},
auto_mixed_precision=kwargs['auto_mixed_precision'],
optimizer_args=optimizer_args,
use_tf_data=kwargs['use_tf_data'],
save_every=kwargs['checkpoints'],
snapshot_every=kwargs['snapshots']) +
# save the passing batch as an HDF5 file for inspection
gp.Snapshot(
snapshot,
output_dir=os.path.join(kwargs['output_folder'], 'snapshots'),
output_filename='batch_{iteration}.hdf',
every=kwargs['snapshots'],
additional_request=snapshot_request,
compression_type='gzip') +
# show a summary of time spend in each node every 10 iterations
gp.PrintProfilingStats(every=kwargs['profiling'])
)
#########
# TRAIN #
#########
print("Starting training...")
try:
with gp.build(pipeline):
print(pipeline)
for i in range(trained_until, kwargs['max_iteration']):
start = time.time()
pipeline.request_batch(request)
time_of_iteration = time.time() - start
logger.info("Batch: iteration=%d, time=%f", i,
time_of_iteration)
# exit()
except KeyboardInterrupt:
sys.exit()
print("Training finished")
