def __call__(self, chunk: Chunk): assert isinstance(chunk, Chunk) 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,:,:] imsave(os.path.join(self.output_path, '{:05d}.png'.format(z)), img)
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 __call__(self, input_chunk: Chunk): """ args: input_chunk (Chunk): input chunk with voxel offset and voxel size """ 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: logging.info('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.voxel_offset, voxel_size=input_chunk.voxel_size, ) if np.all(input_chunk == 0): logging.info('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) input_chunk /= dtype_max chunk_time_start = time.time() # iterate the offset list for i in tqdm(range(0, len(self.patch_slices_list), self.batch_size), disable=(self.verbose <= 0), desc='ConvNet inference for patches: '): 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 end = time.time() logging.debug( f'prepare {self.batch_size:d} input patches takes {end-start:.3f} sec' ) 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] if not self.test_time_augmentation: output_patch = self.patch_inferencer(self.input_patch_buffer) else: # test time augmentation pass end = time.time() logging.debug( f'run inference for {self.batch_size:d} patch takes {end-start:.3f} sec' ) 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 = tuple(s.start for s in slices[1]) output_patch_chunk = Chunk(output_patch[batch_idx, :, :, :, :], voxel_offset=offset, voxel_size=input_chunk.voxel_size) ## 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) # output_chunk.to_tif() # #input_chunk.cutout(slices[0]).to_tif() output_buffer.blend(output_patch_chunk) end = time.time() logging.debug('blend patch takes {:.3f} sec'.format(end - start)) logging.debug("Inference of whole chunk takes {:.3f} sec".format( 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) if output_chunk.dtype == np.dtype('<f4'): output_chunk = output_chunk.astype('float32') return output_chunk else: return output_buffer