def create_body_tarball_from_sv_tarball(instance_info, body_id):
"""
Download a supervoxel mesh tarball from the given key-value instance,
concatenate together the component meshes into a single body tarball,
and upload it.
"""
keyEncodeLevel0 = 10000000000000
keyEncodeLevel1 = 10100000000000
encoded_sv = str(body_id + keyEncodeLevel0)
sv_tarball_path = f'/tmp/{encoded_sv}.tar'
logger.info(f'Fetching {encoded_sv}.tar')
tarball_contents = fetch_key(*instance_info, f'{encoded_sv}.tar')
with open(sv_tarball_path, 'wb') as f:
f.write(tarball_contents)
logger.info(f'Unpacking {encoded_sv}.tar')
sv_dir = f'/tmp/{encoded_sv}'
os.makedirs(sv_dir, exist_ok=True)
os.chdir(sv_dir)
subprocess.check_call(f'tar -xf {sv_tarball_path}', shell=True)
encoded_body = str(body_id + keyEncodeLevel1)
body_tarball_path = f'/tmp/{encoded_body}.tar'
logger.info(f"Constructing {encoded_body}.drc")
mesh = Mesh.from_directory(sv_dir)
mesh.serialize(f'/tmp/{encoded_body}.drc')
subprocess.check_call(f'tar -cf {body_tarball_path} /tmp/{encoded_body}.drc', shell=True)
with open(body_tarball_path, 'rb') as f:
logger.info(f'Posting {encoded_body}.tar')
post_key(*instance_info, f'{encoded_body}.tar', f)
def test_createmeshes_to_keyvalue(setup_createmeshes_config,
disable_auto_retry):
template_dir, config, _dvid_address, _repo_uuid, object_boxes, _object_sizes = setup_createmeshes_config
kv_instance = 'test_createmeshes_to_keyvalue'
config['output'] = {
'keyvalue': {
'instance': kv_instance,
'create-if-necessary': True
}
}
YAML().dump(config, open(f"{template_dir}/workflow.yaml", 'w'))
_execution_dir, _workflow = launch_flow(template_dir, 1)
server = config['input']['dvid']['server']
uuid = config['input']['dvid']['uuid']
for sv in [100, 200, 300]:
mesh_bytes = fetch_key(server, uuid, kv_instance, f'{sv}.obj')
mesh = Mesh.from_buffer(mesh_bytes, fmt='obj')
assert np.allclose(mesh.vertices_zyx.min(axis=0), object_boxes[sv][0],
1)
assert np.allclose(mesh.vertices_zyx.max(axis=0), object_boxes[sv][1],
1)
def autogen_points(input_seg,
count,
roi,
body,
tbars,
use_skeleton,
random_seed=None,
minimum_distance=0):
"""
Generate a list of points within the input segmentation, based on the given criteria.
See the main help text below for details.
"""
if tbars and not body:
sys.exit(
"If you want to auto-generate tbar points, please specify a body.")
if not tbars and not count:
sys.exit(
"You must supply a --count unless you are generating all tbars of a body."
)
if use_skeleton:
if not body:
sys.exit(
"You must supply a body ID if you want to use a skeleton.")
if tbars:
sys.exit(
"You can't select both tbar points and skeleton points. Pick one or the other."
)
if not count and minimum_distance > 0:
sys.exit(
"You must supply a --count if you want skeleton point samples to respect the minimum distance."
)
if not count and not roi and minimum_distance == 0:
logger.warning(
"You are using all nodes of a skeleton without any ROI filter! Is that what you meant?"
)
rng = default_rng(random_seed)
if tbars:
logger.info(f"Fetching synapses for body {body}")
syn_df = fetch_annotation_label(*input_seg[:2],
'synapses',
body,
format='pandas')
tbars = syn_df.query('kind == "PreSyn"')[[*'zyx']]
if roi:
logger.info(f"Filtering tbars for roi {roi}")
determine_point_rois(*input_seg[:2], [roi], tbars)
tbars = tbars.query('roi == @roi')[[*'zyx']]
if minimum_distance:
logger.info(
f"Pruning close points from {len(tbars)} total tbar points")
tbars = prune_close_pairs(tbars, minimum_distance, rng)
logger.info(f"After pruning, {len(tbars)} tbars remain.")
if count:
count = min(count, len(tbars))
logger.info(f"Sampling {count} tbars")
choices = rng.choice(tbars.index, size=count, replace=False)
tbars = tbars.loc[choices]
logger.info(f"Returning {len(tbars)} tbar points")
return tbars
elif use_skeleton:
assert body
logger.info(f"Fetching skeleton for body {body}")
skeleton_instance = f'{input_seg[2]}_skeletons'
swc = fetch_key(*input_seg[:2], skeleton_instance, f'{body}_swc')
skeleton_df = swc_to_dataframe(swc)
skeleton_df['x'] = skeleton_df['x'].astype(int)
skeleton_df['y'] = skeleton_df['y'].astype(int)
skeleton_df['z'] = skeleton_df['z'].astype(int)
if roi:
logger.info(f"Filtering skeleton for roi {roi}")
determine_point_rois(*input_seg[:2], [roi], skeleton_df)
skeleton_df = skeleton_df.query('roi == @roi')[[*'zyx']]
if minimum_distance:
assert count
# Distance-pruning is very expensive on a huge number of close points.
# If skeleton is large, first reduce the workload by pre-selecting a
# random sample of skeleton points, and prune more from there.
if len(skeleton_df) > 10_000:
# FIXME: random_state can't use rng until I upgrade to pandas 1.0
skeleton_df = skeleton_df.sample(min(4 * count,
len(skeleton_df)),
random_state=None)
logger.info(
f"Pruning close points from {len(skeleton_df)} skeleton points"
)
prune_close_pairs(skeleton_df, minimum_distance, rng)
logger.info(
f"After pruning, {len(skeleton_df)} skeleton points remain.")
if count:
count = min(count, len(skeleton_df))
logger.info(f"Sampling {count} skeleton points")
choices = rng.choice(skeleton_df.index, size=count, replace=False)
skeleton_df = skeleton_df.loc[choices]
logger.info(f"Returning {len(skeleton_df)} skeleton points")
return skeleton_df
elif body:
assert count
if roi:
# TODO: intersect the ranges with the ROI.
raise NotImplementedError(
"Sorry, I haven't yet implemented support for "
"body+roi filtering. Pick one or the other, "
"or ask Stuart to fix this.")
logger.info(f"Fetching sparsevol for body {body}")
ranges = fetch_sparsevol(*input_seg, body, format='ranges')
logger.info("Sampling from sparsevol")
if minimum_distance > 0:
# Sample 4x extra so we still have enough after pruning.
points = sample_points_from_ranges(ranges, 4 * count, rng)
else:
points = sample_points_from_ranges(ranges, count, rng)
points = pd.DataFrame(points, columns=[*'zyx'])
if minimum_distance > 0:
logger.info(f"Pruning close points from {len(points)} body points")
prune_close_pairs(points, minimum_distance, rng)
logger.info(f"After pruning, {len(points)} body points remain")
points = points.iloc[:count]
logger.info(f"Returning {len(points)} body points")
return points
elif roi:
assert count
logger.info(f"Fetching roi {roi}")
roi_ranges = fetch_roi_roi(*input_seg[:2], roi, format='ranges')
logger.info("Sampling from ranges")
if minimum_distance > 0:
# Sample 4x extra so we can prune some out if necessary.
points_s5 = sample_points_from_ranges(roi_ranges, 4 * count, rng)
else:
points_s5 = sample_points_from_ranges(roi_ranges, count, rng)
corners_s0 = points_s5 * (2**5)
points_s0 = rng.integers(corners_s0, corners_s0 + (2**5))
points = pd.DataFrame(points_s0, columns=[*'zyx'])
if minimum_distance > 0:
logger.info(f"Pruning close points from {len(points)} roi points")
prune_close_pairs(points, minimum_distance, rng)
logger.info(
f"After pruning, points from {len(points)} roi points remain")
points = points.iloc[:count]
logger.info(f"Returning {len(points)} roi points")
return points
else:
# No body or roi specified, just sample from the whole non-zero segmentation area
assert count
logger.info("Sampling random points from entire input segmentation")
logger.info("Fetching low-res input volume")
box_s6 = round_box(fetch_volume_box(*input_seg), 2**6, 'out') // 2**6
seg_s6 = fetch_labelmap_voxels(*input_seg, box_s6, scale=6)
mask_s6 = seg_s6.astype(bool)
logger.info("Encoding segmentation as ranges")
seg_ranges = runlength_encode_mask_to_ranges(mask_s6, box_s6)
logger.info("Sampling from ranges")
if minimum_distance > 0:
# Sample 4x extra so we can prune some out if necessary.
points_s6 = sample_points_from_ranges(seg_ranges, 4 * count, rng)
else:
points_s6 = sample_points_from_ranges(seg_ranges, count, rng)
corners_s0 = points_s6 * (2**6)
points_s0 = rng.integers(corners_s0, corners_s0 + (2**6))
points = pd.DataFrame(points_s0, columns=[*'zyx'])
if minimum_distance > 0:
logger.info(
f"Pruning close points from {len(points)} segmentation points")
prune_close_pairs(points, minimum_distance, rng)
logger.info(
f"After pruning, points from {len(points)} segmentation points remain"
)
points = points.iloc[:count]
logger.info(f"Returning {len(points)} segmentation points")
return points