Source code for

import gc
import time
from typing import Optional, Union, List, Tuple
import numpy

from aydin.features.base import FeatureGeneratorBase
from aydin.features.standard_features import StandardFeatureGenerator
from import DataHistogramBalancer
from import ImageTranslatorBase
from aydin.regression.base import RegressorBase
from aydin.regression.cb import CBRegressor
from aydin.util.array.nd import nd_split_slices
from aydin.util.log.log import lprint, lsection
from aydin.util.offcore.offcore import offcore_array

[docs]class ImageTranslatorFGR(ImageTranslatorBase): """ Feature Generation & Regression (FGR) based Image TranslatorFGR Image Translator. """ feature_generator: FeatureGeneratorBase def __init__( self, feature_generator: FeatureGeneratorBase = None, regressor: RegressorBase = None, balance_training_data: bool = False, voxel_keep_ratio: float = 1, max_voxels_for_training: Optional[int] = None, favour_bright_pixels: float = 0, blind_spots: Optional[Union[str, List[Tuple[int]]]] = None, tile_min_margin: int = 8, tile_max_margin: Optional[int] = None, max_memory_usage_ratio: float = 0.9, max_tiling_overhead: float = 0.1, ): """Constructs a FGR image translator. FGR image translators use feature generation and regression learning to acheive image translation. Typically used to denoise images using the Noise 2Self principle. Parameters ---------- feature_generator : FeatureGeneratorBase Feature generator. (advanced) regressor : RegressorBase Regressor. (advanced) balance_training_data : bool Balancing the training data improves image denoising performance by avoiding the over-representation of certain pixel intensities in the training data. For example, for images that are very dark with mostly dark background and only few bright regions, the brightness histogram will be is very skewed. In that case you want to turn on balancing so that the training data better represents the pixel values at all brightness levels. voxel_keep_ratio : float Ratio of the voxels to keep for training. Reduce this number to speed up denoising or if the images are extremely large and repetitive. (advanced) max_voxels_for_training : int, optional Maximum number of the voxels that can be used for training. This option limits the amount of training data and thus also limits the time required for training. In our experience we don't necessarily need more than several million voxels to train an effective denoiser. Different regressors (lgbm, cb, ...) can handle varying sizes of training data. We recommend at least a million (1e6) and at most 40 million if you have a powerfull machine with GPU and if you are armed with patience. (advanced) favour_bright_pixels : float Value within [-1.0, 1.0] that controls whether to favour bright or dark pixels during training data balancing. By default (0) all pixel intensities are treated equally. Positive values favour bright pixels, negative values favour dark pixels. blind_spots : Optional[Union[str,List[Tuple[int]]]] List of voxel coordinates (relative to receptive field center) to be included in the blind-spot. For example, you can enter: '<axis>#<radius>' to extend the blindspot along a given axis by a certain radius. For example, for an image of dimension 3, 'x#1' extends the blind spot to cover voxels of relative coordinates: (0,0,0),(0,1,0), and (0,-1,0). If you want to extend both in x and y, enter: 'x#1,y#1' by comma separating between axis. To specify the axis you can use integer indices, or 'x', 'y', 'z', and 't' (dimension order is tzyx with x being always the last dimension). If None is passed then the blindspots are automatically discovered from the image content. If 'center' is passed then no additional blindspots to the center pixel are considered. If 'center' is passed then only the default single center voxel blind-spot is used. tile_min_margin : int Minimal width of tile margin in voxels. (advanced) tile_max_margin : Optional[int] Maximal width of tile margin in voxels. (advanced) max_memory_usage_ratio : float Maximum allowed memory load, value must be within [0, 1]. Default is 90%. (advanced) max_tiling_overhead : float Maximum allowed margin overhead during tiling. Default is 10%. (advanced) """ super().__init__( blind_spots=blind_spots, tile_min_margin=tile_min_margin, tile_max_margin=tile_max_margin, max_memory_usage_ratio=max_memory_usage_ratio, max_tiling_overhead=max_tiling_overhead, ) self.voxel_keep_ratio = voxel_keep_ratio self.balance_training_data = balance_training_data self.favour_bright_pixels = max(min(favour_bright_pixels, 1.0), -1.0) self.feature_generator = ( StandardFeatureGenerator() if feature_generator is None else feature_generator ) self.regressor = CBRegressor() if regressor is None else regressor self.max_voxels_for_training = ( self.regressor.recommended_max_num_datapoints() if max_voxels_for_training is None else max_voxels_for_training ) # It seems empirically that excluding the central feature incurs no cost in quality, # simply because for every scale the next scale already partly covers these pixels. self.exclude_center_feature = True # Option to exclude the center value during translation. Set to True by default. self.exclude_center_value_when_translating = False # Advanced private functionality: # This field gives the opportunity to specify which channels must be 'passed-through' # directly as features. This is not intended to be used directly by users. self._passthrough_channels = None with lsection("FGR image translator"): lprint(f"balance training data: {self.balance_training_data}") def __repr__(self): return f"<{self.__class__.__name__}, feature_generator={self.feature_generator}, regressor={self.regressor}, balance_training_data={self.balance_training_data}"
[docs] def save(self, path: str): """Saves a 'all-batteries-included' image translation model at a given path (folder). Parameters ---------- path : str path to save to Returns ------- frozen """ with lsection(f"Saving 'fgr' image translator to {path}"): frozen = super().save(path) frozen += + '\n' frozen += + '\n' return frozen
def _load_internals(self, path: str): with lsection(f"Loading 'fgr' image translator from {path}"): self.feature_generator = FeatureGeneratorBase.load(path) self.regressor = RegressorBase.load(path) # We exclude certain fields from saving: def __getstate__(self): state = self.__dict__.copy() del state['feature_generator'] del state['regressor'] return state
[docs] def stop_training(self): """Stops currently running training within the instance by calling the corresponding `stop_fit()` method on the regressor. """ self.regressor.stop_fit()
def _estimate_memory_needed_and_available(self, image): with lsection( "Estimating the amount of memory needed to store feature arrays:" ): num_spatio_temp_dim = len(image.shape[2:]) dummy_image = numpy.zeros( shape=(1, 1) + (4,) * num_spatio_temp_dim, dtype=numpy.float32 ) x = self._compute_features( dummy_image, exclude_center_feature=self.exclude_center_feature, exclude_center_value=False, ) num_features = x.shape[-1] feature_dtype = x.dtype ( memory_needed, memory_available, ) = super()._estimate_memory_needed_and_available(image) """Here, 1.3 correction factor is chosen to have a bigger safety margin as we use more memory while generating features(namely uniform features) than the amount of memory we need after feature computations.""" memory_needed = max( memory_needed, 1.3 * num_features * image.size * feature_dtype.itemsize ) return memory_needed, memory_available def _compute_features( self, image, exclude_center_feature, exclude_center_value, features_last=True, num_reserved_features=0, image_slice=None, whole_image_shape=None, ): """Internal function that computes features for a given image. Parameters ---------- image exclude_center_feature exclude_center_value features_last : bool num_reserved_features : int image_slice whole_image_shape Returns ------- flattened array of features """ with lsection(f"Computing features for image of shape {image.shape}:"): excluded_voxels = ( None if self.blind_spots is None or 'center' in self.blind_spots else list( [ coordinate for coordinate in self.blind_spots if coordinate != (0,) * (image.ndim - 2) ] ) ) lprint(f"exclude_center_feature = {exclude_center_feature}") lprint(f"exclude_center_value = {exclude_center_value}") lprint(f"excluded_voxels = {excluded_voxels}") # If this is a part of a larger image, we can figure out what are the offsets and scales for the spatial features: spatial_feature_scale = ( None if whole_image_shape is None else tuple(1.0 / s for s in whole_image_shape[2:]) ) spatial_feature_offset = ( None if image_slice is None else tuple(s.start for s in image_slice[2:]) ) lprint(f"spatial_feature_scale = {spatial_feature_scale}") lprint(f"spatial_feature_offset = {spatial_feature_offset}") features = self.feature_generator.compute( image, exclude_center_feature=exclude_center_feature, exclude_center_value=exclude_center_value, num_reserved_features=num_reserved_features, passthrough_channels=self._passthrough_channels, feature_last_dim=False, excluded_voxels=excluded_voxels, spatial_feature_offset=spatial_feature_offset, spatial_feature_scale=spatial_feature_scale, ) num_features = features.shape[0] x = features.reshape(num_features, -1) if features_last: x = numpy.moveaxis(x, 0, -1) return x def _train( self, input_image, target_image, train_valid_ratio, callback_period, jinv ): with lsection( f"Training image translator from image of shape {input_image.shape} to image of shape {target_image.shape}:" ): self.prepare_monitoring_images() num_input_channels = input_image.shape[1] num_target_channels = target_image.shape[1] normalised_input_shape = input_image.shape # Deal with jinv parameter; if jinv is None: exclude_center_value = self.self_supervised elif type(jinv) is tuple: exclude_center_value = jinv else: exclude_center_value = jinv # Prepare the splitting of train from valid data as well as balancing and decimation... self._prepare_split_train_val(input_image, target_image) # Tilling strategy is determined here: tilling_strategy, margins = self._get_tilling_strategy_and_margins( input_image, max_voxels_per_tile=self.max_voxels_per_tile, min_margin=self.tile_min_margin, max_margin=self.tile_max_margin, ) lprint(f"Tilling strategy (just batches): {tilling_strategy}") lprint(f"Margins for tiles: {margins} .") # tile slice objects with margins: tile_slices_margins = list( nd_split_slices( normalised_input_shape, nb_slices=tilling_strategy, margins=margins ) ) # Number of tiles: number_of_tiles = len(tile_slices_margins) lprint(f"Number of tiles (slices): {number_of_tiles}") # We initialise the arrays: x_train, x_valid, y_train, y_valid = (None,) * 4 for idx, slice_margin_tuple in enumerate(tile_slices_margins): with lsection( f"Current tile: {idx}/{number_of_tiles}, slice: {slice_margin_tuple} " ): # We first extract the tile image: input_image_tile = input_image[slice_margin_tuple] target_image_tile = target_image[slice_margin_tuple] x_tile = self._compute_features( input_image_tile, exclude_center_feature=self.exclude_center_feature, exclude_center_value=exclude_center_value, num_reserved_features=0, features_last=False, image_slice=slice_margin_tuple, whole_image_shape=input_image.shape, ) y_tile = target_image_tile.reshape(num_target_channels, -1) num_features_tile = x_tile.shape[0] num_entries_tile = y_tile.shape[1] lprint( f"Number of entries: {num_entries_tile}, features: {num_features_tile}, input channels: {num_input_channels}, target channels: {num_target_channels}" ) # We split this tile's data into train and valid sets: ( x_train_tile, x_valid_tile, y_train_tile, y_valid_tile, ) = self._do_split_train_val( num_target_channels, train_valid_ratio, x_tile, y_tile ) # We get rid of x and y to free memory: del x_tile, y_tile gc.collect() # We now put the feature dimension to the back: x_train_tile = numpy.moveaxis(x_train_tile, 0, -1) x_valid_tile = numpy.moveaxis(x_valid_tile, 0, -1) if x_train is None or x_valid is None: x_train, x_valid, y_train, y_valid = ( x_train_tile, x_valid_tile, y_train_tile, y_valid_tile, ) else: x_train = numpy.append(x_train, x_train_tile, axis=0) x_valid = numpy.append(x_valid, x_valid_tile, axis=0) y_train = numpy.append(y_train, y_train_tile, axis=1) y_valid = numpy.append(y_valid, y_valid_tile, axis=1) lprint("Training now...") x_train=x_train, y_train=y_train, x_valid=x_valid, y_valid=y_valid, regressor_callback=self.get_callback() if self.monitor else None, ) self.loss_history = self.regressor.loss_history def prepare_monitoring_images(self): # Compute features for monitoring images: if self.monitor is not None and self.monitor.monitoring_images is not None: # Normalise monitoring images: normalised_monitoring_images = [ self.shape_normaliser.normalise( monitoring_image, batch_dims=None, channel_dims=None ) for monitoring_image in self.monitor.monitoring_images ] # compute features proper: monitoring_images_features = [ self._compute_features( monitoring_image, exclude_center_feature=self.exclude_center_feature, exclude_center_value=False, features_last=True, ) for monitoring_image in normalised_monitoring_images ] else: monitoring_images_features = None # We keep these features handy... self.monitoring_datasets = monitoring_images_features def get_callback(self): # Regressor callback: def regressor_callback(iteration, val_loss, model): if val_loss is None: return current_time_sec = time.time() # Correct for dtype range: if self.feature_generator.dtype == numpy.uint8: val_loss /= 255 elif self.feature_generator.dtype == numpy.uint16: val_loss /= 255 * 255 if current_time_sec > self.last_callback_time_sec + self.callback_period: if self.monitoring_datasets and self.monitor: predicted_monitoring_datasets = [ self.regressor.predict(x_m, models_to_use=[model]) for x_m in self.monitoring_datasets ] inferred_images = [ y_m.reshape(image.shape) for (image, y_m) in zip( self.monitor.monitoring_images, predicted_monitoring_datasets, ) ] # for image in inferred_images: for index in range(len(inferred_images)): ( inferred_images[index], _, _, ) = self.shape_normaliser.shape_normalize( inferred_images[index] ) denormalised_inferred_images = [ self.target_shape_normaliser.denormalise(inferred_image) for inferred_image in inferred_images ] self.monitor.variables = ( iteration, val_loss, denormalised_inferred_images, ) elif self.monitor: self.monitor.variables = (iteration, val_loss, None) self.last_callback_time_sec = current_time_sec else: pass # print(f"Iteration={iteration} metric value: {eval_metric_value} ") return regressor_callback def _prepare_split_train_val(self, input_image, target_image): # the number of voxels: num_of_voxels = input_image.size lprint( f"Image has: {num_of_voxels} voxels, at most: {self.max_voxels_for_training} voxels will be used for training or validation." ) # This is the ratio of pixels to keep: max_voxels_keep_ratio = float(self.max_voxels_for_training) / num_of_voxels effective_keep_ratio = min(self.voxel_keep_ratio, max_voxels_keep_ratio) lprint( f"Given train ratio is: {self.voxel_keep_ratio}, max_voxels induced keep-ratio is: {max_voxels_keep_ratio}" ) # For small images it is not worth having any limit or balance anything: num_voxels_in_small_image = 5 * 1e6 effective_keep_ratio = ( 1.0 if num_of_voxels < num_voxels_in_small_image else effective_keep_ratio ) balance_training_data = ( False if num_of_voxels < num_voxels_in_small_image else self.balance_training_data ) lprint( f"Data histogram balancer is: {'active' if balance_training_data else 'inactive'}" ) lprint(f"Effective keep-ratio is: {effective_keep_ratio}") lprint( f"Favouring bright pixels: {'yes' if self.favour_bright_pixels > 0 else 'no'}" ) if self.favour_bright_pixels != 0: lprint( f"Favouring bright pixels by a linear slope of: {self.favour_bright_pixels}" ) # We decide on a 'batch' length that will be used to shuffle, select and then copy the training data... num_of_voxels_per_stack = input_image[2:].size batch_length = ( 16 if num_of_voxels_per_stack < 1e5 else ( 32 if num_of_voxels_per_stack < 1e6 else ( 128 if num_of_voxels_per_stack < 1e7 else (512 if num_of_voxels_per_stack < 1e8 else 4096) ) ) ) lprint(f"Using contiguous batches of length: {batch_length} ") # We create a balancer common to all tiles: balancer = DataHistogramBalancer( balance=balance_training_data, keep_ratio=effective_keep_ratio, favour_bright_pixels=self.favour_bright_pixels, ) # Calibration of the balancer is done on the entire image: lprint("Calibrating balancer...") balancer.calibrate(target_image.ravel(), batch_length) # Keep both balancer and batch length self.train_val_split_balancer = balancer self.train_val_split_batch_length = batch_length def _do_split_train_val( self, num_target_channels: int, train_valid_ratio: float, x, y ): with lsection( f"Splitting train and test sets (train_test_ratio={train_valid_ratio}) " ): balancer: DataHistogramBalancer = self.train_val_split_balancer batch_length: int = self.train_val_split_batch_length nb_features = x.shape[0] nb_entries = y.shape[1] nb_split_batches = max(nb_entries // batch_length, 64) lprint( f"Creating random indices for train/val split (nb_split_batches={nb_split_batches})" ) nb_split_batches_valid = int(train_valid_ratio * nb_split_batches) nb_split_batches_train = nb_split_batches - nb_split_batches_valid is_train_array = numpy.full(nb_split_batches, False) is_train_array[nb_split_batches_valid:] = True lprint(f"train/valid bool array created (length={is_train_array.shape[0]})") lprint("Shuffling train/valid bool array...") numpy.random.shuffle(is_train_array) lprint("Calculating number of entries for train and validation...") nb_entries_per_split_batch = max(1, nb_entries // nb_split_batches) nb_entries_train = nb_split_batches_train * nb_entries_per_split_batch nb_entries_valid = nb_split_batches_valid * nb_entries_per_split_batch lprint( f"Number of entries for training: {nb_entries_train} = {nb_split_batches_train}*{nb_entries_per_split_batch}, validation: {nb_entries_valid} = {nb_split_batches_valid} * {nb_entries_per_split_batch}" ) lprint("Allocating arrays...") x_train = offcore_array( shape=(nb_features, nb_entries_train), dtype=x.dtype, max_memory_usage_ratio=self.max_memory_usage_ratio, ) y_train = offcore_array( shape=(num_target_channels, nb_entries_train), dtype=y.dtype, max_memory_usage_ratio=self.max_memory_usage_ratio, ) x_valid = offcore_array( shape=(nb_features, nb_entries_valid), dtype=x.dtype, max_memory_usage_ratio=self.max_memory_usage_ratio, ) y_valid = offcore_array( shape=(num_target_channels, nb_entries_valid), dtype=y.dtype, max_memory_usage_ratio=self.max_memory_usage_ratio, ) with lsection("Copying data for training and validation sets..."): # We use a random permutation to avoid having the balancer drop only from the 'end' of the image permutation = numpy.random.permutation(nb_split_batches) i, jt, jv = 0, 0, 0 dst_stop_train = 0 dst_stop_valid = 0 balancer.initialise(nb_split_batches) for is_train in numpy.nditer(is_train_array): if i % (nb_split_batches // 64) == 0: lprint( f"Copying section [{i},{min(nb_split_batches, i + nb_split_batches // 64)}]" ) permutated_i = permutation[i] src_start = permutated_i * nb_entries_per_split_batch src_stop = src_start + nb_entries_per_split_batch i += 1 xsrc = x[:, src_start:src_stop] ysrc = y[:, src_start:src_stop] if balancer.add_entry(ysrc): if is_train: dst_start_train = jt * nb_entries_per_split_batch dst_stop_train = (jt + 1) * nb_entries_per_split_batch jt += 1 xdst = x_train[:, dst_start_train:dst_stop_train] ydst = y_train[:, dst_start_train:dst_stop_train] numpy.copyto(xdst, xsrc) numpy.copyto(ydst, ysrc) else: dst_start_valid = jv * nb_entries_per_split_batch dst_stop_valid = (jv + 1) * nb_entries_per_split_batch jv += 1 xdst = x_valid[:, dst_start_valid:dst_stop_valid] ydst = y_valid[:, dst_start_valid:dst_stop_valid] numpy.copyto(xdst, xsrc) numpy.copyto(ydst, ysrc) # Now we actually truncate out all the zeros at the end of the arrays: x_train = x_train[:, 0:dst_stop_train] y_train = y_train[:, 0:dst_stop_train] x_valid = x_valid[:, 0:dst_stop_valid] y_valid = y_valid[:, 0:dst_stop_valid] lprint(f"Histogram all : {balancer.get_histogram_all_as_string()}") lprint(f"Histogram kept : {balancer.get_histogram_kept_as_string()}") lprint( f"Histogram dropped: {balancer.get_histogram_dropped_as_string()}" ) lprint( f"Number of entries kept: {balancer.total_kept()} out of {balancer.total_entries} total" ) lprint( f"Percentage of data kept: {100 * balancer.percentage_kept():.3f}% (train_data_ratio={balancer.keep_ratio}) " ) if balancer.keep_ratio >= 1 and balancer.percentage_kept() < 1: lprint( "Note: balancer has dropped entries that fell on over-represented histogram bins" ) return x_train, x_valid, y_train, y_valid def _translate(self, input_image, image_slice=None, whole_image_shape=None): """Internal method that translates an input image on the basis of the trained model. Parameters ---------- input_image image_slice whole_image_shape Returns ------- inferred image """ shape = input_image.shape num_batches = shape[0] x = self._compute_features( input_image, exclude_center_feature=self.exclude_center_feature, exclude_center_value=self.exclude_center_value_when_translating, num_reserved_features=0, features_last=True, image_slice=image_slice, whole_image_shape=whole_image_shape, ) with lsection( f"Predict from feature vector of dimension {x.shape} and dtype: {x.dtype}:" ): lprint("Predicting... ") # Predict using regressor: yp = self.regressor.predict(x) # We make sure that we have the result in float type, but make _sure_ to avoid copying data: if yp.dtype != numpy.float32 and yp.dtype != numpy.float64: yp = yp.astype(numpy.float32, copy=False) # We reshape the array: num_target_channels = yp.shape[0] translated_image_shape = (num_batches, num_target_channels) + shape[2:] lprint(f"Reshaping array to {translated_image_shape}... ") inferred_image = yp.reshape(translated_image_shape) return inferred_image