initiate demo

app.py

@@ -0,0 +1,66 @@
+import sys
+sys.path.append("src")
+from interactive_pipe import interactive_pipeline
+from rstor.analyzis.interactive.pipelines import natural_inference_pipeline, morph_canvas, CANVAS
+from rstor.analyzis.interactive.model_selection import get_default_models
+from pathlib import Path
+from rstor.analyzis.parser import get_parser
+import argparse
+from batch_processing import Batch
+from interactive_pipe.data_objects.image import Image
+from rstor.analyzis.interactive.images import image_selector
+from rstor.analyzis.interactive.crop import plug_crop_selector
+from rstor.analyzis.interactive.metrics import plug_configure_metrics
+from interactive_pipe import interactive, KeyboardControl
+
+
+def plug_morph_canvas():
+    interactive(
+        canvas=KeyboardControl(CANVAS[0], CANVAS, name="canvas", keyup="p", modulo=True)
+    )(morph_canvas)
+
+
+def image_loading_batch(input: Path, args: argparse.Namespace) -> dict:
+    """Wrapper to load images files from a directory using batch_processing
+    """
+
+    if not args.disable_preload:
+        img = Image.from_file(input).data
+        return {"name": input.name, "path": input, "buffer": img}
+    else:
+        return {"name": input.name, "path": input, "buffer": None}
+
+
+def main(argv):
+    batch = Batch(argv)
+    batch.set_io_description(
+        input_help='input image files',
+        output_help=argparse.SUPPRESS
+    )
+    parser = get_parser()
+    parser.add_argument("-nop", "--disable-preload", action="store_true", help="Disable images preload")
+    args = batch.parse_args(parser)
+    # batch.set_multiprocessing_enabled(False)
+    img_list = batch.run(image_loading_batch)
+    if args.keyboard:
+        image_control = KeyboardControl(0, [0, len(img_list)-1], keydown="3", keyup="9", modulo=True)
+    else:
+        image_control = (0, [0, len(img_list)-1])
+    interactive(image_index=image_control)(image_selector)
+    plug_crop_selector(num_pad=args.keyboard)
+    plug_configure_metrics(key_shortcut="a")  # "a" if args.keyboard else None)
+    plug_morph_canvas()
+    model_dict = get_default_models(args.experiments, Path(args.models_storage), keyboard_control=args.keyboard)
+    interactive_pipeline(
+        gui=args.backend,
+        cache=True,
+        safe_input_buffer_deepcopy=False
+    )(natural_inference_pipeline)(
+        img_list,
+        model_dict
+    )
+
+
+if __name__ == "__main__":
+    # main(sys.argv[1:])
+    main(["-e", "6002", "-i", "__dataset/sample/*.*g",  "-b","gradio"])

requirements.txt

@@ -0,0 +1,6 @@
+interactive-pipe>=0.7.8
+opencv_python_headless==4.8.0.74
+torch>=2.0.0
+tqdm
+
+

scripts/configuration.py

@@ -0,0 +1,18 @@
+from pathlib import Path
+
+NB_ID = "blind-deblurring-from-synthetic-data"  # This will be the name which appears on Kaggle.
+GIT_USER = "balthazarneveu"  # Your git user name
+GIT_REPO = "blind-deblurring-from-synthetic-data"  # Your current git repo
+# Keep free unless you need to acess kaggle datasets. You'll need to modify the remote_training_template.ipynb.
+KAGGLE_DATASET_LIST = [
+    "balthazarneveu/deadleaves-div2k-512", # Deadleaves classic
+    "balthazarneveu/deadleaves-primitives-div2k-512", # Deadleaves with extra primitives
+    "balthazarneveu/motion-blur-kernels", # Motion blur kernels
+    "joe1995/div2k-dataset",
+]
+WANDBSPACE = "deblur-from-deadleaves"
+TRAIN_SCRIPT = "scripts/train.py"  # Location of the training script
+
+ROOT_DIR = Path(__file__).parent
+OUTPUT_FOLDER_NAME = "__output"
+INFERENCE_FOLDER_NAME = "__inference"

scripts/infer.py

@@ -0,0 +1,205 @@
+from configuration import ROOT_DIR, OUTPUT_FOLDER_NAME, INFERENCE_FOLDER_NAME
+from rstor.analyzis.parser import get_models_parser
+from batch_processing import Batch
+from rstor.properties import (
+    DEVICE, NAME, PRETTY_NAME, DATALOADER, CONFIG_DEAD_LEAVES, VALIDATION,
+    BATCH_SIZE, SIZE,
+    REDUCTION_SKIP,
+    TRACES_TARGET, TRACES_DEGRADED, TRACES_RESTORED, TRACES_METRICS, TRACES_ALL,
+    SAMPLER_SATURATED,
+    CONFIG_DEGRADATION,
+    DATASET_DIV2K,
+    DATASET_DL_DIV2K_512,
+    DATASET_DL_EXTRAPRIMITIVES_DIV2K_512,
+    METRIC_PSNR, METRIC_SSIM, METRIC_LPIPS,
+    DEGRADATION_BLUR_MAT, DEGRADATION_BLUR_NONE
+)
+from rstor.data.dataloader import get_data_loader
+from tqdm import tqdm
+from pathlib import Path
+import torch
+from typing import Optional
+import argparse
+import sys
+from rstor.analyzis.interactive.model_selection import get_default_models
+from rstor.learning.metrics import compute_metrics
+from interactive_pipe.data_objects.image import Image
+from interactive_pipe.data_objects.parameters import Parameters
+from typing import List
+from itertools import product
+import pandas as pd
+ALL_TRACES = [TRACES_TARGET, TRACES_DEGRADED, TRACES_RESTORED, TRACES_METRICS]
+
+
+def parse_int_pairs(s):
+    try:
+        # Split the input string by spaces to separate pairs, then split each pair by ',' and convert to tuple of ints
+        return [tuple(map(int, item.split(','))) for item in s.split()]
+    except ValueError:
+        raise argparse.ArgumentTypeError("Must be a series of pairs 'a,b' separated by spaces.")
+
+
+def get_parser(parser: Optional[argparse.ArgumentParser] = None, batch_mode=False) -> argparse.ArgumentParser:
+    parser = get_models_parser(
+        parser=parser,
+        help="Inference on validation set",
+        default_models_path=ROOT_DIR/OUTPUT_FOLDER_NAME)
+    if not batch_mode:
+        parser.add_argument("-o", "--output-dir", type=str, default=ROOT_DIR /
+                            INFERENCE_FOLDER_NAME, help="Output directory")
+    parser.add_argument("--cpu", action="store_true", help="Force CPU")
+    parser.add_argument("--traces", "-t", nargs="+", type=str, choices=ALL_TRACES+[TRACES_ALL],
+                        help="Traces to be computed", default=TRACES_ALL)
+    parser.add_argument("--size", type=parse_int_pairs,
+                        default=[(256, 256)], help="Size of the images like '256,512 512,512'")
+    parser.add_argument("--std-dev", type=parse_int_pairs, default=[(0, 50)],
+                        help="Noise standard deviation (a, b) as pairs separated by spaces, e.g., '0,50 8,8 6,10'")
+    parser.add_argument("-n", "--number-of-images", type=int, default=None,
+                        required=False, help="Number of images to process")
+    parser.add_argument("-d", "--dataset", type=str,
+                        choices=[None,  DATASET_DL_DIV2K_512, DATASET_DIV2K, DATASET_DL_EXTRAPRIMITIVES_DIV2K_512],
+                        default=None),
+    parser.add_argument("-b", "--blur", action="store_true")
+    parser.add_argument("--blur-index", type=int, nargs="+", default=None)
+    return parser
+
+
+def to_image(img: torch.Tensor):
+    return img.permute(0, 2, 3, 1).cpu().numpy()
+
+
+def infer(model, dataloader, config, device, output_dir: Path, traces: List[str] = ALL_TRACES, number_of_images=None, degradation_key=CONFIG_DEAD_LEAVES,
+          chosen_metrics=[METRIC_PSNR, METRIC_SSIM]):  # add METRIC_LPIPS here!
+    img_index = 0
+    if TRACES_ALL in traces:
+        traces = ALL_TRACES
+    if TRACES_METRICS in traces:
+        all_metrics = {}
+    else:
+        all_metrics = None
+    with torch.no_grad():
+        model.eval()
+        for img_degraded, img_target in tqdm(dataloader):
+            img_degraded = img_degraded.to(device)
+            img_target = img_target.to(device)
+            img_restored = model(img_degraded)
+            if TRACES_METRICS in traces:
+                metrics_input_per_image = compute_metrics(
+                    img_degraded, img_target, reduction=REDUCTION_SKIP, chosen_metrics=chosen_metrics)
+                metrics_per_image = compute_metrics(
+                    img_restored, img_target, reduction=REDUCTION_SKIP, chosen_metrics=chosen_metrics)
+                # print(metrics_per_image)
+            img_degraded = to_image(img_degraded)
+            img_target = to_image(img_target)
+            img_restored = to_image(img_restored)
+            for idx in range(img_restored.shape[0]):
+                degradation_parameters = dataloader.dataset.current_degradation[img_index]
+                common_prefix = f"{img_index:05d}_{img_degraded.shape[-3]:04d}x{img_degraded.shape[-2]:04d}"
+                common_prefix += f"_noise=[{config[DATALOADER][degradation_key]['noise_stddev'][0]:02d},{config[DATALOADER][degradation_key]['noise_stddev'][1]:02d}]"
+                suffix_deg = ""
+                if degradation_parameters['noise_stddev'] > 0:
+                    suffix_deg += f"_noise={round(degradation_parameters['noise_stddev']):02d}"
+                suffix_deg += f"_blur={degradation_parameters['blur_kernel_id']:04d}"
+                #if degradation_parameters.get("blur_kernel_id", False) else ""
+                save_path_pred = output_dir/f"{common_prefix}_pred{suffix_deg}_{config[PRETTY_NAME]}.png"
+                save_path_degr = output_dir/f"{common_prefix}_degr{suffix_deg}.png"
+                save_path_targ = output_dir/f"{common_prefix}_targ.png"
+                if TRACES_RESTORED in traces:
+                    Image(img_restored[idx]).save(save_path_pred)
+                if TRACES_DEGRADED in traces:
+                    Image(img_degraded[idx]).save(save_path_degr)
+                if TRACES_TARGET in traces:
+                    Image(img_target[idx]).save(save_path_targ)
+                if TRACES_METRICS in traces:
+                    # current_metrics = {"in": {}, "out": {}}
+                    # for key, value in metrics_per_image.items():
+                    #     print(f"{key}: {value[idx]:.3f}")
+                    #     current_metrics["in"][key] = metrics_input_per_image[key][idx].item()
+                    #     current_metrics["out"][key] = metrics_per_image[key][idx].item()
+                    current_metrics = {}
+                    for key, value in metrics_per_image.items():
+                        current_metrics["in_"+key] = metrics_input_per_image[key][idx].item()
+                        current_metrics["out_"+key] = metrics_per_image[key][idx].item()
+                    current_metrics["degradation"] = degradation_parameters
+                    current_metrics["size"] = (img_degraded.shape[-3], img_degraded.shape[-2])
+                    current_metrics["deadleaves_config"] = config[DATALOADER][degradation_key]
+                    current_metrics["restored"] = save_path_pred.relative_to(output_dir).as_posix()
+                    current_metrics["degraded"] = save_path_degr.relative_to(output_dir).as_posix()
+                    current_metrics["target"] = save_path_targ.relative_to(output_dir).as_posix()
+                    current_metrics["model"] = config[PRETTY_NAME]
+                    current_metrics["model_id"] = config[NAME]
+                    Parameters(current_metrics).save(output_dir/f"{common_prefix}_metrics.json")
+                    # for key, value in all_metrics.items():
+
+                    all_metrics[img_index] = current_metrics
+                img_index += 1
+                if number_of_images is not None and img_index > number_of_images:
+                    return all_metrics
+    return all_metrics
+
+
+def infer_main(argv, batch_mode=False):
+    parser = get_parser(batch_mode=batch_mode)
+    if batch_mode:
+        batch = Batch(argv)
+        batch.set_io_description(
+            input_help='input image files',
+            output_help=f'output directory {str(ROOT_DIR/INFERENCE_FOLDER_NAME)}',
+        )
+        batch.parse_args(parser)
+    else:
+        args = parser.parse_args(argv)
+    device = "cpu" if args.cpu else DEVICE
+    dataset = args.dataset
+    blur_flag = args.blur
+    for exp in args.experiments:
+        model_dict = get_default_models([exp], Path(args.models_storage), interactive_flag=False)
+        # print(list(model_dict.keys()))
+        current_model_dict = model_dict[list(model_dict.keys())[0]]
+        model = current_model_dict["model"]
+        config = current_model_dict["config"]
+        for std_dev, size, blur_index in product(args.std_dev, args.size, args.blur_index):
+            if dataset is None:
+                config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+                    blur_kernel_half_size=[0, 0],
+                    ds_factor=1,
+                    noise_stddev=list(std_dev),
+                    sampler=SAMPLER_SATURATED
+                )
+                config[DATALOADER]["gpu_gen"] = True
+                config[DATALOADER][SIZE] = size
+                config[DATALOADER][BATCH_SIZE][VALIDATION] = 1 if size[0] > 512 else 4
+            else:
+                config[DATALOADER][CONFIG_DEGRADATION] = dict(
+                    noise_stddev=list(std_dev),
+                    degradation_blur=DEGRADATION_BLUR_MAT if blur_flag else DEGRADATION_BLUR_NONE,
+                    blur_index=blur_index
+                )
+                config[DATALOADER][NAME] = dataset
+                config[DATALOADER][SIZE] = size
+                config[DATALOADER][BATCH_SIZE][VALIDATION] = 1 if size[0] > 512 else 4
+            dataloader = get_data_loader(config, frozen_seed=42)
+            # print(config)
+            output_dir = Path(args.output_dir)/(config[NAME] + "_" +
+                                                config[PRETTY_NAME])  # + "_" + f"{size[0]:04d}x{size[1]:04d}")
+            output_dir.mkdir(parents=True, exist_ok=True)
+
+            all_metrics = infer(model, dataloader[VALIDATION], config, device, output_dir,
+                                traces=args.traces, number_of_images=args.number_of_images,
+                                degradation_key=CONFIG_DEAD_LEAVES if dataset is None else CONFIG_DEGRADATION)
+            if all_metrics is not None:
+                # print(all_metrics)
+                df = pd.DataFrame(all_metrics).T
+                prefix = f"{size[0]:04d}x{size[1]:04d}"
+                if not (std_dev[0] == 0 and std_dev[1] == 0):
+                    prefix += f"_noise=[{std_dev[0]:02d},{std_dev[1]:02d}]"
+                if blur_index is not None:
+                    prefix += f"_blur={blur_index:02d}"
+                prefix += f"_{config[PRETTY_NAME]}"
+                df.to_csv(output_dir/f"__{prefix}_metrics_.csv", index=False)
+                # Normally this could go into another script to handle the metrics analyzis
+                # print(df)
+
+
+if __name__ == "__main__":
+    infer_main(sys.argv[1:])

scripts/interactive_inference_natural.py

@@ -0,0 +1,65 @@
+import sys
+sys.path.append("src")
+from interactive_pipe import interactive_pipeline
+from rstor.analyzis.interactive.pipelines import natural_inference_pipeline, morph_canvas, CANVAS
+from rstor.analyzis.interactive.model_selection import get_default_models
+from pathlib import Path
+from rstor.analyzis.parser import get_parser
+import argparse
+from batch_processing import Batch
+from interactive_pipe.data_objects.image import Image
+from rstor.analyzis.interactive.images import image_selector
+from rstor.analyzis.interactive.crop import plug_crop_selector
+from rstor.analyzis.interactive.metrics import plug_configure_metrics
+from interactive_pipe import interactive, KeyboardControl
+
+
+def plug_morph_canvas():
+    interactive(
+        canvas=KeyboardControl(CANVAS[0], CANVAS, name="canvas", keyup="p", modulo=True)
+    )(morph_canvas)
+
+
+def image_loading_batch(input: Path, args: argparse.Namespace) -> dict:
+    """Wrapper to load images files from a directory using batch_processing
+    """
+
+    if not args.disable_preload:
+        img = Image.from_file(input).data
+        return {"name": input.name, "path": input, "buffer": img}
+    else:
+        return {"name": input.name, "path": input, "buffer": None}
+
+
+def main(argv):
+    batch = Batch(argv)
+    batch.set_io_description(
+        input_help='input image files',
+        output_help=argparse.SUPPRESS
+    )
+    parser = get_parser()
+    parser.add_argument("-nop", "--disable-preload", action="store_true", help="Disable images preload")
+    args = batch.parse_args(parser)
+    # batch.set_multiprocessing_enabled(False)
+    img_list = batch.run(image_loading_batch)
+    if args.keyboard:
+        image_control = KeyboardControl(0, [0, len(img_list)-1], keydown="3", keyup="9", modulo=True)
+    else:
+        image_control = (0, [0, len(img_list)-1])
+    interactive(image_index=image_control)(image_selector)
+    plug_crop_selector(num_pad=args.keyboard)
+    plug_configure_metrics(key_shortcut="a")  # "a" if args.keyboard else None)
+    plug_morph_canvas()
+    model_dict = get_default_models(args.experiments, Path(args.models_storage), keyboard_control=args.keyboard)
+    interactive_pipeline(
+        gui=args.backend,
+        cache=True,
+        safe_input_buffer_deepcopy=False
+    )(natural_inference_pipeline)(
+        img_list,
+        model_dict
+    )
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

scripts/interactive_inference_synthetic.py

@@ -0,0 +1,25 @@
+import sys
+sys.path.append("src")
+from interactive_pipe import interactive_pipeline
+from rstor.synthetic_data.interactive.interactive_dead_leaves import dead_leave_plugin
+from rstor.analyzis.interactive.pipelines import deadleave_inference_pipeline
+from rstor.analyzis.interactive.model_selection import get_default_models
+from rstor.analyzis.interactive.crop import plug_crop_selector
+from rstor.analyzis.interactive.metrics import plug_configure_metrics
+from pathlib import Path
+from rstor.analyzis.parser import get_parser
+
+
+def main(argv):
+    parser = get_parser()
+    args = parser.parse_args(argv)
+    plug_crop_selector(num_pad=args.keyboard)
+    model_dict = get_default_models(args.experiments, Path(args.models_storage))
+    dead_leave_plugin(ds=1)
+    plug_configure_metrics(key_shortcut="a")
+    interactive_pipeline(gui="auto", cache=True, safe_input_buffer_deepcopy=False)(
+        deadleave_inference_pipeline)(model_dict)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

scripts/remote_training.py

