def read_write_h5(chunk): assert isinstance(chunk, np.ndarray) file_name = 'test.h5' if os.path.exists(file_name): os.remove(file_name) chunk.to_h5(file_name) chunk2 = Chunk.from_h5(file_name) assert np.alltrue(chunk == chunk2) assert chunk.global_offset == chunk2.global_offset os.remove(file_name)
def array_to_chunk(arr: Union[np.ndarray, Chunk], voxel_offset: Cartesian, voxel_size: Cartesian, shape: tuple): if isinstance(arr, np.ndarray): # in case the plugin did some symmetric cropping offset = tuple( vo + (ins - outs) // 2 for vo, ins, outs in zip(voxel_offset, shape[-3:], arr.shape[-3:])) return Chunk(arr, voxel_offset=offset, voxel_size=voxel_size) elif isinstance(arr, Chunk): return arr else: raise TypeError(f'only support ndarray and Chunk, but got {type(arr)}')
def read_write_h5(chunk): assert isinstance(chunk, Chunk) file_name = 'test.h5' if os.path.exists(file_name): os.remove(file_name) chunk.to_h5(file_name, chunk_size=None) chunk2 = Chunk.from_h5(file_name) assert chunk == chunk2 assert chunk.voxel_offset == chunk2.voxel_offset os.remove(file_name)
def test_aligned_patch_num(): print('\ntest block inference with patch number...') patch_size = (32, 256, 256) patch_overlap = (4, 64, 64) patch_num = (2, 2, 2) num_output_channels = 2 dtype = 'float16' image = np.random.randint(1, 255, size=(28 * 2 + 4, (256 - 64) * 2 + 64, (256 - 64) * 2 + 64), dtype=np.uint8) image = Chunk(image) with Inferencer(None, None, patch_size, output_patch_overlap=patch_overlap, num_output_channels=num_output_channels, patch_num=patch_num, framework='identity', dtype=dtype, batch_size=5, mask_output_chunk=False) as block_inference_engine: output = block_inference_engine(image) # only use the first channel to check correctness output = output[0, :, :, :] # we need to crop the patch overlap since the values were changed image = image[patch_overlap[0]:-patch_overlap[0], patch_overlap[1]:-patch_overlap[1], patch_overlap[2]:-patch_overlap[2]] image = image.astype(dtype) / 255 print('maximum difference: ', np.max(image - output)) # some of the image voxel is 0, the test can only work with rtol=1 np.testing.assert_allclose(image, output, rtol=1e-3, atol=1e-3)
def test_aligned_input_chunk_with_croped_patch(): print('\ntest block inference engine...') input_patch_size = (20, 256, 256) output_patch_size = (16, 192, 192) output_patch_crop_margin = (2, 32, 32) output_patch_overlap = (2, 32, 32) input_patch_overlap = (6, 96, 96) input_patch_stride = (14, 160, 160) input_size = (2 * 14 + 6, 2 * 160 + 96, 2 * 160 + 96) num_output_channels = 1 image = np.random.randint(1, 255, size=input_size, dtype=np.uint8) # make sure that it works with arbitrary global offset image = Chunk(image, global_offset=(123, 345, 567)) with Inferencer(None, None, input_patch_size, output_patch_size=output_patch_size, output_patch_overlap=output_patch_overlap, patch_num=(2, 2, 2), num_output_channels=num_output_channels, framework='identity', batch_size=5, mask_output_chunk=False) as chunk_inferencer: output = chunk_inferencer(image) # only use the first channel to check correctness output = output[0, :, :, :] assert output.ndim == 3 # we need to crop the patch overlap since the values were changed image = image[4:-4, 64:-64, 64:-64] image = image.astype(np.float32) / 255 print('maximum difference: ', np.max(image - output)) # some of the image voxel is 0, the test can only work with rtol=1 np.testing.assert_allclose(image, output, rtol=1e-5, atol=1e-5)
def test_aligned_input_chunk(): print('\ntest block inference engine...') patch_size = (32, 256, 256) patch_overlap = (4, 64, 64) num_output_channels = 2 block_inference_engine = Engine(None, None, patch_size=patch_size, patch_overlap=patch_overlap, num_output_channels=num_output_channels, framework='identity', batch_size=5, mask_output_chunk=False) image = np.random.randint(1, 255, size=(28 * 2 + 4, (256 - 64) * 2 + 64, (256 - 64) * 2 + 64), dtype=np.uint8) image = Chunk(image) output = block_inference_engine(image) # only use the first channel to check correctness output = output[0, :, :, :] output = np.reshape(output, image.shape) # we need to crop the patch overlap since the values were changed image = image[patch_overlap[0]:-patch_overlap[0], patch_overlap[1]:-patch_overlap[1], patch_overlap[2]:-patch_overlap[2]] output = output[patch_overlap[0]:-patch_overlap[0], patch_overlap[1]:-patch_overlap[1], patch_overlap[2]:-patch_overlap[2]] image = image.astype(np.float32) / 255 print('maximum difference: ', np.max(image - output)) # some of the image voxel is 0, the test can only work with rtol=1 np.testing.assert_allclose(image, output, rtol=1e-5, atol=1e-5)
def _construct_output_chunk_mask(self, input_chunk): if not self.mask_output_chunk: return if self.verbose: print('creating output chunk mask...') if self.output_chunk_mask is None or not np.array_equal( input_chunk.shape, self.output_chunk_mask.shape): # To-Do: clean up extra temporal files if we created # multiple mmap files #output_mask_mmap_file = mktemp(suffix='.dat') ## the memory map is initialized with 0 in default #output_mask_array = np.memmap(output_mask_mmap_file, # dtype=self.dtype, mode='w+', # shape=self.output_size) output_mask_array = np.zeros(self.output_size, self.dtype) else: output_chunk_mask_array = self.output_chunk_mask.array output_chunk_mask_array.fill(0) output_global_offset = tuple( io + ocso for io, ocso in zip(input_chunk.global_offset, self.output_offset)) self.output_chunk_mask = Chunk(output_mask_array, global_offset=output_global_offset) assert len(self.patch_slices_list) > 0 for _, output_patch_slice in self.patch_slices_list: # accumulate weights using the patch mask in RAM patch_global_offset = tuple(s.start for s in output_patch_slice) patch_mask = Chunk(self.patch_inferencer.output_patch_mask_numpy, global_offset=patch_global_offset) self.output_chunk_mask.blend(patch_mask) # normalize weight, so accumulated inference result multiplies # this mask will result in 1 self.output_chunk_mask.array = 1.0 / self.output_chunk_mask.array
def test_mesh(): ck = Chunk.create(dtype=np.float32) # segment it with threshold ck = ck < 0.5 ck = ck.astype(np.uint32) tempdir = tempfile.mkdtemp() #mesher = MeshOperator('file://' + tempdir, 'precomputed', manifest=True) mesher = MeshOperator('file://' + tempdir, 'obj') mesher(ck) mesher = MeshOperator('file://' + tempdir, 'ply') mesher(ck)
def test_normalize_section_contrast(): DIR = os.path.join(os.path.dirname(__file__), '../data/levels/1/') #histogram = np.random.randint(10000, size=256, dtype=np.uint32) image = np.arange(256, dtype=np.uint8).reshape(1, 16, 16) image = Chunk(image, voxel_offset=(26774, 234, 456)) levels_path = 'precomputed://file://' + DIR operator = NormalizeSectionContrastOperator(levels_path) normalized = operator(image) normalized_filename = os.path.join(DIR, '../normalized.npz') #np.savez( normalized_filename , normalized=normalized) groundtruth = np.load(normalized_filename)['normalized'] np.testing.assert_array_equal(normalized, groundtruth)
def read_write_tif(chunk): """We'll lost global offset information using tif format!""" assert isinstance(chunk, Chunk) file_name = 'test.tif' if os.path.exists(file_name): os.remove(file_name) chunk.to_tif(file_name) chunk2 = Chunk.from_tif(file_name) # we can not preserve the global offset here # so chunk2's global offset will all be 0 np.testing.assert_array_equal(chunk.array, chunk2.array) os.remove(file_name)
def normalize_section_shang(image: np.ndarray, nominalmin: float, nominalmax: float, clipvalues: bool): """ Parameters ------------ image: image volume. nominalmin: min threshold nominalmax: max threshold clipvalues: clip values or not. """ assert nominalmin < nominalmax assert image.ndim == 3 global_offset = image.global_offset originaltype = image.dtype arr = image.astype(np.float32) # number of bits per voxel nbits = np.dtype(originaltype).itemsize * 8 default_nominalmax = float(2**nbits - 1) nominalmin = nominalmin if nominalmin is not None else 0.0 nominalmax = nominalmax if nominalmax is not None else default_nominalmax normalization = 'fill' # stack/chunk-wise normalization first if necessary (for blank slices within a valid stack) #arr = normalize(arr, normalization, target_scale = [-1,1], min_max_invalid = [True]*2, make_copy=False) # slice-wise normalization # Note in chunkflow the first dim is z/slice for ii in range(arr.shape[0]): normalize(arr[ii, :, :], normalization, target_scale=[nominalmin, nominalmax], min_max_invalid=[True, True], do_clipping=clipvalues, make_copy=False) # cast to original data type if necessary #arr = np.round(arr) #arr = arr.astype(originaltype) return Chunk(arr, global_offset=global_offset)
def test_maskout(): mask = np.ones((8, 16, 4), dtype=np.uint32) mask[:4, :8, :2] = 0 mask = Chunk(mask, voxel_offset=(-2, -3, -4), voxel_size=(2, 4, 8)) image = create_chunk(size=(16, 64, 32), voxel_size=(1, 1, 1), dtype=np.uint8) mask.maskout(image) np.testing.assert_array_equal(image.array[:8, :32, :16], 0) affs = create_chunk(size=(3, 16, 64, 32), voxel_size=(1, 1, 1), dtype=np.float32) mask.maskout(affs) np.testing.assert_array_equal(affs.array[:, :8, :32, :16], 0)
def __call__(self, chunk): # this is a image, not affinitymap assert chunk.ndim == 3 # we have to use np.dtype function to make it match # https://stackoverflow.com/questions/26921836/correct-way-to-test-for-numpy-dtype assert chunk.dtype is np.dtype(np.uint8) for z in range(chunk.bbox.minpt[-3], chunk.bbox.maxpt[-3]): lookup_table = self.fetch_lookup_table(z) slices = (slice(z, z + 1), *chunk.slices[-2:]) image = chunk.cutout(slices) image_global_offset = image.global_offset image = lookup_table[image] image = Chunk(image, global_offset=image_global_offset) chunk.save(image) return chunk
def create_chunk(size:tuple = (7, 8, 9), voxel_offset=(-2, -3, -4), dtype=np.float32, voxel_size=(1,1,1)): # make the tests consistent in multiple runs np.random.seed(1) if len(size)==4 and len(voxel_offset)==3: voxel_offset = tuple((0, *voxel_offset)) if np.issubdtype(dtype, np.float32): arr = (np.random.rand(*size)).astype(dtype) else: arr = (np.random.rand(*size) * (np.iinfo(dtype).max)).astype(dtype) chunk = Chunk(arr, voxel_offset=voxel_offset, voxel_size=voxel_size) # test chunk properties chunk.voxel_stop return chunk
def _append_probability_map_layer( self, viewer_state: ng.viewer_state.ViewerState, chunk_name: str, chunk: Chunk): if chunk.dtype == np.dtype('<f4') or chunk.dtype == np.dtype( 'float16'): chunk = chunk.astype(np.float32) voxel_size = self._get_voxel_size(chunk) # chunk = np.ascontiguousarray(chunk) if chunk.shape[0] == 1: shader = """void main() { emitGrayscale(toNormalized(getDataValue(0))); } """ elif chunk.shape[0] == 2: shader = """void main() { emitRGB(vec3(toNormalized(getDataValue(0)), toNormalized(getDataValue(1)), 0.)); } """ else: shader = """void main() { emitRGB(vec3(toNormalized(getDataValue(0)), toNormalized(getDataValue(1)), toNormalized(getDataValue(2)))); } """ dimensions = ng.CoordinateSpace(scales=(1, ) + voxel_size, units=['', 'nm', 'nm', 'nm'], names=['c^', 'z', 'y', 'x']) viewer_state.layers.append( name=chunk_name, layer=ng.LocalVolume( data=chunk.array, dimensions=dimensions, # offset is in nm, not voxels voxel_offset=(0, *chunk.voxel_offset), ), shader=shader)
def _get_output_buffer(self, input_chunk): output_buffer_size = ( self.patch_inferencer.num_output_channels, ) + self.output_size #if self.mask_myelin_threshold is None: # a random temporal file. will be removed later. #output_buffer_mmap_file = mktemp(suffix='.dat') ## the memory map is initialized with 0 in default #output_buffer_array = np.memmap(output_buffer_mmap_file, # dtype=self.dtype, mode='w+', # shape=output_buffer_size) ##else: # # when we use myelin mask, the masking computation will create a full array in RAM! # # and it will duplicate the array! thus, we should use normal array in this case. output_buffer_array = np.zeros(output_buffer_size, dtype=self.dtype) output_global_offset = tuple( io + ocso for io, ocso in zip(input_chunk.global_offset, self.output_offset)) output_buffer = Chunk(output_buffer_array, global_offset=(0, ) + output_global_offset) assert output_buffer == 0 return output_buffer
def test_inference(): # compute parameters input_size = (18, 224, 224) patch_overlap = (2, 32, 32) patch_size = (10, 128, 128) image = Chunk.create(size=input_size, dtype='uint8') inference_operator = InferenceOperator(None, None, patch_size=patch_size, output_key='affinity', num_output_channels=3, patch_overlap=patch_overlap, framework='identity') output = inference_operator(image) # ignore the cropping region output = output[0, 2:-2, 32:-32, 32:-32] image = image[2:-2, 32:-32, 32:-32] output = output * 255 output = output.astype(np.uint8) assert np.alltrue(np.isclose(image, output, atol=1))
def test_read_write_aff(self): print('test affinitymap io...') arr = np.random.rand(3, 8, 16, 16).astype(np.float32) chunk = Chunk(arr, voxel_offset=Cartesian(1, 2, 3)) read_write_h5(chunk)
class Inferencer(object): """ Inferencer convnet inference for a whole chunk. if the patches align with the input chunk size, we do not need chunk mask. if the patches do not align, we'll create a chunk mask to make sure that the output have the same size with input. The output buffer is smaller than the input chunk size, and the cropped margin area is not allocated. This will save about 20% of memory usage. what's more, the output buffer is formated as memory map and was mapped to disk. This is particularly useful for multiple channel output with large chunk size. """ def __init__(self, convnet_model: str, convnet_weight_path: str, input_patch_size: Union[tuple, list], output_patch_size: Union[tuple, list] = None, patch_num: Union[tuple, list] = None, num_output_channels: int = 3, output_patch_overlap: Union[tuple, list] = (4, 64, 64), output_crop_margin: Union[tuple, list] = None, dtype='float32', framework: str = 'identity', batch_size: int = 1, bump: str = 'wu', input_size: tuple = None, mask_output_chunk: bool = False, mask_myelin_threshold=None, dry_run: bool = False, verbose: int = 1): assert input_size is None or patch_num is None if output_patch_size is None: output_patch_size = input_patch_size self.input_patch_size = input_patch_size self.output_patch_size = output_patch_size self.output_patch_overlap = output_patch_overlap self.patch_num = patch_num self.batch_size = batch_size self.input_size = input_size if mask_output_chunk: # the chunk mask will handle the boundaries self.output_crop_margin = (0, 0, 0) else: if output_crop_margin is None: self.output_crop_margin = self.output_patch_overlap else: self.output_crop_margin = output_crop_margin # we should always crop more than the patch overlap # since the overlap region is reweighted by patch mask # To-Do: equal should also be OK assert np.alltrue([ v <= m for v, m in zip(self.output_patch_overlap, self.output_crop_margin) ]) self.output_patch_crop_margin = tuple( (ips - ops) // 2 for ips, ops in zip(input_patch_size, output_patch_size)) self.output_offset = tuple(opcm + ocm for opcm, ocm in zip( self.output_patch_crop_margin, self.output_crop_margin)) self.output_patch_stride = tuple( s - o for s, o in zip(output_patch_size, output_patch_overlap)) self.input_patch_overlap = tuple(opcm * 2 + oo for opcm, oo in zip( self.output_patch_crop_margin, self.output_patch_overlap)) self.input_patch_stride = tuple( ps - po for ps, po in zip(input_patch_size, self.input_patch_overlap)) # no chunk wise mask, the patches should be aligned inside chunk if not mask_output_chunk: assert (self.input_size is not None) or (self.patch_num is not None) if patch_num is None: assert input_size is not None self.patch_num = tuple((isz - o) // s for isz, o, s in zip( self.input_size, self.input_patch_overlap, self.input_patch_stride)) if self.input_size is None: assert self.patch_num is not None self.input_size = tuple(pst * pn + po for pst, pn, po in zip( self.input_patch_stride, self.patch_num, self.input_patch_overlap)) self.output_size = tuple( pst * pn + po - 2 * ocm for pst, pn, po, ocm in zip( self.output_patch_stride, self.patch_num, self.output_patch_overlap, self.output_crop_margin)) else: # we can handle arbitrary input and output size self.input_size = None self.output_size = None self.num_output_channels = num_output_channels self.verbose = verbose self.mask_output_chunk = mask_output_chunk self.output_chunk_mask = None self.dtype = dtype self.mask_myelin_threshold = mask_myelin_threshold self.dry_run = dry_run # allocate a buffer to avoid redundant memory allocation self.input_patch_buffer = np.zeros((batch_size, 1, *input_patch_size), dtype=dtype) self.patch_slices_list = [] if isinstance(convnet_model, str): convnet_model = os.path.expanduser(convnet_model) if isinstance(convnet_weight_path, str): convnet_weight_path = os.path.expanduser(convnet_weight_path) self._prepare_patch_inferencer(framework, convnet_model, convnet_weight_path, bump) @property def compute_device(self): return self.patch_inferencer.compute_device def __enter__(self): return self def __exit__(self, exception_type, exception_value, traceback): pass def _update_parameters_for_input_chunk(self, input_chunk): """ if the input size is consistent with old one, reuse the patch offset list and output chunk mask. Otherwise, recompute them. """ if np.array_equal(self.input_size, input_chunk.shape): print('reusing output chunk mask.') assert self.patch_slices_list is not None else: if self.input_size is not None: warn('the input size has changed, using new intput size.') self.input_size = input_chunk.shape if not self.mask_output_chunk: self._check_alignment() self.output_size = tuple( isz - 2 * ocso for isz, ocso in zip(self.input_size, self.output_offset)) self.output_patch_stride = tuple( s - o for s, o in zip(self.output_patch_size, self.output_patch_overlap)) self._construct_patch_slices_list(input_chunk.global_offset) self._construct_output_chunk_mask(input_chunk) def _prepare_patch_inferencer(self, framework, convnet_model, convnet_weight_path, bump): # prepare for inference if framework == 'pznet': from .patch.pznet import PZNet as PatchInferencer elif framework == 'pytorch': # pytorch will not output consistent result if we use batch size > 1 # https://discuss.pytorch.org/t/solved-inconsistent-results-during-test-using-different-batch-size/2265 assert self.batch_size == 1 from .patch.pytorch import PyTorch as PatchInferencer # currently, we do not support pytorch backend with different # input and output patch size and overlap. elif framework == 'pytorch-multitask': # currently only this type of task support mask in device from .patch.pytorch_multitask import PyTorchMultitask as PatchInferencer elif framework == 'identity': from .patch.identity import Identity as PatchInferencer elif framework == 'general': from .patch.general import General as PatchInferencer else: raise Exception(f'invalid inference backend: {self.framework}') self.patch_inferencer = PatchInferencer( convnet_model, convnet_weight_path, input_patch_size=self.input_patch_size, output_patch_size=self.output_patch_size, output_patch_overlap=self.output_patch_overlap, num_output_channels=self.num_output_channels, dtype=self.dtype, bump=bump) def _check_alignment(self): is_align = tuple((i - o) % s == 0 for i, s, o in zip( self.input_size, self.patch_inferencer.input_patch_stride, self.patch_inferencer.input_patch_overlap)) # all axis should be aligned # the patches should aligned with input size in case # we will not mask the output chunk assert np.all(is_align) if self.verbose: print('great! patches aligns in chunk.') def _construct_patch_slices_list(self, input_chunk_offset): """ create the normalization mask and patch bounding box list """ self.patch_slices_list = [] # the step is the stride, so the end of aligned patch is # input_size - patch_overlap input_patch_size = self.input_patch_size output_patch_size = self.output_patch_size input_patch_overlap = self.input_patch_overlap input_patch_stride = self.input_patch_stride print('Construct patch slices list...') for iz in range(0, self.input_size[0] - input_patch_overlap[0], input_patch_stride[0]): if iz + input_patch_size[0] > self.input_size[0]: iz = self.input_size[0] - input_patch_size[0] assert iz >= 0 iz += input_chunk_offset[-3] oz = iz + self.output_patch_crop_margin[0] for iy in range(0, self.input_size[1] - input_patch_overlap[1], input_patch_stride[1]): if iy + input_patch_size[1] > self.input_size[1]: iy = self.input_size[1] - input_patch_size[1] assert iy >= 0 iy += input_chunk_offset[-2] oy = iy + self.output_patch_crop_margin[1] for ix in range(0, self.input_size[2] - input_patch_overlap[2], input_patch_stride[2]): if ix + input_patch_size[2] > self.input_size[2]: ix = self.input_size[2] - input_patch_size[2] assert ix >= 0 ix += input_chunk_offset[-1] ox = ix + self.output_patch_crop_margin[2] input_patch_slice = (slice(iz, iz + input_patch_size[0]), slice(iy, iy + input_patch_size[1]), slice(ix, ix + input_patch_size[2])) output_patch_slice = (slice(oz, oz + output_patch_size[0]), slice(oy, oy + output_patch_size[1]), slice(ox, ox + output_patch_size[2])) self.patch_slices_list.append( (input_patch_slice, output_patch_slice)) def _construct_output_chunk_mask(self, input_chunk): if not self.mask_output_chunk: return if self.verbose: print('creating output chunk mask...') if self.output_chunk_mask is None or not np.array_equal( input_chunk.shape, self.output_chunk_mask.shape): # To-Do: clean up extra temporal files if we created # multiple mmap files #output_mask_mmap_file = mktemp(suffix='.dat') ## the memory map is initialized with 0 in default #output_mask_array = np.memmap(output_mask_mmap_file, # dtype=self.dtype, mode='w+', # shape=self.output_size) output_mask_array = np.zeros(self.output_size, self.dtype) else: output_chunk_mask_array = self.output_chunk_mask.array output_chunk_mask_array.fill(0) output_global_offset = tuple( io + ocso for io, ocso in zip(input_chunk.global_offset, self.output_offset)) self.output_chunk_mask = Chunk(output_mask_array, global_offset=output_global_offset) assert len(self.patch_slices_list) > 0 for _, output_patch_slice in self.patch_slices_list: # accumulate weights using the patch mask in RAM patch_global_offset = tuple(s.start for s in output_patch_slice) patch_mask = Chunk(self.patch_inferencer.output_patch_mask_numpy, global_offset=patch_global_offset) self.output_chunk_mask.blend(patch_mask) # normalize weight, so accumulated inference result multiplies # this mask will result in 1 self.output_chunk_mask.array = 1.0 / self.output_chunk_mask.array def _get_output_buffer(self, input_chunk): output_buffer_size = ( self.patch_inferencer.num_output_channels, ) + self.output_size #if self.mask_myelin_threshold is None: # a random temporal file. will be removed later. #output_buffer_mmap_file = mktemp(suffix='.dat') ## the memory map is initialized with 0 in default #output_buffer_array = np.memmap(output_buffer_mmap_file, # dtype=self.dtype, mode='w+', # shape=output_buffer_size) ##else: # # when we use myelin mask, the masking computation will create a full array in RAM! # # and it will duplicate the array! thus, we should use normal array in this case. output_buffer_array = np.zeros(output_buffer_size, dtype=self.dtype) output_global_offset = tuple( io + ocso for io, ocso in zip(input_chunk.global_offset, self.output_offset)) output_buffer = Chunk(output_buffer_array, global_offset=(0, ) + output_global_offset) assert output_buffer == 0 return output_buffer def __call__(self, input_chunk: np.ndarray): """ args: input_chunk (Chunk): input chunk with global offset """ assert isinstance(input_chunk, Chunk) self._update_parameters_for_input_chunk(input_chunk) output_buffer = self._get_output_buffer(input_chunk) if not self.mask_output_chunk: self._check_alignment() if self.dry_run: print('dry run, return a special artifical chunk.') size = output_buffer.shape if self.mask_myelin_threshold: # eleminate the myelin channel size = (size[0] - 1, *size[1:]) return Chunk.create(size=size, dtype=output_buffer.dtype, voxel_offset=output_buffer.global_offset) if input_chunk == 0: print('input is all zero, return zero buffer directly') if self.mask_myelin_threshold: assert output_buffer.shape[0] == 4 return output_buffer[:-1, ...] else: return output_buffer if np.issubdtype(input_chunk.dtype, np.integer): # normalize to 0-1 value range dtype_max = np.iinfo(input_chunk.dtype).max input_chunk = input_chunk.astype(self.dtype) / dtype_max if self.verbose: chunk_time_start = time.time() # iterate the offset list for i in tqdm(range(0, len(self.patch_slices_list), self.batch_size), disable=not self.verbose, desc='ConvNet inference for patches: '): if self.verbose: start = time.time() batch_slices = self.patch_slices_list[i:i + self.batch_size] for batch_idx, slices in enumerate(batch_slices): self.input_patch_buffer[batch_idx, 0, :, :, :] = input_chunk.cutout( slices[0]).array if self.verbose > 1: end = time.time() print('prepare %d input patches takes %3f sec' % (self.batch_size, end - start)) start = end # the input and output patch is a 5d numpy array with # datatype of float32, the dimensions are batch/channel/z/y/x. # the input image should be normalized to [0,1] output_patch = self.patch_inferencer(self.input_patch_buffer) if self.verbose > 1: assert output_patch.ndim == 5 end = time.time() print('run inference for %d patch takes %3f sec' % (self.batch_size, end - start)) start = end for batch_idx, slices in enumerate(batch_slices): # only use the required number of channels # the remaining channels are dropped # the slices[0] is for input patch slice # the slices[1] is for output patch slice offset = (0, ) + tuple(s.start for s in slices[1]) output_chunk = Chunk(output_patch[batch_idx, :, :, :, :], global_offset=offset) ## save some patch for debug #bbox = output_chunk.bbox #if bbox.minpt[-1] < 94066 and bbox.maxpt[-1] > 94066 and \ # bbox.minpt[-2]<81545 and bbox.maxpt[-2]>81545 and \ # bbox.minpt[-3]<17298 and bbox.maxpt[-3]>17298: # print('save patch: ', output_chunk.bbox) # breakpoint() # output_chunk.to_tif() # #input_chunk.cutout(slices[0]).to_tif() output_buffer.blend(output_chunk) if self.verbose > 1: end = time.time() print('blend patch takes %3f sec' % (end - start)) if self.verbose: print("Inference of whole chunk takes %3f sec" % (time.time() - chunk_time_start)) if self.mask_output_chunk: output_buffer *= self.output_chunk_mask # theoretically, all the value of output_buffer should not be greater than 1 # we use a slightly higher value here to accomondate numerical precision issue np.testing.assert_array_less( output_buffer, 1.0001, err_msg='output buffer should not be greater than 1') if self.mask_myelin_threshold: # currently only for masking out affinity map assert output_buffer.shape[0] == 4 output_chunk = output_buffer.mask_using_last_channel( threshold=self.mask_myelin_threshold) # currently neuroglancer only support float32, not float16 if output_chunk.dtype == np.dtype('float16'): output_chunk = output_chunk.astype('float32') return output_chunk else: return output_buffer
def _create_output_buffer(self, input_chunk): output_buffer = np.zeros( (self.num_output_channels, ) + input_chunk.shape, dtype=np.float32) return Chunk(output_buffer, global_offset=(0, ) + input_chunk.global_offset)
def test_create_from_bounding_box(self): bbox = BoundingBox.from_delta(self.voxel_offset, self.size) bbox = BoundingBox.from_bbox(bbox) Chunk.from_bbox(bbox)
def setUp(self): self.size = (7, 8, 9) self.voxel_offset = (-1, -1, -1) arr = np.random.rand(*self.size).astype('float32') self.chunk = Chunk(arr, self.voxel_offset)
class Test3DChunk(unittest.TestCase): def setUp(self): self.size = (7, 8, 9) self.voxel_offset = (-1, -1, -1) arr = np.random.rand(*self.size).astype('float32') self.chunk = Chunk(arr, self.voxel_offset) #def test_math(self): # self.assertEqual(np.max(self.chunk), np.max(self.chunk.array)) # self.assertEqual(np.min(self.chunk), np.min(self.chunk.array)) def test_create_from_bounding_box(self): bbox = BoundingBox.from_delta(self.voxel_offset, self.size) bbox = BoundingBox.from_bbox(bbox) Chunk.from_bbox(bbox) def test_computation(self): self.chunk /= 127.5 self.chunk -= 1 def test_bbox(self): self.assertEqual(self.chunk.bbox, BoundingBox.from_delta(self.voxel_offset, self.size)) def test_slices(self): self.assertEqual(self.chunk.slices, (slice(-1, 1), slice(-1, 1), slice(-1, 1))) def test_initialization(self): arr = np.ones((3, 3, 3), dtype='float32') chunk = Chunk(arr) (self.assertEqual(o, 0) for o in chunk.voxel_offset) def test_attr(self): """ attribute should not change after numpy operation """ chunk2 = np.sin(self.chunk) self.assertIsInstance(chunk2, Chunk) self.assertEqual(chunk2.voxel_offset, self.chunk.voxel_offset) chunk2 = self.chunk * 255 self.assertIsInstance(chunk2, Chunk) self.assertEqual(chunk2.voxel_offset, self.chunk.voxel_offset) chunk2 = chunk2.astype(np.uint8) self.assertIsInstance(chunk2, Chunk) self.assertEqual(chunk2.voxel_offset, self.chunk.voxel_offset) chunk2 = self.chunk.transpose() print('type of chunk after transpose: {}'.format(type(chunk2))) self.assertIsInstance(chunk2, Chunk) self.assertEqual(chunk2.shape, self.chunk.shape[::-1]) self.assertEqual(chunk2.voxel_offset, self.chunk.voxel_offset[::-1]) self.assertTrue( np.array_equal(chunk2.array, np.transpose(self.chunk.array))) #chunk2 = np.ascontiguousarray(chunk2) #self.assertIsInstance(chunk2, Chunk) #self.assertEqual(chunk2.voxel_offset, self.chunk.voxel_offset) #chunk2 = np.pad(chunk2, ((0,2),(0,2),(0,2)), 'reflect') #self.assertIsInstance(chunk2, Chunk) #self.assertEqual(chunk2.voxel_offset, self.chunk.voxel_offset) def test_normalize(self): print('\ntest chunk normalization...') arr = np.ones((3, 3, 3), dtype='float32') voxel_offset = (-1, -1, -1) chunk = Chunk(arr, voxel_offset=voxel_offset) mask = np.ones((3, 3, 3), dtype='float32') * 0.5 mask = Chunk(mask, voxel_offset=voxel_offset) chunk *= mask self.assertTrue(chunk == mask) def test_item(self): arr = np.ones((1, 3, 3, 3), dtype='float32') chunk = Chunk(arr, (0, -1, -1, -1)) chunk[:, :, :, :] = 0 np.testing.assert_array_equal(chunk, 0) def test_slices(self): arr = np.ones((1, 3, 3, 3), dtype='float32') chunk = Chunk(arr, (-1, -1, -1)) self.assertEqual( chunk.slices, (slice(0, 1), slice(-1, 2), slice(-1, 2), slice(-1, 2))) def test_where(self): pass
def test_slices(self): arr = np.ones((1, 3, 3, 3), dtype='float32') chunk = Chunk(arr, (-1, -1, -1)) self.assertEqual( chunk.slices, (slice(0, 1), slice(-1, 2), slice(-1, 2), slice(-1, 2)))
def test_item(self): arr = np.ones((1, 3, 3, 3), dtype='float32') chunk = Chunk(arr, (0, -1, -1, -1)) chunk[:, :, :, :] = 0 np.testing.assert_array_equal(chunk, 0)
def test_initialization(self): arr = np.ones((3, 3, 3), dtype='float32') chunk = Chunk(arr) (self.assertEqual(o, 0) for o in chunk.voxel_offset)
def setUp(self): self.size = (3, 3, 3) self.global_offset = (-1, -1, -1) arr = np.ones(self.size, dtype='float32') self.chunk = Chunk(arr, self.global_offset)
def test_read_write_image(self): print('test image io...') arr = np.random.randint(0, 256, size=(8, 16, 16), dtype=np.uint8) chunk = Chunk(arr, voxel_offset=Cartesian(1, 2, 3)) read_write_h5(chunk) read_write_tif(chunk)
def __call__(self, input_chunk: np.ndarray): """ args: input_chunk (Chunk): input chunk with global offset """ assert isinstance(input_chunk, Chunk) self._update_parameters_for_input_chunk(input_chunk) output_buffer = self._get_output_buffer(input_chunk) if not self.mask_output_chunk: self._check_alignment() if self.dry_run: print('dry run, return a special artifical chunk.') size = output_buffer.shape if self.mask_myelin_threshold: # eleminate the myelin channel size = (size[0] - 1, *size[1:]) return Chunk.create(size=size, dtype=output_buffer.dtype, voxel_offset=output_buffer.global_offset) if input_chunk == 0: print('input is all zero, return zero buffer directly') if self.mask_myelin_threshold: assert output_buffer.shape[0] == 4 return output_buffer[:-1, ...] else: return output_buffer if np.issubdtype(input_chunk.dtype, np.integer): # normalize to 0-1 value range dtype_max = np.iinfo(input_chunk.dtype).max input_chunk = input_chunk.astype(self.dtype) / dtype_max if self.verbose: chunk_time_start = time.time() # iterate the offset list for i in tqdm(range(0, len(self.patch_slices_list), self.batch_size), disable=not self.verbose, desc='ConvNet inference for patches: '): if self.verbose: start = time.time() batch_slices = self.patch_slices_list[i:i + self.batch_size] for batch_idx, slices in enumerate(batch_slices): self.input_patch_buffer[batch_idx, 0, :, :, :] = input_chunk.cutout( slices[0]).array if self.verbose > 1: end = time.time() print('prepare %d input patches takes %3f sec' % (self.batch_size, end - start)) start = end # the input and output patch is a 5d numpy array with # datatype of float32, the dimensions are batch/channel/z/y/x. # the input image should be normalized to [0,1] output_patch = self.patch_inferencer(self.input_patch_buffer) if self.verbose > 1: assert output_patch.ndim == 5 end = time.time() print('run inference for %d patch takes %3f sec' % (self.batch_size, end - start)) start = end for batch_idx, slices in enumerate(batch_slices): # only use the required number of channels # the remaining channels are dropped # the slices[0] is for input patch slice # the slices[1] is for output patch slice offset = (0, ) + tuple(s.start for s in slices[1]) output_chunk = Chunk(output_patch[batch_idx, :, :, :, :], global_offset=offset) ## save some patch for debug #bbox = output_chunk.bbox #if bbox.minpt[-1] < 94066 and bbox.maxpt[-1] > 94066 and \ # bbox.minpt[-2]<81545 and bbox.maxpt[-2]>81545 and \ # bbox.minpt[-3]<17298 and bbox.maxpt[-3]>17298: # print('save patch: ', output_chunk.bbox) # breakpoint() # output_chunk.to_tif() # #input_chunk.cutout(slices[0]).to_tif() output_buffer.blend(output_chunk) if self.verbose > 1: end = time.time() print('blend patch takes %3f sec' % (end - start)) if self.verbose: print("Inference of whole chunk takes %3f sec" % (time.time() - chunk_time_start)) if self.mask_output_chunk: output_buffer *= self.output_chunk_mask # theoretically, all the value of output_buffer should not be greater than 1 # we use a slightly higher value here to accomondate numerical precision issue np.testing.assert_array_less( output_buffer, 1.0001, err_msg='output buffer should not be greater than 1') if self.mask_myelin_threshold: # currently only for masking out affinity map assert output_buffer.shape[0] == 4 output_chunk = output_buffer.mask_using_last_channel( threshold=self.mask_myelin_threshold) # currently neuroglancer only support float32, not float16 if output_chunk.dtype == np.dtype('float16'): output_chunk = output_chunk.astype('float32') return output_chunk else: return output_buffer
def test_item(self): arr = np.ones((1, 3, 3, 3), dtype='float32') chunk = Chunk(arr, (0, -1, -1, -1)) chunk[:, :,:,:] = 0 self.assertTrue( chunk == 0 )