Restoration variant option bug fix and refactor

aydin/cli/cli.py

@@ -69,7 +69,7 @@ def exception_hook(exctype, value, traceback):
 @click.option(
     '-ca', '--channel-axes', type=str, help='only pass while denoising a single image'
 )
-@click.option('-v', '--variant', default='noise2selffgr-cb')
+@click.option('-d', '--denoiser', default='noise2selffgr-cb')
 @click.option('--use-model/--save-model', default=False)
 @click.option('--model-path', default=None)
 @click.option('--lower-level-args', default=None)
@@ -89,10 +89,10 @@ def denoise(files, **kwargs):
     # Check whether a filename is provided for lower-level-args json
     if kwargs["lower_level_args"]:
         lower_level_args = load_any_json(kwargs['lower_level_args'])
-        backend = lower_level_args["variant"]
+        denoiser = lower_level_args["variant"]
     else:
         lower_level_args = None
-        backend = kwargs["variant"]
+        denoiser = kwargs["denoiser"]
 
     filenames = []
     for filename in files:
@@ -144,7 +144,7 @@ def denoise(files, **kwargs):
             kwargs_to_pass.pop("channel_axes")
 
             denoiser = get_denoiser_class_instance(
-                lower_level_args=lower_level_args, variant=backend
+                lower_level_args=lower_level_args, variant=denoiser
             )
 
             denoiser.train(

aydin/restoration/denoise/classic.py

@@ -74,35 +74,43 @@
 
 
 class Classic(DenoiseRestorationBase):
-    """Classic Image Denoising.
-
-    Parameters
-    ----------
-    variant : str
-        Opacity of the layer visual, between 0.0 and 1.0.
-    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 : string
-        Path to model that is desired to be used for inference.
-        By default it is None.
-    """
+    """Classic Image Denoising"""
 
     disabled_modules = ["bilateral", "bmnd", "_defaults"]
 
     def __init__(
         self,
         *,
-        variant: str = 'butterworth',
+        variant: str = None,
         use_model=None,
         input_model_path=None,
         lower_level_args=None,
         it_transforms=None,
     ):
+        """
+
+        Parameters
+        ----------
+        variant : str
+            Variant of the Classic 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 'Classic-`) that can be used
+            on a given installation. Example variants: `butterworth`,
+            `gaussian`, `lipschitz`, `nlm`, ...
+        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 : string
+            Path to model that is desired to be used for inference.
+            By default it is None.
+        lower_level_args
+        it_transforms
+        """
         super().__init__()
         self.lower_level_args = lower_level_args
 
-        self.backend = variant
+        self.variant = variant
 
         self.input_model_path = input_model_path
         self.use_model_flag = use_model
@@ -248,6 +256,9 @@ def get_translator(self):
         it : ImageTranslatorBase
 
         """
+        if self.variant:
+            return ImageDenoiserClassic(method=self.variant)
+
         # 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
@@ -256,19 +267,19 @@ def get_translator(self):
             )
             it = ImageTranslatorBase.load(self.input_model_path[:-4])
         else:
-            if self.lower_level_args is not None:
-                method = (
-                    self.backend
-                    if self.lower_level_args["variant"] is None
-                    else self.lower_level_args["variant"].split("-")[1]
-                )
+            if (
+                self.lower_level_args is not None
+                and self.lower_level_args["variant"] is not None
+            ):
+                method = self.lower_level_args["variant"].split("-")[1]
+
                 it = ImageDenoiserClassic(
                     method=method,
                     calibration_kwargs=self.lower_level_args["calibration"]["kwargs"],
                     **self.lower_level_args["it"]["kwargs"],
                 )
             else:
-                it = ImageDenoiserClassic(method=self.backend)
+                it = ImageDenoiserClassic()
 
         return it
 

aydin/restoration/denoise/noise2selfcnn.py

@@ -2,6 +2,7 @@
 import inspect
 import os
 import shutil
+from typing import Optional
 
 from aydin.it.base import ImageTranslatorBase
 from aydin.it.cnn import ImageTranslatorCNN
@@ -23,6 +24,7 @@ class Noise2SelfCNN(DenoiseRestorationBase):
     def __init__(
         self,
         *,
+        variant: Optional[str] = None,
         use_model=None,
         input_model_path=None,
         lower_level_args=None,
@@ -33,6 +35,10 @@ def __init__(
 
         Parameters
         ----------
+        variant : str, optional
+            Variant of CNN denoiser to be used. Variant would supersede
+            the denoiser option passed in lower_level_args. Currently, we
+            support only two variants: `unet` and `jinet`.
         use_model : bool
             Flag to choose to train a new model or infer from a
             previously trained model. By default it is None.
@@ -41,6 +47,7 @@ def __init__(
             By default it is None.
         """
         super().__init__()
+        self.variant = variant
         self.use_model_flag = use_model
         self.input_model_path = input_model_path
         self.lower_level_args = lower_level_args
@@ -134,6 +141,9 @@ def get_translator(self):
         it : ImageTranslatorBase
 
         """
+        if self.variant:
+            return ImageTranslatorCNN(model_architecture=self.variant)
+
         # 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
@@ -162,15 +172,15 @@ def train(self, noisy_image, *, batch_axes=None, chan_axes=None, **kwargs):
 
         Parameters
         ----------
-        noisy_image : numpy.ndarray
+        noisy_image : numpy.ArrayLike
         batch_axes : array_like, optional
             Indices of batch axes.
         chan_axes : array_like, optional
             Indices of channel axes.
 
         Returns
         -------
-        response : numpy.ndarray
+        response : numpy.ArrayLike
 
         """
         with lsection("Noise2Self train is starting..."):

aydin/restoration/denoise/noise2selffgr.py

@@ -2,6 +2,7 @@
 import inspect
 import os
 import shutil
+from typing import Optional
 
 from aydin import regression
 from aydin.features.standard_features import StandardFeatureGenerator
@@ -11,18 +12,15 @@
 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
 
 
-if os.getenv("BUNDLED_AYDIN") == "1":
-    from aydin.regression.lgbm import LGBMRegressor  # noqa: F401
-    from aydin.regression.linear import LinearRegressor  # noqa: F401
-    from aydin.regression.perceptron import PerceptronRegressor  # noqa: F401
-    from aydin.regression.random_forest import RandomForestRegressor  # noqa: F401
-    from aydin.regression.support_vector import SupportVectorRegressor  # noqa: F401
-
-
 class Noise2SelfFGR(DenoiseRestorationBase):
     """
     Noise2Self image denoising using the "Feature Generation & Regression" (
@@ -33,6 +31,7 @@ class Noise2SelfFGR(DenoiseRestorationBase):
     def __init__(
         self,
         *,
+        variant: Optional[str] = None,
         use_model=None,
         input_model_path=None,
         lower_level_args=None,
@@ -41,6 +40,12 @@ def __init__(
         """
         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.
@@ -56,6 +61,7 @@ def __init__(
         self.use_model_flag = use_model
         self.input_model_path = input_model_path
         self.lower_level_args = lower_level_args
+        self.variant = variant
 
         self.it = None
         self.it_transforms = (
@@ -170,7 +176,6 @@ def get_generator(self):
         generator : FeatureGeneratorBase
 
         """
-        # print(self.lower_level_args)
         if self.lower_level_args is not None:
             generator = self.lower_level_args["feature_generator"]["class"](
                 **self.lower_level_args["feature_generator"]["kwargs"]
@@ -188,13 +193,23 @@ def get_regressor(self):
         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 not None:
+        if self.lower_level_args is None:
+            regressor = CBRegressor()
+        else:
             regressor = self.lower_level_args["regressor"]["class"](
                 **self.lower_level_args["regressor"]["kwargs"]
             )
-        else:
-            regressor = CBRegressor()
 
         return regressor
 
@@ -326,7 +341,7 @@ def noise2self_fgr(noisy, *, batch_axes=None, chan_axes=None, variant=None):
 
     """
     # Run N2S and save the result
-    n2s = Noise2SelfFGR(variant=variant)
+    n2s = Noise2SelfFGR()
 
     # Train
     n2s.train(noisy, batch_axes=batch_axes, chan_axes=chan_axes)

aydin/restoration/denoise/test/test_n2s_fgr.py

@@ -35,7 +35,7 @@ def test_run_n2s_fgr():
         {"class": RangeTransform, "kwargs": {}},
         {"class": PaddingTransform, "kwargs": {}},
     ]
-    n2s = Noise2SelfFGR(variant="fgr-cb", it_transforms=transforms)
+    n2s = Noise2SelfFGR(variant="cb", it_transforms=transforms)
     n2s.train(noisy_image)
     denoised_image = n2s.denoise(noisy_image).clip(0, 1)