def test_non_aligned_input_chunk(): print('\ntest no aligned block inference engine...') patch_size = (32, 256, 256) patch_overlap = (4, 64, 64) num_output_channels = 2 input_size = (28 * 2 + 4 + 6, (256 - 64) * 2 + 64 + 7, (256 - 64) * 2 + 64 + 9) image = np.random.randint(1, 255, size=input_size, dtype=np.uint8) image = Chunk(image) with Inferencer(None, None, patch_size, output_patch_overlap=patch_overlap, num_output_channels=num_output_channels, batch_size=5, framework='identity', mask_output_chunk=True) as inferencer: output = inferencer(image) # only use the first channel to check correctness output = output[0, :, :, :] #output = np.reshape(output, image.shape) image = image.astype(np.float32) / 255 residual = image - output print('maximum difference: ', np.max(residual.array)) # 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 __call__(self, chunk: Chunk): assert isinstance(chunk, Chunk) if chunk.is_affinity_map: properties = chunk.properties chunk = (chunk[1, ...] + chunk[2, ...]) / 2. * 255. chunk = chunk.astype(np.uint8) chunk = Chunk(chunk) chunk.set_properties(properties) assert chunk.ndim == 3 for z in tqdm(range(chunk.voxel_offset[0], chunk.bbox.maxpt[0])): img = chunk.cutout( (slice(z, z + 1), chunk.slices[1], chunk.slices[2])) img = img.array[0, :, :] filename = os.path.join(self.output_path, f"{z:05d}.png") with open(filename, "wb") as f: f.write(pyspng.encode(img))
def _append_segmentation_layer(self, viewer_state: ng.viewer_state.ViewerState, chunk_name: str, chunk: Chunk): if np.issubdtype(chunk.dtype, np.int64): assert chunk.min() >= 0 chunk = chunk.astype(np.uint64) voxel_size = self._get_voxel_size(chunk) dimensions = ng.CoordinateSpace(scales=voxel_size, units=['nm', 'nm', 'nm'], names=['z', 'y', 'x']) viewer_state.layers.append( name=chunk_name, layer=ng.LocalVolume( data=chunk, dimensions=dimensions, # offset is in nm, not voxels # chunkflow use C order with zyx, # while neuroglancer use F order with xyz voxel_offset=chunk.voxel_offset, ))
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_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 _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 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 __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