@@ -0,0 +1,116 @@
+import kaggle
+from pathlib import Path, PurePosixPath
+import json
+try:
+    from __kaggle_login import kaggle_users
+except ImportError:
+    raise ImportError("Please create a __kaggle_login.py file with a kaggle_users" +
+                      "dict containing your Kaggle credentials.")
+import argparse
+import sys
+import subprocess
+from configuration import ROOT_DIR, OUTPUT_FOLDER_NAME
+from train import get_parser as get_train_parser
+from typing import Optional
+from configuration import KAGGLE_DATASET_LIST, NB_ID, GIT_USER, GIT_REPO, TRAIN_SCRIPT
+
+
+def get_git_branch_name():
+    try:
+        branch_name = subprocess.check_output(["git", "branch", "--show-current"]).strip().decode()
+        return branch_name
+    except subprocess.CalledProcessError:
+        return "Error: Could not determine the Git branch name."
+
+
+def prepare_notebook(
+    output_nb_path: Path,
+    exp: int,
+    branch: str,
+    git_user: str = None,
+    git_repo: str = None,
+    template_nb_path: Path = Path(__file__).parent/"remote_training_template.ipynb",
+    wandb_flag: bool = False,
+    output_dir: Path = "scripts/"+OUTPUT_FOLDER_NAME,
+    dataset_files: Optional[list] = None,
+    train_script: str = TRAIN_SCRIPT
+):
+    assert git_user is not None, "Please provide a git username for the repo"
+    assert git_repo is not None, "Please provide a git repo name for the repo"
+    expressions = [
+        ("exp", f"{exp}"),
+        ("branch", f"\'{branch}\'"),
+        ("git_user", f"\'{git_user}\'"),
+        ("git_repo", f"\'{git_repo}\'"),
+        ("wandb_flag", "True" if wandb_flag else "False"),
+        ("output_dir", "None" if output_dir is None else f"\'{output_dir}\'"),
+        ("dataset_files", "None" if dataset_files is None else f"{dataset_files}"),
+        ("train_script", "\'"+train_script+"\'")
+    ]
+    with open(template_nb_path) as f:
+        template_nb = f.readlines()
+        for line_idx, li in enumerate(template_nb):
+            for expr, expr_replace in expressions:
+                if f"!!!{expr}!!!" in li:
+                    template_nb[line_idx] = template_nb[line_idx].replace(f"!!!{expr}!!!", expr_replace)
+        template_nb = "".join(template_nb)
+    with open(output_nb_path, "w") as w:
+        w.write(template_nb)
+
+
+def main(argv):
+    parser = argparse.ArgumentParser(description="Train a model on Kaggle using a script")
+    parser.add_argument("-n", "--nb_id", type=str, help="Notebook name in kaggle", default=NB_ID)
+    parser.add_argument("-u", "--user", type=str, help="Kaggle user", choices=list(kaggle_users.keys()))
+    parser.add_argument("--branch", type=str, help="Git branch name", default=get_git_branch_name())
+    parser.add_argument("-p", "--push", action="store_true", help="Push")
+    parser.add_argument("-d", "--download", action="store_true", help="Download results")
+    get_train_parser(parser)
+    args = parser.parse_args(argv)
+    nb_id = args.nb_id
+    exp_str = "_".join(f"{exp:04d}" for exp in args.exp)
+    kaggle_user = kaggle_users[args.user]
+    uname_kaggle = kaggle_user["username"]
+    kaggle.api._load_config(kaggle_user)
+    if args.download:
+        tmp_dir = ROOT_DIR/f"__tmp_{exp_str}"
+        tmp_dir.mkdir(exist_ok=True, parents=True)
+        kaggle.api.kernels_output_cli(f"{kaggle_user['username']}/{nb_id}", path=str(tmp_dir))
+        subprocess.run(["tar", "-xzf", tmp_dir/"output.tgz"])
+        # @FIXME: windows probably does not have tar command
+        import shutil
+        shutil.rmtree(tmp_dir, ignore_errors=True)
+        return
+    kernel_root = ROOT_DIR/f"__nb_{uname_kaggle}"
+    kernel_root.mkdir(exist_ok=True, parents=True)
+
+    kernel_path = kernel_root/exp_str
+    kernel_path.mkdir(exist_ok=True, parents=True)
+    branch = args.branch
+    config = {
+        "id": str(PurePosixPath(f"{kaggle_user['username']}")/nb_id),
+        "title": nb_id.lower(),
+        "code_file": f"{nb_id}.ipynb",
+        "language": "python",
+        "kernel_type": "notebook",
+        "is_private": "true",
+        "enable_gpu": "true" if not args.cpu else "false",
+        "enable_tpu": "false",
+        "enable_internet": "true",
+        "dataset_sources": KAGGLE_DATASET_LIST,
+        "competition_sources": [],
+        "kernel_sources": [],
+        "model_sources": []
+    }
+    prepare_notebook((kernel_path/nb_id).with_suffix(".ipynb"), args.exp, branch,
+                     git_user=GIT_USER, git_repo=GIT_REPO, wandb_flag=not args.no_wandb)
+    assert (kernel_path/nb_id).with_suffix(".ipynb").exists()
+    with open(kernel_path/"kernel-metadata.json", "w") as f:
+        json.dump(config, f, indent=4)
+
+    if args.push:
+        kaggle.api.kernels_push_cli(str(kernel_path))
+
+
+if __name__ == '__main__':
+    main(sys.argv[1:])

scripts/remote_training_template.ipynb

@@ -0,0 +1 @@
+{"cells":[{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["exp = !!!exp!!!\n","branch = !!!branch!!!\n","git_user = !!!git_user!!!\n","git_repo = !!!git_repo!!!\n","wandb_flag = !!!wandb_flag!!!\n","output_dir = !!!output_dir!!!\n","dataset_files = !!!dataset_files!!!\n","train_script = !!!train_script!!!"]},{"cell_type":"markdown","metadata":{},"source":["# Clone git repo"]},{"cell_type":"code","execution_count":null,"metadata":{"trusted":true},"outputs":[],"source":["%cd ~\n","!git clone https://github.com/$git_user/$git_repo >/dev/null\n","%cd $git_repo\n","!git checkout $branch\n","!pip install -e ."]},{"cell_type":"markdown","metadata":{},"source":["# Load Kaggle datasets"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["!mkdir __dataset/deadleaves_div2k_512\n","!cp /kaggle/input/deadleaves-div2k-512/deadleaves_div2k_512/* \"__dataset/deadleaves_div2k_512\""]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["!mkdir __dataset/deadleaves_primitives_div2k_512\n","!cp /kaggle/input/deadleaves-primitives-div2k-512/deadleaves_primitives_div2k_512/* \"__dataset/deadleaves_primitives_div2k_512\""]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["!mkdir __dataset/div2k\n","!cp -r \"/kaggle/input/div2k-dataset/DIV2K_train_HR\" \"__dataset/div2k/\"\n","!cp -r \"/kaggle/input/div2k-dataset/DIV2K_valid_HR/\" \"__dataset/div2k/\""]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["!mkdir __dataset/kernels\n","!cp \"/kaggle/input/motion-blur-kernels/custom_blur_centered.mat\" __dataset/kernels/"]},{"cell_type":"markdown","metadata":{},"source":["# Setup weights and biases"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["if wandb_flag:\n","    from kaggle_secrets import UserSecretsClient\n","    user_secrets = UserSecretsClient()\n","    wandb_api_key = user_secrets.get_secret(\"wandb_api_key\")\n","\n","    !pip install wandb >/dev/null\n","    !wandb login $wandb_api_key"]},{"cell_type":"markdown","metadata":{},"source":["# Launch training"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["exp_str = ' '.join([str(e) for e in exp])\n","wb_ext = \"-nowb\" if not wandb_flag else \"\"\n","!python $train_script -e $exp_str $wb_ext"]},{"cell_type":"markdown","metadata":{},"source":["# Prepare outputs"]},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["if output_dir is not None:\n","    !tar -cvzf /kaggle/working/output.tgz $output_dir"]}],"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"datasetId":4234777,"sourceId":7299921,"sourceType":"datasetVersion"}],"dockerImageVersionId":30626,"isGpuEnabled":true,"isInternetEnabled":true,"language":"python","sourceType":"notebook"},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.12"}},"nbformat":4,"nbformat_minor":4}

scripts/save_deadleaves.py

@@ -0,0 +1,172 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Sat Mar 23 15:38:28 2024
+
+@author: jamyl
+"""
+import cv2
+from pathlib import Path
+from time import perf_counter
+import matplotlib.pyplot as plt
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+from numba import cuda
+from tqdm import tqdm
+import argparse
+from rstor.synthetic_data.dead_leaves_gpu import gpu_dead_leaves_chart
+from rstor.utils import DEFAULT_TORCH_FLOAT_TYPE
+from rstor.properties import DATASET_PATH, DATASET_DL_RANDOMRGB_1024, DATASET_DL_DIV2K_1024, SAMPLER_NATURAL, SAMPLER_UNIFORM, DATASET_DL_DIV2K_512, DATASET_DL_EXTRAPRIMITIVES_DIV2K_512
+
+
+class DeadLeavesDatasetGPU(Dataset):
+    def __init__(
+        self,
+        size: Tuple[int, int] = (128, 128),
+        length: int = 1000,
+        frozen_seed: int = None,  # useful for validation set!
+        ds_factor: int = 5,
+        **config_dead_leaves
+    ):
+
+        self.frozen_seed = frozen_seed
+        self.ds_factor = ds_factor
+        self.size = (size[0]*ds_factor, size[1]*ds_factor)
+        self.length = length
+        self.config_dead_leaves = config_dead_leaves
+
+        # downsample kernel
+        sigma = 3/5
+        k_size = 5  # This fits with sigma = 3/5, the cutoff value is 0.0038 (neglectable)
+        x = (torch.arange(k_size) - 2).to('cuda')
+        kernel = torch.stack(torch.meshgrid((x, x), indexing='ij'))
+        dist_sq = kernel[0]**2 + kernel[1]**2
+        kernel = (-dist_sq.square()/(2*sigma**2)).exp()
+        kernel = kernel / kernel.sum()
+        self.downsample_kernel = kernel.repeat(3, 1, 1, 1)  # shape [3, 1, k_size, k_size]
+
+    def __len__(self) -> int:
+        return self.length
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Get a single deadleave chart and its degraded version.
+
+        Args:
+            idx (int): index of the item to retrieve
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: degraded chart, target chart
+        """
+        seed = self.frozen_seed + idx if self.frozen_seed is not None else None
+
+        # Return numba device array
+        numba_chart = gpu_dead_leaves_chart(self.size, seed=seed, **self.config_dead_leaves)
+        if self.ds_factor > 1:
+            # print(f"Downsampling {chart.shape} with factor {self.ds_factor}...")
+
+            # Downsample using strided gaussian conv (sigma=3/5)
+            th_chart = torch.as_tensor(numba_chart, dtype=DEFAULT_TORCH_FLOAT_TYPE,
+                                       device="cuda").permute(2, 0, 1)[None]  # [b, c, h, w]
+            th_chart = F.pad(th_chart,
+                             pad=(2, 2, 0, 0),
+                             mode="replicate")
+            th_chart = F.conv2d(th_chart,
+                                self.downsample_kernel,
+                                padding='valid',
+                                groups=3,
+                                stride=self.ds_factor)
+
+            # Convert back to numba
+            numba_chart = cuda.as_cuda_array(th_chart.permute(0, 2, 3, 1))  # [b, h, w, c]
+
+        # convert back to numpy (temporary for legacy)
+        chart = numba_chart.copy_to_host()[0]
+
+        return chart
+
+
+def generate_images(path: Path, dataset: Dataset, imin=0):
+    for i in tqdm(range(imin, dataset.length)):
+        img = dataset[i]
+        img = (img * 255).astype(np.uint8)
+        out_path = path / "{:04d}.png".format(i)
+        cv2.imwrite(out_path.as_posix(), img)
+
+
+def bench(dataset):
+
+    print("dataset initialised")
+    t1 = perf_counter()
+    chart = dataset[0]
+
+    d = (perf_counter()-t1)
+    print(f"generation done {d}")
+    print(f"{d*1_000/60} min for 1_000")
+    plt.imshow(chart)
+    plt.show()
+
+
+if __name__ == "__main__":
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("-o", "--output-dir", type=str, default=str(DATASET_PATH))
+    argparser.add_argument(
+        "-n", "--name", type=str,
+        choices=[DATASET_DL_RANDOMRGB_1024, DATASET_DL_DIV2K_1024,
+                 DATASET_DL_DIV2K_512, DATASET_DL_EXTRAPRIMITIVES_DIV2K_512],
+        default=DATASET_DL_RANDOMRGB_1024
+    )
+    argparser.add_argument("-b", "--benchmark", action="store_true")
+    default_config = dict(
+        size=(1_024, 1_024),
+        length=1_000,
+        frozen_seed=42,
+        background_color=(0.2, 0.4, 0.6),
+        colored=True,
+        radius_min=5,
+        radius_max=2_000,
+        ds_factor=5,
+    )
+
+    args = argparser.parse_args()
+    dataset_dir = args.output_dir
+    name = args.name
+    path = Path(dataset_dir)/name
+    # print(path)
+    path.mkdir(parents=True, exist_ok=True)
+    if name == DATASET_DL_RANDOMRGB_1024:
+        config = default_config
+        config["sampler"] = SAMPLER_UNIFORM
+    elif name == DATASET_DL_DIV2K_1024:
+        config = default_config
+        config["sampler"] = SAMPLER_NATURAL
+        config["natural_image_list"] = sorted(
+            list((DATASET_PATH / "div2k" / "DIV2K_train_HR" / "DIV2K_train_HR").glob("*.png"))
+        )
+    elif name == DATASET_DL_DIV2K_512:
+        config = default_config
+        config["size"] = (512, 512)
+        config["rmin"] = 3
+        config["length"] = 4000
+        config["sampler"] = SAMPLER_NATURAL
+        config["natural_image_list"] = sorted(
+            list((DATASET_PATH / "div2k" / "DIV2K_train_HR" / "DIV2K_train_HR").glob("*.png"))
+        )
+    elif name == DATASET_DL_EXTRAPRIMITIVES_DIV2K_512:
+        config = default_config
+        config["size"] = (512, 512)
+        config["sampler"] = SAMPLER_NATURAL
+        config["circle_primitives"] = False
+        config["length"] = 4000
+        config["natural_image_list"] = sorted(
+            list((DATASET_PATH / "div2k" / "DIV2K_train_HR" / "DIV2K_train_HR").glob("*.png"))
+        )
+    else:
+        raise NotImplementedError
+    dataset = DeadLeavesDatasetGPU(**config)
+    if args.benchmark:
+        bench(dataset)
+    else:
+        generate_images(path, dataset)

scripts/train.py

@@ -0,0 +1,171 @@
+import sys
+import argparse
+from typing import Optional
+import torch
+import logging
+from pathlib import Path
+import json
+from tqdm import tqdm
+from rstor.properties import (
+    ID, NAME, NB_EPOCHS,
+    TRAIN, VALIDATION, LR,
+    LOSS_MSE, METRIC_PSNR, METRIC_SSIM,
+    DEVICE, SCHEDULER_CONFIGURATION, SCHEDULER, REDUCELRONPLATEAU,
+    REDUCTION_SUM,
+    SELECTED_METRICS,
+    LOSS
+)
+from rstor.learning.metrics import compute_metrics
+from rstor.learning.loss import compute_loss
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from configuration import WANDBSPACE, ROOT_DIR, OUTPUT_FOLDER_NAME
+from rstor.learning.experiments import get_training_content
+from rstor.learning.experiments_definition import get_experiment_config
+WANDB_AVAILABLE = False
+try:
+    WANDB_AVAILABLE = True
+    import wandb
+except ImportError:
+    logging.warning("Could not import wandb. Disabling wandb.")
+    pass
+
+
+def get_parser(parser: Optional[argparse.ArgumentParser] = None) -> argparse.ArgumentParser:
+    if parser is None:
+        parser = argparse.ArgumentParser(description="Train a model")
+    parser.add_argument("-e", "--exp", nargs="+", type=int, required=True, help="Experiment id")
+    parser.add_argument("-o", "--output-dir", type=str, default=ROOT_DIR/OUTPUT_FOLDER_NAME, help="Output directory")
+    parser.add_argument("-nowb", "--no-wandb", action="store_true", help="Disable weights and biases")
+    parser.add_argument("--cpu", action="store_true", help="Force CPU")
+    return parser
+
+
+def training_loop(
+    model,
+    optimizer,
+    dl_dict: dict,
+    config: dict,
+    scheduler=None,
+    device: str = DEVICE,
+    wandb_flag: bool = False,
+    output_dir: Path = None,
+):
+    best_accuracy = 0.
+    chosen_metrics = config.get(SELECTED_METRICS, [METRIC_PSNR, METRIC_SSIM])
+    for n_epoch in tqdm(range(config[NB_EPOCHS])):
+        current_metrics = {
+            TRAIN: 0.,
+            VALIDATION: 0.,
+            LR: optimizer.param_groups[0]['lr'],
+        }
+        for met in chosen_metrics:
+            current_metrics[met] = 0.
+        for phase in [TRAIN, VALIDATION]:
+            total_elements = 0
+            if phase == TRAIN:
+                model.train()
+            else:
+                model.eval()
+            for x, y in tqdm(dl_dict[phase], desc=f"{phase} - Epoch {n_epoch}"):
+                x, y = x.to(device), y.to(device)
+                optimizer.zero_grad()
+                with torch.set_grad_enabled(phase == TRAIN):
+                    y_pred = model(x)
+                    loss = compute_loss(y_pred, y, mode=config.get(LOSS, LOSS_MSE))
+                    if torch.isnan(loss):
+                        print(f"Loss is NaN at epoch {n_epoch} and phase {phase}!")
+                        continue
+                    if phase == TRAIN:
+                        loss.backward()
+                        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
+                        optimizer.step()
+                current_metrics[phase] += loss.item()
+                if phase == VALIDATION:
+                    metrics_on_batch = compute_metrics(
+                        y_pred,
+                        y,
+                        chosen_metrics=chosen_metrics,
+                        reduction=REDUCTION_SUM
+                    )
+                    total_elements += y_pred.shape[0]
+                    for k, v in metrics_on_batch.items():
+                        current_metrics[k] += v
+
+            current_metrics[phase] /= (len(dl_dict[phase]))
+            if phase == VALIDATION:
+                for k, v in metrics_on_batch.items():
+                    current_metrics[k] /= total_elements
+                    try:
+                        current_metrics[k] = current_metrics[k].item()
+                    except AttributeError:
+                        pass
+        debug_print = f"{phase}: Epoch {n_epoch} - Loss: {current_metrics[phase]:.3e} "
+        for k, v in current_metrics.items():
+            if k not in [TRAIN, VALIDATION, LR]:
+                debug_print += f"{k}: {v:.3} |"
+        print(debug_print)
+        if scheduler is not None and isinstance(scheduler, ReduceLROnPlateau):
+            scheduler.step(current_metrics[VALIDATION])
+        if output_dir is not None:
+            with open(output_dir/f"metrics_{n_epoch}.json", "w") as f:
+                json.dump(current_metrics, f)
+        if wandb_flag:
+            wandb.log(current_metrics)
+        if best_accuracy < current_metrics[METRIC_PSNR]:
+            best_accuracy = current_metrics[METRIC_PSNR]
+            if output_dir is not None:
+                print("new best model saved!")
+                torch.save(model.state_dict(), output_dir/"best_model.pt")
+    if output_dir is not None:
+        torch.save(model.cpu().state_dict(), output_dir/"last_model.pt")
+    return model
+
+
+def train(config: dict, output_dir: Path, device: str = DEVICE, wandb_flag: bool = False):
+    logging.basicConfig(level=logging.INFO)
+    logging.info(f"Training experiment {config[ID]} on device {device}...")
+    output_dir.mkdir(parents=True, exist_ok=True)
+    with open(output_dir/"config.json", "w") as f:
+        json.dump(config, f)
+    model, optimizer, dl_dict = get_training_content(config, training_mode=True, device=device)
+    model.to(device)
+    if wandb_flag:
+        import wandb
+        wandb.init(
+            project=WANDBSPACE,
+            entity="balthazarneveu",
+            name=config[NAME],
+            tags=["debug"],
+            # tags=["base"],
+            config=config
+        )
+    scheduler = None
+    if config.get(SCHEDULER, False):
+        scheduler_config = config[SCHEDULER_CONFIGURATION]
+        if config[SCHEDULER] == REDUCELRONPLATEAU:
+            scheduler = ReduceLROnPlateau(optimizer, mode='min', verbose=True, **scheduler_config)
+        else:
+            raise NameError(f"Scheduler {config[SCHEDULER]} not implemented")
+    model = training_loop(model, optimizer, dl_dict, config, scheduler=scheduler, device=device,
+                          wandb_flag=wandb_flag, output_dir=output_dir)
+
+    if wandb_flag:
+        wandb.finish()
+
+
+def train_main(argv):
+    parser = get_parser()
+    args = parser.parse_args(argv)
+    if not WANDB_AVAILABLE:
+        args.no_wandb = True
+    device = "cpu" if args.cpu else DEVICE
+    for exp in args.exp:
+        config = get_experiment_config(exp)
+        print(config)
+        output_dir = Path(args.output_dir)/config[NAME]
+        logging.info(f"Training experiment {config[ID]} on device {device}...")
+        train(config, device=device, output_dir=output_dir, wandb_flag=not args.no_wandb)
+
+
+if __name__ == "__main__":
+    train_main(sys.argv[1:])

src/rstor/analyzis/interactive/crop.py

