def test_extract_edges_multithreaded(labelmap_setup):
"""
Make sure the extract_edges() can be used from multiple threads without deadlocking.
"""
dvid_server, dvid_repo, merge_table_path, mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
orig_merge_table = load_merge_table(merge_table_path,
mapping_path,
normalize=True)
merge_graph = LabelmapMergeGraph(merge_table_path)
merge_graph.apply_mapping(mapping_path)
def _test(force_dirty):
if force_dirty:
# A little white-box manipulation here to ensure that the cache is out-of-date.
with merge_graph._edge_cache_main_lock:
merge_graph._edge_cache.clear()
# Extraction should still work.
mutid, dvid_supervoxels, edges, _scores = merge_graph.extract_edges(
*instance_info, 1)
assert (dvid_supervoxels == [1, 2, 3, 4, 5]).all()
assert (orig_merge_table[['id_a', 'id_b']] == edges).all().all(), \
f"Original merge table doesn't match fetched:\n{orig_merge_table}\n\n{edges}\n"
# Quickly check the test function before loading it into a pool.
# (If it's going to fail in a trivial way, let's see it in the main thread.)
_test(False)
_test(True)
with ThreadPoolExecutor(max_workers=11) as executor:
list(executor.map(_test, 300 * [[True], [False], [False]]))
def test_fetch_labelmap_voxels_chunkwise(labelmap_setup):
dvid_server, dvid_repo, _merge_table_path, _mapping_path, supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
# Use a ridiculously small chunk shape, since the test data is quite tiny
CHUNK_SHAPE = (1, 1, 10)
# Test raw supervoxels
voxels = fetch_labelmap_voxels_chunkwise(*instance_info,
[(0, 0, 0), supervoxel_vol.shape],
supervoxels=True,
chunk_shape=CHUNK_SHAPE)
assert (voxels == supervoxel_vol).all()
# Test mapped bodies
voxels = fetch_labelmap_voxels_chunkwise(*instance_info,
[(0, 0, 0), supervoxel_vol.shape],
supervoxels=False,
chunk_shape=CHUNK_SHAPE)
assert (voxels == 1).all()
# Test uninflated mode
voxels_proxy = fetch_labelmap_voxels_chunkwise(
*instance_info, [(0, 0, 0), supervoxel_vol.shape],
supervoxels=True,
chunk_shape=CHUNK_SHAPE,
format='lazy-array')
assert (voxels_proxy() == supervoxel_vol).all()
def test_fetch_supervoxel_sizes_for_body(labelmap_setup):
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
sv_sizes = fetch_supervoxel_sizes_for_body(*instance_info, 1)
assert isinstance(sv_sizes, pd.Series)
assert (sv_sizes.index == [1, 2, 3, 4, 5]).all()
assert (sv_sizes == 9).all() # see initialization in conftest.py
def test_fetch_and_apply_mapping(labelmap_setup):
dvid_server, dvid_repo, merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
# Don't give mapping, ensure it's loaded from dvid.
merge_graph = LabelmapMergeGraph(merge_table_path, no_kafka=True)
merge_graph.fetch_and_apply_mapping(*instance_info)
assert (merge_graph.merge_table_df['body'] == 1).all()
def test_fetch_label(labelmap_setup):
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
label = fetch_label(*instance_info, (0, 1, 3), supervoxels=False)
assert isinstance(label, np.uint64)
assert label == 1
label = fetch_label(*instance_info, (0, 1, 3), supervoxels=True)
assert isinstance(label, np.uint64)
assert label == 2
def test_generate_sample_coordinate(labelmap_setup):
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
coord_zyx = generate_sample_coordinate(*instance_info, 1)
label = fetch_label(*instance_info, coord_zyx.tolist())
assert label == 1
coord_zyx = generate_sample_coordinate(*instance_info, 2, supervoxels=True)
label = fetch_label(*instance_info, coord_zyx.tolist(), supervoxels=True)
assert label == 2
def test_fetch_complete_mappings(labelmap_setup):
"""
Very BASIC test for fetch_complete_mappings().
Does not verify features related to split supervoxels
"""
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
mapping = fetch_complete_mappings(*instance_info, kafka_msgs=[])
assert isinstance(mapping, pd.Series)
assert mapping.index.name == 'sv'
assert mapping.name == 'body'
assert (sorted(mapping.index) == [1, 2, 3, 4, 5])
assert (mapping == 1).all() # see initialization in conftest.py
def test_fetch_mappings(labelmap_setup):
"""
Test the wrapper function for the /mappings DVID API.
"""
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
mapping = fetch_mappings(*instance_info)
assert isinstance(mapping, pd.Series)
assert mapping.index.name == 'sv'
assert mapping.name == 'body'
assert (sorted(mapping.index) == [2, 3, 4, 5]
) # Does not include 'identity' row for SV 1. See docstring.
assert (mapping == 1).all() # see initialization in conftest.py
def test_fetch_labels(labelmap_setup):
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
coords = [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3], [0, 0, 4], [0, 0, 4]]
labels = fetch_labels(*instance_info, coords, supervoxels=False)
assert labels.dtype == np.uint64
assert (labels == 1).all() # See init_labelmap_nodes() in conftest.py
labels = fetch_labels(*instance_info, coords, supervoxels=True)
assert labels.dtype == np.uint64
assert (labels == [1, 1, 1, 2, 2,
2]).all() # See init_labelmap_nodes() in conftest.py
def test_post_labelmap_voxels(labelmap_setup):
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo,
'segmentation-scratch')
# Write some random data and read it back.
vol = np.random.randint(10, size=(128, 128, 128), dtype=np.uint64)
offset = (64, 64, 64)
post_labelmap_voxels(dvid_server, dvid_repo, 'segmentation-scratch',
offset, vol, 0)
complete_voxels = fetch_labelmap_voxels(*instance_info, [(0, 0, 0),
(256, 256, 256)],
supervoxels=True)
assert (complete_voxels[64:192, 64:192, 64:192] == vol).all()
def test_post_mappings(labelmap_setup):
"""
Test the wrapper function for the /mappings DVID API.
"""
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
# Fetch the original mapping
orig_mapping = fetch_mappings(*instance_info).sort_index()
assert (orig_mapping.index == [2, 3, 4, 5]).all()
assert (orig_mapping == 1).all() # see initialization in conftest.py
# Make sure post_mappings does not REQUIRE the Series to be named
orig_mapping.index.name = 'barfoo'
orig_mapping.name = 'foobar'
# Now post a new mapping and read it back.
new_mapping = orig_mapping.copy()
new_mapping[:] = 2
new_mapping.sort_index(inplace=True)
# Post all but sv 5
post_mappings(*instance_info, new_mapping.iloc[:-1], mutid=1)
fetched_mapping = fetch_mappings(*instance_info).sort_index()
assert (fetched_mapping.index == [3, 4, 5]).all()
assert (fetched_mapping.iloc[:-1] == 2).all()
assert (fetched_mapping.iloc[-1:] == 1).all()
# Now post sv 5, too
post_mappings(*instance_info, new_mapping.iloc[-1:], mutid=1)
fetched_mapping = fetch_mappings(*instance_info).sort_index()
assert (fetched_mapping.index == [3, 4, 5]).all()
assert (fetched_mapping == 2).all()
# Try batched
new_mapping = pd.Series(index=[1, 2, 3, 4, 5], data=[1, 1, 1, 2, 2])
post_mappings(*instance_info, new_mapping, mutid=1, batch_size=4)
fetched_mapping = fetch_mappings(*instance_info).sort_index()
assert (fetched_mapping.index == [2, 3, 4, 5]).all()
assert (fetched_mapping == [1, 1, 2,
2]).all(), f"{fetched_mapping} != {[1,1,2,2]}"
# Restore the original mapping
post_mappings(*instance_info, orig_mapping, 1)
fetched_mapping = fetch_mappings(*instance_info).sort_index()
assert (fetched_mapping.index == [2, 3, 4, 5]).all()
assert (fetched_mapping == 1).all()
def test_post_labelmap_blocks(labelmap_setup):
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo,
'segmentation-scratch')
# Write some random data and read it back.
blocks = np.random.randint(10, size=(3, 64, 64, 64), dtype=np.uint64)
corners_zyx = [[0, 0, 0], [0, 64, 0], [0, 0, 64]]
post_labelmap_blocks(dvid_server, dvid_repo, 'segmentation-scratch',
corners_zyx, blocks, 0)
complete_voxels = fetch_labelmap_voxels(*instance_info, [(0, 0, 0),
(128, 128, 128)],
supervoxels=True)
assert (complete_voxels[0:64, 0:64, 0:64] == blocks[0]).all()
assert (complete_voxels[0:64, 64:128, 0:64] == blocks[1]).all()
assert (complete_voxels[0:64, 0:64, 64:128] == blocks[2]).all()
def test_fetch_labelmap_voxels(labelmap_setup):
dvid_server, dvid_repo, _merge_table_path, _mapping_path, supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
# Test raw supervoxels
voxels = fetch_labelmap_voxels(*instance_info,
[(0, 0, 0), supervoxel_vol.shape],
supervoxels=True)
assert (voxels == supervoxel_vol).all()
# Test mapped bodies
voxels = fetch_labelmap_voxels(*instance_info,
[(0, 0, 0), supervoxel_vol.shape],
supervoxels=False)
assert (voxels == 1).all()
# Test uninflated mode
voxels_proxy = fetch_labelmap_voxels(*instance_info,
[(0, 0, 0), supervoxel_vol.shape],
supervoxels=True,
format='lazy-array')
assert len(voxels_proxy.content) < supervoxel_vol.nbytes, \
"Fetched data was apparently not compressed"
assert (voxels_proxy() == supervoxel_vol).all()
def test_fetch_supervoxels(labelmap_setup):
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
supervoxels = fetch_supervoxels(*instance_info, 1)
assert (supervoxels == [1, 2, 3, 4, 5]).all()
def test_fetch_mutation_id(labelmap_setup):
dvid_server, dvid_repo, _merge_table_path, _mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
mut_id = fetch_mutation_id(*instance_info, 1)
assert isinstance(mut_id, int)
def _test_extract_edges(labelmap_setup, force_dirty_mapping):
"""
Implementation for testing extract_edges(), starting either with a "clean" mapping
(in which the body column is already correct beforehand),
or a "dirty" mapping (in which the body column is not correct beforehand).
"""
dvid_server, dvid_repo, merge_table_path, mapping_path, _supervoxel_vol = labelmap_setup
instance_info = DvidInstanceInfo(dvid_server, dvid_repo, 'segmentation')
orig_merge_table = load_merge_table(merge_table_path,
mapping_path,
normalize=True)
merge_graph = LabelmapMergeGraph(merge_table_path)
merge_graph.apply_mapping(mapping_path)
if force_dirty_mapping:
# A little white-box manipulation here to ensure that the mapping is dirty
merge_graph.merge_table_df['body'] = np.uint64(0)
merge_graph.mapping[:] = np.uint64(0)
# First test: If nothing has changed in DVID, we get all rows.
# We should be able to repeat this with the same results
# (Make sure the cache is repopulated correctly.)
_mutid, dvid_supervoxels, edges, _scores = merge_graph.extract_edges(
*instance_info, 1)
assert (dvid_supervoxels == [1, 2, 3, 4, 5]).all()
assert (orig_merge_table[['id_a', 'id_b']].values == edges).all().all(), \
f"Original merge table doesn't match fetched:\n{orig_merge_table}\n\n{edges}\n"
# Now change the mapping in DVID and verify it is reflected in the extracted rows.
# For this test, we'll cleave supervoxels [4,5] from the rest of the body.
uuid = post_branch(dvid_server, dvid_repo,
f'extract-rows-test-{force_dirty_mapping}', '')
cleaved_body = post_cleave(dvid_server, uuid, 'segmentation', 1, [4, 5])
cleaved_mutid = fetch_mutation_id(dvid_server, uuid, 'segmentation', 1)
if force_dirty_mapping:
# A little white-box manipulation here to ensure that the mapping is dirty
merge_graph.mapping.loc[2] = 0
merge_graph.merge_table_df['body'].values[0:2] = np.uint64(0)
mutid, dvid_supervoxels, edges, _scores = merge_graph.extract_edges(
dvid_server, uuid, 'segmentation', 1)
assert (dvid_supervoxels == [1, 2, 3]).all()
_cleaved_svs = set([4, 5])
assert (edges == orig_merge_table[[
'id_a', 'id_b'
]].query('id_a not in @_cleaved_svs and id_b not in @_cleaved_svs')
).all().all()
assert mutid == cleaved_mutid, "Expected cached mutation ID to match DVID"
cleaved_mutid = fetch_mutation_id(dvid_server, uuid, 'segmentation',
cleaved_body)
# Check the other body
mutid, dvid_supervoxels, edges, _scores = merge_graph.extract_edges(
dvid_server, uuid, 'segmentation', cleaved_body)
assert (edges == orig_merge_table[[
'id_a', 'id_b'
]].query('id_a in @_cleaved_svs and id_b in @_cleaved_svs')).all().all()
assert mutid == cleaved_mutid, "Expected cached mutation ID to match DVID"