def test_ensure_center_non_zero(self):
path = Path(self.path_to("test_swc_source.swc"))
# write test swc
self._write_swc(path, self._toy_swc_points().to_nx_graph())
# read arrays
swc = PointsKey("SWC")
img = ArrayKey("IMG")
pipeline = (SwcFileSource(
path, [swc], [PointsSpec(roi=Roi((0, 0, 0), (11, 11, 11)))]) +
RandomLocation(ensure_nonempty=swc, ensure_centered=True) +
RasterizeSkeleton(
points=swc,
array=img,
array_spec=ArraySpec(
interpolatable=False,
dtype=np.uint32,
voxel_size=Coordinate((1, 1, 1)),
),
))
request = BatchRequest()
request.add(img, Coordinate((5, 5, 5)))
request.add(swc, Coordinate((5, 5, 5)))
with build(pipeline):
batch = pipeline.request_batch(request)
data = batch[img].data
g = batch[swc]
assert g.num_vertices() > 0
self.assertNotEqual(data[tuple(np.array(data.shape) // 2)], 0)
def test_without_placeholder(self):
test_labels = ArrayKey("TEST_LABELS")
test_points = GraphKey("TEST_POINTS")
pipeline = (
PointTestSource3D() + RandomLocation(ensure_nonempty=test_points) +
ElasticAugment([10, 10, 10], [0.1, 0.1, 0.1], [0, 2.0 * math.pi]) +
Snapshot(
{test_labels: "volumes/labels"},
output_dir=self.path_to(),
output_filename="elastic_augment_test{id}-{iteration}.hdf",
))
with build(pipeline):
for i in range(2):
request_size = Coordinate((40, 40, 40))
request_a = BatchRequest(random_seed=i)
request_a.add(test_points, request_size)
request_b = BatchRequest(random_seed=i)
request_b.add(test_points, request_size)
request_b.add(test_labels, request_size)
# No array to provide a voxel size to ElasticAugment
with pytest.raises(PipelineRequestError):
pipeline.request_batch(request_a)
batch_b = pipeline.request_batch(request_b)
self.assertIn(test_labels, batch_b)
def test_3d(self):
test_graph = GraphKey("TEST_GRAPH")
graph_spec = GraphSpec(roi=Roi((0, 0, 0), (5, 5, 5)))
test_array = ArrayKey("TEST_ARRAY")
array_spec = ArraySpec(
roi=Roi((0, 0, 0), (5, 5, 5)), voxel_size=Coordinate((1, 1, 1))
)
test_array2 = ArrayKey("TEST_ARRAY2")
array2_spec = ArraySpec(
roi=Roi((0, 0, 0), (5, 5, 5)), voxel_size=Coordinate((1, 1, 1))
)
snapshot_request = BatchRequest()
snapshot_request.add(test_graph, Coordinate((5, 5, 5)))
pipeline = ExampleSource(
[test_graph, test_array, test_array2], [graph_spec, array_spec, array2_spec]
) + Snapshot(
{
test_graph: "graphs/graph",
test_array: "volumes/array",
test_array2: "volumes/array2",
},
output_dir=str(self.test_dir),
every=2,
additional_request=snapshot_request,
output_filename="snapshot.hdf",
)
snapshot_file_path = Path(self.test_dir, "snapshot.hdf")
with build(pipeline):
request = BatchRequest()
roi = Roi((0, 0, 0), (5, 5, 5))
request[test_array] = ArraySpec(roi=roi)
request[test_array2] = ArraySpec(roi=roi)
pipeline.request_batch(request)
assert snapshot_file_path.exists()
f = h5py.File(snapshot_file_path)
assert f["volumes/array"] is not None
assert f["graphs/graph-ids"] is not None
snapshot_file_path.unlink()
pipeline.request_batch(request)
assert not snapshot_file_path.exists()
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 visualize_embedding_pipeline(fusion_pipeline, train_embedding):
setup_config = DEFAULT_CONFIG
setup_config["FUSION_PIPELINE"] = fusion_pipeline
setup_config["TRAIN_EMBEDDING"] = train_embedding
voxel_size = Coordinate(setup_config["VOXEL_SIZE"])
output_size = Coordinate(setup_config["OUTPUT_SHAPE"]) * voxel_size
input_size = Coordinate(setup_config["INPUT_SHAPE"]) * voxel_size
pipeline, raw, output = embedding_pipeline(setup_config,
get_test_data_sources)
request = BatchRequest()
request.add(raw, input_size)
request.add(output, output_size)
with build(pipeline):
pipeline.request_batch(request)
visualize_hdf5(Path("snapshots/snapshot_1.hdf"), tuple(voxel_size))
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 test_grow_labels_speed(self):
bb = Roi(Coordinate([0, 0, 0]), ([256, 256, 256]))
voxel_size = Coordinate([1, 1, 1])
swc_file = "test_swc.swc"
swc_path = Path(self.path_to(swc_file))
swc_points = self._get_points(np.array([1, 1, 1]), np.array([1, 1, 1]),
bb)
self._write_swc(swc_path, swc_points.graph)
# create swc sources
swc_key = PointsKey("SWC")
labels_key = ArrayKey("LABELS")
# add request
request = BatchRequest()
request.add(labels_key, bb.get_shape())
request.add(swc_key, bb.get_shape())
# data source for swc a
data_source = tuple()
data_source = (data_source + SwcFileSource(
swc_path, [swc_key], [PointsSpec(roi=bb)]) + RasterizeSkeleton(
points=swc_key,
array=labels_key,
array_spec=ArraySpec(interpolatable=False,
dtype=np.uint32,
voxel_size=voxel_size),
) + GrowLabels(array=labels_key, radius=3))
pipeline = data_source
num_repeats = 10
t1 = time.time()
with build(pipeline):
for i in range(num_repeats):
pipeline.request_batch(request)
t2 = time.time()
self.assertLess((t2 - t1) / num_repeats, 0.1)
def test_embedding_pipeline(
tmpdir, aux_task, blend_mode, fusion_pipeline, train_embedding, snapshot_every
):
setup_config = DEFAULT_CONFIG
setup_config["FUSION_PIPELINE"] = fusion_pipeline
setup_config["TRAIN_EMBEDDING"] = train_embedding
setup_config["SNAPSHOT_EVERY"] = snapshot_every
setup_config["TENSORBOARD_LOG_DIR"] = tmpdir
setup_config["SNAPSHOT_DIR"] = tmpdir
setup_config["SNAPSHOT_FILE_NAME"] = "test_snapshot"
setup_config["MATCHING_FAILURES_DIR"] = None
setup_config["BLEND_MODE"] = blend_mode
setup_config["AUX_TASK"] = aux_task
voxel_size = Coordinate(setup_config["VOXEL_SIZE"])
output_size = Coordinate(setup_config["OUTPUT_SHAPE"]) * voxel_size
input_size = Coordinate(setup_config["INPUT_SHAPE"]) * voxel_size
pipeline, raw, output, inputs = embedding_pipeline(
setup_config, get_test_data_sources
)
request = BatchRequest()
request.add(raw, input_size)
request.add(output, output_size)
for key in inputs:
request.add(key, output_size)
with build(pipeline):
batch = pipeline.request_batch(request)
assert output in batch
assert raw in batch
def test_merge_basics(self):
voxel_size = (1, 1, 1)
GraphKey("PRESYN")
ArrayKey("GT_LABELS")
graphsource = GraphTestSource(voxel_size)
arraysource = ArrayTestSoure(voxel_size)
pipeline = (graphsource, arraysource) + MergeProvider() + RandomLocation()
window_request = Coordinate((50, 50, 50))
with build(pipeline):
# Check basic merging.
request = BatchRequest()
request.add((GraphKeys.PRESYN), window_request)
request.add((ArrayKeys.GT_LABELS), window_request)
batch_res = pipeline.request_batch(request)
self.assertTrue(ArrayKeys.GT_LABELS in batch_res.arrays)
self.assertTrue(GraphKeys.PRESYN in batch_res.graphs)
# Check that request of only one source also works.
request = BatchRequest()
request.add((GraphKeys.PRESYN), window_request)
batch_res = pipeline.request_batch(request)
self.assertFalse(ArrayKeys.GT_LABELS in batch_res.arrays)
self.assertTrue(GraphKeys.PRESYN in batch_res.graphs)
# Check that it fails, when having two sources that provide the same type.
arraysource2 = ArrayTestSoure(voxel_size)
pipeline_fail = (arraysource, arraysource2) + MergeProvider() + RandomLocation()
with self.assertRaises(PipelineSetupError):
with build(pipeline_fail):
pass
def test_placeholder(self):
test_labels = ArrayKey("TEST_LABELS")
test_points = GraphKey("TEST_POINTS")
pipeline = (
PointTestSource3D() + RandomLocation(ensure_nonempty=test_points) +
ElasticAugment([10, 10, 10], [0.1, 0.1, 0.1], [0, 2.0 * math.pi]) +
Snapshot(
{test_labels: "volumes/labels"},
output_dir=self.path_to(),
output_filename="elastic_augment_test{id}-{iteration}.hdf",
))
with build(pipeline):
for i in range(2):
request_size = Coordinate((40, 40, 40))
request_a = BatchRequest(random_seed=i)
request_a.add(test_points, request_size)
request_a.add(test_labels, request_size, placeholder=True)
request_b = BatchRequest(random_seed=i)
request_b.add(test_points, request_size)
request_b.add(test_labels, request_size)
batch_a = pipeline.request_batch(request_a)
batch_b = pipeline.request_batch(request_b)
points_a = batch_a[test_points].nodes
points_b = batch_b[test_points].nodes
for a, b in zip(points_a, points_b):
assert all(np.isclose(a.location, b.location))
def visualize_foreground_pipeline(fusion_pipeline,
train_foreground,
distances,
test_sources=True):
setup_config = DEFAULT_CONFIG
setup_config["FUSION_PIPELINE"] = fusion_pipeline
setup_config["TRAIN_FOREGROUND"] = train_foreground
setup_config["DISTANCES"] = distances
voxel_size = Coordinate(setup_config["VOXEL_SIZE"])
output_size = Coordinate(setup_config["OUTPUT_SHAPE"]) * voxel_size
input_size = Coordinate(setup_config["INPUT_SHAPE"]) * voxel_size
if test_sources:
pipeline, raw, output = foreground_pipeline(setup_config,
get_test_data_sources)
else:
pipeline, raw, output = foreground_pipeline(setup_config)
request = BatchRequest()
request.add(raw, input_size)
request.add(output, output_size)
with build(pipeline):
pipeline.request_batch(request)
visualize_hdf5(Path("snapshots/snapshot_1.hdf"), tuple(voxel_size))
def test_foreground_pipeline(tmpdir, fusion_pipeline, train_foreground,
distance_loss, snapshot_every):
setup_config = DEFAULT_CONFIG
setup_config["FUSION_PIPELINE"] = fusion_pipeline
setup_config["TRAIN_FOREGROUND"] = train_foreground
setup_config["SNAPSHOT_EVERY"] = snapshot_every
setup_config["TENSORBOARD_LOG_DIR"] = tmpdir
setup_config["SNAPSHOT_DIR"] = tmpdir
setup_config["SNAPSHOT_FILE_NAME"] = "test_snapshot"
setup_config["MATCHING_FAILURES_DIR"] = None
voxel_size = Coordinate(setup_config["VOXEL_SIZE"])
output_size = Coordinate(setup_config["OUTPUT_SHAPE"]) * voxel_size
input_size = Coordinate(setup_config["INPUT_SHAPE"]) * voxel_size
pipeline, raw, output, inputs = foreground_pipeline(
setup_config, get_test_data_sources)
request = BatchRequest()
request.add(raw, input_size)
request.add(output, output_size)
for key in inputs:
request.add(key, output_size)
with build(pipeline):
batch = pipeline.request_batch(request)
assert output in batch
assert raw in batch
def test_output_min_distance(self):
voxel_size = Coordinate((20, 2, 2))
ArrayKey("GT_VECTORS_MAP_PRESYN")
GraphKey("PRESYN")
GraphKey("POSTSYN")
arraytypes_to_source_target_pointstypes = {
ArrayKeys.GT_VECTORS_MAP_PRESYN: (GraphKeys.PRESYN, GraphKeys.POSTSYN)
}
arraytypes_to_stayinside_arraytypes = {
ArrayKeys.GT_VECTORS_MAP_PRESYN: ArrayKeys.GT_LABELS
}
# test for partner criterion 'min_distance'
radius_phys = 30
pipeline_min_distance = AddVectorMapTestSource() + AddVectorMap(
src_and_trg_points=arraytypes_to_source_target_pointstypes,
voxel_sizes={ArrayKeys.GT_VECTORS_MAP_PRESYN: voxel_size},
radius_phys=radius_phys,
partner_criterion="min_distance",
stayinside_array_keys=arraytypes_to_stayinside_arraytypes,
pad_for_partners=(0, 0, 0),
)
with build(pipeline_min_distance):
request = BatchRequest()
raw_roi = pipeline_min_distance.spec[ArrayKeys.RAW].roi
gt_labels_roi = pipeline_min_distance.spec[ArrayKeys.GT_LABELS].roi
presyn_roi = pipeline_min_distance.spec[GraphKeys.PRESYN].roi
request.add(ArrayKeys.RAW, raw_roi.get_shape())
request.add(ArrayKeys.GT_LABELS, gt_labels_roi.get_shape())
request.add(GraphKeys.PRESYN, presyn_roi.get_shape())
request.add(GraphKeys.POSTSYN, presyn_roi.get_shape())
request.add(ArrayKeys.GT_VECTORS_MAP_PRESYN, presyn_roi.get_shape())
for identifier, spec in request.items():
spec.roi = spec.roi.shift((1000, 1000, 1000))
batch = pipeline_min_distance.request_batch(request)
presyn_locs = {n.id: n for n in batch.graphs[GraphKeys.PRESYN].nodes}
postsyn_locs = {n.id: n for n in batch.graphs[GraphKeys.POSTSYN].nodes}
vector_map_presyn = batch.arrays[ArrayKeys.GT_VECTORS_MAP_PRESYN].data
offset_vector_map_presyn = request[
ArrayKeys.GT_VECTORS_MAP_PRESYN
].roi.get_offset()
self.assertTrue(len(presyn_locs) > 0)
self.assertTrue(len(postsyn_locs) > 0)
for loc_id, point in presyn_locs.items():
if request[ArrayKeys.GT_VECTORS_MAP_PRESYN].roi.contains(
Coordinate(point.location)
):
self.assertTrue(
batch.arrays[ArrayKeys.GT_VECTORS_MAP_PRESYN].spec.roi.contains(
Coordinate(point.location)
)
)
dist_to_loc = {}
for partner_id in point.attrs["partner_ids"]:
if partner_id in postsyn_locs.keys():
partner_location = postsyn_locs[partner_id].location
dist_to_loc[
np.linalg.norm(partner_location - point.location)
] = partner_location
min_dist = np.min(list(dist_to_loc.keys()))
relevant_partner_loc = dist_to_loc[min_dist]
presyn_loc_shifted_vx = (
point.location - offset_vector_map_presyn
) // voxel_size
radius_vx = [(radius_phys // vx_dim) for vx_dim in voxel_size]
region_to_check = np.clip(
[
(presyn_loc_shifted_vx - radius_vx),
(presyn_loc_shifted_vx + radius_vx),
],
a_min=(0, 0, 0),
a_max=vector_map_presyn.shape[-3:],
)
for x, y, z in itertools.product(
range(int(region_to_check[0][0]), int(region_to_check[1][0])),
range(int(region_to_check[0][1]), int(region_to_check[1][1])),
range(int(region_to_check[0][2]), int(region_to_check[1][2])),
):
if (
np.linalg.norm(
(np.array((x, y, z)) - np.asarray(point.location))
)
< radius_phys
):
vector = [
vector_map_presyn[dim][x, y, z]
for dim in range(vector_map_presyn.shape[0])
]
if not np.sum(vector) == 0:
trg_loc_of_vector_phys = (
np.asarray(offset_vector_map_presyn)
+ (voxel_size * np.array([x, y, z]))
+ np.asarray(vector)
)
self.assertTrue(
np.array_equal(
trg_loc_of_vector_phys, relevant_partner_loc
)
)
# test for partner criterion 'all'
pipeline_all = AddVectorMapTestSource() + AddVectorMap(
src_and_trg_points=arraytypes_to_source_target_pointstypes,
voxel_sizes={ArrayKeys.GT_VECTORS_MAP_PRESYN: voxel_size},
radius_phys=radius_phys,
partner_criterion="all",
stayinside_array_keys=arraytypes_to_stayinside_arraytypes,
pad_for_partners=(0, 0, 0),
)
with build(pipeline_all):
batch = pipeline_all.request_batch(request)
presyn_locs = {n.id: n for n in batch.graphs[GraphKeys.PRESYN].nodes}
postsyn_locs = {n.id: n for n in batch.graphs[GraphKeys.POSTSYN].nodes}
vector_map_presyn = batch.arrays[ArrayKeys.GT_VECTORS_MAP_PRESYN].data
offset_vector_map_presyn = request[
ArrayKeys.GT_VECTORS_MAP_PRESYN
].roi.get_offset()
self.assertTrue(len(presyn_locs) > 0)
self.assertTrue(len(postsyn_locs) > 0)
for loc_id, point in presyn_locs.items():
if request[ArrayKeys.GT_VECTORS_MAP_PRESYN].roi.contains(
Coordinate(point.location)
):
self.assertTrue(
batch.arrays[ArrayKeys.GT_VECTORS_MAP_PRESYN].spec.roi.contains(
Coordinate(point.location)
)
)
partner_ids_to_locs_per_src, count_vectors_per_partner = {}, {}
for partner_id in point.attrs["partner_ids"]:
if partner_id in postsyn_locs.keys():
partner_ids_to_locs_per_src[partner_id] = postsyn_locs[
partner_id
].location.tolist()
count_vectors_per_partner[partner_id] = 0
presyn_loc_shifted_vx = (
point.location - offset_vector_map_presyn
) // voxel_size
radius_vx = [(radius_phys // vx_dim) for vx_dim in voxel_size]
region_to_check = np.clip(
[
(presyn_loc_shifted_vx - radius_vx),
(presyn_loc_shifted_vx + radius_vx),
],
a_min=(0, 0, 0),
a_max=vector_map_presyn.shape[-3:],
)
for x, y, z in itertools.product(
range(int(region_to_check[0][0]), int(region_to_check[1][0])),
range(int(region_to_check[0][1]), int(region_to_check[1][1])),
range(int(region_to_check[0][2]), int(region_to_check[1][2])),
):
if (
np.linalg.norm(
(np.array((x, y, z)) - np.asarray(point.location))
)
< radius_phys
):
vector = [
vector_map_presyn[dim][x, y, z]
for dim in range(vector_map_presyn.shape[0])
]
if not np.sum(vector) == 0:
trg_loc_of_vector_phys = (
np.asarray(offset_vector_map_presyn)
+ (voxel_size * np.array([x, y, z]))
+ np.asarray(vector)
)
self.assertTrue(
trg_loc_of_vector_phys.tolist()
in partner_ids_to_locs_per_src.values()
)
for (
partner_id,
partner_loc,
) in partner_ids_to_locs_per_src.items():
if np.array_equal(
np.asarray(trg_loc_of_vector_phys), partner_loc
):
count_vectors_per_partner[partner_id] += 1
self.assertTrue(
(
list(count_vectors_per_partner.values())
- np.min(list(count_vectors_per_partner.values()))
<= len(count_vectors_per_partner.keys())
).all()
)
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")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
setup_config = DEFAULT_CONFIG
setup_config["FUSION_PIPELINE"] = True
setup_config["TRAIN_EMBEDDING"] = True
setup_config["SNAPSHOT_EVERY"] = 0
setup_config["SNAPSHOT_FILE_NAME"] = None
setup_config["MATCHING_FAILURES_DIR"] = None
setup_config["PROFILE_EVERY"] = 1
setup_config["CLAHE"] = False
setup_config["NUM_FMAPS_EMBEDDING"] = 12
setup_config["FMAP_INC_FACTORS_EMBEDDING"] = 5
voxel_size = Coordinate(setup_config["VOXEL_SIZE"])
output_size = Coordinate(setup_config["OUTPUT_SHAPE"]) * voxel_size
input_size = Coordinate(setup_config["INPUT_SHAPE"]) * voxel_size
pipeline, raw, output, inputs = embedding_pipeline(setup_config,
get_test_data_sources)
request = BatchRequest()
request.add(raw, input_size)
request.add(output, output_size)
for key in inputs:
request.add(key, output_size)
with build(pipeline):
for i in range(1):
batch = pipeline.request_batch(request)
assert output in batch
assert raw in batch
masked_add_b: "volumes/masked_add_b",
softmask: "volumes/softmask",
softmask_b: "volumes/softmask_b",
mask_maxed: "volumes/mask_maxed",
mask_maxed_b: "volumes/mask_maxed_b",
},
every=1,
))
with build(pipeline):
for i in range(1):
request = BatchRequest(random_seed=i)
# add request
request = gp.BatchRequest()
request.add(raw_fused, input_size)
request.add(labels_fused, input_size)
request.add(raw_fused_b, input_size)
request.add(labels_fused_b, input_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)
# debugging
request.add(masked_base, input_size)
def test_output(self):
source = ScanTestSource()
chunk_request = BatchRequest()
chunk_request.add(ArrayKeys.RAW, (400, 30, 34))
chunk_request.add(ArrayKeys.GT_LABELS, (200, 10, 14))
chunk_request.add(GraphKeys.GT_GRAPH, (400, 30, 34))
pipeline = ScanTestSource() + Scan(chunk_request, num_workers=10)
with build(pipeline):
raw_spec = pipeline.spec[ArrayKeys.RAW]
labels_spec = pipeline.spec[ArrayKeys.GT_LABELS]
graph_spec = pipeline.spec[GraphKeys.GT_GRAPH]
full_request = BatchRequest({
ArrayKeys.RAW: raw_spec,
ArrayKeys.GT_LABELS: labels_spec,
GraphKeys.GT_GRAPH: graph_spec,
})
batch = pipeline.request_batch(full_request)
voxel_size = pipeline.spec[ArrayKeys.RAW].voxel_size
# assert that pixels encode their position
for (array_key, array) in batch.arrays.items():
# the z,y,x coordinates of the ROI
roi = array.spec.roi
meshgrids = np.meshgrid(range(roi.get_begin()[0] // voxel_size[0],
roi.get_end()[0] // voxel_size[0]),
range(roi.get_begin()[1] // voxel_size[1],
roi.get_end()[1] // voxel_size[1]),
range(roi.get_begin()[2] // voxel_size[2],
roi.get_end()[2] // voxel_size[2]),
indexing='ij')
data = meshgrids[0] + meshgrids[1] + meshgrids[2]
self.assertTrue((array.data == data).all())
for (graph_key, graph) in batch.graphs.items():
roi = graph.spec.roi
for i, j, k in itertools.product(range(20000, 22000, 100),
range(2000, 2200, 10),
range(2000, 2200, 10)):
assert all(
np.isclose(
graph.node(coordinate_to_id(i, j, k)).location,
np.array([i, j, k])))
assert (batch.arrays[ArrayKeys.RAW].spec.roi.get_offset() == (20000,
2000,
2000))
# test scanning with empty request
pipeline = ScanTestSource() + Scan(chunk_request, num_workers=1)
with build(pipeline):
batch = pipeline.request_batch(BatchRequest())
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 test_get_neuron_pair(self):
path = Path(self.path_to("test_swc_source.swc"))
# write test swc
self._write_swc(path, self._toy_swc_points().to_nx_graph())
# read arrays
swc_source = PointsKey("SWC_SOURCE")
ensure_nonempty = PointsKey("ENSURE_NONEMPTY")
labels_source = ArrayKey("LABELS_SOURCE")
img_source = ArrayKey("IMG_SOURCE")
img_swc = PointsKey("IMG_SWC")
label_swc = PointsKey("LABEL_SWC")
imgs = ArrayKey("IMGS")
labels = ArrayKey("LABELS")
points_a = PointsKey("SKELETON_A")
points_b = PointsKey("SKELETON_B")
img_a = ArrayKey("VOLUME_A")
img_b = ArrayKey("VOLUME_B")
labels_a = ArrayKey("LABELS_A")
labels_b = ArrayKey("LABELS_B")
data_shape = 5
output_shape = Coordinate((data_shape, data_shape, data_shape))
# Get points from test swc
swc_file_source = SwcFileSource(
path,
[swc_source, ensure_nonempty],
[
PointsSpec(roi=Roi((-10, -10, -10), (31, 31, 31))),
PointsSpec(roi=Roi((-10, -10, -10), (31, 31, 31))),
],
)
# Create an artificial image source by rasterizing the points
image_source = (SwcFileSource(
path, [img_swc],
[PointsSpec(roi=Roi((-10, -10, -10), (31, 31, 31)))]) +
RasterizeSkeleton(
points=img_swc,
array=img_source,
array_spec=ArraySpec(
interpolatable=True,
dtype=np.uint32,
voxel_size=Coordinate((1, 1, 1)),
),
) +
BinarizeLabels(labels=img_source, labels_binary=imgs) +
GrowLabels(array=imgs, radius=0))
# Create an artificial label source by rasterizing the points
label_source = (SwcFileSource(path, [label_swc], [
PointsSpec(roi=Roi((-10, -10, -10), (31, 31, 31)))
]) + RasterizeSkeleton(
points=label_swc,
array=labels_source,
array_spec=ArraySpec(
interpolatable=True,
dtype=np.uint32,
voxel_size=Coordinate((1, 1, 1)),
),
) + BinarizeLabels(labels=labels_source, labels_binary=labels) +
GrowLabels(array=labels, radius=1))
skeleton = tuple()
skeleton += ((swc_file_source, image_source, label_source) +
MergeProvider() +
RandomLocation(ensure_nonempty=ensure_nonempty,
ensure_centered=True))
pipeline = skeleton + GetNeuronPair(
point_source=swc_source,
nonempty_placeholder=ensure_nonempty,
array_source=imgs,
label_source=labels,
points=(points_a, points_b),
arrays=(img_a, img_b),
labels=(labels_a, labels_b),
seperate_by=(1, 3),
shift_attempts=100,
request_attempts=10,
output_shape=output_shape,
)
request = BatchRequest()
request.add(points_a, output_shape)
request.add(points_b, output_shape)
request.add(img_a, output_shape)
request.add(img_b, output_shape)
request.add(labels_a, output_shape)
request.add(labels_b, output_shape)
with build(pipeline):
for i in range(10):
batch = pipeline.request_batch(request)
assert all([
x in batch for x in
[points_a, points_b, img_a, img_b, labels_a, labels_b]
])
min_dist = 5
for a, b in itertools.product(
batch[points_a].nodes,
batch[points_b].nodes,
):
min_dist = min(
min_dist,
np.linalg.norm(a.location - b.location),
)
self.assertLessEqual(min_dist, 3)
self.assertGreaterEqual(min_dist, 1)
def test_recenter():
path = Path(self.path_to("test_swc_source.swc"))
# write test swc
self._write_swc(path, self._toy_swc_points())
# read arrays
swc_source = PointsKey("SWC_SOURCE")
labels_source = ArrayKey("LABELS_SOURCE")
img_source = ArrayKey("IMG_SOURCE")
img_swc = PointsKey("IMG_SWC")
label_swc = PointsKey("LABEL_SWC")
imgs = ArrayKey("IMGS")
labels = ArrayKey("LABELS")
points_a = PointsKey("SKELETON_A")
points_b = PointsKey("SKELETON_B")
img_a = ArrayKey("VOLUME_A")
img_b = ArrayKey("VOLUME_B")
labels_a = ArrayKey("LABELS_A")
labels_b = ArrayKey("LABELS_B")
# Get points from test swc
swc_file_source = SwcFileSource(
path, [swc_source], [PointsSpec(roi=Roi((-10, -10, -10), (31, 31, 31)))]
)
# Create an artificial image source by rasterizing the points
image_source = (
SwcFileSource(
path, [img_swc], [PointsSpec(roi=Roi((-10, -10, -10), (31, 31, 31)))]
)
+ RasterizeSkeleton(
points=img_swc,
array=img_source,
array_spec=ArraySpec(
interpolatable=True, dtype=np.uint32, voxel_size=Coordinate((1, 1, 1))
),
)
+ BinarizeLabels(labels=img_source, labels_binary=imgs)
+ GrowLabels(array=imgs, radius=0)
)
# Create an artificial label source by rasterizing the points
label_source = (
SwcFileSource(
path, [label_swc], [PointsSpec(roi=Roi((-10, -10, -10), (31, 31, 31)))]
)
+ RasterizeSkeleton(
points=label_swc,
array=labels_source,
array_spec=ArraySpec(
interpolatable=True, dtype=np.uint32, voxel_size=Coordinate((1, 1, 1))
),
)
+ BinarizeLabels(labels=labels_source, labels_binary=labels)
+ GrowLabels(array=labels, radius=1)
)
skeleton = tuple()
skeleton += (
(swc_file_source, image_source, label_source)
+ MergeProvider()
+ RandomLocation(ensure_nonempty=swc_source, ensure_centered=True)
)
pipeline = (
skeleton
+ GetNeuronPair(
point_source=swc_source,
array_source=imgs,
label_source=labels,
points=(points_a, points_b),
arrays=(img_a, img_b),
labels=(labels_a, labels_b),
seperate_by=4,
shift_attempts=100,
)
+ Recenter(points_a, img_a, max_offset=4)
+ Recenter(points_b, img_b, max_offset=4)
)
request = BatchRequest()
data_shape = 9
request.add(points_a, Coordinate((data_shape, data_shape, data_shape)))
request.add(points_b, Coordinate((data_shape, data_shape, data_shape)))
request.add(img_a, Coordinate((data_shape, data_shape, data_shape)))
request.add(img_b, Coordinate((data_shape, data_shape, data_shape)))
request.add(labels_a, Coordinate((data_shape, data_shape, data_shape)))
request.add(labels_b, Coordinate((data_shape, data_shape, data_shape)))
with build(pipeline):
batch = pipeline.request_batch(request)
data_a = batch[img_a].data
assert data_a[tuple(np.array(data_a.shape) // 2)] == 1
data_a = np.pad(data_a, (1,), "constant", constant_values=(0,))
data_b = batch[img_b].data
assert data_b[tuple(np.array(data_b.shape) // 2)] == 1
data_b = np.pad(data_b, (1,), "constant", constant_values=(0,))
data_c = data_a + data_b
data = np.array((data_a, data_b, data_c))
for _, point in batch[points_a].data.items():
assert (
data[(0,) + tuple(int(x) + 1 for x in point.location)] == 1
), "data at {} is not 1, its {}".format(
point.location, data[(0,) + tuple(int(x) for x in point.location)]
)
for _, point in batch[points_b].data.items():
assert (
data[(1,) + tuple(int(x) + 1 for x in point.location)] == 1
), "data at {} is not 1".format(point.location)
def test_rasterize_speed(self):
# This is worryingly slow for such a small volume (256**3) and only 2
# straight lines for skeletons.
LABEL_RADIUS = 3
bb = Roi(Coordinate([0, 0, 0]), ([256, 256, 256]))
voxel_size = Coordinate([1, 1, 1])
swc_files = ("test_line_a.swc", "test_line_b.swc")
swc_paths = tuple(Path(self.path_to(file_name)) for file_name in swc_files)
# create two lines seperated by a given distance and write them to swc files
intercepts, slopes = self._get_line_pair(roi=bb, dist=3 * LABEL_RADIUS)
for intercept, slope, swc_path in zip(intercepts, slopes, swc_paths):
swc_points = self._get_points(intercept, slope, bb)
self._write_swc(swc_path, swc_points)
# create swc sources
swc_key_names = ("SWC_A", "SWC_B")
labels_key_names = ("LABELS_A", "LABELS_B")
swc_keys = tuple(PointsKey(name) for name in swc_key_names)
labels_keys = tuple(ArrayKey(name) for name in labels_key_names)
# add request
request = BatchRequest()
request.add(labels_keys[0], bb.get_shape())
request.add(labels_keys[1], bb.get_shape())
request.add(swc_keys[0], bb.get_shape())
request.add(swc_keys[1], bb.get_shape())
# data source for swc a
data_sources_a = tuple()
data_sources_a = (
data_sources_a
+ SwcFileSource(swc_paths[0], swc_keys[0], PointsSpec(roi=bb))
+ RasterizeSkeleton(
points=swc_keys[0],
array=labels_keys[0],
array_spec=ArraySpec(
interpolatable=False, dtype=np.uint32, voxel_size=voxel_size
),
radius=LABEL_RADIUS,
)
)
# data source for swc b
data_sources_b = tuple()
data_sources_b = (
data_sources_b
+ SwcFileSource(swc_paths[1], swc_keys[1], PointsSpec(roi=bb))
+ RasterizeSkeleton(
points=swc_keys[1],
array=labels_keys[1],
array_spec=ArraySpec(
interpolatable=False, dtype=np.uint32, voxel_size=voxel_size
),
radius=LABEL_RADIUS,
)
)
data_sources = tuple([data_sources_a, data_sources_b]) + MergeProvider()
pipeline = data_sources
t1 = time.time()
with build(pipeline):
batch = pipeline.request_batch(request)
a_data = batch[labels_keys[0]].data
a_data = np.pad(a_data, (1,), "constant", constant_values=(0,))
b_data = batch[labels_keys[1]].data
b_data = np.pad(b_data, (1,), "constant", constant_values=(0,))
t2 = time.time()
self.assertLess(t2 - t1, 0.1)
def test_two_disjoint_lines_softmask(self):
LABEL_RADIUS = 3
RAW_RADIUS = 3
# exagerated to show problem
BLEND_SMOOTHNESS = 10
bb = Roi(Coordinate([0, 0, 0]), ([256, 256, 256]))
voxel_size = Coordinate([1, 1, 1])
swc_files = ("test_line_a.swc", "test_line_b.swc")
swc_paths = tuple(
Path(self.path_to(file_name)) for file_name in swc_files)
# create two lines seperated by a given distance and write them to swc files
intercepts, slopes = self._get_line_pair(roi=bb, dist=3 * LABEL_RADIUS)
for intercept, slope, swc_path in zip(intercepts, slopes, swc_paths):
swc_points = self._get_points(intercept, slope, bb)
self._write_swc(swc_path, swc_points)
# create swc sources
fused = ArrayKey("FUSED")
fused_labels = ArrayKey("FUSED_LABELS")
swc_key_names = ("SWC_A", "SWC_B")
labels_key_names = ("LABELS_A", "LABELS_B")
raw_key_names = ("RAW_A", "RAW_B")
swc_keys = tuple(PointsKey(name) for name in swc_key_names)
labels_keys = tuple(ArrayKey(name) for name in labels_key_names)
raw_keys = tuple(ArrayKey(name) for name in raw_key_names)
# add request
request = BatchRequest()
request.add(fused, bb.get_shape())
request.add(fused_labels, bb.get_shape())
request.add(labels_keys[0], bb.get_shape())
request.add(labels_keys[1], bb.get_shape())
request.add(raw_keys[0], bb.get_shape())
request.add(raw_keys[1], bb.get_shape())
request.add(swc_keys[0], bb.get_shape())
request.add(swc_keys[1], bb.get_shape())
# data source for swc a
data_sources_a = tuple()
data_sources_a = (data_sources_a + SwcFileSource(
swc_paths[0], swc_keys[0], PointsSpec(roi=bb)) + RasterizeSkeleton(
points=swc_keys[0],
array=labels_keys[0],
array_spec=ArraySpec(interpolatable=False,
dtype=np.uint32,
voxel_size=voxel_size),
radius=LABEL_RADIUS,
) + RasterizeSkeleton(
points=swc_keys[0],
array=raw_keys[0],
array_spec=ArraySpec(interpolatable=False,
dtype=np.uint32,
voxel_size=voxel_size),
radius=RAW_RADIUS,
))
# data source for swc b
data_sources_b = tuple()
data_sources_b = (data_sources_b + SwcFileSource(
swc_paths[1], swc_keys[1], PointsSpec(roi=bb)) + RasterizeSkeleton(
points=swc_keys[1],
array=labels_keys[1],
array_spec=ArraySpec(interpolatable=False,
dtype=np.uint32,
voxel_size=voxel_size),
radius=LABEL_RADIUS,
) + RasterizeSkeleton(
points=swc_keys[1],
array=raw_keys[1],
array_spec=ArraySpec(interpolatable=False,
dtype=np.uint32,
voxel_size=voxel_size),
radius=RAW_RADIUS,
))
data_sources = tuple([data_sources_a, data_sources_b
]) + MergeProvider()
pipeline = data_sources + FusionAugment(
raw_keys[0],
raw_keys[1],
labels_keys[0],
labels_keys[1],
fused,
fused_labels,
blend_mode="labels_mask",
blend_smoothness=BLEND_SMOOTHNESS,
num_blended_objects=0,
)
with build(pipeline):
batch = pipeline.request_batch(request)
fused_data = batch[fused].data
fused_data = np.pad(fused_data, (1, ),
"constant",
constant_values=(0, ))
a_data = batch[raw_keys[0]].data
a_data = np.pad(a_data, (1, ), "constant", constant_values=(0, ))
b_data = batch[raw_keys[1]].data
b_data = np.pad(b_data, (1, ), "constant", constant_values=(0, ))
all_data = np.zeros((5, ) + fused_data.shape)
all_data[0, :, :, :] = fused_data
all_data[1, :, :, :] = a_data + b_data
all_data[2, :, :, :] = fused_data - a_data - b_data
all_data[3, :, :, :] = a_data
all_data[4, :, :, :] = b_data
# Uncomment to visualize problem
if imported_volshow:
volshow(all_data)
# input("Press enter when you are done viewing the data: ")
diff = np.linalg.norm(fused_data - a_data - b_data)
self.assertAlmostEqual(diff, 0)
def test_two_disjoint_lines_intensity(self):
LABEL_RADIUS = 3
RAW_RADIUS = 3
BLEND_SMOOTHNESS = 3
bb = Roi(Coordinate([0, 0, 0]), ([256, 256, 256]))
voxel_size = Coordinate([1, 1, 1])
swc_files = ("test_line_a.swc", "test_line_b.swc")
swc_paths = tuple(Path(self.path_to(file_name)) for file_name in swc_files)
# create two lines seperated by a given distance and write them to swc files
intercepts, slopes = self._get_line_pair(roi=bb, dist=3 * LABEL_RADIUS)
for intercept, slope, swc_path in zip(intercepts, slopes, swc_paths):
swc_points = self._get_points(intercept, slope, bb)
self._write_swc(swc_path, swc_points.to_nx_graph())
# create swc sources
fused = ArrayKey("FUSED")
fused_labels = ArrayKey("FUSED_LABELS")
fused_swc = PointsKey("FUSED_SWC")
swc_key_names = ("SWC_A", "SWC_B")
labels_key_names = ("LABELS_A", "LABELS_B")
raw_key_names = ("RAW_A", "RAW_B")
swc_keys = tuple(PointsKey(name) for name in swc_key_names)
labels_keys = tuple(ArrayKey(name) for name in labels_key_names)
raw_keys = tuple(ArrayKey(name) for name in raw_key_names)
# add request
request = BatchRequest()
request.add(fused, bb.get_shape())
request.add(fused_labels, bb.get_shape())
request.add(fused_swc, bb.get_shape())
request.add(labels_keys[0], bb.get_shape())
request.add(labels_keys[1], bb.get_shape())
request.add(raw_keys[0], bb.get_shape())
request.add(raw_keys[1], bb.get_shape())
request.add(swc_keys[0], bb.get_shape())
request.add(swc_keys[1], bb.get_shape())
# data source for swc a
data_sources_a = tuple()
data_sources_a = (
data_sources_a
+ SwcFileSource(swc_paths[0], [swc_keys[0]], [PointsSpec(roi=bb)])
+ RasterizeSkeleton(
points=swc_keys[0],
array=labels_keys[0],
array_spec=ArraySpec(
interpolatable=False, dtype=np.uint16, voxel_size=voxel_size
),
)
+ GrowLabels(array=labels_keys[0], radii=[LABEL_RADIUS])
+ RasterizeSkeleton(
points=swc_keys[0],
array=raw_keys[0],
array_spec=ArraySpec(
interpolatable=False, dtype=np.uint16, voxel_size=voxel_size
),
)
+ GrowLabels(array=raw_keys[0], radii=[RAW_RADIUS])
+ Normalize(raw_keys[0])
)
# data source for swc b
data_sources_b = tuple()
data_sources_b = (
data_sources_b
+ SwcFileSource(swc_paths[1], [swc_keys[1]], [PointsSpec(roi=bb)])
+ RasterizeSkeleton(
points=swc_keys[1],
array=labels_keys[1],
array_spec=ArraySpec(
interpolatable=False, dtype=np.uint16, voxel_size=voxel_size
),
)
+ GrowLabels(array=labels_keys[1], radii=[LABEL_RADIUS])
+ RasterizeSkeleton(
points=swc_keys[1],
array=raw_keys[1],
array_spec=ArraySpec(
interpolatable=False, dtype=np.uint16, voxel_size=voxel_size
),
)
+ GrowLabels(array=raw_keys[1], radii=[RAW_RADIUS])
+ Normalize(raw_keys[1])
)
data_sources = tuple([data_sources_a, data_sources_b]) + MergeProvider()
pipeline = data_sources + FusionAugment(
raw_keys[0],
raw_keys[1],
labels_keys[0],
labels_keys[1],
swc_keys[0],
swc_keys[1],
fused,
fused_labels,
fused_swc,
blend_mode="intensity",
blend_smoothness=BLEND_SMOOTHNESS,
num_blended_objects=0,
)
with build(pipeline):
batch = pipeline.request_batch(request)
fused_data = batch[fused].data
fused_data = np.pad(fused_data, (1,), "constant", constant_values=(0,))
a_data = batch[raw_keys[0]].data
a_data = np.pad(a_data, (1,), "constant", constant_values=(0,))
b_data = batch[raw_keys[1]].data
b_data = np.pad(b_data, (1,), "constant", constant_values=(0,))
diff = np.linalg.norm(fused_data - a_data - b_data)
self.assertAlmostEqual(diff, 0)
labels_fused: "volumes/labels_fused",
labels_base: "volumes/labels_base",
labels_add: "volumes/labels_add",
raw_fused_b: "volumes/raw_fused_b",
labels_fused_b: "volumes/labels_fused_b",
},
every=1,
))
with build(pipeline):
for i in range(1):
request = BatchRequest(random_seed=i)
# 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(raw_fused_b, input_size)
request.add(labels_fused_b, input_size)
request.add(swc_fused_b, input_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)