def test_square(self):
test_graph = GraphKey("TEST_GRAPH")
test_array1 = ArrayKey("TEST_ARRAY1")
test_array2 = ArrayKey("TEST_ARRAY2")
pipeline = ((ArrayTestSource(), TestSource()) + MergeProvider() +
Pad(test_array1, None) + Pad(test_array2, None) + Pad(test_graph, None)
+ SimpleAugment(
mirror_only=[1,2], transpose_only=[1,2]
))
request = BatchRequest()
request[GraphKeys.TEST_GRAPH] = GraphSpec(roi=Roi((0, 50, 65), (100, 100, 100)))
request[ArrayKeys.TEST_ARRAY1] = ArraySpec(roi=Roi((0, 0, 15), (100, 200, 200)))
request[ArrayKeys.TEST_ARRAY2] = ArraySpec(roi=Roi((0, 50, 65), (100, 100, 100)))
with build(pipeline):
for i in range(100):
batch = pipeline.request_batch(request)
assert len(list(batch[GraphKeys.TEST_GRAPH].nodes)) == 1
for (array_key, array) in batch.arrays.items():
assert batch.arrays[array_key].data.shape == batch.arrays[array_key].spec.roi.get_shape()
def test_output(self):
a = ArrayKey("A")
b = ArrayKey("B")
source_a = TestSourceRandomLocation(a)
source_b = TestSourceRandomLocation(b)
pipeline = (source_a, source_b) + \
MergeProvider() + CustomRandomLocation()
with build(pipeline):
for i in range(10):
batch = pipeline.request_batch(
BatchRequest({
a: ArraySpec(roi=Roi((0, 0, 0), (20, 20, 20))),
b: ArraySpec(roi=Roi((0, 0, 0), (20, 20, 20)))
}))
self.assertTrue(np.sum(batch.arrays[a].data) > 0)
self.assertTrue(np.sum(batch.arrays[b].data) > 0)
# Request a ROI with the same shape as the entire ROI
full_roi_a = Roi((0, 0, 0), source_a.roi.get_shape())
full_roi_b = Roi((0, 0, 0), source_b.roi.get_shape())
batch = pipeline.request_batch(
BatchRequest({
a: ArraySpec(roi=full_roi_a),
b: ArraySpec(roi=full_roi_b)
}))
def test_precache(self):
logging.getLogger("gunpowder.torch.nodes.predict").setLevel(
logging.INFO)
a = ArrayKey("A")
pred = ArrayKey("PRED")
model = ExampleModel()
reference_request = BatchRequest()
reference_request[a] = ArraySpec(roi=Roi((0, 0), (7, 7)))
reference_request[pred] = ArraySpec(roi=Roi((1, 1), (5, 5)))
source = ExampleTorchTrain2DSource()
predict = Predict(
model=model,
inputs={"a": a},
outputs={0: pred},
array_specs={pred: ArraySpec()},
)
pipeline = source + predict + PreCache(cache_size=3, num_workers=2)
request = BatchRequest({
a: ArraySpec(roi=Roi((0, 0), (17, 17))),
pred: ArraySpec(roi=Roi((0, 0), (15, 15))),
})
# train for a couple of iterations
with build(pipeline):
batch = pipeline.request_batch(request)
assert pred in batch
def test_6_neighborhood():
# array keys
graph = GraphKey("GRAPH")
neighborhood = ArrayKey("NEIGHBORHOOD")
neighborhood_mask = ArrayKey("NEIGHBORHOOD_MASK")
distance = 1
pipeline = TestSource(graph) + Neighborhood(
graph,
neighborhood,
neighborhood_mask,
distance,
array_specs={
neighborhood: ArraySpec(voxel_size=Coordinate((1, 1, 1))),
neighborhood_mask: ArraySpec(voxel_size=Coordinate((1, 1, 1))),
},
k=6,
)
request = BatchRequest()
request[neighborhood] = ArraySpec(roi=Roi((0, 0, 0), (10, 10, 10)))
request[neighborhood_mask] = ArraySpec(roi=Roi((0, 0, 0), (10, 10, 10)))
with build(pipeline):
batch = pipeline.request_batch(request)
n_data = batch[neighborhood].data
n_mask = batch[neighborhood_mask].data
masked_ind = list(
set([(0, i, 0) for i in range(10) if i not in [0, 4]] +
[(i, 5, 0)
for i in range(10)] + [(i, 4, 0)
for i in range(10) if i not in [0]]))
assert all(n_mask[tuple(zip(*masked_ind))]
), f"expected {masked_ind} but saw {np.where(n_mask==1)}"
def test_3d_basics(self):
test_labels = ArrayKey("TEST_LABELS")
test_points = GraphKey("TEST_POINTS")
test_raster = ArrayKey("TEST_RASTER")
pipeline = (
PointTestSource3D() + ElasticAugment(
[10, 10, 10],
[0.1, 0.1, 0.1],
# [0, 0, 0], # no jitter
[0, 2.0 * math.pi],
) + RasterizeGraph(
test_points,
test_raster,
settings=RasterizationSettings(radius=2, mode="peak"),
) + Snapshot(
{
test_labels: "volumes/labels",
test_raster: "volumes/raster"
},
dataset_dtypes={test_raster: np.float32},
output_dir=self.path_to(),
output_filename="elastic_augment_test{id}-{iteration}.hdf",
))
for _ in range(5):
with build(pipeline):
request_roi = Roi((-20, -20, -20), (40, 40, 40))
request = BatchRequest()
request[test_labels] = ArraySpec(roi=request_roi)
request[test_points] = GraphSpec(roi=request_roi)
request[test_raster] = ArraySpec(roi=request_roi)
batch = pipeline.request_batch(request)
labels = batch[test_labels]
points = batch[test_points]
# the point at (0, 0, 0) should not have moved
self.assertTrue(points.contains(0))
labels_data_roi = (
labels.spec.roi -
labels.spec.roi.get_begin()) / labels.spec.voxel_size
# points should have moved together with the voxels
for point in points.nodes:
loc = point.location - labels.spec.roi.get_begin()
loc = loc / labels.spec.voxel_size
loc = Coordinate(int(round(x)) for x in loc)
if labels_data_roi.contains(loc):
self.assertEqual(labels.data[loc], point.id)
def test_output(self):
meta_base = self.path_to('tf_graph')
ArrayKey('A')
ArrayKey('B')
ArrayKey('C')
ArrayKey('GRADIENT_A')
# create model meta graph file and get input/output names
(a, b, c, optimizer, loss) = self.create_meta_graph(meta_base)
source = ExampleTensorflowTrainSource()
train = Train(
meta_base,
optimizer=optimizer,
loss=loss,
inputs={a: ArrayKeys.A, b: ArrayKeys.B},
outputs={c: ArrayKeys.C},
gradients={a: ArrayKeys.GRADIENT_A},
save_every=100)
pipeline = source + train
request = BatchRequest({
ArrayKeys.A: ArraySpec(roi=Roi((0, 0), (2, 2))),
ArrayKeys.B: ArraySpec(roi=Roi((0, 0), (2, 2))),
ArrayKeys.C: ArraySpec(roi=Roi((0, 0), (2, 2))),
ArrayKeys.GRADIENT_A: ArraySpec(roi=Roi((0, 0), (2, 2))),
})
# train for a couple of iterations
with build(pipeline):
batch = pipeline.request_batch(request)
self.assertAlmostEqual(batch.loss, 9.8994951)
gradient_a = batch.arrays[ArrayKeys.GRADIENT_A].data
self.assertTrue(gradient_a[0, 0] < gradient_a[0, 1])
self.assertTrue(gradient_a[0, 1] < gradient_a[1, 0])
self.assertTrue(gradient_a[1, 0] < gradient_a[1, 1])
for i in range(200-1):
loss1 = batch.loss
batch = pipeline.request_batch(request)
loss2 = batch.loss
self.assertLess(loss2, loss1)
# resume training
with build(pipeline):
for i in range(100):
loss1 = batch.loss
batch = pipeline.request_batch(request)
loss2 = batch.loss
self.assertLess(loss2, loss1)
def get_test_data_sources(setup_config):
input_shape = Coordinate(setup_config["INPUT_SHAPE"])
voxel_size = Coordinate(setup_config["VOXEL_SIZE"])
input_size = input_shape * voxel_size
micron_scale = voxel_size[0]
# New array keys
# Note: These are intended to be requested with size input_size
raw = ArrayKey("RAW")
matched = GraphKey("MATCHED")
nonempty_placeholder = GraphKey("NONEMPTY")
labels = ArrayKey("LABELS")
ensure_nonempty = matched
data_sources = ((
TestImageSource(
array=raw,
array_specs={
raw:
ArraySpec(interpolatable=True,
voxel_size=voxel_size,
dtype=np.uint16)
},
size=input_size * 3,
voxel_size=voxel_size,
),
TestPointSource(
points=[matched, nonempty_placeholder],
directed=False,
size=input_size * 3,
num_points=333,
),
) + MergeProvider() + RandomLocation(
ensure_nonempty=ensure_nonempty,
ensure_centered=True,
point_balance_radius=10 * micron_scale,
) + RasterizeSkeleton(
points=matched,
array=labels,
array_spec=ArraySpec(
interpolatable=False, voxel_size=voxel_size, dtype=np.uint64),
) + Normalize(raw))
return (
data_sources,
raw,
labels,
nonempty_placeholder,
matched,
)
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_multi_transpose(self):
test_graph = GraphKey("TEST_GRAPH")
test_array1 = ArrayKey("TEST_ARRAY1")
test_array2 = ArrayKey("TEST_ARRAY2")
point = np.array([50, 70, 100])
transpose_dims = [0, 1, 2]
pipeline = (ArrayTestSource(),
ExampleSource()) + MergeProvider() + SimpleAugment(
mirror_only=[], transpose_only=transpose_dims)
request = BatchRequest()
offset = (0, 20, 33)
request[GraphKeys.TEST_GRAPH] = GraphSpec(
roi=Roi(offset, (100, 100, 120)))
request[ArrayKeys.TEST_ARRAY1] = ArraySpec(
roi=Roi((0, 0, 0), (100, 200, 300)))
request[ArrayKeys.TEST_ARRAY2] = ArraySpec(
roi=Roi((0, 100, 250), (100, 100, 50)))
# Create all possible permurations of our transpose dims
transpose_combinations = list(permutations(transpose_dims, 3))
possible_loc = np.zeros((len(transpose_combinations), 3))
# Transpose points in all possible ways
for i, comb in enumerate(transpose_combinations):
possible_loc[i] = point[np.array(comb)]
with build(pipeline):
seen_transposed = False
seen_node = True
for i in range(100):
batch = pipeline.request_batch(request)
if len(list(batch[GraphKeys.TEST_GRAPH].nodes)) == 1:
seen_node = True
node = list(batch[GraphKeys.TEST_GRAPH].nodes)[0]
assert node.location in possible_loc
seen_transposed = seen_transposed or any(
[node.location[dim] != point[dim] for dim in range(3)])
assert Roi((0, 20, 33), (100, 100, 120)).contains(
batch[GraphKeys.TEST_GRAPH].spec.roi)
assert batch[GraphKeys.TEST_GRAPH].spec.roi.contains(
node.location)
for (array_key, array) in batch.arrays.items():
assert batch.arrays[array_key].data.shape == batch.arrays[
array_key].spec.roi.get_shape()
assert seen_transposed
assert seen_node
def process_function():
scheduler = Client()
worker_id = scheduler.context.worker_id
num_workers = scheduler.context.num_workers
gpu = actor_id_to_gpu_mapping(worker_id)
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
_logger.info("Worker %d uses gpu %d with %d workers", worker_id, gpu,
num_workers)
_logger.info("Environment:")
for name in os.environ.keys():
_logger.info(' %s:%s', name, os.environ[name])
# _logger.info("GPU is available: %s", tf.test.is_gpu_available())
# from tensorflow.python.client import device_lib
# def get_available_gpus():
# local_device_protos = device_lib.list_local_devices()
# return [x.name for x in local_device_protos if x.device_type == 'GPU']
# available_gpus = get_available_gpus()
# for gpu in available_gpus:
# print("Worker %d sees gpus %s" % (actor_id, available_gpus))
import tensorflow as tf
with tf.device('/gpu:%d' % 0):
from gunpowder import ArrayKey, ArraySpec, build, BatchRequest, DaisyRequestBlocks
_RAW = ArrayKey('RAW')
roi_map = {
ArrayKey('OUTPUT_%d' % i): 'write_roi'
for i in range(len(outputs))
}
roi_map[_RAW] = 'read_roi'
reference = BatchRequest()
reference[_RAW] = ArraySpec(roi=None, voxel_size=input_voxel_size)
for i in range(len(outputs)):
reference[ArrayKey('OUTPUT_%d' % i)] = ArraySpec(
roi=None, voxel_size=output_voxel_size)
pipeline = pipeline_factory()
pipeline += DaisyRequestBlocks(reference=reference,
roi_map=roi_map,
num_workers=num_cpu_workers)
with build(pipeline):
pipeline.request_batch(BatchRequest())
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 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_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_with_edge(self):
graph_with_edge = GraphKey("TEST_GRAPH_WITH_EDGE")
array_with_edge = ArrayKey("RASTERIZED_EDGE")
pipeline = GraphTestSourceWithEdge() + RasterizeGraph(
GraphKeys.TEST_GRAPH_WITH_EDGE,
ArrayKeys.RASTERIZED_EDGE,
ArraySpec(voxel_size=(1, 1, 1)),
settings=RasterizationSettings(0.5),
)
with build(pipeline):
request = BatchRequest()
roi = Roi((0, 0, 0), (10, 10, 10))
request[GraphKeys.TEST_GRAPH_WITH_EDGE] = GraphSpec(roi=roi)
request[ArrayKeys.RASTERIZED_EDGE] = ArraySpec(roi=roi)
batch = pipeline.request_batch(request)
rasterized = batch.arrays[ArrayKeys.RASTERIZED_EDGE].data
assert (
rasterized.sum() == 10
), f"rasterized has ones at: {np.where(rasterized==1)}"
def test_get_total_roi_nonspatial_array():
raw = ArrayKey('RAW')
nonspatial = ArrayKey('NONSPATIAL')
voxel_size = Coordinate((1, 2))
roi = Roi((100, 200), (20, 20))
raw_spec = ArraySpec(roi=roi, voxel_size=voxel_size)
nonspatial_spec = ArraySpec(nonspatial=True)
batch = Batch()
batch[raw] = Array(data=np.zeros((20, 10)), spec=raw_spec)
batch[nonspatial] = Array(data=np.zeros((2, 3)), spec=nonspatial_spec)
assert batch.get_total_roi() == roi
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 __init__(
self,
labelname,
labelid,
targetid=None,
thr=128,
scale_loss=True,
scale_key=None,
data_dir="/groups/saalfeld/saalfeldlab/larissa/data/cell/{0:}.n5",
data_sources=None,
):
self.labelname = labelname
if not isinstance(labelid, collections.Iterable):
labelid = (labelid, )
self.labelid = labelid
self.targetid = targetid
self.thr = thr
self.gt_dist_key = ArrayKey("GT_DIST_" + self.labelname.upper())
self.pred_dist_key = ArrayKey("PRED_DIST_" + self.labelname.upper())
self.mask_key = ArrayKey("MASK_" + self.labelname.upper())
self.scale_loss = scale_loss
self.data_dir = data_dir
self.data_sources = data_sources
self.total_voxels = compute_total_voxels(self.data_dir,
self.data_sources)
num = 0
if data_sources is not None:
for ds in data_sources:
zf = z5py.File(ds.full_path, use_zarr_format=False)
for l in labelid:
if l in zf["volumes/labels/all"].attrs["relabeled_ids"]:
num += zf["volumes/labels/all"].attrs[
"relabeled_counts"][zf["volumes/labels/all"].attrs[
"relabeled_ids"].index(l)]
if num > 0:
self.class_weight = float(self.total_voxels) / num
else:
self.class_weight = 0.0
print(labelname, self.class_weight)
if self.scale_loss:
self.scale_key = ArrayKey("SCALE_" + self.labelname.upper())
if scale_key is not None:
self.scale_key = scale_key
if not self.scale_loss and scale_key is None:
self.scale_key = self.mask_key
def factory():
from gunpowder import ArrayKey, Pad, Normalize, IntensityScaleShift, ArraySpec
from gunpowder.tensorflow import Predict
from gunpowder.nodes.hdf5like_source_base import Hdf5LikeSource
from gunpowder.nodes.hdf5like_write_base import Hdf5LikeWrite
from gunpowder.coordinate import Coordinate
from gunpowder.compat import ensure_str
from fuse import Z5Source, Z5Write
_RAW = ArrayKey('RAW')
output_dataset_names = {
ArrayKey('OUTPUT_%d' % i): ds
for i, (ds, _) in enumerate(outputs)
}
output_tensor_to_key = {
tensor: ArrayKey('OUTPUT_%d' % i)
for i, (_, tensor) in enumerate(outputs)
}
output_array_specs = {
ArrayKey('OUTPUT_%d' % i): ArraySpec(voxel_size=output_voxel_size)
for i in range(len(outputs))
}
input_source = Z5Source(
input_container,
datasets={_RAW: input[0]},
array_specs={_RAW: ArraySpec(voxel_size=input_voxel_size)})
output_write = Z5Write(output_filename=output_filename,
output_dir=output_dir,
dataset_names=output_dataset_names,
compression_type=output_compression_type)
return \
input_source + \
Normalize(_RAW) + \
Pad(_RAW, size=None) + \
IntensityScaleShift(_RAW, 2, -1) + \
Predict(
weight_graph,
inputs={input[1]: _RAW},
outputs=output_tensor_to_key,
graph=meta_graph,
array_specs=output_array_specs) + \
output_write
def test_impossible(self):
a = ArrayKey("A")
b = ArrayKey("B")
source_a = TestSourceRandomLocation(a)
source_b = TestSourceRandomLocation(b)
pipeline = (source_a, source_b) + \
MergeProvider() + CustomRandomLocation()
with build(pipeline):
with self.assertRaises(AssertionError):
batch = pipeline.request_batch(
BatchRequest({
a:
ArraySpec(roi=Roi((0, 0, 0), (200, 20, 20))),
b:
ArraySpec(roi=Roi((1000, 100, 100), (220, 22, 22))),
}))
def test_two_transpose(self):
test_graph = GraphKey("TEST_GRAPH")
test_array1 = ArrayKey("TEST_ARRAY1")
test_array2 = ArrayKey("TEST_ARRAY2")
transpose_dims = [1, 2]
pipeline = (ArrayTestSource(),
ExampleSource()) + MergeProvider() + SimpleAugment(
mirror_only=[], transpose_only=transpose_dims)
request = BatchRequest()
request[GraphKeys.TEST_GRAPH] = GraphSpec(
roi=Roi((0, 20, 33), (100, 100, 120)))
request[ArrayKeys.TEST_ARRAY1] = ArraySpec(
roi=Roi((0, 0, 0), (100, 200, 300)))
request[ArrayKeys.TEST_ARRAY2] = ArraySpec(
roi=Roi((0, 100, 250), (100, 100, 50)))
possible_loc = [[50, 50], [70, 100], [100, 70]]
with build(pipeline):
seen_transposed = False
for i in range(100):
batch = pipeline.request_batch(request)
assert len(list(batch[GraphKeys.TEST_GRAPH].nodes)) == 1
node = list(batch[GraphKeys.TEST_GRAPH].nodes)[0]
logging.debug(node.location)
assert all([
node.location[dim] in possible_loc[dim] for dim in range(3)
])
seen_transposed = seen_transposed or any([
node.location[dim] != possible_loc[dim][0]
for dim in range(3)
])
assert Roi((0, 20, 33), (100, 100, 120)).contains(
batch[GraphKeys.TEST_GRAPH].spec.roi)
assert batch[GraphKeys.TEST_GRAPH].spec.roi.contains(
node.location)
for (array_key, array) in batch.arrays.items():
assert batch.arrays[array_key].data.shape == batch.arrays[
array_key].spec.roi.get_shape()
assert seen_transposed
def test_impossible(self):
a = ArrayKey("A")
b = ArrayKey("B")
null_key = ArrayKey("NULL")
source_a = ExampleSourceRandomLocation(a)
source_b = ExampleSourceRandomLocation(b)
pipeline = ((source_a, source_b) + MergeProvider() +
CustomRandomLocation(null_key))
with build(pipeline):
with self.assertRaises(PipelineRequestError):
batch = pipeline.request_batch(
BatchRequest({
a:
ArraySpec(roi=Roi((0, 0, 0), (200, 20, 20))),
b:
ArraySpec(roi=Roi((1000, 100, 100), (220, 22, 22))),
}))
def test_transpose():
voxel_size = Coordinate((20, 20))
graph_key = GraphKey("GRAPH")
array_key = ArrayKey("ARRAY")
graph = Graph(
[Node(id=1, location=np.array([450, 550]))],
[],
GraphSpec(roi=Roi((100, 200), (800, 600))),
)
data = np.zeros([40, 30])
data[17, 17] = 1
array = Array(
data, ArraySpec(roi=Roi((100, 200), (800, 600)),
voxel_size=voxel_size))
default_pipeline = (
(GraphSource(graph_key, graph), ArraySource(array_key, array)) +
MergeProvider() + SimpleAugment(
mirror_only=[], transpose_only=[0, 1], transpose_probs=[0, 0]))
transpose_pipeline = (
(GraphSource(graph_key, graph), ArraySource(array_key, array)) +
MergeProvider() + SimpleAugment(
mirror_only=[], transpose_only=[0, 1], transpose_probs=[1, 1]))
request = BatchRequest()
request[graph_key] = GraphSpec(roi=Roi((400, 500), (200, 300)))
request[array_key] = ArraySpec(roi=Roi((400, 500), (200, 300)))
with build(default_pipeline):
expected_location = [450, 550]
batch = default_pipeline.request_batch(request)
assert len(list(batch[graph_key].nodes)) == 1
node = list(batch[graph_key].nodes)[0]
assert all(np.isclose(node.location, expected_location))
node_voxel_index = Coordinate(
(node.location - batch[array_key].spec.roi.get_offset()) /
voxel_size)
assert (
batch[array_key].data[node_voxel_index] == 1
), f"Node at {np.where(batch[array_key].data == 1)} not {node_voxel_index}"
with build(transpose_pipeline):
expected_location = [410, 590]
batch = transpose_pipeline.request_batch(request)
assert len(list(batch[graph_key].nodes)) == 1
node = list(batch[graph_key].nodes)[0]
assert all(np.isclose(node.location, expected_location))
node_voxel_index = Coordinate(
(node.location - batch[array_key].spec.roi.get_offset()) /
voxel_size)
assert (
batch[array_key].data[node_voxel_index] == 1
), f"Node at {np.where(batch[array_key].data == 1)} not {node_voxel_index}"
def test_pipeline3(self):
array_key = ArrayKey("TEST_ARRAY")
points_key = GraphKey("TEST_POINTS")
voxel_size = Coordinate((1, 1))
spec = ArraySpec(voxel_size=voxel_size, interpolatable=True)
hdf5_source = Hdf5Source(self.fake_data_file, {array_key: "testdata"},
array_specs={array_key: spec})
csv_source = CsvPointsSource(
self.fake_points_file,
points_key,
GraphSpec(roi=Roi(shape=Coordinate((100, 100)), offset=(0, 0))),
)
request = BatchRequest()
shape = Coordinate((60, 60))
request.add(array_key, shape, voxel_size=Coordinate((1, 1)))
request.add(points_key, shape)
shift_node = ShiftAugment(prob_slip=0.2,
prob_shift=0.2,
sigma=5,
shift_axis=0)
pipeline = ((hdf5_source, csv_source) + MergeProvider() +
RandomLocation(ensure_nonempty=points_key) + shift_node)
with build(pipeline) as b:
request = b.request_batch(request)
# print(request[points_key])
target_vals = [
self.fake_data[point[0]][point[1]] for point in self.fake_points
]
result_data = request[array_key].data
result_points = list(request[points_key].nodes)
result_vals = [
result_data[int(point.location[0])][int(point.location[1])]
for point in result_points
]
for result_val in result_vals:
self.assertTrue(
result_val in target_vals,
msg=
"result value {} at points {} not in target values {} at points {}"
.format(
result_val,
list(result_points),
target_vals,
self.fake_points,
),
)
def test_mismatched_voxel_multiples():
"""
Ensure we don't shift by half a voxel when transposing 2 axes.
If voxel_size = [2, 2], and we transpose array of shape [4, 6]:
center = total_roi.get_center() -> [2, 3]
# Get distance from center, then transpose
dist_to_center = center - roi.get_offset() -> [2, 3]
dist_to_center = transpose(dist_to_center) -> [3, 2]
# Using the transposed distance to center, get the offset.
new_offset = center - dist_to_center -> [-1, 1]
shape = transpose(shape) -> [6, 4]
original = ((0, 0), (4, 6))
transposed = ((-1, 1), (6, 4))
This result is what we would expect from tranposing, but no longer fits the voxel grid.
dist_to_center should be limited to multiples of the lcm_voxel_size.
instead we should get:
original = ((0, 0), (4, 6))
transposed = ((0, 0), (6, 4))
"""
test_array = ArrayKey("TEST_ARRAY")
data = np.zeros([3, 3])
data[
2,
1] = 1 # voxel has Roi((4, 2) (2, 2)). Contained in Roi((0, 0), (6, 4)). at 2, 1
source = ArraySource(
test_array,
Array(
data,
ArraySpec(roi=Roi((0, 0), (6, 6)), voxel_size=(2, 2)),
),
)
pipeline = source + SimpleAugment(
mirror_only=[], transpose_only=[0, 1], transpose_probs={(1, 0): 1})
with build(pipeline):
request = BatchRequest()
request[test_array] = ArraySpec(roi=Roi((0, 0), (4, 6)))
batch = pipeline.request_batch(request)
data = batch[test_array].data
assert data[1, 2] == 1, f"{data}"
def test_precache(self):
if torch.cuda.is_initialized():
raise RuntimeError(
"Cuda is already initialized in the main process! Will not be able "
"to reinitialize in forked subprocesses.")
logging.getLogger("gunpowder.torch.nodes.predict").setLevel(
logging.INFO)
a = ArrayKey("A")
pred = ArrayKey("PRED")
model = TestModel()
reference_request = BatchRequest()
reference_request[a] = ArraySpec(roi=Roi((0, 0), (7, 7)))
reference_request[pred] = ArraySpec(roi=Roi((1, 1), (5, 5)))
source = TestTorchTrain2DSource()
predict = Predict(
model=model,
inputs={"a": a},
outputs={0: pred},
array_specs={pred: ArraySpec()},
)
pipeline = source + predict + PreCache(cache_size=3, num_workers=2)
request = BatchRequest({
a: ArraySpec(roi=Roi((0, 0), (17, 17))),
pred: ArraySpec(roi=Roi((0, 0), (15, 15))),
})
# train for a couple of iterations
with build(pipeline):
batch = pipeline.request_batch(request)
assert pred in batch
def setUp(self):
super(ProviderTest, self).setUp()
# create some common array keys to be used by concrete tests
ArrayKey("RAW")
ArrayKey("GT_LABELS")
ArrayKey("GT_AFFINITIES")
ArrayKey("GT_AFFINITIES_MASK")
ArrayKey("GT_MASK")
ArrayKey("GT_IGNORE")
ArrayKey("LOSS_SCALE")
self.test_source = TestSource()
self.test_request = BatchRequest()
self.test_request[ArrayKeys.RAW] = ArraySpec(
roi=Roi((20, 20, 20), (10, 10, 10)))
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_mismatched_voxel_multiples():
"""
Ensure we don't shift by half a voxel when transposing 2 axes.
If voxel_size = [2, 2], and we transpose array of shape [4, 6]:
center = total_roi.get_center() -> [2, 3]
# Get distance from center, then transpose
dist_to_center = center - roi.get_offset() -> [2, 3]
dist_to_center = transpose(dist_to_center) -> [3, 2]
# Using the tranposed distance to center, get the correct offset.
new_offset = center - dist_to_center -> [-1, 1]
shape = transpose(shape) -> [6, 4]
original = ((0, 0), (4, 6))
transposed = ((-1, 1), (6, 4))
This result is what we would expect from tranposing, but no longer fits the voxel grid.
dist_to_center should be limited to multiples of the lcm_voxel_size.
"""
test_array = ArrayKey("TEST_ARRAY")
pipeline = (CornerSource(test_array, voxel_size=(2, 2)) +
SimpleAugment(transpose_only=[0, 1]))
request = BatchRequest()
request[test_array] = ArraySpec(roi=Roi((0, 0), (4, 6)))
with build(pipeline):
loop = 100
while loop > 0:
loop -= 1
batch = pipeline.request_batch(request)
data = batch[test_array].data
if data.sum(axis=1)[0] == 1:
loop = -1
assert loop < 0, "Data was never transposed!"
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)
