import importlib
import inspect
import os
import shutil
from typing import Optional
from aydin import regression
from aydin.features.standard_features import StandardFeatureGenerator
from aydin.it.base import ImageTranslatorBase
from aydin.it.fgr import ImageTranslatorFGR
from aydin.it.transforms.padding import PaddingTransform
from aydin.it.transforms.range import RangeTransform
from aydin.it.transforms.variance_stabilisation 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="https://arxiv.org/abs/1901.11365">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, module.name.replace("_", "") + "Regressor"
)
arguments["Noise2SelfFGR-" + module.name] = {
"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-" + x.name
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__ + '.' + module.name)
elem = [
x
for x in dir(response)
if (module.name.replace("_", "") + "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"""
self.it.stop_training()
[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 self.it
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"]
self.it.add_transform(transform_class(**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.it = self.get_translator(
feature_generator=self.get_generator(), regressor=self.get_regressor()
)
self.add_transforms()
# Train a new model
self.it.train(
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 = self.it.translate(
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