def test_bbox_grow():
bbox = Bbox((1, 1, 1), (10, 10, 10), dtype=np.float32)
bbox2 = bbox.clone()
bbox2.adjust(1)
assert np.all(bbox2.minpt == bbox.minpt - 1)
assert np.all(bbox2.maxpt == bbox.maxpt + 1)
bbox3 = bbox.clone()
bbox3.adjust(-1)
assert np.all(bbox3.minpt == bbox.minpt + 1)
assert np.all(bbox3.maxpt == bbox.maxpt - 1)
bbox4 = bbox.clone()
bbox4.adjust((1, 1, 1))
assert np.all(bbox2.minpt == bbox.minpt - 1)
assert np.all(bbox2.maxpt == bbox.maxpt + 1)
def MeshSpatialIndex(
cloudpath:str,
shape:Tuple[int,int,int],
offset:Tuple[int,int,int],
mip:int = 0,
fill_missing:bool=False,
compress:Optional[Union[str,bool]] = 'gzip',
mesh_dir:Optional[str] = None
) -> None:
"""
The main way to add a spatial index is to use the MeshTask,
but old datasets or broken datasets may need it to be
reconstituted. An alternative use is create the spatial index
over a different area size than the mesh task.
"""
cv = CloudVolume(
cloudpath, mip=mip,
bounded=False, fill_missing=fill_missing
)
cf = CloudFiles(cloudpath)
bounds = Bbox(Vec(*offset), Vec(*shape) + Vec(*offset))
bounds = Bbox.clamp(bounds, cv.bounds)
data_bounds = bounds.clone()
data_bounds.maxpt += 1 # match typical Marching Cubes overlap
precision = cv.mesh.spatial_index.precision
resolution = cv.resolution
if not mesh_dir:
mesh_dir = cv.info["mesh"]
# remap: old img -> img
img, remap = cv.download(data_bounds, renumber=True)
img = img[...,0]
slcs = find_objects(img)
del img
reverse_map = { v:k for k,v in remap.items() } # img -> old img
bboxes = {}
for label, slc in enumerate(slcs):
if slc is None:
continue
mesh_bounds = Bbox.from_slices(slc)
mesh_bounds += Vec(*offset)
mesh_bounds *= Vec(*resolution, dtype=np.float32)
bboxes[str(reverse_map[label+1])] = \
mesh_bounds.astype(resolution.dtype).to_list()
bounds = bounds.astype(resolution.dtype) * resolution
cf.put_json(
f"{mesh_dir}/{bounds.to_filename(precision)}.spatial",
bboxes,
compress=compress,
cache_control=False,
)
def child_upload_process(
meta,
cache,
img_shape,
offset,
mip,
compress,
cdn_cache,
progress,
location,
location_bbox,
location_order,
delete_black_uploads,
background_color,
green,
chunk_ranges,
compress_level=None,
):
global fs_lock
reset_connection_pools()
shared_shape = img_shape
if location_bbox:
shared_shape = list(location_bbox.size3()) + [meta.num_channels]
array_like, renderbuffer = shm.ndarray(shape=shared_shape,
dtype=meta.dtype,
location=location,
order=location_order,
lock=fs_lock,
readonly=True)
if location_bbox:
cutout_bbox = Bbox(offset, offset + img_shape[:3])
delta_box = cutout_bbox.clone() - location_bbox.minpt
renderbuffer = renderbuffer[delta_box.to_slices()]
threaded_upload_chunks(
meta,
cache,
renderbuffer,
mip,
chunk_ranges,
compress=compress,
cdn_cache=cdn_cache,
progress=progress,
delete_black_uploads=delete_black_uploads,
background_color=background_color,
green=green,
compress_level=compress_level,
)
array_like.close()
def test_bbox_division():
bbox = Bbox((1, 1, 1), (10, 10, 10), dtype=np.float32)
bbox2 = bbox.clone()
bbox /= 3.0
bbx3 = bbox2 / 3.0
point333 = np.float32(1) / np.float32(3)
assert np.all(np.abs(bbx3.minpt - point333) < 1e-6)
assert np.all(bbx3.maxpt == np.float32(3) + point333)
assert bbox == bbx3
bbox = Bbox((1, 1, 1), (10, 10, 10), dtype=np.float32)
x = bbox.minpt
bbox /= 3.0
assert np.all(x == point333)
class MeshTask(RegisteredTask):
def __init__(self, shape, offset, layer_path, **kwargs):
"""
Convert all labels in the specified bounding box into meshes
via marching cubes and quadratic edge collapse (github.com/seung-lab/zmesh).
Required:
shape: (sx,sy,sz) size of task
offset: (x,y,z) offset from (0,0,0)
layer_path: neuroglancer/cloudvolume dataset path
Optional:
lod: (uint) level of detail to record these meshes at
mip: (uint) level of the resolution pyramid to download segmentation from
simplification_factor: (uint) try to reduce the number of triangles in the
mesh by this factor (but constrained by max_simplification_error)
max_simplification_error: The maximum physical distance that
simplification is allowed to move a triangle vertex by.
mesh_dir: which subdirectory to write the meshes to (overrides info file location)
remap_table: agglomerate segmentation before meshing using { orig_id: new_id }
generate_manifests: (bool) if it is known that the meshes generated by this
task will not be cropped by the bounding box, avoid needing to run a seperate
MeshManifestTask pass by generating manifests on the spot.
These two options are used to allow sufficient overlap for trivial mesh stitching
between adjacent tasks.
low_padding: (uint) expand the bounding box by this many pixels by subtracting
this padding from the minimum point of the bounding box on all axes.
high_padding: (uint) expand the bounding box by this many pixels adding
this padding to the maximum point of the bounding box on all axes.
parallel_download: (uint: 1) number of processes to use during the segmentation download
cache_control: (str: None) specify the cache-control header when uploading mesh files
dust_threshold: (uint: None) don't bother meshing labels strictly smaller than this number of voxels.
encoding: (str) 'precomputed' (default) or 'draco'
draco_compression_level: (uint: 1) only applies to draco encoding
draco_create_metadata: (bool: False) only applies to draco encoding
progress: (bool: False) show progress bars for meshing
object_ids: (list of ints) if specified, only mesh these ids
fill_missing: (bool: False) replace missing segmentation files with zeros instead of erroring
spatial_index: (bool: False) generate a JSON spatial index of which meshes are available in
a given bounding box.
sharded: (bool: False) If True, upload all meshes together as a single mapbuffer
fragment file.
timestamp: (int: None) (graphene only) use the segmentation existing at this
UNIX timestamp.
"""
super(MeshTask, self).__init__(shape, offset, layer_path, **kwargs)
self.shape = Vec(*shape)
self.offset = Vec(*offset)
self.layer_path = layer_path
self.options = {
'cache_control': kwargs.get('cache_control', None),
'draco_compression_level': kwargs.get('draco_compression_level',
1),
'draco_create_metadata': kwargs.get('draco_create_metadata',
False),
'dust_threshold': kwargs.get('dust_threshold', None),
'encoding': kwargs.get('encoding', 'precomputed'),
'fill_missing': kwargs.get('fill_missing', False),
'generate_manifests': kwargs.get('generate_manifests', False),
'high_padding': kwargs.get('high_padding', 1),
'low_padding': kwargs.get('low_padding', 0),
'lod': kwargs.get('lod', 0),
'max_simplification_error': kwargs.get('max_simplification_error',
40),
'simplification_factor': kwargs.get('simplification_factor', 100),
'mesh_dir': kwargs.get('mesh_dir', None),
'mip': kwargs.get('mip', 0),
'object_ids': kwargs.get('object_ids', None),
'parallel_download': kwargs.get('parallel_download', 1),
'progress': kwargs.get('progress', False),
'remap_table': kwargs.get('remap_table', None),
'spatial_index': kwargs.get('spatial_index', False),
'sharded': kwargs.get('sharded', False),
'timestamp': kwargs.get('timestamp', None),
'agglomerate': kwargs.get('agglomerate', True),
'stop_layer': kwargs.get('stop_layer', 2),
'compress': kwargs.get('compress', 'gzip'),
'closed_dataset_edges': kwargs.get('closed_dataset_edges', True),
}
supported_encodings = ['precomputed', 'draco']
if not self.options['encoding'] in supported_encodings:
raise ValueError(
'Encoding {} is not supported. Options: {}'.format(
self.options['encoding'], ', '.join(supported_encodings)))
self._encoding_to_compression_dict = {
'precomputed': self.options['compress'],
'draco': False,
}
def execute(self):
self._volume = CloudVolume(self.layer_path,
self.options['mip'],
bounded=False,
parallel=self.options['parallel_download'],
fill_missing=self.options['fill_missing'])
self._bounds = Bbox(self.offset, self.shape + self.offset)
self._bounds = Bbox.clamp(self._bounds, self._volume.bounds)
self.progress = bool(self.options['progress'])
self._mesher = zmesh.Mesher(self._volume.resolution)
# Marching cubes loves its 1vx overlaps.
# This avoids lines appearing between
# adjacent chunks.
data_bounds = self._bounds.clone()
data_bounds.minpt -= self.options['low_padding']
data_bounds.maxpt += self.options['high_padding']
self._mesh_dir = self.get_mesh_dir()
if self.options['encoding'] == 'draco':
self.draco_encoding_settings = draco_encoding_settings(
shape=(self.shape + self.options['low_padding'] +
self.options['high_padding']),
offset=self.offset,
resolution=self._volume.resolution,
compression_level=self.options["draco_compression_level"],
create_metadata=self.options['draco_create_metadata'],
uses_new_draco_bin_size=False,
)
# chunk_position includes the overlap specified by low_padding/high_padding
# agglomerate, timestamp, stop_layer only applies to graphene volumes,
# no-op for precomputed
data = self._volume.download(data_bounds,
agglomerate=self.options['agglomerate'],
timestamp=self.options['timestamp'],
stop_layer=self.options['stop_layer'])
if not np.any(data):
if self.options['spatial_index']:
self._upload_spatial_index(self._bounds, {})
return
left_offset = Vec(0, 0, 0)
if self.options["closed_dataset_edges"]:
data, left_offset = self._handle_dataset_boundary(
data, data_bounds)
data = self._remove_dust(data, self.options['dust_threshold'])
data = self._remap(data)
if self.options['object_ids']:
data = fastremap.mask_except(data,
self.options['object_ids'],
in_place=True)
data, renumbermap = fastremap.renumber(data, in_place=True)
renumbermap = {v: k for k, v in renumbermap.items()}
self._mesher.mesh(data[..., 0].T)
del data
self.compute_meshes(renumbermap, left_offset)
def _handle_dataset_boundary(self, data, bbox):
"""
This logic is used to add a black border along sides
of the image that touch the dataset boundary which
results in the closure of the mesh faces on that side.
"""
if ((not np.any(bbox.minpt == self._volume.bounds.minpt))
and (not np.any(bbox.maxpt == self._volume.bounds.maxpt))):
return data, Vec(0, 0, 0)
shape = Vec(*data.shape, dtype=np.int64)
offset = Vec(0, 0, 0, 0)
for i in range(3):
if bbox.minpt[i] == self._volume.voxel_offset[i]:
offset[i] += 1
shape[i] += 1
if bbox.maxpt[i] == self._volume.bounds.maxpt[i]:
shape[i] += 1
slices = (
slice(offset.x, offset.x + data.shape[0]),
slice(offset.y, offset.y + data.shape[1]),
slice(offset.z, offset.z + data.shape[2]),
)
mirror_data = np.zeros(shape, dtype=data.dtype, order="F")
mirror_data[slices] = data
if offset[0]:
mirror_data[0, :, :] = 0
if offset[1]:
mirror_data[:, 0, :] = 0
if offset[2]:
mirror_data[:, :, 0] = 0
return mirror_data, offset[:3]
def get_mesh_dir(self):
if self.options['mesh_dir'] is not None:
return self.options['mesh_dir']
elif 'mesh' in self._volume.info:
return self._volume.info['mesh']
else:
raise ValueError(
"The mesh destination is not present in the info file.")
def _remove_dust(self, data, dust_threshold):
if dust_threshold:
segids, pxct = fastremap.unique(data, return_counts=True)
dust_segids = [
sid for sid, ct in zip(segids, pxct)
if ct < int(dust_threshold)
]
data = fastremap.mask(data, dust_segids, in_place=True)
return data
def _remap(self, data):
if self.options['remap_table'] is None:
return data
self.options['remap_table'] = {
int(k): int(v)
for k, v in self.options['remap_table'].items()
}
remap = self.options['remap_table']
remap[0] = 0
data = fastremap.mask_except(data, list(remap.keys()), in_place=True)
return fastremap.remap(data, remap, in_place=True)
def compute_meshes(self, renumbermap, offset):
bounding_boxes = {}
meshes = {}
for obj_id in tqdm(self._mesher.ids(),
disable=(not self.progress),
desc="Mesh"):
remapped_id = renumbermap[obj_id]
mesh_binary, mesh_bounds = self._create_mesh(obj_id, offset)
bounding_boxes[remapped_id] = mesh_bounds.to_list()
meshes[remapped_id] = mesh_binary
if self.options['sharded']:
self._upload_batch(meshes, self._bounds)
else:
self._upload_individuals(meshes,
self.options['generate_manifests'])
if self.options['spatial_index']:
self._upload_spatial_index(self._bounds, bounding_boxes)
def _upload_batch(self, meshes, bbox):
cf = CloudFiles(self.layer_path, progress=self.options['progress'])
mbuf = MapBuffer(meshes, compress="br")
cf.put(
f"{self._mesh_dir}/{bbox.to_filename()}.frags",
content=mbuf.tobytes(),
compress=None,
content_type="application/x.mapbuffer",
cache_control=False,
)
def _upload_individuals(self, mesh_binaries, generate_manifests):
cf = CloudFiles(self.layer_path)
content_type = "model/mesh"
if self.options["encoding"] == "draco":
content_type = "model/x.draco"
cf.puts(
((f"{self._mesh_dir}/{segid}:{self.options['lod']}:{self._bounds.to_filename()}",
mesh_binary) for segid, mesh_binary in mesh_binaries.items()),
compress=self._encoding_to_compression_dict[
self.options['encoding']],
cache_control=self.options['cache_control'],
content_type=content_type,
)
if generate_manifests:
cf.put_jsons(
((f"{self._mesh_dir}/{segid}:{self.options['lod']}", {
"fragments": [
f"{segid}:{self.options['lod']}:{self._bounds.to_filename()}"
]
}) for segid, mesh_binary in mesh_binaries.items()),
compress=None,
cache_control=self.options['cache_control'],
)
def _create_mesh(self, obj_id, left_bound_offset):
mesh = self._mesher.get_mesh(
obj_id,
simplification_factor=self.options['simplification_factor'],
max_simplification_error=self.options['max_simplification_error'],
voxel_centered=True,
)
self._mesher.erase(obj_id)
resolution = self._volume.resolution
offset = (self._bounds.minpt - self.options['low_padding']).astype(
np.float32)
mesh.vertices[:] += (offset - left_bound_offset) * resolution
mesh_bounds = Bbox(np.amin(mesh.vertices, axis=0),
np.amax(mesh.vertices, axis=0))
if self.options['encoding'] == 'draco':
mesh_binary = DracoPy.encode(mesh.vertices, mesh.faces,
**self.draco_encoding_settings)
elif self.options['encoding'] == 'precomputed':
mesh_binary = mesh.to_precomputed()
return mesh_binary, mesh_bounds
def _upload_spatial_index(self, bbox, mesh_bboxes):
cf = CloudFiles(self.layer_path, progress=self.options['progress'])
precision = self._volume.mesh.spatial_index.precision
resolution = self._volume.resolution
bbox = bbox.astype(resolution.dtype) * resolution
cf.put_json(
f"{self._mesh_dir}/{bbox.to_filename(precision)}.spatial",
mesh_bboxes,
compress=self.options['compress'],
cache_control=False,
)
class HyperSquareTask(RegisteredTask):
def __init__(self, bucket_name, dataset_name, layer_name,
volume_dir, layer_type, overlap, resolution):
self.bucket_name = bucket_name
self.dataset_name = dataset_name
self.layer_name = layer_name
self.volume_dir = volume_dir
self.layer_type = layer_type
self.overlap = Vec(*overlap)
self.resolution = Vec(*resolution)
self._volume_cloudpath = 'gs://{}/{}'.format(self.bucket_name, self.volume_dir)
self._bucket = None
self._metadata = None
self._bounds = None
def execute(self):
client = storage.Client.from_service_account_json(
lib.credentials_path(), project=lib.GCLOUD_PROJECT_NAME
)
self._bucket = client.get_bucket(self.bucket_name)
self._metadata = meta = self._download_metadata()
self._bounds = Bbox(
meta['physical_offset_min'], # in voxels
meta['physical_offset_max']
)
shape = Vec(*meta['chunk_voxel_dimensions'])
shape = Vec(shape.x, shape.y, shape.z, 1)
if self.layer_type == 'image':
dtype = meta['image_type'].lower()
cube = self._materialize_images(shape, dtype)
elif self.layer_type == 'segmentation':
dtype = meta['segment_id_type'].lower()
cube = self._materialize_segmentation(shape, dtype)
else:
dtype = meta['affinity_type'].lower()
return NotImplementedError("Don't know how to get the images for this layer.")
self._upload_chunk(cube, dtype)
def _download_metadata(self):
cloudpath = '{}/metadata.json'.format(self.volume_dir)
metadata = self._bucket.get_blob(cloudpath).download_as_string()
return json.loads(metadata)
def _materialize_segmentation(self, shape, dtype):
segmentation_path = '{}/segmentation.lzma'.format(self.volume_dir)
seg_blob = self._bucket.get_blob(segmentation_path)
return self._decode_lzma(seg_blob.download_as_string(), shape, dtype)
def _materialize_images(self, shape, dtype):
cloudpaths = [ '{}/jpg/{}.jpg'.format(self.volume_dir, i) for i in xrange(shape.z) ]
datacube = np.zeros(shape=shape, dtype=np.uint8) # x,y,z,channels
prefix = '{}/jpg/'.format(self.volume_dir)
blobs = self._bucket.list_blobs(prefix=prefix)
for blob in blobs:
z = int(re.findall(r'(\d+)\.jpg', blob.name)[0])
imgdata = blob.download_as_string()
# Hypersquare images are each situated in the xy plane
# so the shape should be (width,height,1)
shape = self._bounds.size3()
shape.z = 1
datacube[:,:,z,:] = chunks.decode_jpeg(imgdata, shape=tuple(shape))
return datacube
def _decode_lzma(self, string_data, shape, dtype):
arr = lzma.decompress(string_data)
arr = np.fromstring(arr, dtype=dtype)
return arr.reshape(shape[::-1]).T
def _upload_chunk(self, datacube, dtype):
vol = CloudVolume(self.dataset_name, self.layer_name, mip=0)
hov = self.overlap / 2 # half overlap, e.g. 32 -> 16 in e2198
img = datacube[ hov.x:-hov.x, hov.y:-hov.y, hov.z:-hov.z, : ] # e.g. 256 -> 224
bounds = self._bounds.clone()
# the boxes are offset left of zero by half overlap, so no need to
# compensate for weird shifts. only upload the non-overlap region.
downsample_and_upload(image, bounds, vol, ds_shape=img.shape)
vol[ bounds.to_slices() ] = img
class MeshTask(RegisteredTask):
def __init__(self, shape, offset, layer_path, mip=0, simplification_factor=100, max_simplification_error=40):
super(MeshTask, self).__init__(shape, offset, layer_path, mip, simplification_factor, max_simplification_error)
self.shape = Vec(*shape)
self.offset = Vec(*offset)
self.mip = mip
self.layer_path = layer_path
self.lod = 0 # level of detail -- to be implemented
self.simplification_factor = simplification_factor
self.max_simplification_error = max_simplification_error
def execute(self):
self._mesher = Mesher()
self._volume = CloudVolume(self.layer_path, self.mip, bounded=False)
self._bounds = Bbox( self.offset, self.shape + self.offset )
self._bounds = Bbox.clamp(self._bounds, self._volume.bounds)
# Marching cubes loves its 1vx overlaps.
# This avoids lines appearing between
# adjacent chunks.
data_bounds = self._bounds.clone()
data_bounds.minpt -= 1
data_bounds.maxpt += 1
self._mesh_dir = None
if 'meshing' in self._volume.info:
self._mesh_dir = self._volume.info['meshing']
elif 'mesh' in self._volume.info:
self._mesh_dir = self._volume.info['mesh']
if not self._mesh_dir:
raise ValueError("The mesh destination is not present in the info file.")
self._data = self._volume[data_bounds.to_slices()] # chunk_position includes a 1 pixel overlap
self._compute_meshes()
def _compute_meshes(self):
with Storage(self.layer_path) as storage:
data = self._data[:,:,:,0].T
self._mesher.mesh(data.flatten(), *data.shape[:3])
for obj_id in self._mesher.ids():
storage.put_file(
file_path='{}/{}:{}:{}'.format(self._mesh_dir, obj_id, self.lod, self._bounds.to_filename()),
content=self._create_mesh(obj_id),
compress=True,
)
def _create_mesh(self, obj_id):
mesh = self._mesher.get_mesh(obj_id,
simplification_factor=self.simplification_factor,
max_simplification_error=self.max_simplification_error
)
vertices = self._update_vertices(np.array(mesh['points'], dtype=np.float32))
vertex_index_format = [
np.uint32(len(vertices) / 3), # Number of vertices ( each vertex is three numbers (x,y,z) )
vertices,
np.array(mesh['faces'], dtype=np.uint32)
]
return b''.join([ array.tobytes() for array in vertex_index_format ])
def _update_vertices(self, points):
# zlib meshing multiplies verticies by two to avoid working with floats like 1.5
# but we need to recover the exact position for display
points /= 2.0
resolution = self._volume.resolution
xmin, ymin, zmin = self._bounds.minpt
points[0::3] = (points[0::3] + xmin) * resolution.x
points[1::3] = (points[1::3] + ymin) * resolution.y
points[2::3] = (points[2::3] + zmin) * resolution.z
return points
class MeshTask(RegisteredTask):
def __init__(self, shape, offset, layer_path, **kwargs):
"""
Convert all labels in the specified bounding box into meshes
via marching cubes and quadratic edge collapse (github.com/seung-lab/zmesh).
Required:
shape: (sx,sy,sz) size of task
offset: (x,y,z) offset from (0,0,0)
layer_path: neuroglancer/cloudvolume dataset path
Optional:
lod: (uint) level of detail to record these meshes at
mip: (uint) level of the resolution pyramid to download segmentation from
simplification_factor: (uint) try to reduce the number of triangles in the
mesh by this factor (but constrained by max_simplification_error)
max_simplification_error: The maximum physical distance that
simplification is allowed to move a triangle vertex by.
mesh_dir: which subdirectory to write the meshes to (overrides info file location)
remap_table: agglomerate segmentation before meshing using { orig_id: new_id }
generate_manifests: (bool) if it is known that the meshes generated by this
task will not be cropped by the bounding box, avoid needing to run a seperate
MeshManifestTask pass by generating manifests on the spot.
These two options are used to allow sufficient overlap for trivial mesh stitching
between adjacent tasks.
low_padding: (uint) expand the bounding box by this many pixels by subtracting
this padding from the minimum point of the bounding box on all axes.
high_padding: (uint) expand the bounding box by this many pixels adding
this padding to the maximum point of the bounding box on all axes.
parallel_download: (uint: 1) number of processes to use during the segmentation download
cache_control: (str: None) specify the cache-control header when uploading mesh files
dust_threshold: (uint: None) don't bother meshing labels strictly smaller than this number of voxels.
encoding: (str) 'precomputed' (default) or 'draco'
draco_compression_level: (uint: 1) only applies to draco encoding
draco_create_metadata: (bool: False) only applies to draco encoding
progress: (bool: False) show progress bars for meshing
object_ids: (list of ints) if specified, only mesh these ids
fill_missing: (bool: False) replace missing segmentation files with zeros instead of erroring
spatial_index: (bool: False) generate a JSON spatial index of which meshes are available in
a given bounding box.
sharded: (bool: False) If True, upload all meshes together as a single pickled
fragment file.
timestamp: (int: None) (graphene only) use the segmentation existing at this
UNIX timestamp.
"""
super(MeshTask, self).__init__(shape, offset, layer_path, **kwargs)
self.shape = Vec(*shape)
self.offset = Vec(*offset)
self.layer_path = layer_path
self.options = {
'cache_control': kwargs.get('cache_control', None),
'draco_compression_level': kwargs.get('draco_compression_level',
1),
'draco_create_metadata': kwargs.get('draco_create_metadata',
False),
'dust_threshold': kwargs.get('dust_threshold', None),
'encoding': kwargs.get('encoding', 'precomputed'),
'fill_missing': kwargs.get('fill_missing', False),
'generate_manifests': kwargs.get('generate_manifests', False),
'high_padding': kwargs.get('high_padding', 1),
'low_padding': kwargs.get('low_padding', 0),
'lod': kwargs.get('lod', 0),
'max_simplification_error': kwargs.get('max_simplification_error',
40),
'simplification_factor': kwargs.get('simplification_factor', 100),
'mesh_dir': kwargs.get('mesh_dir', None),
'mip': kwargs.get('mip', 0),
'object_ids': kwargs.get('object_ids', None),
'parallel_download': kwargs.get('parallel_download', 1),
'progress': kwargs.get('progress', False),
'remap_table': kwargs.get('remap_table', None),
'spatial_index': kwargs.get('spatial_index', False),
'sharded': kwargs.get('sharded', False),
'timestamp': kwargs.get('timestamp', None),
'agglomerate': kwargs.get('agglomerate', True),
'stop_layer': kwargs.get('stop_layer', 2),
'compress': kwargs.get('compress', 'gzip'),
}
supported_encodings = ['precomputed', 'draco']
if not self.options['encoding'] in supported_encodings:
raise ValueError(
'Encoding {} is not supported. Options: {}'.format(
self.options['encoding'], ', '.join(supported_encodings)))
self._encoding_to_compression_dict = {
'precomputed': self.options['compress'],
'draco': False,
}
def execute(self):
self._volume = CloudVolume(self.layer_path,
self.options['mip'],
bounded=False,
parallel=self.options['parallel_download'],
fill_missing=self.options['fill_missing'])
self._bounds = Bbox(self.offset, self.shape + self.offset)
self._bounds = Bbox.clamp(self._bounds, self._volume.bounds)
self.progress = bool(self.options['progress'])
self._mesher = zmesh.Mesher(self._volume.resolution)
# Marching cubes loves its 1vx overlaps.
# This avoids lines appearing between
# adjacent chunks.
data_bounds = self._bounds.clone()
data_bounds.minpt -= self.options['low_padding']
data_bounds.maxpt += self.options['high_padding']
self._mesh_dir = self.get_mesh_dir()
if self.options['encoding'] == 'draco':
self.draco_encoding_settings = self._compute_draco_encoding_settings(
)
# chunk_position includes the overlap specified by low_padding/high_padding
# agglomerate, timestamp, stop_layer only applies to graphene volumes,
# no-op for precomputed
data = self._volume.download(data_bounds,
agglomerate=self.options['agglomerate'],
timestamp=self.options['timestamp'],
stop_layer=self.options['stop_layer'])
if not np.any(data):
return
data = self._remove_dust(data, self.options['dust_threshold'])
data = self._remap(data)
if self.options['object_ids']:
data = fastremap.mask_except(data,
self.options['object_ids'],
in_place=True)
data, renumbermap = fastremap.renumber(data, in_place=True)
renumbermap = {v: k for k, v in renumbermap.items()}
self.compute_meshes(data, renumbermap)
def get_mesh_dir(self):
if self.options['mesh_dir'] is not None:
return self.options['mesh_dir']
elif 'mesh' in self._volume.info:
return self._volume.info['mesh']
else:
raise ValueError(
"The mesh destination is not present in the info file.")
def _compute_draco_encoding_settings(self):
min_quantization_range = max(
(self.shape + self.options['low_padding'] +
self.options['high_padding']) * self._volume.resolution)
max_draco_bin_size = np.floor(
min(self._volume.resolution) / np.sqrt(2))
draco_quantization_bits, draco_quantization_range, draco_bin_size = \
calculate_draco_quantization_bits_and_range(min_quantization_range, max_draco_bin_size)
draco_quantization_origin = self.offset - (self.offset %
draco_bin_size)
return {
'quantization_bits': draco_quantization_bits,
'compression_level': self.options['draco_compression_level'],
'quantization_range': draco_quantization_range,
'quantization_origin': draco_quantization_origin,
'create_metadata': self.options['draco_create_metadata']
}
def _remove_dust(self, data, dust_threshold):
if dust_threshold:
segids, pxct = fastremap.unique(data, return_counts=True)
dust_segids = [
sid for sid, ct in zip(segids, pxct)
if ct < int(dust_threshold)
]
data = fastremap.mask(data, dust_segids, in_place=True)
return data
def _remap(self, data):
if self.options['remap_table'] is None:
return data
self.options['remap_table'] = {
int(k): int(v)
for k, v in self.options['remap_table'].items()
}
remap = self.options['remap_table']
remap[0] = 0
data = fastremap.mask_except(data, list(remap.keys()), in_place=True)
return fastremap.remap(data, remap, in_place=True)
def compute_meshes(self, data, renumbermap):
data = data[:, :, :, 0].T
self._mesher.mesh(data)
del data
bounding_boxes = {}
meshes = {}
for obj_id in tqdm(self._mesher.ids(),
disable=(not self.progress),
desc="Mesh"):
remapped_id = renumbermap[obj_id]
mesh_binary, mesh_bounds = self._create_mesh(obj_id)
bounding_boxes[remapped_id] = mesh_bounds.to_list()
meshes[remapped_id] = mesh_binary
if self.options['sharded']:
self._upload_batch(meshes, self._bounds)
else:
self._upload_individuals(meshes,
self.options['generate_manifests'])
if self.options['spatial_index']:
self._upload_spatial_index(self._bounds, bounding_boxes)
def _upload_batch(self, meshes, bbox):
with SimpleStorage(self.layer_path,
progress=self.options['progress']) as stor:
# Create mesh batch for postprocessing later
stor.put_file(
file_path="{}/{}.frags".format(self._mesh_dir,
bbox.to_filename()),
content=pickle.dumps(meshes),
compress=self.options['compress'],
content_type="application/python-pickle",
cache_control=False,
)
def _upload_individuals(self, mesh_binaries, generate_manifests):
with Storage(self.layer_path) as storage:
for segid, mesh_binary in mesh_binaries.items():
storage.put_file(file_path='{}/{}:{}:{}'.format(
self._mesh_dir, segid, self.options['lod'],
self._bounds.to_filename()),
content=mesh_binary,
compress=self._encoding_to_compression_dict[
self.options['encoding']],
cache_control=self.options['cache_control'])
if generate_manifests:
fragments = []
fragments.append('{}:{}:{}'.format(
segid, self.options['lod'],
self._bounds.to_filename()))
storage.put_file(
file_path='{}/{}:{}'.format(self._mesh_dir, segid,
self.options['lod']),
content=json.dumps({"fragments": fragments}),
content_type='application/json',
cache_control=self.options['cache_control'])
def _create_mesh(self, obj_id):
mesh = self._mesher.get_mesh(
obj_id,
simplification_factor=self.options['simplification_factor'],
max_simplification_error=self.options['max_simplification_error'])
self._mesher.erase(obj_id)
resolution = self._volume.resolution
offset = self._bounds.minpt - self.options['low_padding']
mesh.vertices[:] += offset * resolution
mesh_bounds = Bbox(np.amin(mesh.vertices, axis=0),
np.amax(mesh.vertices, axis=0))
if self.options['encoding'] == 'draco':
mesh_binary = DracoPy.encode_mesh_to_buffer(
mesh.vertices.flatten('C'), mesh.faces.flatten('C'),
**self.draco_encoding_settings)
elif self.options['encoding'] == 'precomputed':
mesh_binary = mesh.to_precomputed()
return mesh_binary, mesh_bounds
def _upload_spatial_index(self, bbox, mesh_bboxes):
with SimpleStorage(self.layer_path,
progress=self.options['progress']) as stor:
stor.put_file(
file_path="{}/{}.spatial".format(self._mesh_dir,
bbox.to_filename()),
content=jsonify(mesh_bboxes).encode('utf8'),
compress=self.options['compress'],
content_type="application/json",
cache_control=False,
)
class MeshTask(RegisteredTask):
def __init__(self, shape, offset, layer_path, **kwargs):
super(MeshTask, self).__init__(shape, offset, layer_path, **kwargs)
self.shape = Vec(*shape)
self.offset = Vec(*offset)
self.layer_path = layer_path
self.options = {
'lod': kwargs.get('lod', 0),
'mip': kwargs.get('mip', 0),
'simplification_factor': kwargs.get('simplification_factor', 100),
'max_simplification_error': kwargs.get('max_simplification_error', 40),
'mesh_dir': kwargs.get('mesh_dir', None),
'remap_table': kwargs.get('remap_table', None),
'generate_manifests': kwargs.get('generate_manifests', False),
'low_padding': kwargs.get('low_padding', 0),
'high_padding': kwargs.get('high_padding', 1),
'parallel_download': kwargs.get('parallel_download', 1),
'cache_control': kwargs.get('cache_control', None)
}
def execute(self):
self._volume = CloudVolume(
self.layer_path, self.options['mip'], bounded=False,
parallel=self.options['parallel_download'])
self._bounds = Bbox(self.offset, self.shape + self.offset)
self._bounds = Bbox.clamp(self._bounds, self._volume.bounds)
self._mesher = Mesher(self._volume.resolution)
# Marching cubes loves its 1vx overlaps.
# This avoids lines appearing between
# adjacent chunks.
data_bounds = self._bounds.clone()
data_bounds.minpt -= self.options['low_padding']
data_bounds.maxpt += self.options['high_padding']
self._mesh_dir = None
if self.options['mesh_dir'] is not None:
self._mesh_dir = self.options['mesh_dir']
elif 'mesh' in self._volume.info:
self._mesh_dir = self._volume.info['mesh']
if not self._mesh_dir:
raise ValueError("The mesh destination is not present in the info file.")
# chunk_position includes the overlap specified by low_padding/high_padding
self._data = self._volume[data_bounds.to_slices()]
self._remap()
self._compute_meshes()
def _remap(self):
if self.options['remap_table'] is not None:
actual_remap = {
int(k): int(v) for k, v in self.options['remap_table'].items()
}
self._remap_list = [0] + list(actual_remap.values())
enumerated_remap = {int(v): i for i, v in enumerate(self._remap_list)}
do_remap = lambda x: enumerated_remap[actual_remap.get(x, 0)]
self._data = np.vectorize(do_remap)(self._data)
def _compute_meshes(self):
with Storage(self.layer_path) as storage:
data = self._data[:, :, :, 0].T
self._mesher.mesh(data)
for obj_id in self._mesher.ids():
if self.options['remap_table'] is None:
remapped_id = obj_id
else:
remapped_id = self._remap_list[obj_id]
storage.put_file(
file_path='{}/{}:{}:{}'.format(
self._mesh_dir, remapped_id, self.options['lod'],
self._bounds.to_filename()
),
content=self._create_mesh(obj_id),
compress=True,
cache_control=self.options['cache_control']
)
if self.options['generate_manifests']:
fragments = []
fragments.append('{}:{}:{}'.format(remapped_id, self.options['lod'],
self._bounds.to_filename()))
storage.put_file(
file_path='{}/{}:{}'.format(
self._mesh_dir, remapped_id, self.options['lod']),
content=json.dumps({"fragments": fragments}),
content_type='application/json',
cache_control=self.options['cache_control']
)
def _create_mesh(self, obj_id):
mesh = self._mesher.get_mesh(
obj_id,
simplification_factor=self.options['simplification_factor'],
max_simplification_error=self.options['max_simplification_error']
)
vertices = self._update_vertices(
np.array(mesh['points'], dtype=np.float32))
vertex_index_format = [
np.uint32(len(vertices) / 3), # Number of vertices (3 coordinates)
vertices,
np.array(mesh['faces'], dtype=np.uint32)
]
return b''.join([array.tobytes() for array in vertex_index_format])
def _update_vertices(self, points):
# zi_lib meshing multiplies vertices by 2.0 to avoid working with floats,
# but we need to recover the exact position for display
# Note: points are already multiplied by resolution, but missing the offset
points /= 2.0
resolution = self._volume.resolution
xmin, ymin, zmin = self._bounds.minpt - self.options['low_padding']
points[0::3] = points[0::3] + xmin * resolution.x
points[1::3] = points[1::3] + ymin * resolution.y
points[2::3] = points[2::3] + zmin * resolution.z
return points
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)
def child_upload_process(meta,
cache,
img_shape,
offset,
mip,
compress,
cdn_cache,
progress,
location,
location_bbox,
location_order,
delete_black_uploads,
background_color,
green,
chunk_ranges,
compress_level=None,
secrets=None):
global fs_lock
reset_connection_pools()
shared_shape = img_shape
if location_bbox:
shared_shape = list(location_bbox.size3()) + [meta.num_channels]
array_like, renderbuffer = shm.ndarray(shape=shared_shape,
dtype=meta.dtype,
location=location,
order=location_order,
lock=fs_lock,
readonly=True)
def updatefn():
if progress:
# This is not good programming practice, but
# I could not find a clean way to do this that
# did not result in warnings about leaked semaphores.
# progress_queue is created in common.py:initialize_progress_queue
# as a global for this module.
progress_queue.put(1)
try:
if location_bbox:
cutout_bbox = Bbox(offset, offset + img_shape[:3])
delta_box = cutout_bbox.clone() - location_bbox.minpt
renderbuffer = renderbuffer[delta_box.to_slices()]
return threaded_upload_chunks(
meta,
cache,
None,
renderbuffer,
mip,
chunk_ranges,
compress=compress,
cdn_cache=cdn_cache,
progress=updatefn,
delete_black_uploads=delete_black_uploads,
background_color=background_color,
green=green,
compress_level=compress_level,
secrets=secrets,
)
finally:
array_like.close()
def create_contrast_normalization_tasks(task_queue,
src_path,
dest_path,
shape=None,
mip=0,
clip_fraction=0.01,
fill_missing=False,
translate=(0, 0, 0)):
srcvol = CloudVolume(src_path, mip=mip)
try:
dvol = CloudVolume(dest_path, mip=mip)
except Exception: # no info file
info = copy.deepcopy(srcvol.info)
dvol = CloudVolume(dest_path, mip=mip, info=info)
dvol.info['scales'] = dvol.info['scales'][:mip + 1]
dvol.commit_info()
if shape == None:
shape = Bbox((0, 0, 0), (2048, 2048, 64))
shape = shape.shrink_to_chunk_size(dvol.underlying).size3()
shape = Vec(*shape)
create_downsample_scales(dest_path,
mip=mip,
ds_shape=shape,
preserve_chunk_size=True)
dvol.refresh_info()
bounds = srcvol.bounds.clone()
for startpt in tqdm(xyzrange(bounds.minpt, bounds.maxpt, shape),
desc="Inserting Contrast Normalization Tasks"):
task_shape = min2(shape.clone(), srcvol.bounds.maxpt - startpt)
task = ContrastNormalizationTask(
src_path=src_path,
dest_path=dest_path,
shape=task_shape,
offset=startpt.clone(),
clip_fraction=clip_fraction,
mip=mip,
fill_missing=fill_missing,
translate=translate,
)
task_queue.insert(task)
task_queue.wait('Uploading Contrast Normalization Tasks')
dvol.provenance.processing.append({
'method': {
'task': 'ContrastNormalizationTask',
'src_path': src_path,
'dest_path': dest_path,
'shape': Vec(*shape).tolist(),
'clip_fraction': clip_fraction,
'mip': mip,
'translate': Vec(*translate).tolist(),
},
'by': USER_EMAIL,
'date': strftime('%Y-%m-%d %H:%M %Z'),
})
dvol.commit_provenance()
