def ImageShardTransferTask(
src_path: str,
dst_path: str,
shape: ShapeType,
offset: ShapeType,
mip: int = 0,
fill_missing: bool = False,
translate: ShapeType = (0, 0, 0),
agglomerate: bool = False,
timestamp: Optional[int] = None,
):
"""
Generates a sharded image volume from
a preexisting CloudVolume readable data
source. Downsamples are not generated.
The sharded specification can be read here:
Shard Container:
https://github.com/google/neuroglancer/blob/056a3548abffc3c76c93c7a906f1603ce02b5fa3/src/neuroglancer/datasource/precomputed/sharded.md
Sharded Images:
https://github.com/google/neuroglancer/blob/056a3548abffc3c76c93c7a906f1603ce02b5fa3/src/neuroglancer/datasource/precomputed/volume.md#unsharded-chunk-storage
"""
shape = Vec(*shape)
offset = Vec(*offset)
mip = int(mip)
fill_missing = bool(fill_missing)
translate = Vec(*translate)
src_vol = CloudVolume(
src_path, fill_missing=fill_missing,
mip=mip, bounded=False
)
dst_vol = CloudVolume(
dst_path,
fill_missing=fill_missing,
mip=mip,
compress=None
)
dst_bbox = Bbox(offset, offset + shape)
dst_bbox = Bbox.clamp(dst_bbox, dst_vol.meta.bounds(mip))
dst_bbox = dst_bbox.expand_to_chunk_size(
dst_vol.meta.chunk_size(mip),
offset=dst_vol.meta.voxel_offset(mip)
)
src_bbox = dst_bbox - translate
img = src_vol.download(
src_bbox, agglomerate=agglomerate, timestamp=timestamp
)
(filename, shard) = dst_vol.image.make_shard(
img, dst_bbox, mip, progress=False
)
del img
basepath = dst_vol.meta.join(
dst_vol.cloudpath, dst_vol.meta.key(mip)
)
CloudFiles(basepath).put(filename, shard)
def create_blackout_tasks(cloudpath: str,
bounds: Bbox,
mip: int = 0,
shape: ShapeType = (2048, 2048, 64),
value: int = 0,
non_aligned_writes: bool = False):
vol = CloudVolume(cloudpath, mip=mip)
shape = Vec(*shape)
bounds = Bbox.create(bounds)
bounds = vol.bbox_to_mip(bounds, mip=0, to_mip=mip)
if not non_aligned_writes:
bounds = bounds.expand_to_chunk_size(vol.chunk_size, vol.voxel_offset)
bounds = Bbox.clamp(bounds, vol.mip_bounds(mip))
class BlackoutTaskIterator(FinelyDividedTaskIterator):
def task(self, shape, offset):
bounded_shape = min2(shape, vol.bounds.maxpt - offset)
return partial(
igneous.tasks.BlackoutTask,
cloudpath=cloudpath,
mip=mip,
shape=shape.clone(),
offset=offset.clone(),
value=value,
non_aligned_writes=non_aligned_writes,
)
def on_finish(self):
vol.provenance.processing.append({
'method': {
'task': 'BlackoutTask',
'cloudpath': cloudpath,
'mip': mip,
'non_aligned_writes': non_aligned_writes,
'value': value,
'shape': shape.tolist(),
'bounds': [
bounds.minpt.tolist(),
bounds.maxpt.tolist(),
],
},
'by':
operator_contact(),
'date':
strftime('%Y-%m-%d %H:%M %Z'),
})
return BlackoutTaskIterator(bounds, shape)
def pull_vertex_list(
self,
seg_id: int,
v_id_list: List[int],
buffer: int = 1,
expand: bool = False,
) -> Tuple[np.ndarray, Bbox, List[Tuple[int, int, int]]]:
"""Pull a subvolume containing all listed vertices.
Arguments:
seg_id: ID of the segment to use, depends on data in s3.
v_id_list: list of vertex IDs to use.
buffer: Buffer around the bounding box (in voxels). Default 1, set to 0 if expand is True.
expand: Flag whether to expand subvolume to closest set of chunks.
Returns:
img: The image volume containing all vertices.
bounds: Bounding box object which contains the bounds of the volume.
vox_in_img_list: List of coordinates which locate the vertices in the volume.
"""
check_type(seg_id, (int, np.integer))
check_iterable_type(v_id_list, (int, np.integer))
check_type(buffer, (int, np.integer))
if buffer < 0:
raise ValueError(f"Buffer {buffer} shouild not be negative.")
check_type(expand, bool)
if expand:
buffer = 0
buffer = [buffer] * 3
voxel_list = [self._get_voxel(seg_id, i) for i in v_id_list]
if len(voxel_list) == 1: # edge case of 1 vertex
bounds = Bbox(voxel_list[0] - buffer, voxel_list[0] + buffer + 1)
else:
voxel_list = np.array(voxel_list)
lower = list(np.min(voxel_list, axis=0) - buffer)
higher = list(np.max(voxel_list, axis=0) + buffer + 1)
bounds = Bbox(lower, higher)
if expand:
bounds = bounds.expand_to_chunk_size(self.chunk_size)
vox_in_img_list = np.array(voxel_list) - bounds.to_list()[:3]
img = self.pull_bounds_img(bounds)
return img, bounds, vox_in_img_list
def pull_vertex_list(self,
seg_id,
v_id_list,
buffer=[0, 0, 0],
expand=False):
"""
Pull a region containing all listed vertices.
Parameters
----------
seg_id : int
ID of the segment to use, depends on data in s3.
v_id_list : list of ints
list of vertex IDs to use.
buffer : list of ints, optional (default=[0, 0, 0])
Buffer around the bounding box of seed vertices (on lower and higher bound).
expand : bool, optional (default=False)
Flag whether to expand region to closest combination of chunks.
Returns
-------
img : ndarray
The image volume containing all vertices.
bounds : Bbox object
Bounding box object which contains the bounds of the volume.
vox_in_img_list : ndarray, shape nx3
List of coordinates which locate the vertices in the volume.
"""
voxel_list = np.array([self._get_voxel(seg_id, i) for i in v_id_list])
lower = list(np.min(voxel_list, axis=0) - buffer)
higher = list(np.max(voxel_list, axis=0) + buffer)
bounds = Bbox(lower, higher)
if expand:
bounds = bounds.expand_to_chunk_size(self.chunk_size)
lower = bounds.to_list()[:3]
img = self.pull_bounds_img(bounds)
vox_in_img_list = voxel_list - lower
return img, bounds, vox_in_img_list
def get_chunk_aligned_bcube(self, bcube, mip, chunk_xy, chunk_z):
cv_chunk = self.cv[mip].chunk_size
# Expand bbox to be difizible by chunk_size
bbox = Bbox(
(bcube.x_range(mip)[0], bcube.y_range(mip)[0], bcube.z_range()[0]),
(bcube.x_range(mip)[1], bcube.y_range(mip)[1], bcube.z_range()[1]))
aligned_bbox = bbox.expand_to_chunk_size(self.cv[mip].chunk_size,
self.cv[mip].voxel_offset)
aligned_bcube = copy.deepcopy(bcube)
aligned_bcube.reset_coords(aligned_bbox.minpt[0],
aligned_bbox.maxpt[0],
aligned_bbox.minpt[1],
aligned_bbox.maxpt[1],
aligned_bbox.minpt[2],
aligned_bbox.maxpt[2],
mip=mip)
if chunk_xy is not None:
if chunk_xy % cv_chunk[0] != 0:
raise exceptions.ChunkingError(
self, f"Processing chunk_xy {chunk_xy} is not"
f"divisible by MIP{mip} CV chunk {cv_chunk[0]}")
if chunk_z % cv_chunk[2] != 0:
raise exceptions.ChunkingError(
self, f"Processing chunk_z {chunk_z} is not"
f"divisible by MIP{mip} CV chunk {cv_chunk[2]}")
if chunk_xy > aligned_bcube.x_size(mip):
x_adj = chunk_xy - aligned_bcube.x_size(mip)
else:
x_rem = aligned_bcube.x_size(mip) % chunk_xy
if x_rem == 0:
x_adj = 0
else:
x_adj = chunk_xy - x_rem
if chunk_xy > aligned_bcube.y_size(mip):
y_adj = chunk_xy - aligned_bcube.y_size(mip)
else:
y_rem = aligned_bcube.y_size(mip) % chunk_xy
if y_rem == 0:
y_adj = 0
else:
y_adj = chunk_xy - y_rem
if chunk_z > aligned_bcube.z_size():
z_adj = chunk_z - aligned_bcube.z_size()
else:
rem = aligned_bcube.z_size() % chunk_z
if rem == 0:
z_adj = 0
else:
z_adj = chunk_z - z_rem
if x_adj != 0:
xe = aligned_bcube.x_range(mip)[1]
aligned_bcube.reset_coords(xe=xe + x_adj, mip=mip)
if y_adj != 0:
ye = aligned_bcube.y_range(mip)[1]
aligned_bcube.reset_coords(ye=ye + y_adj, mip=mip)
if z_adj != 0:
ze = aligned_bcube.z_range()[1]
aligned_bcube.reset_coords(ze=ze + z_adj)
return aligned_bcube
def upload(meta,
cache,
image,
offset,
mip,
compress=None,
compress_level=None,
cdn_cache=None,
parallel=1,
progress=False,
delete_black_uploads=False,
background_color=0,
non_aligned_writes=False,
location=None,
location_bbox=None,
location_order='F',
use_shared_memory=False,
use_file=False,
green=False,
fill_missing=False):
"""Upload img to vol with offset. This is the primary entry point for uploads."""
if not np.issubdtype(image.dtype, np.dtype(meta.dtype).type):
raise ValueError("""
The uploaded image data type must match the volume data type.
Volume: {}
Image: {}
""".format(meta.dtype, image.dtype))
shape = Vec(*image.shape)[:3]
offset = Vec(*offset)[:3]
bounds = Bbox(offset, shape + offset)
is_aligned = check_grid_aligned(meta,
image,
bounds,
mip,
throw_error=(non_aligned_writes == False))
options = {
"compress": compress,
"compress_level": compress_level,
"cdn_cache": cdn_cache,
"parallel": parallel,
"progress": progress,
"location": location,
"location_bbox": location_bbox,
"location_order": location_order,
"use_shared_memory": use_shared_memory,
"use_file": use_file,
"delete_black_uploads": delete_black_uploads,
"background_color": background_color,
"green": green,
}
if is_aligned:
upload_aligned(meta, cache, image, offset, mip, **options)
return
# Upload the aligned core
expanded = bounds.expand_to_chunk_size(meta.chunk_size(mip),
meta.voxel_offset(mip))
retracted = bounds.shrink_to_chunk_size(meta.chunk_size(mip),
meta.voxel_offset(mip))
core_bbox = retracted.clone() - bounds.minpt
if not core_bbox.subvoxel():
core_img = image[core_bbox.to_slices()]
upload_aligned(meta, cache, core_img, retracted.minpt, mip, **options)
# Download the shell, paint, and upload
all_chunks = set(
chunknames(expanded, meta.bounds(mip), meta.key(mip),
meta.chunk_size(mip)))
core_chunks = set(
chunknames(retracted, meta.bounds(mip), meta.key(mip),
meta.chunk_size(mip)))
shell_chunks = all_chunks.difference(core_chunks)
def shade_and_upload(img3d, bbox):
# decode is returning non-writable chunk
# we're throwing them away so safe to write
img3d.setflags(write=1)
shade(img3d, bbox, image, bounds)
threaded_upload_chunks(
meta,
cache,
img3d,
mip,
((Vec(0, 0, 0), Vec(*img3d.shape[:3]), bbox.minpt, bbox.maxpt), ),
compress=compress,
cdn_cache=cdn_cache,
progress=progress,
n_threads=0,
delete_black_uploads=delete_black_uploads,
green=green,
)
compress_cache = should_compress(meta.encoding(mip),
compress,
cache,
iscache=True)
download_chunks_threaded(meta,
cache,
mip,
shell_chunks,
fn=shade_and_upload,
fill_missing=fill_missing,
progress=progress,
compress_cache=compress_cache,
green=green)
def ImageShardDownsampleTask(
src_path: str,
shape: ShapeType,
offset: ShapeType,
mip: int = 0,
fill_missing: bool = False,
sparse: bool = False,
agglomerate: bool = False,
timestamp: Optional[int] = None,
factor: ShapeType = (2,2,1)
):
"""
Generate a single downsample level for a shard.
Shards are usually hundreds of megabytes to several
gigabyte of data, so it is usually unrealistic from a
memory perspective to make more than one mip at a time.
"""
shape = Vec(*shape)
offset = Vec(*offset)
mip = int(mip)
fill_missing = bool(fill_missing)
src_vol = CloudVolume(
src_path, fill_missing=fill_missing,
mip=mip, bounded=False, progress=False
)
chunk_size = src_vol.meta.chunk_size(mip)
bbox = Bbox(offset, offset + shape)
bbox = Bbox.clamp(bbox, src_vol.meta.bounds(mip))
bbox = bbox.expand_to_chunk_size(
chunk_size, offset=src_vol.meta.voxel_offset(mip)
)
shard_shape = igneous.shards.image_shard_shape_from_spec(
src_vol.scales[mip + 1]["sharding"],
src_vol.meta.volume_size(mip + 1),
src_vol.meta.chunk_size(mip + 1)
)
upper_offset = offset // Vec(*factor)
shape_bbox = Bbox(upper_offset, upper_offset + shard_shape)
shape_bbox = shape_bbox.astype(np.int64)
shape_bbox = Bbox.clamp(shape_bbox, src_vol.meta.bounds(mip + 1))
shape_bbox = shape_bbox.expand_to_chunk_size(src_vol.meta.chunk_size(mip + 1))
if shape_bbox.subvoxel():
return
shard_shape = list(shape_bbox.size3()) + [ 1 ]
output_img = np.zeros(shard_shape, dtype=src_vol.dtype)
nz = int(math.ceil(bbox.dz / chunk_size.z))
dsfn = tinybrain.downsample_with_averaging
if src_vol.layer_type == "segmentation":
dsfn = tinybrain.downsample_segmentation
zbox = bbox.clone()
zbox.maxpt.z = zbox.minpt.z + chunk_size.z
for z in range(nz):
img = src_vol.download(
zbox, agglomerate=agglomerate, timestamp=timestamp
)
(ds_img,) = dsfn(img, factor, num_mips=1, sparse=sparse)
# ds_img[slc] b/c sometimes the size round up in tinybrain
# makes this too large by one voxel on an axis
output_img[:,:,(z*chunk_size.z):(z+1)*chunk_size.z] = ds_img
del img
del ds_img
zbox.minpt.z += chunk_size.z
zbox.maxpt.z += chunk_size.z
(filename, shard) = src_vol.image.make_shard(
output_img, shape_bbox, (mip + 1), progress=False
)
basepath = src_vol.meta.join(
src_vol.cloudpath, src_vol.meta.key(mip + 1)
)
CloudFiles(basepath).put(filename, shard)