@@ -0,0 +1,75 @@
+import cv2
+import numpy as np
+from interactive_pipe import interactive
+
+
+def get_color_channel_offset(image):
+    # size is defined in power of 2
+    if len(image.shape) == 2:
+        offset = 0
+    elif len(image.shape) == 3:
+        channel_guesser_max_size = 4
+        if image.shape[0] <= channel_guesser_max_size:  # channel first C,H,W
+            offset = 0
+        elif image.shape[-1] <= channel_guesser_max_size:  # channel last or numpy H,W,C
+            offset = 1
+    else:
+        raise NameError(f"Not supported shape {image.shape}")
+    return offset
+
+
+def crop_selector(image, center_x=0.5, center_y=0.5, size=9., global_params={}):
+    offset = get_color_channel_offset(image)
+    crop_size_pixels = int(2.**(size)/2.)
+    h, w = image.shape[-2-offset], image.shape[-1-offset]
+    ar = w/h
+    half_crop_h, half_crop_w = crop_size_pixels, int(ar*crop_size_pixels)
+
+    def round(val):
+        return int(np.round(val))
+    center_x_int = round(half_crop_w + center_x*(w-2*half_crop_w))
+    center_y_int = round(half_crop_h + center_y*(h-2*half_crop_h))
+    start_x = max(0, center_x_int-half_crop_w)
+    start_y = max(0, center_y_int-half_crop_h)
+    end_x = min(start_x+2*half_crop_w, w-1)
+    end_y = min(start_y+2*half_crop_h, h-1)
+    start_x = max(0, end_x-2*half_crop_w)
+    start_y = max(0, end_y-2*half_crop_h)
+    MAX_ALLOWED_SIZE = 512
+    w_resize = int(min(MAX_ALLOWED_SIZE, w))
+    h_resize = int(w_resize/w*h)
+    h_resize = int(min(MAX_ALLOWED_SIZE, h_resize))
+    w_resize = int(h_resize/h*w)
+    global_params["crop"] = (start_x, start_y, end_x, end_y)
+    global_params["resize"] = (w_resize, h_resize)
+    return
+
+
+def plug_crop_selector(num_pad: bool = False):
+    interactive(
+        center_x=(0.5, [0., 1.], "cx", ["4" if num_pad else "left", "6" if num_pad else "right"]),
+        center_y=(0.5, [0., 1.], "cy", ["8" if num_pad else "up", "2" if num_pad else "down"]),
+        size=(9., [6., 13., 0.3], "crop size", ["+", "-"])
+    )(crop_selector)
+
+
+def crop(*images, global_params={}):
+    images_resized = []
+    for image in images:
+        offset = get_color_channel_offset(image)
+        start_x, start_y, end_x, end_y = global_params["crop"]
+        w_resize, h_resize = global_params["resize"]
+        if offset == 0:
+            crop = image[..., start_y:end_y, start_x:end_x]
+        if offset == 1:
+            crop = image[..., start_y:end_y, start_x:end_x, :]
+        image_resized = cv2.resize(crop, (w_resize, h_resize), interpolation=cv2.INTER_NEAREST)
+        images_resized.append(image_resized)
+    return tuple(images_resized)
+
+
+def rescale_thumbnail(image, global_params={}):
+    if image is None:  # support no blur kernel!
+        return None
+    resize_dim = max(global_params.get("resize", (512, 512)))
+    return cv2.resize(image, (resize_dim, resize_dim), interpolation=cv2.INTER_NEAREST)

src/rstor/analyzis/interactive/degradation.py

