ENH new tasks

.gitignore

@@ -16,3 +16,6 @@ benchopt.ini
 
 .DS_Store
 coverage.xml
+
+tmp
+data/

benchmark_utils/denoiser_2c.py

@@ -0,0 +1,26 @@
+import torch
+from deepinv.models import DRUNet
+from deepinv.models import Denoiser
+
+
+class Denoiser_2c(Denoiser):
+    def __init__(self, device):
+        super(Denoiser_2c, self).__init__()
+        self.model_c1 = DRUNet(
+            in_channels=1, out_channels=1,
+            pretrained="download", device=device
+        )
+        self.model_c2 = DRUNet(
+            in_channels=1, out_channels=1,
+            pretrained="download", device=device
+        )
+
+    def forward(self, y, sigma):
+        y1, y2 = torch.split(y, 1, dim=1)
+
+        x_hat_1 = self.model_c1(y1, sigma=sigma)
+        x_hat_2 = self.model_c2(y2, sigma=sigma)
+
+        x_hat = torch.cat([x_hat_1, x_hat_2], dim=1)
+
+        return x_hat

benchmark_utils/helper.py

@@ -0,0 +1,78 @@
+from benchopt import safe_import_context
+
+with safe_import_context() as import_ctx:
+    import deepinv as dinv
+    from deepinv.physics import (
+        Denoising,
+        GaussianNoise,
+        Downsampling,
+        Demosaicing,
+        BlurFFT,
+        Inpainting
+    )
+    from deepinv.physics.generator import MotionBlurGenerator
+
+
+DEVICE = None
+
+
+def get_device():
+    global DEVICE
+    if DEVICE is not None:
+        return DEVICE
+    if dinv.torch.cuda.is_available():
+        DEVICE = dinv.utils.get_freer_gpu()
+    else:
+        DEVICE = "cpu"
+    return DEVICE
+
+
+def get_task_physic(task, img_size, device):
+    if task == "denoising":
+        noise_level_img = 0.1
+        physics = Denoising(GaussianNoise(sigma=noise_level_img))
+    elif task == "gaussian-debluring":
+        filter_torch = dinv.physics.blur.gaussian_blur(sigma=(3, 3))
+        noise_level_img = 0.03
+
+        physics = BlurFFT(
+            img_size=img_size,
+            filter=filter_torch,
+            noise_model=GaussianNoise(sigma=noise_level_img),
+            device=device
+        )
+    elif task == "motion-debluring":
+        psf_size = 31
+        motion_generator = MotionBlurGenerator(
+            (psf_size, psf_size),
+            device=device
+        )
+
+        filters = motion_generator.step(batch_size=1)
+
+        physics = BlurFFT(
+            img_size=img_size,
+            filter=filters["filter"],
+            device=device
+        )
+    elif task == "SRx4":
+        physics = Downsampling(
+            img_size=img_size,
+            filter="bicubic",
+            factor=4,
+            device=device
+        )
+    elif task == "inpainting":
+        physics = Inpainting(
+            img_size,
+            mask=0.7,
+            device=device
+        )
+    elif task == "demosaicing":
+        physics = Demosaicing(
+            img_size=img_size,
+            device=device
+        )
+    else:
+        raise Exception("Unknown task")
+    return physics

benchmark_utils/hugging_face_torch_dataset.py

@@ -2,19 +2,26 @@
 
 
 class HuggingFaceTorchDataset(torch.utils.data.Dataset):
-    def __init__(self, hf_dataset, key, transform=None):
+    def __init__(self, hf_dataset, key, physics, device, transform=None):
         self.hf_dataset = hf_dataset
         self.transform = transform
         self.key = key
+        self.device = device
+        self.physics = physics
 
     def __len__(self):
         return len(self.hf_dataset)
 
     def __getitem__(self, idx):
         sample = self.hf_dataset[idx]
-        image = sample[self.key]  # Image PIL
+        x = sample[self.key]  # Image PIL
 
         if self.transform:
-            image = self.transform(image)
+            x = self.transform(x)
 
-        return image
+        x = x.to(self.device)
+
+        y = self.physics(x.unsqueeze(0))
+        y = y.squeeze(0)
+
+        return x, y

datasets/bsd500_cbsd68.py

@@ -3,116 +3,73 @@
 
 with safe_import_context() as import_ctx:
     import deepinv as dinv
-    import torch
     from torchvision import transforms
     from datasets import load_dataset
     from benchmark_utils.hugging_face_torch_dataset import (
         HuggingFaceTorchDataset
     )
