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()