@@ -0,0 +1,71 @@
+import numpy as np
+from interactive_pipe import interactive
+from skimage.filters import gaussian
+from rstor.properties import DATASET_BLUR_KERNEL_PATH
+from scipy.io import loadmat
+import cv2
+
+
+@interactive(
+    sigma=(3/5, [0., 2.])
+)
+def downsample(chart: np.ndarray, sigma=3/5, global_params={}):
+    ds_factor = global_params.get("ds_factor", 5)
+    if sigma > 0.:
+        ds_chart = gaussian(chart, sigma=(sigma, sigma, 0), mode='nearest', cval=0, preserve_range=True, truncate=4.0)
+    else:
+        ds_chart = chart.copy()
+    ds_chart = ds_chart[ds_factor//2::ds_factor, ds_factor//2::ds_factor]
+    return ds_chart
+
+
+@interactive(
+    k_size_x=(0, [0, 10]),
+    k_size_y=(0, [0, 10]),
+)
+def degrade_blur_gaussian(chart: np.ndarray, k_size_x: int = 1, k_size_y: int = 1):
+    if k_size_x == 0 and k_size_y == 0:
+        blurred = chart
+    blurred = cv2.GaussianBlur(chart, (2*k_size_x+1, 2*k_size_y+1), 0)
+    return blurred
+
+
+@interactive(
+    noise_stddev=(0., [0., 50.])
+)
+def degrade_noise(img: np.ndarray, noise_stddev=0., global_params={}):
+    seed = global_params.get("seed", 42)
+    np.random.seed(seed)
+    if noise_stddev > 0.:
+        noise = np.random.normal(0, noise_stddev/255., img.shape)
+        img = img.copy()+noise
+    return img
+
+
+@interactive(
+    ksize=(3, [1, 10])
+)
+def get_blur_kernel_box(ksize=3):
+    return np.ones((ksize, ksize), dtype=np.float32) / (1.*ksize**2)
+
+
+@interactive(
+    blur_index=(-1, [-1, 1000])
+)
+def get_blur_kernel(blur_index: int = -1, global_params={}):
+    if blur_index == -1:
+        return None
+    blur_mat = global_params.get("blur_mat", False)
+    if blur_mat is False:
+        blur_mat = loadmat(DATASET_BLUR_KERNEL_PATH)["kernels"].squeeze()
+        global_params["blur_mat"] = blur_mat
+    blur_k = blur_mat[blur_index]
+    blur_k = blur_k/blur_k.sum()
+    return blur_k
+
+
+def degrade_blur(img: np.ndarray, blur_kernel: np.ndarray, global_params={}):
+    if blur_kernel is None:
+        return img
+    img_blur = cv2.filter2D(img, -1, blur_kernel)
+    return img_blur

src/rstor/analyzis/interactive/images.py

@@ -0,0 +1,10 @@
+from interactive_pipe.data_objects.image import Image
+from typing import List
+
+
+def image_selector(image_list: List[dict], image_index: int = 0) -> dict:
+    current_image = image_list[image_index % len(image_list)]
+    img = current_image.get("buffer", None)
+    if img is None:
+        img = Image.from_file(current_image["path"]).data
+    return img

src/rstor/analyzis/interactive/inference.py

@@ -0,0 +1,12 @@
+import torch
+import numpy as np
+from rstor.properties import DEVICE
+
+
+def infer(degraded: np.ndarray, model: torch.nn.Module):
+    degraded_tensor = torch.from_numpy(degraded).permute(-1, 0, 1).float().unsqueeze(0)
+    model.eval()
+    with torch.no_grad():
+        output = model(degraded_tensor.to(DEVICE))
+    output = output.squeeze().permute(1, 2, 0).cpu().numpy()
+    return np.ascontiguousarray(output)

src/rstor/analyzis/interactive/metrics.py

@@ -0,0 +1,36 @@
+from rstor.learning.metrics import compute_metrics, ALL_METRICS
+import torch
+import numpy as np
+from rstor.properties import METRIC_PSNR, METRIC_SSIM
+from interactive_pipe import interactive, KeyboardControl
+from typing import Optional
+
+
+def plug_configure_metrics(key_shortcut: Optional[str] = None) -> None:
+    interactive(
+        advanced_metrics=KeyboardControl(False, keydown=key_shortcut) if key_shortcut is not None else (True,)
+    )(configure_metrics)
+
+
+def configure_metrics(advanced_metrics=False, global_params={}) -> None:
+    chosen_metrics = ALL_METRICS if advanced_metrics else [METRIC_PSNR, METRIC_SSIM]
+    global_params["chosen_metrics"] = chosen_metrics
+
+
+def get_metrics(prediction: torch.Tensor, target: torch.Tensor,
+                image_name: str,  # use functools.partial to root where you want the title to appear
+                global_params: dict = {}) -> None:
+    if isinstance(prediction, np.ndarray):
+        prediction_ = torch.from_numpy(prediction).permute(-1, 0, 1).float().unsqueeze(0)
+    else:
+        prediction_ = prediction
+    if isinstance(target, np.ndarray):
+        target_ = torch.from_numpy(target).permute(-1, 0, 1).float().unsqueeze(0)
+    else:
+        target_ = target
+    chosen_metrics = global_params.get("chosen_metrics", [METRIC_PSNR])
+    metrics = compute_metrics(prediction_, target_, chosen_metrics=chosen_metrics)
+    global_params["metrics"] = metrics
+    title = f"{image_name}: "
+    title += " ".join([f"{key}: {value:.4f}" for key, value in metrics.items()])
+    global_params["__output_styles"][image_name] = {"title": title, "image_name": image_name}

src/rstor/analyzis/interactive/model_selection.py

@@ -0,0 +1,58 @@
+import torch
+from interactive_pipe import KeyboardControl
+from rstor.learning.experiments import get_training_content
+from rstor.learning.experiments_definition import get_experiment_config
+from rstor.properties import DEVICE, PRETTY_NAME
+from tqdm import tqdm
+from pathlib import Path
+from typing import List, Tuple
+
+from interactive_pipe import interactive
+MODELS_PATH = Path("scripts")/"__output"
+
+
+def model_selector(models_dict: dict, global_params={}, model_name="vanilla"):
+    if isinstance(model_name, str):
+        current_model = models_dict[model_name]
+    elif isinstance(model_name, int):
+        model_names = [name for name in models_dict.keys()]
+        current_model = models_dict[model_names[model_name % len(model_names)]]
+    else:
+        raise ValueError(f"Model name {model_name} not understood")
+    global_params["model_config"] = current_model["config"]
+    return current_model["model"]
+
+
+def get_model_from_exp(exp: int, model_storage: Path = MODELS_PATH, device=DEVICE) -> Tuple[torch.nn.Module, dict]:
+    config = get_experiment_config(exp)
+    model, _, _ = get_training_content(config, training_mode=False)
+    model_path = torch.load(model_storage/f"{exp:04d}"/"best_model.pt")
+    assert model_path is not None, f"Model {exp} not found"
+    model.load_state_dict(model_path)
+    model = model.to(device)
+    return model, config
+
+
+def get_default_models(
+    exp_list: List[int] = [1000, 1001],
+    model_storage: Path = MODELS_PATH,
+    keyboard_control: bool = False,
+    interactive_flag: bool = True
+) -> dict:
+    model_dict = {}
+    assert model_storage.exists(), f"Model storage {model_storage} does not exist"
+    for exp in tqdm(exp_list, desc="Loading models"):
+        model, config = get_model_from_exp(exp, model_storage=model_storage)
+        name = config.get(PRETTY_NAME, f"{exp:04d}")
+        model_dict[name] = {
+            "model": model,
+            "config": config
+        }
+    exp_names = [name for name in model_dict.keys()]
+    if interactive_flag:
+        if keyboard_control:
+            model_control = KeyboardControl(0, [0, len(exp_names)-1], keydown="pagedown", keyup="pageup", modulo=True)
+        else:
+            model_control = (exp_names[0], exp_names)
+        interactive(model_name=model_control)(model_selector)  # Create the model dialog
+    return model_dict

src/rstor/analyzis/interactive/pipelines.py

@@ -0,0 +1,61 @@
+from rstor.synthetic_data.interactive.interactive_dead_leaves import generate_deadleave
+from rstor.analyzis.interactive.crop import crop_selector, crop, rescale_thumbnail
+from rstor.analyzis.interactive.inference import infer
+from rstor.analyzis.interactive.degradation import degrade_noise, degrade_blur, downsample, degrade_blur_gaussian, get_blur_kernel
+from rstor.analyzis.interactive.model_selection import model_selector
+from rstor.analyzis.interactive.images import image_selector
+from rstor.analyzis.interactive.metrics import get_metrics, configure_metrics
+from interactive_pipe import interactive, KeyboardControl
+from typing import Tuple, List
+from functools import partial
+import numpy as np
+
+
+get_metrics_restored = partial(get_metrics, image_name="restored")
+get_metrics_degraded = partial(get_metrics, image_name="degraded")
+
+
+def deadleave_inference_pipeline(models_dict: dict) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    groundtruth = generate_deadleave()
+    groundtruth = downsample(groundtruth)
+    model = model_selector(models_dict)
+    degraded = degrade_blur_gaussian(groundtruth)
+    degraded = degrade_noise(degraded)
+    restored = infer(degraded, model)
+    crop_selector(restored)
+    groundtruth, degraded, restored = crop(groundtruth, degraded, restored)
+    configure_metrics()
+    get_metrics_restored(restored, groundtruth)
+    get_metrics_degraded(degraded, groundtruth)
+    return groundtruth, degraded, restored
+
+
+CANVAS_DICT = {
+    "demo": [["degraded", "restored"]],
+    "landscape_light": [["degraded", "restored", "groundtruth"]],
+    "landscape": [["degraded", "restored", "blur_kernel", "groundtruth"]],
+    "full": [["degraded", "restored"], ["blur_kernel", "groundtruth"]]
+}
+CANVAS = list(CANVAS_DICT.keys())
+
+
+def morph_canvas(canvas=CANVAS[0], global_params={}):
+    global_params["__pipeline"].outputs = CANVAS_DICT[canvas]
+    return None
+
+
+def natural_inference_pipeline(input_image_list: List[np.ndarray], models_dict: dict):
+    model = model_selector(models_dict)
+    img_clean = image_selector(input_image_list)
+    crop_selector(img_clean)
+    groundtruth = crop(img_clean)
+    blur_kernel = get_blur_kernel()
+    degraded = degrade_blur(groundtruth, blur_kernel)
+    degraded = degrade_noise(degraded)
+    blur_kernel = rescale_thumbnail(blur_kernel)
+    restored = infer(degraded, model)
+    configure_metrics()
+    get_metrics_restored(restored, groundtruth)
+    get_metrics_degraded(degraded, groundtruth)
+    morph_canvas()
+    return [[degraded, restored], [blur_kernel, groundtruth]]

src/rstor/analyzis/metrics_plots.py

@@ -0,0 +1,73 @@
+import json
+import matplotlib.pyplot as plt
+import numpy as np
+from pathlib import Path
+import pandas as pd
+
+
+def snr_to_sigma(snr):
+    return 10**(-snr/20.)*255.
+
+
+def sigma_to_snr(sigma):
+    return -20.*np.log10(sigma/255.)
+
+
+def plot_results(selected_paths, title=None, diff=True, ylim=None):
+    # plt.figure(figsize=(10, 10))
+    all_stats = {}
+    fig, ax = plt.subplots(layout='constrained', figsize=(10, 10))
+    for selected_path, selected_regex in selected_paths:
+        selected_path = Path(selected_path)
+        assert selected_path.exists()
+        results_path = sorted(list(selected_path.glob(selected_regex)))
+        stats = []
+        for result_path in results_path:
+            df = pd.read_csv(result_path)
+            in_psnr = df["in_PSNR"].mean()
+            out_psnr = df["out_PSNR"].mean()
+            out_ssim = df["out_SSIM"].mean()
+            noise_stddev = np.array([
+                float(el.replace("}", "").split(":")[1]) for el in df["degradation"]]).mean()
+            stats.append({
+                # "label": label,
+                "in_psnr": in_psnr,
+                "out_psnr": out_psnr,
+                "noise_stddev": noise_stddev,
+                "ssim": out_ssim
+            })
+            label = selected_path.name + " " + df["size"][0]
+        
+        stats_array = pd.DataFrame(stats)
+        all_stats[label] = stats_array
+        x_data = stats_array["in_psnr"].copy()
+        x_data = snr_to_sigma(x_data)
+
+        ax.plot(
+            x_data,
+            stats_array["out_psnr"]-stats_array["in_psnr"] if diff else stats_array["out_psnr"],
+            "-o",
+            label=label
+        )
+        # label=selected_path.name)
+    if not diff:
+        neutral_sigma = np.linspace(1, 80, 80)
+        ax.plot(neutral_sigma, sigma_to_snr(neutral_sigma), "k--", alpha=0.1, label="Neutral")
+    secax = ax.secondary_xaxis('top', functions=(sigma_to_snr, snr_to_sigma))
+    secax.set_xlabel('PSNR [db]')
+
+    ax.set_xlabel("sigma 255")
+    ax.set_ylabel("PSNR improvement" if diff else "PSNR out")
+    plt.xlim(1., 50.)
+    if diff:
+        plt.ylim(0, 15)
+    else:
+        if ylim is not None:
+            plt.ylim(*ylim)
+    if title is not None:
+        plt.title(title)
+    plt.legend()
+    plt.grid()
+    plt.show()
+
+    return all_stats

src/rstor/analyzis/parser.py

@@ -0,0 +1,26 @@
+from rstor.analyzis.interactive.model_selection import MODELS_PATH
+import argparse
+
+
+def get_models_parser(parser: argparse.ArgumentParser = None, help: str = "Inference",
+                      default_models_path: str = MODELS_PATH) -> argparse.ArgumentParser:
+    if parser is None:
+        parser = argparse.ArgumentParser(description=help)
+    parser.add_argument("-e", "--experiments", type=int, nargs="+", required=True,
+                        help="Experience indexes to be used at inference time")
+    parser.add_argument("-m", "--models-storage", type=str, help="Model storage path", default=default_models_path)
+    return parser
+
+
+def get_parser(
+    parser: argparse.ArgumentParser = None,
+    help: str = "Live inference pipeline"
+) -> argparse.ArgumentParser:
+    """Generic parser for live interactive inference
+    """
+    if parser is None:
+        parser = argparse.ArgumentParser(description=help)
+    get_models_parser(parser=parser, help=help)
+    parser.add_argument("-k", "--keyboard", action="store_true", help="Keyboard control - less sliders")
+    parser.add_argument("-b", "--backend", default="gradio", help="Backend to use for the GUI", choices=["gradio", "qt"])
+    return parser

src/rstor/architecture/base.py

@@ -0,0 +1,56 @@
+import torch
+from rstor.properties import LEAKY_RELU, RELU, SIMPLE_GATE
+from typing import Optional, Tuple
+
+
+class SimpleGate(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x.chunk(2, dim=1)
+        return x1 * x2
+
+
+def get_non_linearity(activation: str):
+    if activation == LEAKY_RELU:
+        non_linearity = torch.nn.LeakyReLU()
+    elif activation == RELU:
+        non_linearity = torch.nn.ReLU()
+    elif activation is None:
+        non_linearity = torch.nn.Identity()
+    elif activation == SIMPLE_GATE:
+        non_linearity = SimpleGate()
+    else:
+        raise ValueError(f"Unknown activation {activation}")
+    return non_linearity
+
+
+class BaseModel(torch.nn.Module):
+    """Base class for all restoration models with additional useful methods"""
+
+    def count_parameters(self):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+    def receptive_field(
+        self,
+        channels: Optional[int] = 3,
+        size: Optional[int] = 256,
+        device: Optional[str] = None
+    ) -> Tuple[int, int]:
+        """Compute the receptive field of the model
+
+        Returns:
+            int: receptive field
+        """
+        input_tensor = torch.ones(1, channels, size, size, requires_grad=True)
+        if device is not None:
+            input_tensor = input_tensor.to(device)
+        out = self.forward(input_tensor)
+        grad = torch.zeros_like(out)
+        grad[..., out.shape[-2]//2, out.shape[-1]//2] = torch.nan  # set NaN gradient at the middle of the output
+        out.backward(gradient=grad)
+        self.zero_grad()
+        receptive_field_mask = input_tensor.grad.isnan()[0, 0]
+        receptive_field_indexes = torch.where(receptive_field_mask)
+        # Count NaN in the input
+        receptive_x = 1+receptive_field_indexes[-1].max() - receptive_field_indexes[-1].min()  # Horizontal x
+        receptive_y = 1+receptive_field_indexes[-2].max() - receptive_field_indexes[-2].min()  # Vertical y
+        return receptive_x.item(), receptive_y.item()

src/rstor/architecture/convolution_blocks.py

@@ -0,0 +1,47 @@
+import torch
+from rstor.properties import LEAKY_RELU
+from rstor.architecture.base import get_non_linearity
+
+
+class BaseConvolutionBlock(torch.nn.Module):
+    def __init__(
+        self,
+        ch_in: int,
+        ch_out: int,
+        k_size: int,
+        activation=LEAKY_RELU,
+        bias: bool = True
+    ) -> None:
+        super().__init__()
+        self.conv = torch.nn.Conv2d(ch_in, ch_out, k_size, padding=k_size//2, bias=bias)
+        self.non_linearity = get_non_linearity(activation)
+        self.conv_non_lin = torch.nn.Sequential(self.conv, self.non_linearity)
+
+    def forward(self, x_in: torch.Tensor) -> torch.Tensor:
+        return self.conv_non_lin(x_in)
+
+
+class ResConvolutionBlock(torch.nn.Module):
+    def __init__(
+        self,
+        ch_in: int,
+        ch_out: int,
+        k_size: int,
+        activation=LEAKY_RELU,
+        bias: bool = True,
+        residual: bool = True
+    ) -> None:
+        super().__init__()
+        self.conv1 = torch.nn.Conv2d(ch_in, ch_out, k_size, padding=k_size//2, bias=bias)
+        self.non_linearity = get_non_linearity(activation)
+        self.conv2 = torch.nn.Conv2d(ch_out, ch_out, k_size, padding=k_size//2, bias=bias)
+        self.residual = residual
+
+    def forward(self, x_in: torch.Tensor) -> torch.Tensor:
+        y = self.conv1(x_in)
+        y = self.non_linearity(y)
+        y = self.conv2(y)
+        if self.residual:
+            y = x_in + y
+        y = self.non_linearity(y)
+        return y

src/rstor/architecture/nafnet.py

@@ -0,0 +1,299 @@
+"""
+NAFNet: Non linear activation free neural network
+Architecture adapted from Simple Baselines for Image Restoration
+https://github.com/megvii-research/NAFNet/tree/main
+"""
+from torch import nn
+import torch.nn.functional as F
+import torch
+from rstor.architecture.base import BaseModel, get_non_linearity
+from typing import Optional, List
+from rstor.properties import RELU, SIMPLE_GATE
+
+
+class LayerNormFunction(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, weight, bias, eps):
+        ctx.eps = eps
+        N, C, H, W = x.size()
+        mu = x.mean(1, keepdim=True)
+        var = (x - mu).pow(2).mean(1, keepdim=True)
+        y = (x - mu) / (var + eps).sqrt()
+        ctx.save_for_backward(y, var, weight)
+        y = weight.view(1, C, 1, 1) * y + bias.view(1, C, 1, 1)
+        return y
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        eps = ctx.eps
+
+        N, C, H, W = grad_output.size()
+        y, var, weight = ctx.saved_variables
+        g = grad_output * weight.view(1, C, 1, 1)
+        mean_g = g.mean(dim=1, keepdim=True)
+
+        mean_gy = (g * y).mean(dim=1, keepdim=True)
+        gx = 1. / torch.sqrt(var + eps) * (g - y * mean_gy - mean_g)
+        return gx, (grad_output * y).sum(dim=3).sum(dim=2).sum(dim=0), grad_output.sum(dim=3).sum(dim=2).sum(
+            dim=0), None
+
+
+class LayerNorm2d(nn.Module):
+    def __init__(self, channels, eps=1e-6):
+        super(LayerNorm2d, self).__init__()
+        self.register_parameter('weight', nn.Parameter(torch.ones(channels)))
+        self.register_parameter('bias', nn.Parameter(torch.zeros(channels)))
+        self.eps = eps
+
+    def forward(self, x):
+        return LayerNormFunction.apply(x, self.weight, self.bias, self.eps)
+
+
+class NAFBlock(nn.Module):
+    def __init__(
+        self,
+        c, DW_Expand=2, FFN_Expand=2, drop_out_rate=0.,
+        activation: Optional[str] = SIMPLE_GATE,
+        layer_norm_flag: Optional[bool] = True,
+        channel_attention_flag: Optional[bool] = True,
+    ):
+        super().__init__()
+        self.layer_norm_flag = layer_norm_flag
+        self.channel_attention_flag = channel_attention_flag
+        dw_channel = c * DW_Expand
+        half_dw_channel = dw_channel // 2
+        self.conv1 = nn.Conv2d(in_channels=c, out_channels=dw_channel, kernel_size=1,
+                               padding=0, stride=1, groups=1, bias=True)
+        self.conv2 = nn.Conv2d(
+            in_channels=dw_channel,
+            out_channels=dw_channel if activation == SIMPLE_GATE else half_dw_channel,
+            kernel_size=3,
+            padding=1, stride=1,
+            groups=dw_channel if activation == SIMPLE_GATE else half_dw_channel,
+            bias=True
+        )
+        # To grand the same amount of parameters between Simple Gate and ReLU versions...
+        # Conv2 has to reduce the number of channels to half but... using grouped convolution
+        # w -> w/2 ... not really a depthwise convolution but rather by channels of 2!
+        self.conv3 = nn.Conv2d(in_channels=half_dw_channel, out_channels=c,
+                               kernel_size=1, padding=0, stride=1, groups=1, bias=True)
+
+        # Simplified Channel Attention
+        if self.channel_attention_flag:
+            self.sca = nn.Sequential(
+                nn.AdaptiveAvgPool2d(1),
+                nn.Conv2d(in_channels=half_dw_channel, out_channels=half_dw_channel, kernel_size=1,
+                          padding=0, stride=1,
+                          groups=1, bias=True),
+            )
+
+        # SimpleGate
+        self.sg = get_non_linearity(activation)
+        ffn_channel = FFN_Expand
+        half_ffn_channel = ffn_channel // 2 if activation == SIMPLE_GATE else ffn_channel
+        self.conv4 = nn.Conv2d(
+            in_channels=c,
+            out_channels=ffn_channel if activation == SIMPLE_GATE else half_ffn_channel,
+            kernel_size=1,
+            padding=0, stride=1, groups=1, bias=True)
+        self.conv5 = nn.Conv2d(in_channels=half_ffn_channel, out_channels=c,
+                               kernel_size=1, padding=0, stride=1, groups=1, bias=True)
+        if self.layer_norm_flag:
+            self.norm1 = LayerNorm2d(c)
+            self.norm2 = LayerNorm2d(c)
+
+        self.dropout1 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity()
+        self.dropout2 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity()
+
+        self.beta = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
+        self.gamma = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
+
+    def forward(self, inp):
+        x = inp
+        if self.layer_norm_flag:
+            x = self.norm1(x)
+
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.sg(x)
+        if self.channel_attention_flag:
+            x = x * self.sca(x)
+        x = self.conv3(x)
+
+        x = self.dropout1(x)
+
+        y = inp + x * self.beta
+
+        x = self.conv4(self.norm2(y) if self.layer_norm_flag else y)
+        x = self.sg(x)
+        x = self.conv5(x)
+
+        x = self.dropout2(x)
+
+        return y + x * self.gamma
+
+
+class NAFNet(BaseModel):
+    def __init__(
+            self,
+            img_channel: Optional[int] = 3,
+            width: Optional[int] = 16,
+            middle_blk_num: Optional[int] = 1,
+            enc_blk_nums: List[int] = [],
+            dec_blk_nums: List[int] = [],
+            activation: Optional[bool] = SIMPLE_GATE,
+            layer_norm_flag: Optional[bool] = True,
+            channel_attention_flag: Optional[bool] = True,
+    ) -> None:
+        super().__init__()
+
+        self.intro = nn.Conv2d(
+            in_channels=img_channel,
+            out_channels=width,
+            kernel_size=3,
+            padding=1, stride=1,
+            groups=1,
+            bias=True
+        )
+        config_block = {
+            "activation": activation,
+            "layer_norm_flag": layer_norm_flag,
+            "channel_attention_flag": channel_attention_flag
+        }
+        self.ending = nn.Conv2d(
+            in_channels=width, out_channels=img_channel, kernel_size=3,
+            padding=1, stride=1, groups=1,
+            bias=True)
+
+        self.encoders = nn.ModuleList()
+        self.decoders = nn.ModuleList()
+        self.middle_blks = nn.ModuleList()
+        self.ups = nn.ModuleList()
+        self.downs = nn.ModuleList()
+
+        chan = width
+        for num in enc_blk_nums:
+            self.encoders.append(
+                nn.Sequential(
+                    *[NAFBlock(chan, **config_block) for _ in range(num)]
+                )
+            )
+            self.downs.append(
+                nn.Conv2d(chan, 2*chan, 2, 2)
+            )
+            chan = chan * 2
+
+        self.middle_blks = \
+            nn.Sequential(
+                *[NAFBlock(chan, **config_block) for _ in range(middle_blk_num)]
+            )
+
+        for num in dec_blk_nums:
+            self.ups.append(
+                nn.Sequential(
+                    nn.Conv2d(chan, chan * 2, 1, bias=False),
+                    nn.PixelShuffle(2)
+                )
+            )
+            chan = chan // 2
+            self.decoders.append(
+                nn.Sequential(
+                    *[NAFBlock(chan, **config_block) for _ in range(num)]
+                )
+            )
+
+        self.padder_size = 2 ** len(self.encoders)
+
+    def forward(self, inp: torch.Tensor) -> torch.Tensor:
+        B, C, H, W = inp.shape
+        inp = self.sanitize_image_size(inp)
+
+        x = self.intro(inp)
+
+        encs = []
+
+        for encoder, down in zip(self.encoders, self.downs):
+            x = encoder(x)
+            encs.append(x)
+            x = down(x)
+
+        x = self.middle_blks(x)
+
+        for decoder, up, enc_skip in zip(self.decoders, self.ups, encs[::-1]):
+            x = up(x)
+            x = x + enc_skip
+            x = decoder(x)
+
+        x = self.ending(x)
+        x = x + inp
+
+        return x[:, :, :H, :W]
+
+    def sanitize_image_size(self, x: torch.Tensor) -> torch.Tensor:
+        _, _, h, w = x.size()
+        mod_pad_h = (self.padder_size - h % self.padder_size) % self.padder_size
+        mod_pad_w = (self.padder_size - w % self.padder_size) % self.padder_size
+        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h))
+        return x
+
+
+class UNet(NAFNet):
+    def __init__(
+            self,
+            activation: Optional[bool] = RELU,
+            layer_norm_flag: Optional[bool] = False,
+            channel_attention_flag: Optional[bool] = False,
+            **kwargs):
+        super().__init__(
+            activation=activation,
+            layer_norm_flag=layer_norm_flag,
+            channel_attention_flag=channel_attention_flag, **kwargs)
+
+
+if __name__ == '__main__':
+    tiny_recetive_field = True
+    if tiny_recetive_field:
+        enc_blks = [1, 1, 2]
+        middle_blk_num = 1
+        dec_blks = [1, 1, 1]
+        width = 16
+        # Receptive field is 208x208
+    else:
+        enc_blks = [1, 1, 1, 28]
+        middle_blk_num = 1
+        dec_blks = [1, 1, 1, 1]
+        width = 2
+        # Receptive field is 544x544
+    device = "cpu"
+
+    for model_name in ["NAFNet", "UNet"]:
+        if model_name == "NAFNet":
+            model = NAFNet(
+                img_channel=3,
+                width=width,
+                middle_blk_num=middle_blk_num,
+                enc_blk_nums=enc_blks,
+                dec_blk_nums=dec_blks,
+                activation=SIMPLE_GATE,
+                layer_norm_flag=False,
+                channel_attention_flag=False
+            )
+        if model_name == "UNet":
+            model = UNet(
+                img_channel=3,
+                width=width,
+                middle_blk_num=middle_blk_num,
+                enc_blk_nums=enc_blks,
+                dec_blk_nums=dec_blks
+            )
+        model.to(device)
+        with torch.no_grad():
+            x = torch.randn(1, 3, 256, 256).to(device)
+            y = model(x)
+
+            # print(y.shape)
+            # print(y)
+            # print(model)
+            print(f"{model.count_parameters()/1E3:.2f}k parameters")
+        print(model.receptive_field(size=256 if tiny_recetive_field else 1024, device=device))

src/rstor/architecture/selector.py

@@ -0,0 +1,19 @@
+from rstor.properties import MODEL, NAME, N_PARAMS, ARCHITECTURE
+from rstor.architecture.stacked_convolutions import StackedConvolutions
+from rstor.architecture.nafnet import NAFNet, UNet
+import torch
+
+
+def load_architecture(config: dict) -> torch.nn.Module:
+    conf_model = config[MODEL][ARCHITECTURE]
+    if config[MODEL][NAME] == StackedConvolutions.__name__:
+        model = StackedConvolutions(**conf_model)
+    elif config[MODEL][NAME] == NAFNet.__name__:
+        model = NAFNet(**conf_model)
+    elif config[MODEL][NAME] == UNet.__name__:
+        model = UNet(**conf_model)
+    else:
+        raise ValueError(f"Unknown model {config[MODEL][NAME]}")
+    config[MODEL][N_PARAMS] = model.count_parameters()
+    config[MODEL]["receptive_field"] = model.receptive_field()
+    return model

src/rstor/architecture/stacked_convolutions.py

@@ -0,0 +1,30 @@
+from rstor.architecture.base import BaseModel
+from rstor.architecture.convolution_blocks import BaseConvolutionBlock, ResConvolutionBlock
+from rstor.properties import LEAKY_RELU
+import torch
+
+
+class StackedConvolutions(BaseModel):
+    def __init__(self,
+                 ch_in: int = 3,
+                 ch_out: int = 3,
+                 h_dim: int = 64,
+                 num_layers: int = 8,
+                 k_size: int = 3,
+                 activation: str = LEAKY_RELU,
+                 bias: bool = True,
+                 ) -> None:
+        super().__init__()
+        assert num_layers % 2 == 0, "Number of layers should be even"
+        self.conv_in_modality = BaseConvolutionBlock(
+            ch_in, h_dim, k_size, activation=activation, bias=bias)
+        conv_list = []
+        for _i in range(num_layers-2):
+            conv_list.append(ResConvolutionBlock(
+                h_dim, h_dim, k_size, activation=activation, bias=bias, residual=True))
+        self.conv_out_modality = BaseConvolutionBlock(
+            h_dim, ch_out, k_size, activation=None, bias=bias)
+        self.conv_stack = torch.nn.Sequential(self.conv_in_modality, *conv_list, self.conv_out_modality)
+
+    def forward(self, x_in: torch.Tensor) -> torch.Tensor:
+        return self.conv_stack(x_in)

src/rstor/data/augmentation.py

@@ -0,0 +1,27 @@
+import torch
+from typing import Tuple, Optional
+
+
+def augment_flip(
+    img: torch.Tensor,
+    flip: Optional[Tuple[bool, bool]] = None
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Roll pixels horizontally to avoid negative index
+
+    Args:
+        img (torch.Tensor): [N, 3, H, W] image tensor
+        lab (torch.Tensor): [N, 3, H, W] label tensor
+        flip (Optional[bool], optional): forced flip_h, flip_v value. Defaults to None.
+        If not provided, a random flip_h, flip_v values are used
+    Returns:
+        torch.Tensor, torch.Tensor: flipped image, labels
+
+    """
+    if flip is None:
+        flip = torch.randint(0, 2, (2,))
+    flipped_img = img
+    if flip[0] > 0:
+        flipped_img = torch.flip(flipped_img, (-1,))
+    if flip[1] > 0:
+        flipped_img = torch.flip(flipped_img, (-2,))
+    return flipped_img

src/rstor/data/dataloader.py

@@ -0,0 +1,120 @@
+from torch.utils.data import DataLoader
+from rstor.data.synthetic_dataloader import DeadLeavesDataset, DeadLeavesDatasetGPU
+from rstor.data.stored_images_dataloader import RestorationDataset
+from rstor.properties import (
+    DATALOADER, BATCH_SIZE, TRAIN, VALIDATION, LENGTH, CONFIG_DEAD_LEAVES, SIZE, NAME, CONFIG_DEGRADATION,
+    DATASET_SYNTH_LIST, DATASET_DIV2K,
+    DATASET_PATH
+)
+from typing import Optional
+from random import seed, shuffle
+
+
+def get_data_loader_synthetic(config, frozen_seed=42):
+    # print(config[DATALOADER].get(CONFIG_DEAD_LEAVES, {}))
+    if config[DATALOADER].get("gpu_gen", False):
+        print("Using GPU dead leaves generator")
+        ds = DeadLeavesDatasetGPU
+    else:
+        ds = DeadLeavesDataset
+    dl_train = ds(config[DATALOADER][SIZE], config[DATALOADER][LENGTH][TRAIN],
+                  frozen_seed=None, **config[DATALOADER].get(CONFIG_DEAD_LEAVES, {}))
+    dl_valid = ds(config[DATALOADER][SIZE], config[DATALOADER][LENGTH][VALIDATION],
+                  frozen_seed=frozen_seed, **config[DATALOADER].get(CONFIG_DEAD_LEAVES, {}))
+    dl_dict = create_dataloaders(config, dl_train, dl_valid)
+    return dl_dict
+
+
+def create_dataloaders(config, dl_train, dl_valid) -> dict:
+    dl_dict = {
+        TRAIN: DataLoader(
+            dl_train,
+            shuffle=True,
+            batch_size=config[DATALOADER][BATCH_SIZE][TRAIN],
+        ),
+        VALIDATION: DataLoader(
+            dl_valid,
+            shuffle=False,
+            batch_size=config[DATALOADER][BATCH_SIZE][VALIDATION]
+        ),
+        # TEST: DataLoader(dl_test, shuffle=False, batch_size=config[DATALOADER][BATCH_SIZE][TEST])
+    }
+    return dl_dict
+
+
+def get_data_loader_from_disk(config, frozen_seed: Optional[int] = 42) -> dict:
+    ds = RestorationDataset
+    dataset_name = config[DATALOADER][NAME]  # NAME shall be here!
+    if dataset_name == DATASET_DIV2K:
+        dataset_root = DATASET_PATH/DATASET_DIV2K
+        train_root = dataset_root/"DIV2K_train_HR"/"DIV2K_train_HR"
+        valid_root = dataset_root/"DIV2K_valid_HR"/"DIV2K_valid_HR"
+        train_files = sorted(list(train_root.glob("*.png")))
+        train_files = 5*train_files  # Just to get 4000 elements...
+        valid_files = sorted(list(valid_root.glob("*.png")))
+    elif dataset_name in DATASET_SYNTH_LIST:
+        dataset_root = DATASET_PATH/dataset_name
+        all_files = sorted(list(dataset_root.glob("*.png")))
+        seed(frozen_seed)
+        shuffle(all_files)  # Easy way to perform cross validation if neeeded
+        cut_index = int(0.9*len(all_files))
+        train_files = all_files[:cut_index]
+        valid_files = all_files[cut_index:]
+    dl_train = ds(
+        train_files,
+        size=config[DATALOADER][SIZE],
+        frozen_seed=None,
+        **config[DATALOADER].get(CONFIG_DEGRADATION, {})
+    )
+    dl_valid = ds(
+        valid_files,
+        size=config[DATALOADER][SIZE],
+        frozen_seed=frozen_seed,
+        **config[DATALOADER].get(CONFIG_DEGRADATION, {})
+    )
+    dl_dict = create_dataloaders(config, dl_train, dl_valid)
+    return dl_dict
+
+
+def get_data_loader(config, frozen_seed=42):
+    dataset_name = config[DATALOADER].get(NAME, False)
+    if dataset_name:
+        return get_data_loader_from_disk(config, frozen_seed)
+    else:
+        return get_data_loader_synthetic(config, frozen_seed)
+
+
+if __name__ == "__main__":
+    # Example of usage synthetic dataset
+    for dataset_name in [DATASET_DIV2K, None, DATASET_DL_DIV2K_512, DATASET_DL_DIV2K_1024]:
+        if dataset_name is None:
+            dead_leaves_dataset = DeadLeavesDatasetGPU(colored=True)
+            dl = DataLoader(dead_leaves_dataset, batch_size=4, shuffle=True)
+        else:
+            # Example of usage stored images dataset
+            config = {
+                DATALOADER: {
+                    NAME: dataset_name,
+                    SIZE: (128, 128),
+                    BATCH_SIZE: {
+                        TRAIN: 4,
+                        VALIDATION: 4
+                    },
+                }
+            }
+            dl_dict = get_data_loader(config)
+            dl = dl_dict[TRAIN]
+            # dl = dl_dict[VALIDATION]
+        for i, (batch_inp, batch_target) in enumerate(dl):
+            print(batch_inp.shape, batch_target.shape)  # Should print [batch_size, size[0], size[1], 3] for each batch
+            if i == 1:  # Just to break the loop after two batches for demonstration
+                import matplotlib.pyplot as plt
+                plt.subplot(1, 2, 1)
+                plt.imshow(batch_inp.permute(0, 2, 3, 1).reshape(-1, batch_inp.shape[-1], 3).cpu().numpy())
+                plt.title("Degraded")
+                plt.subplot(1, 2, 2)
+                plt.imshow(batch_target.permute(0, 2, 3, 1).reshape(-1, batch_inp.shape[-1], 3).cpu().numpy())
+                plt.title("Target")
+                plt.show()
+                # print(batch_target)
+                break

src/rstor/data/degradation.py

@@ -0,0 +1,156 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Sun Mar 24 01:21:46 2024
+
+@author: jamyl
+"""
+import torch
+from rstor.properties import DATASET_BLUR_KERNEL_PATH
+import random
+from scipy.io import loadmat
+import cv2
+
+
+class Degradation():
+    def __init__(self,
+                 length: int = 1000,
+                 frozen_seed: int = None):
+        self.frozen_seed = frozen_seed
+        self.current_degradation = {}
+
+
+class DegradationNoise(Degradation):
+    def __init__(self,
+                 length: int = 1000,
+                 noise_stddev: float = [0., 50.],
+                 frozen_seed: int = None):
+        super().__init__(length, frozen_seed)
+        self.noise_stddev = noise_stddev
+
+        if frozen_seed is not None:
+            random.seed(frozen_seed)
+            self.noise_stddev = [(self.noise_stddev[1] - self.noise_stddev[0]) *
+                                 random.random() + self.noise_stddev[0] for _ in range(length)]
+
+    def __call__(self, x: torch.Tensor, idx: int):
+        # WARNING! INPLACE OPERATIONS!!!!!
+        # expects x of shape [b, c, h, w]
+        assert x.ndim == 4
+        assert x.shape[1] in [1, 3]
+
+        if self.frozen_seed is not None:
+            std_dev = self.noise_stddev[idx]
+        else:
+            std_dev = (self.noise_stddev[1] - self.noise_stddev[0]) * random.random() + self.noise_stddev[0]
+
+        if std_dev > 0.:
+            # x += (std_dev/255.)*np.random.randn(*x.shape)
+            x += (std_dev/255.)*torch.randn(*x.shape, device=x.device)
+        self.current_degradation[idx] = {
+            "noise_stddev": std_dev
+        }
+        return x
+
+
+class DegradationBlurMat(Degradation):
+    def __init__(self,
+                 length: int = 1000,
+                 frozen_seed: int = None,
+                 blur_index: int = None):
+        super().__init__(length, frozen_seed)
+
+        kernels = loadmat(DATASET_BLUR_KERNEL_PATH)["kernels"].squeeze()
+        # conversion to torch (the shape of the kernel is not constant)
+        self.kernels = tuple([
+            torch.from_numpy(kernel/kernel.sum(keepdims=True)).unsqueeze(0).unsqueeze(0)
+            for kernel in kernels] + [torch.ones((1, 1)).unsqueeze(0).unsqueeze(0)])
+        self.n_kernels = len(self.kernels)
+
+        if frozen_seed is not None:
+            random.seed(frozen_seed)
+            self.kernel_ids = [random.randint(0, self.n_kernels-1) for _ in range(length)]
+        if blur_index is not None:
+            self.frozen_seed = 42
+            self.kernel_ids = [blur_index for _ in range(length)]
+
+    def __call__(self, x: torch.Tensor, idx: int):
+        # expects x of shape [b, c, h, w]
+        assert x.ndim == 4
+        assert x.shape[1] in [1, 3]
+        device = x.device
+
+        if self.frozen_seed is not None:
+            kernel_id = self.kernel_ids[idx]
+        else:
+            kernel_id = random.randint(0, self.n_kernels-1)
+
+        kernel = self.kernels[kernel_id].to(device).repeat(3, 1, 1, 1).float()  # repeat for grouped conv
+        _, _, kh, kw = kernel.shape
+        # We use padding = same to make
+        # sure that the output size does not depend on the kernel.
+
+        # define nn.Conf layer to define both padding mode and padding value...
+        conv_layer = torch.nn.Conv2d(in_channels=x.shape[1],
+                                     out_channels=x.shape[1],
+                                     kernel_size=(kh, kw),
+                                     padding="same",
+                                     padding_mode='replicate',
+                                     groups=3,
+                                     bias=False)
+
+        # Set the predefined kernel as weights and freeze the parameters
+        with torch.no_grad():
+            conv_layer.weight = torch.nn.Parameter(kernel)
+            conv_layer.weight.requires_grad = False
+        # breakpoint()
+        x = conv_layer(x)
+        # Alternative Functional version with 0 padding :
+        # x = F.conv2d(x, kernel, padding="same", groups=3)
+
+        self.current_degradation[idx] = {
+            "blur_kernel_id": kernel_id
+        }
+        return x
+
+
+class DegradationBlurGauss(Degradation):
+    def __init__(self,
+                 length: int = 1000,
+                 blur_kernel_half_size: int = [0, 2],
+                 frozen_seed: int = None):
+        super().__init__(length, frozen_seed)
+
+        self.blur_kernel_half_size = blur_kernel_half_size
+        # conversion to torch (the shape of the kernel is not constant)
+        if frozen_seed is not None:
+            random.seed(self.frozen_seed)
+            self.blur_kernel_half_size = [
+                (
+                    random.randint(self.blur_kernel_half_size[0], self.blur_kernel_half_size[1]),
+                    random.randint(self.blur_kernel_half_size[0], self.blur_kernel_half_size[1])
+                ) for _ in range(length)
+            ]
+
+    def __call__(self, x: torch.Tensor, idx: int):
+        # expects x of shape [b, c, h, w]
+        assert x.ndim == 4
+        assert x.shape[1] in [1, 3]
+        device = x.device
+
+        if self.frozen_seed is not None:
+            k_size_x, k_size_y = self.blur_kernel_half_size[idx]
+        else:
+            k_size_x = random.randint(self.blur_kernel_half_size[0], self.blur_kernel_half_size[1])
+            k_size_y = random.randint(self.blur_kernel_half_size[0], self.blur_kernel_half_size[1])
+
+        k_size_x = 2 * k_size_x + 1
+        k_size_y = 2 * k_size_y + 1
+
+        x = x.squeeze(0).permute(1, 2, 0).cpu().numpy()
+        x = cv2.GaussianBlur(x, (k_size_x, k_size_y), 0)
+        x = torch.from_numpy(x).to(device).permute(2, 0, 1).unsqueeze(0)
+
+        self.current_degradation[idx] = {
+            "blur_kernel_half_size": (k_size_x, k_size_y),
+        }
+        return x

src/rstor/data/stored_images_dataloader.py

@@ -0,0 +1,156 @@
+import torch
+from torch.utils.data import DataLoader, Dataset
+from rstor.data.augmentation import augment_flip
+from rstor.data.degradation import DegradationBlurMat, DegradationBlurGauss, DegradationNoise
+from rstor.properties import DEVICE, AUGMENTATION_FLIP, AUGMENTATION_ROTATE, DEGRADATION_BLUR_NONE, DEGRADATION_BLUR_MAT, DEGRADATION_BLUR_GAUSS
+from rstor.properties import DATALOADER, BATCH_SIZE, TRAIN, VALIDATION, LENGTH, CONFIG_DEAD_LEAVES, SIZE
+from typing import Tuple, Optional, Union
+from torchvision import transforms
+# from torchvision.transforms import RandomCrop
+from pathlib import Path
+from tqdm import tqdm
+from time import time
+from torchvision.io import read_image
+IMAGES_FOLDER = "images"
+
+
+def load_image(path):
+    return read_image(str(path))
+
+
+class RestorationDataset(Dataset):
+    def __init__(
+        self,
+        images_path: Path,
+        size: Tuple[int, int] = (128, 128),
+        device: str = DEVICE,
+        preloaded: bool = False,
+        augmentation_list: Optional[list] = [],
+        frozen_seed: int = None,  # useful for validation set!
+        blur_kernel_half_size: int = [0, 2],
+        noise_stddev: float = [0., 50.],
+        degradation_blur=DEGRADATION_BLUR_NONE,
+        blur_index=None,
+        **_extra_kwargs
+    ):
+        self.preloaded = preloaded
+        self.augmentation_list = augmentation_list
+        self.device = device
+        self.frozen_seed = frozen_seed
+        if not isinstance(images_path, list):
+            self.path_list = sorted(list(images_path.glob("*.png")))
+        else:
+            self.path_list = images_path
+
+        self.length = len(self.path_list)
+
+        self.n_samples = len(self.path_list)
+        # If we can preload everything in memory, we can do it
+        if preloaded:
+            self.data_list = [load_image(pth) for pth in tqdm(self.path_list)]
+        else:
+            self.data_list = self.path_list
+
+        # if AUGMENTATION_FLIP in self.augmentation_list:
+        #     img_data = augment_flip(img_data)
+        # img_data = self.cropper(img_data)
+        self.transforms = []
+
+        if self.frozen_seed is None:
+            if AUGMENTATION_FLIP in self.augmentation_list:
+                self.transforms.append(transforms.RandomHorizontalFlip(p=0.5))
+                self.transforms.append(transforms.RandomVerticalFlip(p=0.5))
+            if AUGMENTATION_ROTATE in self.augmentation_list:
+                self.transforms.append(transforms.RandomRotation(degrees=180))
+
+        crop = transforms.RandomCrop(size) if frozen_seed is None else transforms.CenterCrop(size)
+        self.transforms.append(crop)
+        self.transforms = transforms.Compose(self.transforms)
+
+        # self.cropper = RandomCrop(size=size)
+
+        self.degradation_blur_type = degradation_blur
+        if degradation_blur == DEGRADATION_BLUR_GAUSS:
+            self.degradation_blur = DegradationBlurGauss(self.length,
+                                                         blur_kernel_half_size,
+                                                         frozen_seed)
+            self.blur_deg_str = "blur_kernel_half_size"
+        elif degradation_blur == DEGRADATION_BLUR_MAT:
+            self.degradation_blur = DegradationBlurMat(self.length,
+                                                       frozen_seed,
+                                                       blur_index)
+            self.blur_deg_str = "blur_kernel_id"
+        elif degradation_blur == DEGRADATION_BLUR_NONE:
+            pass
+        else:
+            raise ValueError(f"Unknown degradation blur {degradation_blur}")
+
+        self.degradation_noise = DegradationNoise(self.length,
+                                                  noise_stddev,
+                                                  frozen_seed)
+        self.current_degradation = {}
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, Union[torch.Tensor, None]]:
+        """Access a specific image from dataset and augment
+
+        Args:
+            index (int): access index
+
+        Returns:
+            torch.Tensor: image tensor [C, H, W]
+        """
+        if self.preloaded:
+            img_data = self.data_list[index]
+        else:
+            img_data = load_image(self.data_list[index])
+        img_data = img_data.to(self.device)
+
+        # if AUGMENTATION_FLIP in self.augmentation_list:
+        #     img_data = augment_flip(img_data)
+        # img_data = self.cropper(img_data)
+
+        img_data = self.transforms(img_data)
+        img_data = img_data.float()/255.
+        degraded_img = img_data.clone().unsqueeze(0)
+
+        self.current_degradation[index] = {}
+        if self.degradation_blur_type != DEGRADATION_BLUR_NONE:
+            degraded_img = self.degradation_blur(degraded_img, index)
+            self.current_degradation[index][self.blur_deg_str] = self.degradation_blur.current_degradation[index][self.blur_deg_str]
+
+        degraded_img = self.degradation_noise(degraded_img, index)
+        self.current_degradation[index]["noise_stddev"] = self.degradation_noise.current_degradation[index]["noise_stddev"]
+
+        degraded_img = degraded_img.squeeze(0)
+        self.current_degradation[index] = {
+            "noise_stddev": self.degradation_noise.current_degradation[index]["noise_stddev"]
+        }
+        try:
+            self.current_degradation[index][self.blur_deg_str] = self.degradation_blur.current_degradation[index][self.blur_deg_str]
+        except KeyError:
+            pass
+
+        return degraded_img, img_data
+
+    def __len__(self):
+        return self.n_samples
+
+
+if __name__ == "__main__":
+    dataset_restoration = RestorationDataset(
+        Path("__dataset/div2k/DIV2K_train_HR/DIV2K_train_HR/"),
+        preloaded=True,
+    )
+    dataloader = DataLoader(
+        dataset_restoration,
+        batch_size=16,
+        shuffle=True
+    )
+    start = time()
+    total = 0
+    for batch in tqdm(dataloader):
+        # print(batch.shape)
+        torch.cuda.synchronize()
+        total += batch.shape[0]
+    end = time()
+    print(f"Time elapsed: {(end-start)/total*1000.:.2f}ms/image")

src/rstor/data/synthetic_dataloader.py

@@ -0,0 +1,187 @@
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+from typing import Tuple
+from rstor.data.degradation import DegradationBlurMat, DegradationBlurGauss, DegradationNoise
+from rstor.properties import DEVICE, AUGMENTATION_FLIP, DEGRADATION_BLUR_NONE, DEGRADATION_BLUR_MAT, DEGRADATION_BLUR_GAUSS
+from rstor.synthetic_data.dead_leaves_cpu import cpu_dead_leaves_chart
+from rstor.synthetic_data.dead_leaves_gpu import gpu_dead_leaves_chart
+import cv2
+from skimage.filters import gaussian
+import random
+import numpy as np
+
+from rstor.utils import DEFAULT_TORCH_FLOAT_TYPE
+
+
+class DeadLeavesDataset(Dataset):
+    def __init__(
+        self,
+        size: Tuple[int, int] = (128, 128),
+        length: int = 1000,
+        frozen_seed: int = None,  # useful for validation set!
+        blur_kernel_half_size: int = [0, 2],
+        ds_factor: int = 5,
+        noise_stddev: float = [0., 50.],
+        degradation_blur=DEGRADATION_BLUR_NONE,
+        **config_dead_leaves
+        # number_of_circles: int = -1,
+        # background_color: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5),
+        # colored: Optional[bool] = False,
+        # radius_mean: Optional[int] = -1,
+        # radius_stddev: Optional[int] = -1,
+    ):
+
+        self.frozen_seed = frozen_seed
+        self.ds_factor = ds_factor
+        self.size = (size[0]*ds_factor, size[1]*ds_factor)
+        self.length = length
+        self.config_dead_leaves = config_dead_leaves
+        self.blur_kernel_half_size = blur_kernel_half_size
+        self.noise_stddev = noise_stddev
+        
+
+        self.degradation_blur_type = degradation_blur
+        if degradation_blur == DEGRADATION_BLUR_GAUSS:
+            self.degradation_blur = DegradationBlurGauss(self.length,
+                                                         blur_kernel_half_size,
+                                                         frozen_seed)
+            self.blur_deg_str = "blur_kernel_half_size"
+        elif degradation_blur == DEGRADATION_BLUR_MAT:
+            self.degradation_blur = DegradationBlurMat(self.length,
+                                                       frozen_seed)
+            self.blur_deg_str = "blur_kernel_id"
+        elif degradation_blur == DEGRADATION_BLUR_NONE:
+            pass
+        else:
+            raise ValueError(f"Unknown degradation blur {degradation_blur}")
+            
+        self.degradation_noise = DegradationNoise(self.length,
+                                                  noise_stddev,
+                                                  frozen_seed)
+        self.current_degradation = {}
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        # TODO there is a bug on this cpu version, the dead leaved dont appear ot be right
+        seed = self.frozen_seed + idx if self.frozen_seed is not None else None
+        chart = cpu_dead_leaves_chart(self.size, seed=seed, **self.config_dead_leaves)
+
+        if self.ds_factor > 1:
+            # print(f"Downsampling {chart.shape} with factor {self.ds_factor}...")
+            sigma = 3/5
+            chart = gaussian(
+                chart, sigma=(sigma, sigma, 0), mode='nearest',
+                cval=0, preserve_range=True, truncate=4.0)
+            chart = chart[::self.ds_factor, ::self.ds_factor]
+
+        th_chart = torch.from_numpy(chart).permute(2, 0, 1).unsqueeze(0)
+        degraded_chart = th_chart
+
+        self.current_degradation[idx] = {}
+        if self.degradation_blur_type != DEGRADATION_BLUR_NONE:
+            degraded_chart = self.degradation_blur(degraded_chart, idx)
+            self.current_degradation[idx][self.blur_deg_str] = self.degradation_blur.current_degradation[idx][self.blur_deg_str]
+        
+        degraded_chart = self.degradation_noise(degraded_chart, idx)
+        self.current_degradation[idx]["noise_stddev"] = self.degradation_noise.current_degradation[idx]["noise_stddev"]
+        
+        degraded_chart = degraded_chart.squeeze(0)
+        th_chart = th_chart.squeeze(0)
+
+        return degraded_chart, th_chart
+
+
+class DeadLeavesDatasetGPU(Dataset):
+    def __init__(
+        self,
+        size: Tuple[int, int] = (128, 128),
+        length: int = 1000,
+        frozen_seed: int = None,  # useful for validation set!
+        blur_kernel_half_size: int = [0, 2],
+        ds_factor: int = 5,
+        noise_stddev: float = [0., 50.],
+        use_gaussian_kernel=True,
+        **config_dead_leaves
+        # number_of_circles: int = -1,
+        # background_color: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5),
+        # colored: Optional[bool] = False,
+        # radius_mean: Optional[int] = -1,
+        # radius_stddev: Optional[int] = -1,
+    ):
+        self.frozen_seed = frozen_seed
+        self.ds_factor = ds_factor
+        self.size = (size[0]*ds_factor, size[1]*ds_factor)
+        self.length = length
+        self.config_dead_leaves = config_dead_leaves
+
+        # downsample kernel
+        sigma = 3/5
+        k_size = 5  # This fits with sigma = 3/5, the cutoff value is 0.0038 (neglectable)
+        x = (torch.arange(k_size) - 2).to('cuda')
+        kernel = torch.stack(torch.meshgrid((x, x), indexing='ij'))
+        kernel.requires_grad = False
+        dist_sq = kernel[0]**2 + kernel[1]**2
+        kernel = (-dist_sq.square()/(2*sigma**2)).exp()
+        kernel = kernel / kernel.sum()
+        self.downsample_kernel = kernel.repeat(3, 1, 1, 1)  # shape [3, 1, k_size, k_size]
+        self.downsample_kernel.requires_grad = False
+        self.use_gaussian_kernel = use_gaussian_kernel
+        if use_gaussian_kernel:
+            self.degradation_blur = DegradationBlurGauss(length,
+                                                         blur_kernel_half_size,
+                                                         frozen_seed)
+        else:
+            self.degradation_blur = DegradationBlurMat(length,
+                                                       frozen_seed)
+
+        self.degradation_noise = DegradationNoise(length,
+                                                  noise_stddev,
+                                                  frozen_seed)
+        self.current_degradation = {}
+
+    def __len__(self) -> int:
+        return self.length
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Get a single deadleave chart and its degraded version.
+
+        Args:
+            idx (int): index of the item to retrieve
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: degraded chart, target chart
+        """
+        seed = self.frozen_seed + idx if self.frozen_seed is not None else None
+
+        # Return numba device array
+        numba_chart = gpu_dead_leaves_chart(self.size, seed=seed, **self.config_dead_leaves)
+        th_chart = torch.as_tensor(numba_chart, dtype=DEFAULT_TORCH_FLOAT_TYPE, device="cuda")[
+            None].permute(0, 3, 1, 2)  # [1, c, h, w]
+        if self.ds_factor > 1:
+            # Downsample using strided gaussian conv (sigma=3/5)
+            th_chart = F.pad(th_chart,
+                             pad=(2, 2, 0, 0),
+                             mode="replicate")
+            th_chart = F.conv2d(th_chart,
+                                self.downsample_kernel,
+                                padding='valid',
+                                groups=3,
+                                stride=self.ds_factor)
+
+        degraded_chart = self.degradation_blur(th_chart, idx)
+        degraded_chart = self.degradation_noise(degraded_chart, idx)
+
+        blur_deg_str = "blur_kernel_half_size" if self.use_gaussian_kernel else "blur_kernel_id"
+        self.current_degradation[idx] = {
+            blur_deg_str: self.degradation_blur.current_degradation[idx][blur_deg_str],
+            "noise_stddev": self.degradation_noise.current_degradation[idx]["noise_stddev"]
+        }
+
+        degraded_chart = degraded_chart.squeeze(0)
+        th_chart = th_chart.squeeze(0)
+
+        return degraded_chart, th_chart

src/rstor/learning/experiments.py

@@ -0,0 +1,24 @@
+from rstor.properties import DEVICE, OPTIMIZER, PARAMS
+from rstor.architecture.selector import load_architecture
+from rstor.data.dataloader import get_data_loader
+from typing import Tuple
+import torch
+
+
+def get_training_content(
+        config: dict,
+        training_mode: bool = False,
+        device=DEVICE) -> Tuple[torch.nn.Module, torch.optim.Optimizer, dict]:
+    model = load_architecture(config)
+    optimizer, dl_dict = None, None
+    if training_mode:
+        optimizer = torch.optim.Adam(model.parameters(), **config[OPTIMIZER][PARAMS])
+        dl_dict = get_data_loader(config)
+    return model, optimizer, dl_dict
+
+
+if __name__ == "__main__":
+    from rstor.learning.experiments_definition import default_experiment
+    config = default_experiment(1)
+    model, optimizer, dl_dict = get_training_content(config, training_mode=True)
+    print(config)

src/rstor/learning/experiments_definition.py

@@ -0,0 +1,489 @@
+from rstor.properties import (NB_EPOCHS, DATALOADER, BATCH_SIZE, SIZE, LENGTH,
+                              TRAIN, VALIDATION, SCHEDULER, REDUCELRONPLATEAU,
+                              MODEL, ARCHITECTURE, ID, NAME, SCHEDULER_CONFIGURATION, OPTIMIZER, PARAMS, LR,
+                              LOSS, LOSS_MSE, CONFIG_DEAD_LEAVES,
+                              SELECTED_METRICS, METRIC_PSNR, METRIC_SSIM, METRIC_LPIPS,
+                              DATASET_DL_DIV2K_512, DATASET_DIV2K,
+                              CONFIG_DEGRADATION,
+                              PRETTY_NAME,
+                              DEGRADATION_BLUR_NONE, DEGRADATION_BLUR_MAT, DEGRADATION_BLUR_GAUSS,
+                              AUGMENTATION_FLIP, AUGMENTATION_ROTATE,
+                              DATASET_DL_EXTRAPRIMITIVES_DIV2K_512)
+
+
+from typing import Tuple
+
+
+def model_configurations(config, model_preset="StackedConvolutions", bias: bool = True) -> dict:
+    if model_preset == "StackedConvolutions":
+        config[MODEL] = {
+            ARCHITECTURE: dict(
+                num_layers=8,
+                k_size=3,
+                h_dim=16,
+                bias=bias
+            ),
+            NAME: "StackedConvolutions"
+        }
+    elif model_preset == "NAFNet" or model_preset == "UNet":
+        # https://github.com/megvii-research/NAFNet/blob/main/options/test/GoPro/NAFNet-width64.yml
+        config[MODEL] = {
+            ARCHITECTURE: dict(
+                width=64,
+                enc_blk_nums=[1, 1, 1, 28],
+                middle_blk_num=1,
+                dec_blk_nums=[1, 1, 1, 1],
+            ),
+            NAME: model_preset
+        }
+    else:
+        raise ValueError(f"Unknown model preset {model_preset}")
+
+
+def presets_experiments(
+    exp: int,
+    b: int = 32,
+    n: int = 50,
+    bias: bool = True,
+    length: int = 5000,
+    data_size: Tuple[int, int] = (128, 128),
+    model_preset: str = "StackedConvolutions",
+    lpips: bool = False
+) -> dict:
+    config = {
+        ID: exp,
+        NAME: f"{exp:04d}",
+        NB_EPOCHS: n
+    }
+    config[DATALOADER] = {
+        BATCH_SIZE: {
+            TRAIN: b,
+            VALIDATION: b
+        },
+        SIZE: data_size,  # (width, height)
+        LENGTH: {
+            TRAIN: length,
+            VALIDATION: 800
+        }
+    }
+    config[OPTIMIZER] = {
+        NAME: "Adam",
+        PARAMS: {
+            LR: 1e-3
+        }
+    }
+    model_configurations(config, model_preset=model_preset, bias=bias)
+    config[SCHEDULER] = REDUCELRONPLATEAU
+    config[SCHEDULER_CONFIGURATION] = {
+        "factor": 0.8,
+        "patience": 5
+    }
+    config[LOSS] = LOSS_MSE
+    config[SELECTED_METRICS] = [METRIC_PSNR, METRIC_SSIM]
+    if lpips:
+        config[SELECTED_METRICS].append(METRIC_LPIPS)
+    return config
+
+
+def get_experiment_config(exp: int) -> dict:
+    if exp == -1:
+        config = presets_experiments(exp, length=10, n=2)
+    elif exp == -2:
+        config = presets_experiments(exp, length=10, n=2, lpips=True)
+    elif exp == -3:
+        config = presets_experiments(exp, n=20)
+        config[DATALOADER]["gpu_gen"] = True
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+            blur_kernel_half_size=[0, 0],
+            ds_factor=1,
+            noise_stddev=[0., 50.]
+        )
+        config[PRETTY_NAME] = "Vanilla denoise only - ds=1 - noisy 0-50"
+    elif exp == -4:
+        config = presets_experiments(exp, b=4, n=20)
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.]
+        )
+        config[PRETTY_NAME] = "Vanilla exp from disk - noisy 0-50"
+    elif exp == 1000:
+        config = presets_experiments(exp, n=60)
+        config[PRETTY_NAME] = "Vanilla small blur"
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(blur_kernel_half_size=[0, 2], ds_factor=1, noise_stddev=[0., 0.])
+    elif exp == 1001:
+        config = presets_experiments(exp, n=60)
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(blur_kernel_half_size=[0, 6], ds_factor=1, noise_stddev=[0., 0.])
+        config[PRETTY_NAME] = "Vanilla large blur 0 - 6"
+    elif exp == 1002:
+        config = presets_experiments(exp, n=6)  # Less epochs because of the large downsample factor
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(blur_kernel_half_size=[0, 2], ds_factor=5, noise_stddev=[0., 0.])
+        config[PRETTY_NAME] = "Vanilla small blur - ds=5"
+    elif exp == 1003:
+        config = presets_experiments(exp, n=6)  # Less epochs because of the large downsample factor
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(blur_kernel_half_size=[0, 2], ds_factor=5, noise_stddev=[0., 50.])
+        config[PRETTY_NAME] = "Vanilla small blur - ds=5 - noisy 0-50"
+    elif exp == 1004:
+        config = presets_experiments(exp, n=60)
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(blur_kernel_half_size=[0, 0], ds_factor=1, noise_stddev=[0., 50.])
+        config[PRETTY_NAME] = "Vanilla denoise only - ds=1 - noisy 0-50"
+    elif exp == 1005:
+        config = presets_experiments(exp, bias=False, n=60)
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(blur_kernel_half_size=[0, 0], ds_factor=1, noise_stddev=[0., 50.])
+        config[PRETTY_NAME] = "Vanilla denoise only - ds=1 - noisy 0-50 - bias free"
+    elif exp == 1006:
+        config = presets_experiments(exp, n=60)
+        config[PRETTY_NAME] = "Vanilla small blur - noisy 0-50"
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(blur_kernel_half_size=[0, 2], ds_factor=1, noise_stddev=[0., 50.])
+    elif exp == 1007:
+        config = presets_experiments(exp, n=60)
+        config[PRETTY_NAME] = "Vanilla large blur 0 - 6 - noisy 0-50"
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(blur_kernel_half_size=[0, 6], ds_factor=1, noise_stddev=[0., 50.])
+    elif exp == 2000:
+        config = presets_experiments(exp, n=60,  b=16, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet denoise 0-50"
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+            blur_kernel_half_size=[0, 0],
+            ds_factor=1,
+            noise_stddev=[0., 50.]
+        )
+    elif exp == 2001:
+        config = presets_experiments(exp, n=60,  b=16, model_preset="UNet")
+        config[PRETTY_NAME] = "UNEt denoise 0-50"
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+            blur_kernel_half_size=[0, 0],
+            ds_factor=1,
+            noise_stddev=[0., 50.]
+        )
+    elif exp == 2002:
+        config = presets_experiments(exp, n=20,  b=8, data_size=(256, 256), model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet denoise  0-50 gpu dl 256x256"
+        config[DATALOADER]["gpu_gen"] = True
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+            blur_kernel_half_size=[0, 0],
+            ds_factor=1,
+            noise_stddev=[0., 50.]
+        )
+    elif exp == 2003:
+        config = presets_experiments(exp, n=20,  b=8, data_size=(128, 128), model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet denoise  0-50 gpu dl - 128x128"
+        config[DATALOADER]["gpu_gen"] = True
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+            blur_kernel_half_size=[0, 0],
+            ds_factor=1,
+            noise_stddev=[0., 50.]
+        )
+    elif exp == 2004:
+        config = presets_experiments(exp, n=20,  b=16, data_size=(128, 128), model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet Light denoise  0-50 gpu dl - 128x128"
+        config[DATALOADER]["gpu_gen"] = True
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+            blur_kernel_half_size=[0, 0],
+            ds_factor=1,
+            noise_stddev=[0., 50.]
+        )
+        config[MODEL][ARCHITECTURE] = dict(
+            width=64,
+            enc_blk_nums=[1, 1, 1, 2],
+            middle_blk_num=1,
+            dec_blk_nums=[1, 1, 1, 1],
+        )
+    elif exp == 2005:
+        config = presets_experiments(exp, n=20,  b=16, data_size=(128, 128), model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet TresLight denoise  0-50 gpu dl - 128x128"
+        config[DATALOADER]["gpu_gen"] = True
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+            blur_kernel_half_size=[0, 0],
+            ds_factor=1,
+            noise_stddev=[0., 50.]
+        )
+        config[MODEL][ARCHITECTURE] = dict(
+            width=64,
+            enc_blk_nums=[1, 1, 2],
+            middle_blk_num=1,
+            dec_blk_nums=[1, 1, 1],
+        )
+    elif exp == 2006:
+        config = presets_experiments(exp, n=20,  b=16, data_size=(128, 128), model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet TresLight denoise  0-50 ds=5 gpu dl - 128x128"
+        config[DATALOADER]["gpu_gen"] = True
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+            blur_kernel_half_size=[0, 0],
+            ds_factor=5,
+            noise_stddev=[0., 50.]
+        )
+        config[MODEL][ARCHITECTURE] = dict(
+            width=64,
+            enc_blk_nums=[1, 1, 2],
+            middle_blk_num=1,
+            dec_blk_nums=[1, 1, 1],
+        )
+    elif exp == 2007:
+        config = presets_experiments(exp, n=20,  b=16, data_size=(128, 128), model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet denoise  0-50 gpu dl -ds=5 128x128"
+        config[DATALOADER]["gpu_gen"] = True
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+            blur_kernel_half_size=[0, 0],
+            ds_factor=5,
+            noise_stddev=[0., 50.]
+        )
+    elif exp == 1008:
+        config = presets_experiments(exp, n=20)
+        config[DATALOADER]["gpu_gen"] = True
+        config[DATALOADER][CONFIG_DEAD_LEAVES] = dict(
+            blur_kernel_half_size=[0, 0],
+            ds_factor=5,
+            noise_stddev=[0., 50.]
+        )
+        config[PRETTY_NAME] = "Vanilla denoise only - ds=5 - noisy 0-50"
+    # ---------------------------------
+    # Pure DL DENOISING trainings!
+    # ---------------------------------
+    elif exp == 3000:
+        config = presets_experiments(exp, n=30,  b=4, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet denoise - DL_DIV2K_512 0-50"
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.]
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 3001:  # ENABLE GRADIENT CLIPPING
+        config = presets_experiments(exp, n=30,  b=8, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet41.4M denoise - DL_DIV2K_512 0-50 256x256"
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.]
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 3002:  # ENABLE GRADIENT CLIPPING
+        config = presets_experiments(exp, n=30,  b=16, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet41.4M denoise - DL_DIV2K_512 0-50 128x128"
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.]
+        )
+        config[DATALOADER][SIZE] = (128, 128)
+    elif exp == 3010 or exp == 3011:  # exp 3011 = REDO with Gradient clipping
+        config = presets_experiments(exp, n=50,  b=4, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet3.4M Light denoise - DL_DIV2K_512 0-50 256x256"
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.]
+        )
+        config[MODEL][ARCHITECTURE] = dict(
+            width=64,
+            enc_blk_nums=[1, 1, 2],
+            middle_blk_num=1,
+            dec_blk_nums=[1, 1, 1],
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 3020:
+        config = presets_experiments(exp, b=32, n=50)
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.]
+        )
+        config[PRETTY_NAME] = "Vanilla denoise DL  0-50 - noisy 0-50"
+    # ---------------------------------
+    # Pure DIV2K DENOISING trainings!
+    # ---------------------------------
+    elif exp == 3120:
+        config = presets_experiments(exp, b=32, n=50)
+        config[DATALOADER][NAME] = DATASET_DIV2K
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.]
+        )
+        config[PRETTY_NAME] = "Vanilla DIV2K_512 0-50 - noisy 0-50"
+    elif exp == 3111:
+        config = presets_experiments(exp, n=50,  b=4, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet3.4M Light denoise - DIV2K_512 0-50 256x256"
+        config[DATALOADER][NAME] = DATASET_DIV2K
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(noise_stddev=[0., 50.])
+        config[MODEL][ARCHITECTURE] = dict(
+            width=64,
+            enc_blk_nums=[1, 1, 2],
+            middle_blk_num=1,
+            dec_blk_nums=[1, 1, 1],
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 3101:
+        config = presets_experiments(exp, n=30,  b=8, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet41.4M denoise - DIV2K_512 0-50 256x256"
+        config[DATALOADER][NAME] = DATASET_DIV2K
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.]
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    # ---------------------------------
+    # Pure EXTRA PRIMITIVES
+    # ---------------------------------
+    elif exp == 3030:
+        config = presets_experiments(exp, b=128, n=50)
+        config[DATALOADER][NAME] = DATASET_DL_EXTRAPRIMITIVES_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.]
+        )
+        config[PRETTY_NAME] = "Vanilla DL_PRIMITIVES_512 0-50 - noisy 0-50"
+        # config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 3040:
+        config = presets_experiments(exp, n=50,  b=8, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet3.4M Light denoise - DL_PRIMITIVES_512 0-50 256x256"
+        config[DATALOADER][NAME] = DATASET_DL_EXTRAPRIMITIVES_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(noise_stddev=[0., 50.])
+        config[MODEL][ARCHITECTURE] = dict(
+            width=64,
+            enc_blk_nums=[1, 1, 2],
+            middle_blk_num=1,
+            dec_blk_nums=[1, 1, 1],
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 3050:
+        config = presets_experiments(exp, n=30,  b=8, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet41.4M denoise - DL_PRIMITIVES_512 0-50 256x256"
+        config[DATALOADER][NAME] = DATASET_DL_EXTRAPRIMITIVES_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(noise_stddev=[0., 50.])
+        config[DATALOADER][SIZE] = (256, 256)
+    # ---------------------------------
+    # DEBLURRING
+    # ---------------------------------
+    elif exp == 5000:
+        config = presets_experiments(exp, n=30,  b=8, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet deblur - DL_DIV2K_512 256x256"
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 0.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 5001:
+        config = presets_experiments(exp, n=30,  b=8, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet deblur - DIV2K_512 256x256"
+        config[DATALOADER][NAME] = DATASET_DIV2K
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 0.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 5002:
+        config = presets_experiments(exp, n=30,  b=8, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet deblur - DL_DIV2K_512 256x256"
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 0.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 5003:
+        config = presets_experiments(exp, n=30,  b=8, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet deblur - DIV2K_512 256x256"
+        config[DATALOADER][NAME] = DATASET_DIV2K
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 0.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 5004:
+        config = presets_experiments(exp, n=30,  b=8, model_preset="NAFNet")
+        config[PRETTY_NAME] = "NAFNet deblur - DIV2K_512 256x256"
+        config[DATALOADER][NAME] = DATASET_DIV2K
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 0.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 5005:
+        config = presets_experiments(exp, n=30,  b=8, model_preset="UNet")
+        config[PRETTY_NAME] = "UNet deblur - DL_512 256x256"
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 0.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    # elif exp == 6000:  # -> FAILED, no kernels normalization!
+    #     config = presets_experiments(exp, b=32, n=50)
+    #     config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+    #     config[DATALOADER][CONFIG_DEGRADATION] = dict(
+    #         noise_stddev=[0., 50.],
+    #         degradation_blur=DEGRADATION_BLUR_MAT,  # Deblur = Using .mat kernels
+    #         augmentation_list=[AUGMENTATION_FLIP]
+    #     )
+    #     config[PRETTY_NAME] = "Vanilla deblur DL_DIV2K_512"
+    elif exp == 6002:
+        config = presets_experiments(exp, b=128, n=50)
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Deblur = Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[PRETTY_NAME] = "Vanilla deblur DL_DIV2K_512"
+    # elif exp == 6001:  # -> FAILED, no kernels normalization!
+    #     config = presets_experiments(exp, b=32, n=50)
+    #     config[DATALOADER][NAME] = DATASET_DIV2K
+    #     config[DATALOADER][CONFIG_DEGRADATION] = dict(
+    #         noise_stddev=[0., 50.],
+    #         degradation_blur=DEGRADATION_BLUR_MAT,  # Deblur = Using .mat kernels
+    #         augmentation_list=[AUGMENTATION_FLIP]
+    #     )
+    #     config[PRETTY_NAME] = "Vanilla delbur DIV2K_512"
+    elif exp == 6003:
+        config = presets_experiments(exp, b=128, n=50)
+        config[DATALOADER][NAME] = DATASET_DIV2K
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Deblur = Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[PRETTY_NAME] = "Vanilla delbur DIV2K_512"
+
+    elif exp == 7000:
+        config = presets_experiments(exp, b=16, n=30, model_preset="NAFNet")
+        config[MODEL][ARCHITECTURE] = dict(
+            width=64,
+            enc_blk_nums=[1, 1, 2],
+            middle_blk_num=1,
+            dec_blk_nums=[1, 1, 1],
+        )
+        config[DATALOADER][NAME] = DATASET_DL_DIV2K_512
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Deblur = Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[PRETTY_NAME] = "NafNet Light deblur DL"
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 7001:
+        config = presets_experiments(exp, b=16, n=50, model_preset="NAFNet")
+        config[DATALOADER][NAME] = DATASET_DIV2K
+        config[MODEL][ARCHITECTURE] = dict(
+            width=64,
+            enc_blk_nums=[1, 1, 2],
+            middle_blk_num=1,
+            dec_blk_nums=[1, 1, 1],
+        )
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 50.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Deblur = Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[PRETTY_NAME] = "NafNet Light deblur DIV2K"
+        config[DATALOADER][SIZE] = (256, 256)
+    elif exp == 7002:
+        config = presets_experiments(exp, n=20,  b=8, model_preset="UNet")
+        config[PRETTY_NAME] = "UNET deblur - DIV2K"
+        config[DATALOADER][NAME] = DATASET_DIV2K
+        config[DATALOADER][CONFIG_DEGRADATION] = dict(
+            noise_stddev=[0., 0.],
+            degradation_blur=DEGRADATION_BLUR_MAT,  # Using .mat kernels
+            augmentation_list=[AUGMENTATION_FLIP]
+        )
+        config[DATALOADER][SIZE] = (256, 256)
+    else:
+        raise ValueError(f"Experiment {exp} not found")
+    return config

src/rstor/learning/loss.py

@@ -0,0 +1,25 @@
+import torch
+from typing import Optional
+from rstor.properties import LOSS_MSE
+
+
+def compute_loss(
+    predic: torch.Tensor,
+    target: torch.Tensor,
+    mode: Optional[str] = LOSS_MSE
+) -> torch.Tensor:
+    """
+    Compute loss based on the predicted and true values.
+
+    Args:
+        predic (torch.Tensor): [N, C, H, W] predicted values
+        target (torch.Tensor): [N, C, H, W] target values.
+        mode (Optional[str], optional): mode of loss computation.
+
+    Returns:
+        torch.Tensor: The computed loss.
+    """
+    assert mode in [LOSS_MSE], f"Mode {mode} not supported"
+    if mode == LOSS_MSE:
+        loss = torch.nn.functional.mse_loss(predic, target)
+    return loss

src/rstor/learning/metrics.py

@@ -0,0 +1,140 @@
+import torch
+from rstor.properties import METRIC_PSNR, METRIC_SSIM, METRIC_LPIPS, REDUCTION_AVERAGE, REDUCTION_SKIP, REDUCTION_SUM
+from torchmetrics.image import StructuralSimilarityIndexMeasure as SSIM
+from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from typing import List, Optional
+ALL_METRICS = [METRIC_PSNR, METRIC_SSIM, METRIC_LPIPS]
+
+
+def compute_psnr(
+    predic: torch.Tensor,
+    target: torch.Tensor,
+    clamp_mse=1e-10,
+    reduction: Optional[str] = REDUCTION_AVERAGE
+) -> torch.Tensor:
+    """
+    Compute the average PSNR metric for a batch of predicted and true values.
+
+    Args:
+        predic (torch.Tensor): [N, C, H, W] predicted values.
+        target (torch.Tensor): [N, C, H, W] target values.
+        reduction (str): Reduction method. REDUCTION_AVERAGE/REDUCTION_SKIP/REDUCTION_SUM.
+
+    Returns:
+        torch.Tensor: The average PSNR value for the batch.
+    """
+    with torch.no_grad():
+        mse_per_image = torch.mean((predic - target) ** 2, dim=(-3, -2, -1))
+        mse_per_image = torch.clamp(mse_per_image, min=clamp_mse)
+        psnr_per_image = 10 * torch.log10(1 / mse_per_image)
+        if reduction == REDUCTION_AVERAGE:
+            average_psnr = torch.mean(psnr_per_image)
+        elif reduction == REDUCTION_SUM:
+            average_psnr = torch.sum(psnr_per_image)
+        elif reduction == REDUCTION_SKIP:
+            average_psnr = psnr_per_image
+        else:
+            raise ValueError(f"Unknown reduction {reduction}")
+    return average_psnr
+
+
+def compute_ssim(
+    predic: torch.Tensor,
+    target: torch.Tensor,
+    reduction: Optional[str] = REDUCTION_AVERAGE
+) -> torch.Tensor:
+    """
+    Compute the average SSIM metric for a batch of predicted and true values.
+
+    Args:
+        predic (torch.Tensor): [N, C, H, W] predicted values.
+        target (torch.Tensor): [N, C, H, W] target values.
+        reduction (str): Reduction method. REDUCTION_AVERAGE/REDUCTION_SKIP.
+
+    Returns:
+        torch.Tensor: The average SSIM value for the batch.
+    """
+    with torch.no_grad():
+        reduction_mode = {
+            REDUCTION_SKIP: None,
+            REDUCTION_AVERAGE: "elementwise_mean",
+            REDUCTION_SUM: "sum"
+        }[reduction]
+        ssim = SSIM(data_range=1.0, reduction=reduction_mode).to(predic.device)
+        assert predic.shape == target.shape, f"{predic.shape} != {target.shape}"
+        assert predic.device == target.device, f"{predic.device} != {target.device}"
+        ssim_value = ssim(predic, target)
+    return ssim_value
+
+
+def compute_lpips(
+    predic: torch.Tensor,
+    target: torch.Tensor,
+    reduction: Optional[str] = REDUCTION_AVERAGE,
+) -> torch.Tensor:
+    """
+    Compute the average LPIPS metric for a batch of predicted and true values.
+    https://richzhang.github.io/PerceptualSimilarity/
+
+    Args:
+        predic (torch.Tensor): [N, C, H, W] predicted values.
+        target (torch.Tensor): [N, C, H, W] target values.
+        reduction (str): Reduction method. REDUCTION_AVERAGE/REDUCTION_SKIP.
+
+    Returns:
+        torch.Tensor: The average SSIM value for the batch.
+    """
+    reduction_mode = {
+        REDUCTION_SKIP:  "sum",  # does not really matter
+        REDUCTION_AVERAGE: "mean",
+        REDUCTION_SUM: "sum"
+    }[reduction]
+
+    with torch.no_grad():
+        lpip_metrics = LearnedPerceptualImagePatchSimilarity(
+            reduction=reduction_mode,
+            normalize=True  # If set to True will instead expect input to be in the [0,1] range.
+        ).to(predic.device)
+        assert predic.shape == target.shape, f"{predic.shape} != {target.shape}"
+        assert predic.device == target.device, f"{predic.device} != {target.device}"
+        if reduction == REDUCTION_SKIP:
+            lpip_value = []
+            for idx in range(predic.shape[0]):
+                lpip_value.append(lpip_metrics(
+                    predic[idx, ...].unsqueeze(0).clip(0, 1),
+                    target[idx, ...].unsqueeze(0).clip(0, 1)
+                ))
+            lpip_value = torch.stack(lpip_value)
+        elif reduction in [REDUCTION_SUM, REDUCTION_AVERAGE]:
+            lpip_value = lpip_metrics(predic.clip(0, 1), target.clip(0, 1))
+    return lpip_value
+
+
+def compute_metrics(
+        predic: torch.Tensor,
+        target: torch.Tensor,
+        reduction: Optional[str] = REDUCTION_AVERAGE,
+        chosen_metrics: Optional[List[str]] = ALL_METRICS) -> dict:
+    """
+    Compute the metrics for a batch of predicted and true values.
+
+    Args:
+        predic (torch.Tensor): [N, C, H, W] predicted values.
+        target (torch.Tensor): [N, C, H, W] target values.
+        reduction (str): Reduction method. REDUCTION_AVERAGE/REDUCTION_SKIP/REDUCTION SUM.
+        chosen_metrics (list): List of metrics to compute, default [METRIC_PSNR, METRIC_SSIM]
+
+    Returns:
+        dict: computed metrics.
+    """
+    metrics = {}
+    if METRIC_PSNR in chosen_metrics:
+        average_psnr = compute_psnr(predic, target, reduction=reduction)
+        metrics[METRIC_PSNR] = average_psnr.item() if reduction != REDUCTION_SKIP else average_psnr
+    if METRIC_SSIM in chosen_metrics:
+        ssim_value = compute_ssim(predic, target, reduction=reduction)
+        metrics[METRIC_SSIM] = ssim_value.item() if reduction != REDUCTION_SKIP else ssim_value
+    if METRIC_LPIPS in chosen_metrics:
+        lpip_value = compute_lpips(predic, target, reduction=reduction)
+        metrics[METRIC_LPIPS] = lpip_value.item() if reduction != REDUCTION_SKIP else lpip_value
+    return metrics

src/rstor/properties.py

@@ -0,0 +1,67 @@
+import torch
+from pathlib import Path
+RELU = "ReLU"
+LEAKY_RELU = "LeakyReLU"
+SIMPLE_GATE = "simple_gate"
+LOSS = "loss"
+LOSS_MSE = "MSE"
+METRIC_PSNR = "PSNR"
+METRIC_SSIM = "SSIM"
+METRIC_LPIPS = "LPIPS"
+SELECTED_METRICS = "selected_metrics"
+DATALOADER = "data_loader"
+BATCH_SIZE = "batch_size"
+SIZE = "size"
+TRAIN, VALIDATION, TEST = "train", "validation", "test"
+LENGTH = "length"
+ID = "id"
+NAME = "name"
+PRETTY_NAME = "pretty_name"
+NB_EPOCHS = "nb_epochs"
+ARCHITECTURE = "architecture"
+MODEL = "model"
+NAME = "name"
+N_PARAMS = "n_params"
+OPTIMIZER = "optimizer"
+LR = "lr"
+PARAMS = "parameters"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+SCHEDULER_CONFIGURATION = "scheduler_configuration"
+SCHEDULER = "scheduler"
+REDUCELRONPLATEAU = "ReduceLROnPlateau"
+ARCHITECTURE = "architecture"
+CONFIG_DEAD_LEAVES = "config_dead_leaves"
+CONFIG_DEGRADATION = "config_degradation"
+REDUCTION_SUM = "reduction_sum"
+REDUCTION_AVERAGE = "reduction_average"
+REDUCTION_SKIP = "reduction_skip"
+TRACES_TARGET = "target"
+TRACES_DEGRADED = "degraded"
+TRACES_RESTORED = "restored"
+TRACES_METRICS = "metrics"
+TRACES_ALL = "all"
+
+DEGRADATION_BLUR_NONE = "none"
+DEGRADATION_BLUR_MAT = "mat"
+DEGRADATION_BLUR_GAUSS = "gauss"
+
+
+SAMPLER_SATURATED = "saturated"
+SAMPLER_UNIFORM = "uniform"
+SAMPLER_NATURAL = "natural"
+SAMPLER_DIV2K = "div2k"
+
+DATASET_FOLDER = "__dataset"
+DATASET_PATH = Path(__file__).parent.parent.parent/DATASET_FOLDER
+DATASET_DL_RANDOMRGB_1024 = "deadleaves_randomrgb_1024"
+DATASET_DL_DIV2K_1024 = "deadleaves_div2k_1024"
+DATASET_DL_DIV2K_512 = "deadleaves_div2k_512"
+DATASET_DL_EXTRAPRIMITIVES_DIV2K_512 = "deadleaves_primitives_div2k_512"
+DATASET_SYNTH_LIST = [DATASET_DL_DIV2K_512, DATASET_DL_DIV2K_1024,
+                      DATASET_DL_RANDOMRGB_1024, DATASET_DL_EXTRAPRIMITIVES_DIV2K_512]
+DATASET_BLUR_KERNEL_PATH = DATASET_PATH / "kernels" / "custom_blur_centered.mat"
+AUGMENTATION_FLIP = "flip"
+AUGMENTATION_ROTATE = "rotate"
+
+
+DATASET_DIV2K = "div2k"

src/rstor/synthetic_data/color_sampler.py

@@ -0,0 +1,73 @@
+import numpy as np
+import cv2
+
+from pathlib import Path
+from rstor.properties import DATASET_PATH
+from typing import List
+
+
+def sample_uniform_rgb(size: int, seed: int = None) -> np.ndarray:
+    """
+    Generate n random RGB values.
+
+    Args:
+        n (int): number of colors to sample
+        seed (int, optional): Seed for the random number generator. Defaults to None.
+
+    Returns:
+        np.ndarray: Random RGB values as a numpy array.
+    """
+    # https://github.com/numpy/numpy/issues/17079
+    # https://numpy.org/devdocs/reference/random/new-or-different.html#new-or-different
+    rng = np.random.default_rng(np.random.SeedSequence(seed))
+
+    random_samples = rng.uniform(size=(size, 3))
+    rgb = random_samples
+
+    # Below old version with sturation
+    # lab = (random_samples + np.array([0., -0.5, -0.5])[None]) * np.array([100., 127 * 2, 127 * 2])[None]
+    # rgb = cv2.cvtColor(lab[None, :].astype(np.float32), cv2.COLOR_Lab2RGB)
+    return rgb.squeeze()
+
+
+def sample_saturated_color(size: int, seed: int = None) -> np.ndarray:
+    """
+    Generate n saturated RGB values.
+
+    Args:
+        n (int): number of colors to sample
+        seed (int, optional): Seed for the random number generator. Defaults to None.
+
+    Returns:
+        np.ndarray: Random RGB values as a numpy array.
+    """
+    # https://github.com/numpy/numpy/issues/17079
+    # https://numpy.org/devdocs/reference/random/new-or-different.html#new-or-different
+    rng = np.random.default_rng(np.random.SeedSequence(seed))
+
+    random_samples = rng.uniform(size=(size, 3))
+
+    lab = (random_samples + np.array([0., -0.5, -0.5])[None]) * np.array([100., 127 * 2, 127 * 2])[None]
+    rgb = cv2.cvtColor(lab[None, :].astype(np.float32), cv2.COLOR_Lab2RGB)
+    return rgb.squeeze()
+
+
+def sample_color_from_images(size: int, seed: int = None, path_to_images: List[Path] = []) -> np.ndarray:
+    print("path : ", path_to_images)
+    assert len(path_to_images) > 0, "Please provide a list of images to sample colors from."
+    rng = np.random.default_rng(np.random.SeedSequence(seed))
+
+    # Randomly pick an image and load it
+    img_id = rng.integers(0, len(path_to_images))
+
+    img = cv2.imread(path_to_images[img_id].as_posix())
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255
+
+    pixels = img.reshape(-1, 3)
+    n_pixels = pixels.shape[0]
+
+    # sample a pixel color for each disc
+    pixel_ids = rng.integers(0, n_pixels, size)
+    colors = pixels[pixel_ids, :]
+
+    return colors

src/rstor/synthetic_data/dead_leaves_cpu.py

@@ -0,0 +1,79 @@
+from typing import Tuple, Optional, List
+import numpy as np
+import cv2
+from rstor.synthetic_data.dead_leaves_sampler import define_dead_leaves_chart
+from rstor.properties import SAMPLER_UNIFORM
+from pathlib import Path
+
+
+def cpu_dead_leaves_chart(size: Tuple[int, int] = (100, 100),
+                          number_of_circles: int = -1,
+                          background_color: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5),
+                          colored: Optional[bool] = True,
+                          radius_min: Optional[int] = -1,
+                          radius_max: Optional[int] = -1,
+                          radius_alpha: Optional[int] = 3,
+                          seed: int = None,
+                          reverse: Optional[bool] = True,
+                          sampler=SAMPLER_UNIFORM,
+                          natural_image_list: List[Path] = []) -> np.ndarray:
+    """
+    Generation of a deqqad leaves chart by splatting circles on top of each other.
+
+    Args:
+        size (Tuple[int, int], optional): size of the generated chart. Defaults to (100, 100).
+        number_of_circles (int, optional): number of circles to generate.
+        If negative, it is computed based on the size. Defaults to -1.
+        background_color (Optional[Tuple[float, float, float]], optional):
+        background color of the chart. Defaults to gray (0.5, 0.5, 0.5).
+        colored (Optional[bool], optional): Whether to generate colored circles. Defaults to True.
+        radius_min (Optional[int], optional): minimum radius of the circles. Defaults to -1. (=> 1)
+        radius_max (Optional[int], optional): maximum radius of the circles. Defaults to -1. (=> 2000)
+        radius_alpha (Optional[int], optional): standard deviation of the radius of the circles.
+        If negative, it is calculated based on the size. Defaults to -1.
+        seed (int, optional): seed for the random number generator. Defaults to None
+        reverse: (Optional[bool], optional): View circles from the back view
+        by reversing order. Defaults to True.
+        WARNING: This option is extremely slow on CPU.
+
+    Returns:
+        np.ndarray: generated dead leaves chart as a NumPy array.
+    """
+    center_x, center_y, radius, color = define_dead_leaves_chart(
+        size,
+        number_of_circles,
+        colored,
+        radius_min,
+        radius_max,
+        radius_alpha,
+        seed,
+        sampler=sampler,
+        natural_image_list=natural_image_list
+    )
+    if not colored:
+        color = np.concatenate((color, color, color), axis=1)
+
+    if reverse:
+        chart = np.zeros((size[0], size[1], 3), dtype=np.float32)
+        buffer = np.zeros_like(chart)
+        is_not_covered_mask = np.ones((*chart.shape[:2], 1))
+        for i in range(number_of_circles):
+            cv2.circle(buffer, (center_x[i], center_y[i]), radius[i], color[i], -1)
+            chart += buffer * is_not_covered_mask
+            is_not_covered_mask = cv2.circle(is_not_covered_mask, (center_x[i], center_y[i]), radius[i], 0, -1)
+
+            if not np.any(is_not_covered_mask):
+                break
+
+        chart += np.multiply(background_color, np.ones((size[0], size[1], 3), dtype=np.float32)) * is_not_covered_mask
+    else:
+        chart = np.multiply(background_color, np.ones((size[0], size[1], 3), dtype=np.float32))
+        for i in range(number_of_circles):
+            # circle is inplace
+            cv2.circle(chart, (center_x[i], center_y[i]), radius[i], color[i], -1)
+
+    chart = chart.clip(0, 1)
+
+    if not colored:
+        chart = chart[:, :, 0, None]  # return shape [h, w, 1] in gray mode
+    return chart

src/rstor/synthetic_data/dead_leaves_gpu.py

@@ -0,0 +1,176 @@
+from rstor.utils import DEFAULT_NUMPY_FLOAT_TYPE, THREADS_PER_BLOCK
+from rstor.properties import SAMPLER_UNIFORM
+from typing import Tuple, Optional
+from rstor.synthetic_data.dead_leaves_cpu import define_dead_leaves_chart
+import numpy as np
+from numba import cuda
+import math
+
+
+def gpu_dead_leaves_chart(
+    size: Tuple[int, int] = (100, 100),
+        number_of_circles: int = -1,
+        background_color: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5),
+        colored: Optional[bool] = True,
+        radius_min: Optional[int] = -1,
+        radius_max: Optional[int] = -1,
+        radius_alpha: Optional[int] = 3,
+        seed: int = None,
+        reverse=True,
+        sampler=SAMPLER_UNIFORM,
+        natural_image_list=None,
+        circle_primitives: bool = True,
+        anisotropy: float = 1.,
+        angle: float = 0.
+) -> np.ndarray:
+    center_x, center_y, radius, color = define_dead_leaves_chart(
+        size,
+        number_of_circles,
+        colored,
+        radius_min,
+        radius_max,
+        radius_alpha,
+        seed,
+        sampler=sampler,
+        natural_image_list=natural_image_list
+    )
+
+    # Generate on gpu
+    chart = _generate_dead_leaves(
+        size,
+        centers=np.stack((center_x, center_y), axis=-1),
+        radia=radius,
+        colors=color,
+        background=background_color,
+        reverse=reverse,
+        circle_primitives=circle_primitives,
+        anisotropy=anisotropy,
+        angle=angle
+    )
+
+    return chart
+
+
+def _generate_dead_leaves(size, centers, radia, colors, background, reverse, circle_primitives: bool, anisotropy: float = 1., angle: float=0.):
+    assert centers.ndim == 2
+    ny, nx = size
+    nc = colors.shape[-1]
+
+    # Init empty array on GPU
+    generation_ = cuda.device_array((ny, nx, nc), DEFAULT_NUMPY_FLOAT_TYPE)
+    # Move useful array to GPU
+    centers_ = cuda.to_device(centers)
+    radia_ = cuda.to_device(radia)
+    colors_ = cuda.to_device(colors)
+
+    # Dispatch threads
+    threadsperblock = (THREADS_PER_BLOCK, THREADS_PER_BLOCK)
+    blockspergrid_x = math.ceil(nx/threadsperblock[1])
+    blockspergrid_y = math.ceil(ny/threadsperblock[0])
+    blockspergrid = (blockspergrid_x, blockspergrid_y)
+
+    if reverse:
+        cuda_dead_leaves_gen_reversed[blockspergrid, threadsperblock](
+            generation_,
+            centers_,
+            radia_,
+            colors_,
+            background,
+            circle_primitives,
+            anisotropy,
+            np.deg2rad(angle)
+        )
+    else:
+        cuda_dead_leaves_gen[blockspergrid, threadsperblock](
+            generation_,
+            centers_,
+            radia_,
+            colors_,
+            background)
+
+    return generation_
+
+
+@cuda.jit(cache=False)
+def cuda_dead_leaves_gen_reversed(generation, centers, radia, colors, background, circle_primitives: bool, anisotropy: float, angle: float):
+    idx, idy = cuda.grid(2)
+    ny, nx, nc = generation.shape
+
+    n_discs = centers.shape[0]
+
+    # Out of bound threads
+    if idx >= nx or idy >= ny:
+        return
+
+    for disc_id in range(n_discs):
+        dx_ = idx - centers[disc_id, 0]
+        dy_ = idy - centers[disc_id, 1]
+        dx = math.cos(angle)*dx_ + math.sin(angle)*dy_
+        dy = -math.sin(angle)*dx_ + math.cos(angle)*dy_
+        dx = dx * anisotropy
+        dist_sq = dx*dx + dy*dy
+
+        # Naive thread diverging version
+        r = radia[disc_id]
+        r_sq = r*r
+        if circle_primitives:
+            if dist_sq <= r_sq:
+                # Copy back to global memory
+                for c in range(nc):
+                    generation[idy, idx, c] = colors[disc_id, c]
+                return
+        else:
+            if (disc_id % 4) == 0 and dist_sq <= r_sq:
+                # Copy back to global memory
+                alpha = dist_sq/r_sq
+                for c in range(nc):
+                    generation[idy, idx, c] = colors[disc_id, c] * alpha + colors[disc_id, (c+1) % 3] * (1-alpha)
+                return
+            elif (disc_id % 4) == 1 and (abs(dx)+abs(dy)) <= r:
+                # Copy back to global memory
+                alpha = dist_sq/r_sq
+                for c in range(nc):
+                    generation[idy, idx, c] = colors[disc_id, c] * alpha + colors[disc_id, (c+1) % 3] * (1-alpha)
+                return
+            elif (disc_id % 4) == 2 and abs(dx) <= r and abs(dy) <= r:
+                for c in range(nc):
+                    generation[idy, idx, c] = colors[disc_id, c]
+                return
+            elif (disc_id % 200) == 3 and abs(dy) <= r//5:
+                for c in range(nc):
+                    generation[idy, idx, c] = colors[disc_id, c] * alpha + colors[disc_id, (c+1) % 3] * (1-alpha)
+                return
+            elif (disc_id % 200) == 4 and abs(dx) <= r//5:
+                for c in range(nc):
+                    generation[idy, idx, c] = colors[disc_id, c] * alpha + colors[disc_id, (c+1) % 3] * (1-alpha)
+                return
+    for c in range(nc):
+        generation[idy, idx, c] = background[c]
+
+
+@cuda.jit(cache=False)
+def cuda_dead_leaves_gen(generation, centers, radia, colors, background):
+    idx, idy, c = cuda.grid(3)
+    ny, nx, nc = generation.shape
+
+    n_discs = centers.shape[0]
+
+    # Out of bound threads
+    if idx >= nx or idy >= ny:
+        return
+
+    out = background[c]
+    for disc_id in range(n_discs):
+        dx = idx - centers[disc_id, 0]
+        dy = idy - centers[disc_id, 1]
+        dist_sq = dx*dx + dy*dy
+
+        # Naive thread diverging version
+        r = radia[disc_id]
+        r_sq = r*r
+
+        if dist_sq <= r_sq:
+            out = colors[disc_id, c]
+
+    # Copy back to global memory
+    generation[idy, idx, c] = out

src/rstor/synthetic_data/dead_leaves_sampler.py

@@ -0,0 +1,74 @@
+from typing import Tuple, Optional, List
+from rstor.properties import SAMPLER_SATURATED, SAMPLER_NATURAL, SAMPLER_UNIFORM
+from rstor.synthetic_data.color_sampler import sample_uniform_rgb, sample_saturated_color, sample_color_from_images
+import numpy as np
+from pathlib import Path
+
+
+def define_dead_leaves_chart(
+    size: Tuple[int, int] = (100, 100),
+    number_of_circles: int = -1,
+    colored: Optional[bool] = True,
+    radius_min: Optional[int] = -1,
+    radius_max: Optional[int] = -1,
+    radius_alpha: Optional[int] = 3,
+    seed: int = None,
+    sampler=SAMPLER_UNIFORM,
+    natural_image_list: Optional[List[Path]] = None
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Defines the geometric and color properties of the primitives in the dead leaves chart to later be sampled.
+
+    Args:
+        size (Tuple[int, int], optional): size of the generated chart. Defaults to (100, 100).
+        number_of_circles (int, optional): number of circles to generate.
+        If negative, it is computed based on the size. Defaults to -1.
+        colored (Optional[bool], optional): Whether to generate colored circles. Defaults to True.
+        radius_min (Optional[int], optional): minimum radius of the circles. Defaults to -1. (=> 1)
+        radius_max (Optional[int], optional): maximum radius of the circles. Defaults to -1. (=> 2000)
+        radius_alpha (Optional[int], optional): standard deviation of the radius of the circles.
+        If negative, it is calculated based on the size. Defaults to -1.
+        seed (int, optional): seed for the random number generator. Defaults to None
+
+    Returns:
+        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: center_x, center_y, radius, color
+    """
+    rng = np.random.default_rng(np.random.SeedSequence(seed))
+
+    if number_of_circles < 0:
+        number_of_circles = 30 * max(size)
+    if radius_min < 0.:
+        radius_min = 1.
+    if radius_max < 0.:
+        radius_max = 2000.
+
+    # Pick random circle centers and radii
+    center_x = rng.integers(0, size[1], size=number_of_circles)
+    center_y = rng.integers(0, size[0], size=number_of_circles)
+
+    # Sample from a power law distribution for the p(radius=r) = (r.clip(radius_min, radius_max))^(-alpha)
+
+    radius = rng.uniform(
+        low=radius_max ** (1 - radius_alpha),
+        high=radius_min ** (1 - radius_alpha),
+        size=number_of_circles
+    )
+    # Using the change of variables formula for random variables.
+    radius = radius ** (-1/(radius_alpha - 1))
+    radius = radius.round().astype(int)
+
+    # Pick random colors
+    if colored:
+        if sampler == SAMPLER_UNIFORM:
+            color = sample_uniform_rgb(number_of_circles, seed=rng.integers(0, 1e10)).astype(float)
+        elif sampler == SAMPLER_SATURATED:
+            color = sample_saturated_color(number_of_circles, seed=rng.integers(0, 1e10)).astype(float)
+        elif sampler == SAMPLER_NATURAL:
+            assert natural_image_list is not None, "Please provide a list of images to sample colors from."
+            color = sample_color_from_images(number_of_circles, seed=rng.integers(0, 1e10),
+                                             path_to_images=natural_image_list).astype(float)
+        else:
+            raise NotImplementedError(f"Unknown color sampler {sampler}")
+    else:
+        color = rng.uniform(0.25, 0.75, size=(number_of_circles, 1))
+    return center_x, center_y, radius, color

src/rstor/synthetic_data/interactive/interactive_dead_leaves.py

@@ -0,0 +1,81 @@
+from rstor.synthetic_data.dead_leaves_cpu import cpu_dead_leaves_chart
+from rstor.synthetic_data.dead_leaves_gpu import gpu_dead_leaves_chart
+from rstor.properties import SAMPLER_UNIFORM, SAMPLER_DIV2K, SAMPLER_NATURAL, SAMPLER_SATURATED, DATASET_PATH
+import sys
+import numpy as np
+from interactive_pipe import interactive_pipeline, interactive
+from typing import Optional
+
+
+def dead_leave_plugin(ds=1):
+    interactive(
+        background_intensity=(0.5, [0., 1.]),
+        number_of_circles=(-1, [-1, 10000]),
+        colored=(True,),
+        radius_alpha=(3., [2., 10.]),
+        seed=(0, [-1, 42]),
+        ds=(ds, [1, 5]),
+        numba_flag=(True,),  # Default CPU to avoid issues by default
+        sampler=(SAMPLER_UNIFORM, [SAMPLER_UNIFORM, SAMPLER_DIV2K, SAMPLER_SATURATED]),
+        circle_primitives=(True,),
+        anisotropy=(1., [0.1, 10.]),
+        angle=(0., [-180., 180.])
+        # ds=(ds, [1, 5])
+    )(generate_deadleave)
+
+
+def generate_deadleave(
+    background_intensity: float = 0.5,
+    number_of_circles: int = -1,
+    colored: Optional[bool] = False,
+    radius_alpha: Optional[int] = 3,
+    seed=0,
+    ds=3,
+    numba_flag=True,
+    sampler=SAMPLER_UNIFORM,
+    circle_primitives=True,
+    anisotropy=1.,
+    angle=0.,
+    global_params={}
+) -> np.ndarray:
+    global_params["ds_factor"] = ds
+    bg_color = (background_intensity, background_intensity, background_intensity)
+    natural_image_list = None
+    if sampler == SAMPLER_DIV2K:
+        sampler = SAMPLER_NATURAL
+        div2k_path = DATASET_PATH / "div2k" / "DIV2K_train_HR" / "DIV2K_train_HR"
+        natural_image_list = sorted([file for file in div2k_path.glob("*.png")])
+    if not numba_flag:
+        chart = cpu_dead_leaves_chart((512*ds, 512*ds), number_of_circles, bg_color, colored,
+                                      radius_alpha=radius_alpha,
+                                      seed=None if seed < 0 else seed,
+                                      sampler=sampler,
+                                      reverse=False,
+                                      natural_image_list=natural_image_list)
+    else:
+        chart = gpu_dead_leaves_chart((512*ds, 512*ds), number_of_circles, bg_color, colored,
+                                      radius_alpha=radius_alpha,
+                                      seed=None if seed < 0 else seed,
+                                      sampler=sampler,
+                                      natural_image_list=natural_image_list,
+                                      circle_primitives=circle_primitives,
+                                      anisotropy=anisotropy,
+                                      angle=angle).copy_to_host()
+    if chart.shape[-1] == 1:
+        chart = chart.repeat(3, axis=-1)
+        # Required to switch from colors to gray scale visualization.
+    return chart
+
+
+def deadleave_pipeline():
+    deadleave_chart = generate_deadleave()
+    return deadleave_chart
+
+
+def main(argv):
+    dead_leave_plugin(ds=1)
+    interactive_pipeline(gui="auto")(deadleave_pipeline)()
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

src/rstor/utils.py

@@ -0,0 +1,12 @@
+import numpy as np
+import numba
+import torch
+
+THREADS_PER_BLOCK = 32  # 32 or 16
+DEFAULT_NUMPY_FLOAT_TYPE = np.float32
+DEFAULT_CUDA_FLOAT_TYPE = numba.float32
+DEFAULT_TORCH_FLOAT_TYPE = torch.float32
+
+
+DEFAULT_NUMPY_INT_TYPE = np.int32
+DEFAULT_CUDA_INT_TYPE = numba.int32

test/test_dataloader.py

@@ -0,0 +1,49 @@
+import torch
+from rstor.data.synthetic_dataloader import DeadLeavesDataset
+
+
+def test_dead_leaves_dataset():
+    # Test case 1: Default parameters
+    dataset = DeadLeavesDataset(noise_stddev=(0, 0), ds_factor=1)
+    assert len(dataset) == 1000
+    assert dataset.size == (128, 128)
+    assert dataset.frozen_seed is None
+    assert dataset.config_dead_leaves == {}
+
+    # Test case 2: Custom parameters
+    dataset = DeadLeavesDataset(size=(256, 256), length=500, frozen_seed=42, number_of_circles=5,
+                                background_color=(0.2, 0.4, 0.6), colored=True, radius_min=1, radius_alpha=3,
+                                noise_stddev=(0, 0), ds_factor=1)
+    assert len(dataset) == 500
+    assert dataset.size == (256, 256)
+    assert dataset.frozen_seed == 42
+    assert dataset.config_dead_leaves == {
+        'number_of_circles': 5,
+        'background_color': (0.2, 0.4, 0.6),
+        'colored': True,
+        'radius_min': 1,
+        'radius_alpha': 3
+    }
+
+    # Test case 3: Check item retrieval
+    item, item_tgt = dataset[0]
+    assert isinstance(item, torch.Tensor)
+    assert item.shape == (3, 256, 256)
+
+    # Test case 4: Repeatable results with frozen seed
+    dataset1 = DeadLeavesDataset(frozen_seed=42, noise_stddev=(0, 0), number_of_circles=256)
+    dataset2 = DeadLeavesDataset(frozen_seed=42, noise_stddev=(0, 0), number_of_circles=256)
+    item1, item_tgt1 = dataset1[0]
+    item2, item_tgt2 = dataset2[0]
+    assert torch.all(torch.eq(item1, item2))
+    
+    # Test case 5: Visualize
+    # dataset = DeadLeavesDataset(size=(256, 256), length=500, frozen_seed=43,
+    #                                 background_color=(0.2, 0.4, 0.6), colored=True, radius_min=1, radius_alpha=3,
+    #                                 noise_stddev=(0, 0), ds_factor=1)
+    # item, item_tgt = dataset[0]
+    # import matplotlib.pyplot as plt
+    # plt.figure()
+    # plt.imshow(item.permute(1, 2, 0).detach().cpu())
+    # plt.show()
+    # print("done")

test/test_dataloader_gpu.py

@@ -0,0 +1,56 @@
+import torch
+from rstor.data.synthetic_dataloader import DeadLeavesDatasetGPU
+import numba
+
+
+def test_dead_leaves_dataset_gpu():
+    if not numba.cuda.is_available():
+        return
+
+    # Test case 1: Default parameters
+    dataset = DeadLeavesDatasetGPU(noise_stddev=(0, 0), ds_factor=1)
+    assert len(dataset) == 1000
+    assert dataset.size == (128, 128)
+    assert dataset.frozen_seed is None
+    assert dataset.config_dead_leaves == {}
+
+    # Test case 2: Custom parameters
+    dataset = DeadLeavesDatasetGPU(size=(256, 256), length=500, frozen_seed=42, number_of_circles=5,
+                                   background_color=(0.2, 0.4, 0.6), colored=True, radius_min=1, radius_alpha=3,
+                                   noise_stddev=(0, 0), ds_factor=1)
+    assert len(dataset) == 500
+    assert dataset.size == (256, 256)
+    assert dataset.frozen_seed == 42
+    assert dataset.config_dead_leaves == {
+        'number_of_circles': 5,
+        'background_color': (0.2, 0.4, 0.6),
+        'colored': True,
+        'radius_min': 1,
+        'radius_alpha': 3
+    }
+
+    # Test case 3: Check item retrieval
+    item, item_tgt = dataset[0]
+    assert isinstance(item, torch.Tensor)
+    assert item.shape == (3, 256, 256)
+
+    # Test case 4: Repeatable results with frozen seed
+    dataset1 = DeadLeavesDatasetGPU(frozen_seed=42, noise_stddev=(0, 0), number_of_circles=256)
+    dataset2 = DeadLeavesDatasetGPU(frozen_seed=42, noise_stddev=(0, 0), number_of_circles=256)
+    item1, item_tgt1 = dataset1[0]
+    item2, item_tgt2 = dataset2[0]
+    assert torch.all(torch.eq(item1, item2))
+
+
+
+    # Test case 5: Visualize
+    # dataset = DeadLeavesDatasetGPU(size=(256, 256), length=500, frozen_seed=44, number_of_circles=10_000,
+    #                                 background_color=(0.2, 0.4, 0.6), colored=True, radius_min=1, radius_alpha=3,
+    #                                 noise_stddev=(0, 0), ds_factor=1)
+    # item, item_tgt = dataset[0]
+    # import matplotlib.pyplot as plt
+    # plt.figure()
+    # plt.imshow(item.permute(1, 2, 0).detach().cpu())
+    # plt.show()
+    # print("done")
+

test/test_dataloader_stored.py

@@ -0,0 +1,67 @@
+import torch
+from rstor.data.stored_images_dataloader import RestorationDataset
+from numba import cuda
+from rstor.properties import DATASET_PATH, AUGMENTATION_FLIP, AUGMENTATION_ROTATE
+
+
+def test_dataloader_stored():
+    if not cuda.is_available():
+        print("cuda unavailable, exiting")
+        return
+
+    # Test case 1: Default parameters
+    dataset = RestorationDataset(noise_stddev=(0, 0),
+                                 images_path=DATASET_PATH/"sample")
+    assert len(dataset) == 2
+    assert dataset.frozen_seed is None
+
+    # Test case 2: Custom parameters
+    dataset = RestorationDataset(images_path=DATASET_PATH/"sample",
+                                 size=(64, 64),
+                                 frozen_seed=42,
+                                 noise_stddev=(0, 0))
+    assert len(dataset) == 2
+    assert dataset.frozen_seed == 42
+
+    # Test case 3: Check item retrieval
+    item, item_tgt = dataset[0]
+    assert isinstance(item, torch.Tensor)
+    assert item.shape == item_tgt.shape
+    assert item.shape == (3, 64, 64)
+
+    # Test case 4: Repeatable results with frozen seed
+    dataset1 = RestorationDataset(images_path=DATASET_PATH/"sample", 
+                                  frozen_seed=42, noise_stddev=(0, 0))
+    dataset2 = RestorationDataset(images_path=DATASET_PATH/"sample",
+                                  frozen_seed=42, noise_stddev=(0, 0))
+    item1, item_tgt1 = dataset1[0]
+    item2, item_tgt2 = dataset2[0]
+    
+    assert torch.all(torch.eq(item1, item2))
+    
+    # Test case 4: Repeatable results with frozen seed and augmentation
+    augmentation_list = [AUGMENTATION_FLIP, AUGMENTATION_ROTATE]
+    dataset1 = RestorationDataset(images_path=DATASET_PATH/"sample", 
+                                  frozen_seed=42, noise_stddev=(0, 0),
+                                  augmentation_list=augmentation_list)
+    dataset2 = RestorationDataset(images_path=DATASET_PATH/"sample",
+                                  frozen_seed=42, noise_stddev=(0, 0),
+                                  augmentation_list=augmentation_list)
+    item1, item_tgt1 = dataset1[0]
+    item2, item_tgt2 = dataset2[0]
+    assert torch.all(torch.eq(item1, item2))
+
+
+
+    # Test case 5: Visualize
+    # dataset = RestorationDataset(images_path=DATASET_PATH/"sample",
+    #                              noise_stddev=(0, 0),
+    #                              augmentation_list=augmentation_list)
+    # item, item_tgt = dataset[0]
+    # import matplotlib.pyplot as plt
+    # plt.figure()
+    # plt.imshow(item.permute(1, 2, 0).detach().cpu())
+    # plt.show()
+    # breakpoint()
+    print("done")
+

test/test_dead_leaves.py

@@ -0,0 +1,44 @@
+import numpy as np
+from rstor.synthetic_data.dead_leaves_cpu import cpu_dead_leaves_chart
+from rstor.properties import SAMPLER_NATURAL, DATASET_PATH
+
+
+def test_dead_leaves_chart():
+    # Test case 1: Default parameters
+    chart = cpu_dead_leaves_chart()
+    assert isinstance(chart, np.ndarray)
+    assert chart.shape == (100, 100, 3)
+
+    # Test case 2: Custom size and number of circles
+    chart = cpu_dead_leaves_chart(size=(200, 150), number_of_circles=10)
+    assert isinstance(chart, np.ndarray)
+    assert chart.shape == (200, 150, 3)
+
+    # Test case 3: Colored circles
+    chart = cpu_dead_leaves_chart(colored=True, number_of_circles=300)
+    assert isinstance(chart, np.ndarray)
+    assert chart.shape == (100, 100, 3)
+
+    # Test case 4: Custom radius mean and stddev
+    chart = cpu_dead_leaves_chart(radius_min=5, radius_alpha=2, number_of_circles=300)
+    assert isinstance(chart, np.ndarray)
+    assert chart.shape == (100, 100, 3)
+
+    # Test case 5: Custom background color
+    chart = cpu_dead_leaves_chart(background_color=(0.2, 0.4, 0.6), number_of_circles=300)
+    assert isinstance(chart, np.ndarray)
+    assert chart.shape == (100, 100, 3)
+
+    # Test case 6: Custom seed
+    chart1 = cpu_dead_leaves_chart(seed=42, number_of_circles=300)
+    chart2 = cpu_dead_leaves_chart(seed=42, number_of_circles=300)
+    assert np.array_equal(chart1, chart2)
+
+
+def test_dead_leaves_color_sampler():
+    img_list = sorted(
+        list((DATASET_PATH / "sample").glob("*.png"))
+    )
+    _gen = cpu_dead_leaves_chart(number_of_circles=300, sampler=SAMPLER_NATURAL, natural_image_list=img_list)
+    # from interactive_pipe.data_objects.image import Image
+    # Image(_gen).show()