Source code for aydin.restoration.denoise.noise2selffgr

import importlib
import inspect
import os
import shutil
from typing import Optional

from aydin import regression
from aydin.features.standard_features import StandardFeatureGenerator
from import ImageTranslatorBase
from import ImageTranslatorFGR
from import PaddingTransform
from import RangeTransform
from import VarianceStabilisationTransform
from aydin.regression.cb import CBRegressor
from aydin.regression.lgbm import LGBMRegressor
from aydin.regression.linear import LinearRegressor
from aydin.regression.perceptron import PerceptronRegressor
from aydin.regression.random_forest import RandomForestRegressor
from aydin.regression.support_vector import SupportVectorRegressor
from aydin.restoration.denoise.base import DenoiseRestorationBase
from aydin.util.log.log import lsection

[docs]class Noise2SelfFGR(DenoiseRestorationBase): """ Noise2Self image denoising using the "Feature Generation & Regression" ( FGR) approach. Follows from the theory exposed in the <a href="">Noise2Self paper</a>. """ def __init__( self, *, variant: Optional[str] = None, use_model=None, input_model_path=None, lower_level_args=None, it_transforms=None, ): """ Parameters ---------- variant : str, optional Variant of FGR denoiser to be used. Variant would supersede the denoiser option passed in lower_level_args. `implementations` property would return a complete list of variants (with a prefix of 'Noise2SelfFGR-`) that can be used on a given installation. Example variants: `cb`, `lgbm`, ... use_model : bool Flag to choose to train a new model or infer from a previously trained model. By default it is None. input_model_path : str Path to model that is desired to be used for inference. By default it is None. lower_level_args : args Additional 'low-level' arguments to be passed. it_transforms : Transforms to be applied. """ super().__init__(variant=variant) self.use_model_flag = use_model self.input_model_path = input_model_path self.lower_level_args = lower_level_args self.it_transforms = ( [ {"class": RangeTransform, "kwargs": {}}, {"class": PaddingTransform, "kwargs": {}}, {"class": VarianceStabilisationTransform, "kwargs": {}}, ] if it_transforms is None else it_transforms ) @property def configurable_arguments(self): """Returns the configurable arguments that will be exposed on GUI and CLI. """ arguments = {} # Feature Generator feature_generator = StandardFeatureGenerator fullargspec2 = inspect.getfullargspec(feature_generator.__init__) feature_generator_args = { "arguments": fullargspec2.args[4:], "defaults": fullargspec2.defaults[3:], "annotations": fullargspec2.annotations, "reference_class": feature_generator, } # IT FGR it = ImageTranslatorFGR fullargspec3 = inspect.getfullargspec(ImageTranslatorFGR.__init__) it_args = { "arguments": fullargspec3.args[3:], "defaults": fullargspec3.defaults[2:], "annotations": fullargspec3.annotations, "reference_class": it, } # Regressor regression_modules = DenoiseRestorationBase.get_implementations_in_a_module( regression ) for module in regression_modules: regressor_args = self.get_class_implementation_kwargs( regression, module,"_", "") + "Regressor" ) arguments["Noise2SelfFGR-" +] = { "feature_generator": feature_generator_args, "regressor": regressor_args, "it": it_args, } return arguments @property def implementations(self): """Returns the list of discovered implementations for given method.""" return [ "Noise2SelfFGR-" + for x in self.get_implementations_in_a_module(regression) ] @property def implementations_description(self): fgr_description = Noise2SelfFGR.__doc__.strip() feature_generator_name = StandardFeatureGenerator.__name__.replace( "FeatureGenerator", "" ) feature_generator_description = StandardFeatureGenerator.__doc__.strip() descriptions = [] for module in self.get_implementations_in_a_module(regression): response = importlib.import_module(regression.__name__ + '.' + elem = [ x for x in dir(response) if ("_", "") + "Regressor").lower() in x.lower() ][ 0 ] # class name elem_class = response.__getattribute__(elem) regressor_name = elem_class.__name__.replace("Regressor", "") regressor_description = elem_class.__doc__.replace("\n\n", "<br><br>") descriptions.append( fgr_description + f" Uses the {feature_generator_name} feature generator and {regressor_name} regressor. " + f"<br><br>About the feature generator: {feature_generator_description}" + f"<br><br>About the regressor: {regressor_description}" ) return descriptions
[docs] def stop_running(self): """Method to stop running N2S instance"""
[docs] def get_generator(self): """Returns the corresponding generator instance for given selections. Returns ------- generator : FeatureGeneratorBase """ if self.lower_level_args is not None: generator = self.lower_level_args["feature_generator"]["class"]( **self.lower_level_args["feature_generator"]["kwargs"] ) else: generator = StandardFeatureGenerator() return generator
[docs] def get_regressor(self): """Returns the corresponding regressor instance for given selections. Returns ------- regressor : RegressorBase """ if self.variant: regressors = { "cb": CBRegressor, "lgbm": LGBMRegressor, "linear": LinearRegressor, "perceptron": PerceptronRegressor, "random_forest": RandomForestRegressor, "support_vector": SupportVectorRegressor, } return regressors[self.variant]() if self.lower_level_args is None: regressor = CBRegressor() else: regressor = self.lower_level_args["regressor"]["class"]( **self.lower_level_args["regressor"]["kwargs"] ) return regressor
[docs] def get_translator(self, feature_generator, regressor): """Returns the corresponding translator instance for given selections. Parameters ---------- feature_generator : FeatureGeneratorBase regressor : RegressorBase Returns ------- it : ImageTranslatorBase """ # Use a pre-saved model or train a new one from scratch and save it if self.use_model_flag: # Unarchive the model file and load its ImageTranslator object into shutil.unpack_archive( self.input_model_path, os.path.dirname(self.input_model_path), "zip" ) it = ImageTranslatorBase.load(self.input_model_path[:-4]) else: it = ImageTranslatorFGR( feature_generator=feature_generator, regressor=regressor, **self.lower_level_args["it"]["kwargs"] if self.lower_level_args is not None else {}, ) return it
def add_transforms(self): if self.it_transforms is not None: for transform in self.it_transforms: transform_class = transform["class"] transform_kwargs = transform["kwargs"]**transform_kwargs))
[docs] def train(self, noisy_image, *, batch_axes=None, chan_axes=None, **kwargs): """Method to run training for Noise2Self FGR. Parameters ---------- noisy_image : numpy.ndarray batch_axes : array_like, optional Indices of batch axes. chan_axes : array_like, optional Indices of channel axes. Returns ------- response : numpy.ndarray """ with lsection("Noise2Self train is starting..."): = self.get_translator( feature_generator=self.get_generator(), regressor=self.get_regressor() ) self.add_transforms() # Train a new model noisy_image, noisy_image, batch_axes=batch_axes, channel_axes=chan_axes, train_valid_ratio=kwargs['train_valid_ratio'] if 'train_valid_ratio' in kwargs else 0.1, callback_period=kwargs['callback_period'] if 'callback_period' in kwargs else 3, jinv=kwargs['jinv'] if 'jinv' in kwargs else None, )
[docs] def denoise(self, noisy_image, *, batch_axes=None, chan_axes=None, **kwargs): """Method to denoise an image with trained Noise2Self FGR. Parameters ---------- batch_axes : array_like, optional Indices of batch axes. chan_axes : array_like, optional Indices of channel axes. noisy_image : numpy.ndarray Returns ------- response : numpy.ndarray """ with lsection("Noise2Self denoise is starting..."): # Predict the resulting image response = noisy_image, batch_axes=batch_axes, channel_axes=chan_axes, tile_size=kwargs['tile_size'] if 'tile_size' in kwargs else None, ) response = response.astype(noisy_image.dtype, copy=False) return response
[docs]def noise2self_fgr(noisy, *, batch_axes=None, chan_axes=None, variant=None): """Method to denoise an image with trained Noise2Self FGR. Parameters ---------- noisy : numpy.ndarray Image to denoise batch_axes : array_like, optional Indices of batch axes. chan_axes : array_like, optional Indices of channel axes. variant : str Algorithm variant. Returns ------- Denoised image : numpy.ndarray """ # Run N2S and save the result n2s = Noise2SelfFGR() # Train n2s.train(noisy, batch_axes=batch_axes, chan_axes=chan_axes) # Denoise denoised = n2s.denoise(noisy, batch_axes=batch_axes, chan_axes=chan_axes) return denoised