def multiprocess_download(requested_bbox,
mip,
cloudpaths,
meta,
cache,
compress_cache,
fill_missing,
progress,
parallel,
location,
retain,
use_shared_memory,
order,
green,
secrets=None,
background_color=0):
cloudpaths_by_process = []
length = int(math.ceil(len(cloudpaths) / float(parallel)) or 1)
for i in range(0, len(cloudpaths), length):
cloudpaths_by_process.append(cloudpaths[i:i + length])
cpd = partial(child_process_download, meta, cache, mip, compress_cache,
requested_bbox, fill_missing, progress, location,
use_shared_memory, green, secrets, background_color)
parallel_execution(cpd,
cloudpaths_by_process,
parallel,
cleanup_shm=location)
shape = list(requested_bbox.size3()) + [meta.num_channels]
if use_shared_memory:
mmap_handle, renderbuffer = shm.ndarray(shape,
dtype=meta.dtype,
order=order,
location=location,
lock=fs_lock)
else:
handle, renderbuffer = shm.ndarray_fs(shape,
dtype=meta.dtype,
order=order,
location=location,
lock=fs_lock,
emulate_shm=False)
if not retain:
if use_shared_memory:
shm.unlink(location)
else:
os.unlink(location)
return mmap_handle, renderbuffer
def multiprocess_download(
requested_bbox, mip, cloudpaths,
meta, cache, lru, compress_cache,
fill_missing, progress,
parallel, location,
retain, use_shared_memory, order,
green, secrets=None, background_color=0,
):
cpd = partial(child_process_download,
meta, cache,
mip, compress_cache, requested_bbox,
fill_missing, progress,
location, use_shared_memory,
green, secrets, background_color
)
if lru.size > 0:
for path in cloudpaths:
lru.pop(path, None)
parallel_execution(
cpd, cloudpaths, parallel,
progress=progress,
desc="Download",
cleanup_shm=location,
block_size=750,
)
shape = list(requested_bbox.size3()) + [ meta.num_channels ]
if use_shared_memory:
mmap_handle, renderbuffer = shm.ndarray(
shape, dtype=meta.dtype, order=order,
location=location, lock=fs_lock
)
else:
handle, renderbuffer = shm.ndarray_fs(
shape, dtype=meta.dtype, order=order,
location=location, lock=fs_lock,
emulate_shm=False
)
if meta.encoding(mip) == "raw":
repopulate_lru_from_shm(meta, mip, lru, renderbuffer, requested_bbox)
if not retain:
if use_shared_memory:
shm.unlink(location)
else:
os.unlink(location)
return mmap_handle, renderbuffer
def test_parallel_shared_memory_write():
delete_layer()
cv, _ = create_layer(size=(256, 256, 128, 1), offset=(0, 0, 0))
shm_location = 'cloudvolume-test-shm-parallel-write'
mmapfh, shareddata = shm.ndarray(shape=(256, 256, 128),
dtype=np.uint8,
location=shm_location)
shareddata[:] = 1
cv.parallel = 1
cv.upload_from_shared_memory(shm_location, Bbox((0, 0, 0),
(256, 256, 128)))
assert np.all(cv[:] == 1)
shareddata[:] = 2
cv.parallel = 2
cv.upload_from_shared_memory(shm_location, Bbox((0, 0, 0),
(256, 256, 128)))
assert np.all(cv[:] == 2)
shareddata[:, :, :64] = 3
cv.upload_from_shared_memory(shm_location,
bbox=Bbox((0, 0, 0), (256, 256, 128)),
cutout_bbox=Bbox((0, 0, 0), (256, 256, 64)))
assert np.all(cv[:, :, :64] == 3)
assert np.all(cv[:, :, 64:128] == 2)
shareddata[:, :, :69] = 4
cv.autocrop = True
cv.upload_from_shared_memory(shm_location,
bbox=Bbox((-5, -5, -5), (251, 251, 123)),
cutout_bbox=Bbox((-5, -5, -5),
(128, 128, 64)))
assert np.all(cv[:128, :128, :63] == 4)
assert np.all(cv[128:, 128:, :64] == 3)
assert np.all(cv[:, :, 64:128] == 2)
shareddata[:] = 0
shareddata[:, 0, 0] = 1
cv.upload_from_shared_memory(shm_location,
bbox=Bbox((0, 0, 0), (256, 256, 128)),
order='C')
assert np.all(cv[0, 0, :] == 1)
assert np.all(cv[1, 0, :] == 0)
mmapfh.close()
shm.unlink(shm_location)
def download(requested_bbox,
mip,
meta,
cache,
fill_missing,
progress,
parallel,
location,
retain,
use_shared_memory,
use_file,
compress,
order='F',
green=False,
secrets=None,
renumber=False,
background_color=0):
"""Cutout a requested bounding box from storage and return it as a numpy array."""
full_bbox = requested_bbox.expand_to_chunk_size(
meta.chunk_size(mip), offset=meta.voxel_offset(mip))
full_bbox = Bbox.clamp(full_bbox, meta.bounds(mip))
cloudpaths = list(
chunknames(full_bbox,
meta.bounds(mip),
meta.key(mip),
meta.chunk_size(mip),
protocol=meta.path.protocol))
shape = list(requested_bbox.size3()) + [meta.num_channels]
compress_cache = should_compress(meta.encoding(mip),
compress,
cache,
iscache=True)
handle = None
if renumber and (parallel != 1):
raise ValueError("renumber is not supported for parallel operation.")
if use_shared_memory and use_file:
raise ValueError(
"use_shared_memory and use_file are mutually exclusive arguments.")
dtype = np.uint16 if renumber else meta.dtype
if parallel == 1:
if use_shared_memory: # write to shared memory
handle, renderbuffer = shm.ndarray(shape,
dtype=dtype,
order=order,
location=location,
lock=fs_lock)
if not retain:
shm.unlink(location)
elif use_file: # write to ordinary file
handle, renderbuffer = shm.ndarray_fs(shape,
dtype=dtype,
order=order,
location=location,
lock=fs_lock,
emulate_shm=False)
if not retain:
os.unlink(location)
else:
renderbuffer = np.full(shape=shape,
fill_value=background_color,
dtype=dtype,
order=order)
def process(img3d, bbox):
shade(renderbuffer, requested_bbox, img3d, bbox)
remap = {background_color: background_color}
lock = threading.Lock()
N = 1
def process_renumber(img3d, bbox):
nonlocal N
nonlocal lock
nonlocal remap
nonlocal renderbuffer
img_labels = fastremap.unique(img3d)
with lock:
for lbl in img_labels:
if lbl not in remap:
remap[lbl] = N
N += 1
if N > np.iinfo(renderbuffer.dtype).max:
renderbuffer = fastremap.refit(renderbuffer,
value=N,
increase_only=True)
fastremap.remap(img3d, remap, in_place=True)
shade(renderbuffer, requested_bbox, img3d, bbox)
fn = process
if renumber and not (use_file or use_shared_memory):
fn = process_renumber
download_chunks_threaded(meta,
cache,
mip,
cloudpaths,
fn=fn,
fill_missing=fill_missing,
progress=progress,
compress_cache=compress_cache,
green=green,
secrets=secrets,
background_color=background_color)
else:
handle, renderbuffer = multiprocess_download(
requested_bbox,
mip,
cloudpaths,
meta,
cache,
compress_cache,
fill_missing,
progress,
parallel,
location,
retain,
use_shared_memory=(use_file == False),
order=order,
green=green,
secrets=secrets,
background_color=background_color)
out = VolumeCutout.from_volume(meta,
mip,
renderbuffer,
requested_bbox,
handle=handle)
if renumber:
return (out, remap)
return out
def upload_aligned(
meta,
cache,
img,
offset,
mip,
compress=None,
compress_level=None,
cdn_cache=None,
progress=False,
parallel=1,
location=None,
location_bbox=None,
location_order='F',
use_shared_memory=False,
use_file=False,
delete_black_uploads=False,
background_color=0,
green=False,
):
global fs_lock
chunk_ranges = list(generate_chunks(meta, img, offset, mip))
if parallel == 1:
threaded_upload_chunks(
meta,
cache,
img,
mip,
chunk_ranges,
progress=progress,
compress=compress,
cdn_cache=cdn_cache,
delete_black_uploads=delete_black_uploads,
background_color=background_color,
green=green,
compress_level=compress_level,
)
return
length = (len(chunk_ranges) // parallel) or 1
chunk_ranges_by_process = []
for i in range(0, len(chunk_ranges), length):
chunk_ranges_by_process.append(chunk_ranges[i:i + length])
# use_shared_memory means use a predetermined
# shared memory location, not no shared memory
# at all.
if not use_shared_memory:
array_like, renderbuffer = shm.ndarray(shape=img.shape,
dtype=img.dtype,
location=location,
order=location_order,
lock=fs_lock)
renderbuffer[:] = img
cup = partial(child_upload_process,
meta,
cache,
img.shape,
offset,
mip,
compress,
cdn_cache,
progress,
location,
location_bbox,
location_order,
delete_black_uploads,
background_color,
green,
compress_level=compress_level)
parallel_execution(cup,
chunk_ranges_by_process,
parallel,
cleanup_shm=location)
# If manual mode is enabled, it's the
# responsibilty of the user to clean up
if not use_shared_memory:
array_like.close()
shm.unlink(location)
def download(requested_bbox,
mip,
meta,
cache,
fill_missing,
progress,
parallel,
location,
retain,
use_shared_memory,
use_file,
compress,
order='F',
green=False):
"""Cutout a requested bounding box from storage and return it as a numpy array."""
full_bbox = requested_bbox.expand_to_chunk_size(
meta.chunk_size(mip), offset=meta.voxel_offset(mip))
full_bbox = Bbox.clamp(full_bbox, meta.bounds(mip))
cloudpaths = list(
chunknames(full_bbox,
meta.bounds(mip),
meta.key(mip),
meta.chunk_size(mip),
protocol=meta.path.protocol))
shape = list(requested_bbox.size3()) + [meta.num_channels]
compress_cache = should_compress(meta.encoding(mip),
compress,
cache,
iscache=True)
handle = None
if use_shared_memory and use_file:
raise ValueError(
"use_shared_memory and use_file are mutually exclusive arguments.")
if parallel == 1:
if use_shared_memory: # write to shared memory
handle, renderbuffer = shm.ndarray(shape,
dtype=meta.dtype,
order=order,
location=location,
lock=fs_lock)
if not retain:
shm.unlink(location)
elif use_file: # write to ordinary file
handle, renderbuffer = shm.ndarray_fs(shape,
dtype=meta.dtype,
order=order,
location=location,
lock=fs_lock,
emulate_shm=False)
if not retain:
os.unlink(location)
else:
renderbuffer = np.zeros(shape=shape, dtype=meta.dtype, order=order)
def process(img3d, bbox):
shade(renderbuffer, requested_bbox, img3d, bbox)
download_chunks_threaded(meta,
cache,
mip,
cloudpaths,
fn=process,
fill_missing=fill_missing,
progress=progress,
compress_cache=compress_cache,
green=green)
else:
handle, renderbuffer = multiprocess_download(
requested_bbox,
mip,
cloudpaths,
meta,
cache,
compress_cache,
fill_missing,
progress,
parallel,
location,
retain,
use_shared_memory=(use_file == False),
order=order,
green=green,
)
return VolumeCutout.from_volume(meta,
mip,
renderbuffer,
requested_bbox,
handle=handle)
def cleanup(signum, frame):
shm.unlink(dbf_shm_location)
shm.unlink(cc_shm_location)
executor.terminate()
def skeletonize_parallel(all_dbf_shm, dbf_shm_location, cc_labels_shm,
cc_shm_location, remapping, teasar_params, anisotropy,
all_slices, border_targets, extra_targets_before,
extra_targets_after, progress, fix_borders,
fix_branching, cc_segids, parallel, chunk_size):
prevsigint = signal.getsignal(signal.SIGINT)
prevsigterm = signal.getsignal(signal.SIGTERM)
executor = pathos.pools.ProcessPool(parallel)
def cleanup(signum, frame):
shm.unlink(dbf_shm_location)
shm.unlink(cc_shm_location)
executor.terminate()
signal.signal(signal.SIGINT, cleanup)
signal.signal(signal.SIGTERM, cleanup)
skeletonizefn = partial(
parallel_skeletonize_subset,
dbf_shm_location,
all_dbf_shm.shape,
all_dbf_shm.dtype,
cc_shm_location,
cc_labels_shm.shape,
cc_labels_shm.dtype,
remapping,
teasar_params,
anisotropy,
all_slices,
border_targets,
extra_targets_before,
extra_targets_after,
False, # progress, use our own progress bar below
fix_borders,
fix_branching)
ccids = []
if chunk_size < len(cc_segids) // parallel:
for i in range(0, len(cc_segids), chunk_size):
ccids.append(cc_segids[i:i + chunk_size])
else:
for i in range(parallel):
ccids.append(cc_segids[i::parallel])
skeletons = defaultdict(list)
with tqdm(total=len(cc_segids),
disable=(not progress),
desc="Skeletonizing Labels") as pbar:
for skels in executor.uimap(skeletonizefn, ccids):
for segid, skel in skels.items():
skeletons[segid].append(skel)
pbar.update(len(skels))
executor.close()
executor.join()
executor.clear()
signal.signal(signal.SIGINT, prevsigint)
signal.signal(signal.SIGTERM, prevsigterm)
shm.unlink(dbf_shm_location)
shm.unlink(cc_shm_location)
return merge(skeletons)