-    from deepinv.physics import Denoising, GaussianNoise, Downsampling
-    from deepinv.physics.generator import MotionBlurGenerator
+    from benchmark_utils.helper import get_task_physic, get_device
 
 
 class Dataset(BaseDataset):
 
     name = "BSD500_CBSD68"
 
     parameters = {
-        'task': ['denoising',
-                 'gaussian-debluring',
-                 'motion-debluring',
-                 'SRx4'],
+        'task': [
+            'denoising',
+            'gaussian-debluring',
+            'motion-debluring',
+            'SRx4',
+            'inpainting',
+            'demosaicing'
+        ],
         'img_size': [256],
+        'batch_size': [2]
     }
 
     requirements = ["datasets"]
 
     def get_data(self):
-        # TODO: Remove
-        device = (
-            dinv.utils.get_freer_gpu()) if torch.cuda.is_available() else "cpu"
+        device = get_device()
 
-        if self.task == "denoising":
-            noise_level_img = 0.03
-            physics = Denoising(GaussianNoise(sigma=noise_level_img))
-        elif self.task == "gaussian-debluring":
-            filter_torch = dinv.physics.blur.gaussian_blur(sigma=(3, 3))
-            noise_level_img = 0.03
-            n_channels = 3
+        n_channels = 3
+        img_size = (n_channels, self.img_size, self.img_size)
 
-            physics = dinv.physics.BlurFFT(
-                img_size=(n_channels, self.img_size, self.img_size),
-                filter=filter_torch,
-                noise_model=dinv.physics.GaussianNoise(sigma=noise_level_img),
-                device=device
-            )
-        elif self.task == "motion-debluring":
-            psf_size = 31
-            n_channels = 3
-            motion_generator = MotionBlurGenerator(
-                (psf_size, psf_size),
-                device=device
-            )
-
-            filters = motion_generator.step(batch_size=1)
-
-            physics = dinv.physics.BlurFFT(
-                img_size=(n_channels, self.img_size, self.img_size),
-                filter=filters["filter"],
-                device=device
-            )
-        elif self.task == "SRx4":
-            n_channels = 3
-            physics = Downsampling(img_size=(n_channels,
-                                             self.img_size,
-                                             self.img_size),
-                                   filter="bicubic",
-                                   factor=4,
-                                   device=device)
-        else:
-            raise Exception("Unknown task")
+        physics = get_task_physic(self.task, img_size, device)
 
         transform = transforms.Compose([
             transforms.Resize((self.img_size, self.img_size)),
             transforms.ToTensor()
         ])
 
         path = get_data_path("BSD500")
-        train_dataset = dinv.datasets.BSDS500(
+        bsd500_dataset = dinv.datasets.BSDS500(
             path, download=True, transform=transform
         )
+        train_dataset = HuggingFaceTorchDataset(
+            bsd500_dataset,
+            key=...,
+            physics=physics,
+            device=device,
+            transform=transforms.Resize((self.img_size, self.img_size))
+        )
 
         dataset_cbsd68 = load_dataset("deepinv/CBSD68")
         test_dataset = HuggingFaceTorchDataset(
-            dataset_cbsd68["train"], key="png", transform=transform
-        )
-
-        dinv_dataset_path = dinv.datasets.generate_dataset(
-            train_dataset=train_dataset,
-            test_dataset=test_dataset,
+            dataset_cbsd68["train"],
+            key="png",
             physics=physics,
-            save_dir=get_data_path("bsd500_cbsd68"),
-            dataset_filename=self.task,
-            device=device
-        )
-
-        train_dataset = dinv.datasets.HDF5Dataset(
-            path=dinv_dataset_path, train=True
+            device=device,
+            transform=transform
         )
-        test_dataset = dinv.datasets.HDF5Dataset(
-            path=dinv_dataset_path, train=False
-        )
-
-        x, y = train_dataset[0]
-        dinv.utils.plot([x.unsqueeze(0), y.unsqueeze(0)])
-
-        x, y = test_dataset[0]
-        dinv.utils.plot([x.unsqueeze(0), y.unsqueeze(0)])
 
         return dict(
             train_dataset=train_dataset,
             test_dataset=test_dataset,
             physics=physics,
             dataset_name="BSD68",
-            task_name=self.task
+            task_name=self.task,
+            image_size=img_size,
+            batch_size=self.batch_size
         )