black

Author: tachella

deepinv_bench/__init__.py

@@ -1 +1 @@
-from .run import run_benchmark
+from .run import run_benchmark

deepinv_bench/base/runner_solver.py

@@ -13,14 +13,13 @@ class Solver(BaseSolver):
     To use this class, set the `model` and `name` attributes to run
     it in the model.
     """
+
     name = "runner-solver"
     model = None
     parameters = {}
 
     def set_objective(self, train_dataset=None, physics=None):
-        device = (
-            dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
-        )
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
         # Make sure the model is on the correct device and has a device
         # attribute to we can properly move the test data later.
         self.model.to(device)

deepinv_bench/benchmark_template/datasets/dataset.py

@@ -14,29 +14,24 @@ class Dataset(BaseDataset):
     # IMPORTANT: the names of physics and noise should match
     # the ones defined in deepinv.physics with exact same spelling
     parameters = {
-        'physics': ['physics_name'],
-        'noise': ['noise_name'],
-        'img_size': [256],
+        "physics": ["physics_name"],
+        "noise": ["noise_name"],
+        "img_size": [256],
         # add any other parameter you might need
     }
 
     def get_data(self):
         root = get_data_path("Set14_HR")
 
-        transform = transforms.Compose([
-            transforms.Resize((self.img_size, self.img_size)),
-            transforms.ToTensor()
-        ])
+        transform = transforms.Compose(
+            [transforms.Resize((self.img_size, self.img_size)), transforms.ToTensor()]
+        )
 
         # load the dataset
-        dataset = dinv.datasets.Set14HR(
-            root, download=True, transform=transform
-        )
+        dataset = dinv.datasets.Set14HR(root, download=True, transform=transform)
 
         # define the physics according to the parameters
-        physics = dinv.physics.Denoising(
-            noise_model=dinv.physics.GaussianNoise()
-        )
+        physics = dinv.physics.Denoising(noise_model=dinv.physics.GaussianNoise())
 
         return dict(
             dataset=dataset,

deepinv_bench/benchmark_template/objective.py

@@ -17,21 +17,18 @@ class Objective(BaseObjective):
     # Bump it up if the benchmark depends on a new feature of benchopt.
     min_benchopt_version = "1.8"
 
-    sampling_strategy = 'run_once'
+    sampling_strategy = "run_once"
 
     def set_data(self, dataset, physics):
         self.dataset = dataset
         self.physics = physics
 
     def evaluate_result(self, model):
-        device = getattr(model, 'device', None)
+        device = getattr(model, "device", None)
         self.physics = self.physics.to(device)
 
         # change metrics if needed
-        metrics = [
-            dinv.loss.PSNR(),
-            dinv.loss.NIQE(device=device)
-        ]
+        metrics = [dinv.loss.PSNR(), dinv.loss.NIQE(device=device)]
 
         results = dinv.test(
             model,
@@ -40,14 +37,15 @@ def evaluate_result(self, model):
             online_measurements=True,
             device=device,
             metrics=metrics,
-            compare_no_learning=False
+            compare_no_learning=False,
         )
 
         return results
 
     def get_one_result(self):
         class DummyModel:
-            def eval(self): pass
+            def eval(self):
+                pass
 
             def __call__(self, x, physics=None):
                 return physics.A_adjoint(x)

deepinv_bench/benchmark_template/solvers/solver1.py

@@ -5,15 +5,13 @@
 
 
 class Solver(BaseSolver):
-    name = 'solver_name'
+    name = "solver_name"
 
     # add any hyper-parameters here
     parameters = {}
 
     def set_objective(self, train_dataset=None, physics=None):
-        device = (
-            dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
-        )
+        device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
 
         # replace by your model, should take (y, physics) as input
         self.model = dinv.models.DnCNN(device=device)

deepinv_bench/cbsd500_gaussian_denoising/datasets/cbsd68.py

@@ -14,27 +14,22 @@ class Dataset(BaseDataset):
 
     name = "CBSD68"
     parameters = {
-        'physics': ['Denoising'],
-        'noise': ['GaussianNoise'],
-        'sigma': [0.1],
-        'img_size': [256],
-        'debug': [False],
+        "physics": ["Denoising"],
+        "noise": ["GaussianNoise"],
+        "sigma": [0.1],
+        "img_size": [256],
+        "debug": [False],
     }
 
-    test_parameters = {
-        "debug": [True]
-    }
+    test_parameters = {"debug": [True]}
 
     def get_data(self):
 
         root = get_data_path("CBSD68")
-        transform = transforms.Compose([
-            transforms.Resize((self.img_size, self.img_size)),
-            transforms.ToTensor()
-        ])
-        dataset = CBSD68(
-            root, download=True, transform=transform
+        transform = transforms.Compose(
+            [transforms.Resize((self.img_size, self.img_size)), transforms.ToTensor()]
         )
+        dataset = CBSD68(root, download=True, transform=transform)
 
         if self.debug:
             dataset = dinv.torch.utils.data.Subset(dataset, [0, 1, 2])

deepinv_bench/cbsd500_gaussian_denoising/objective.py

@@ -25,13 +25,10 @@ def set_data(self, dataset, physics):
         self.physics = physics
 
     def evaluate_result(self, model):
-        device = getattr(model, 'device', None)
+        device = getattr(model, "device", None)
         self.physics = self.physics.to(device)
 
-        metrics = [
-            dinv.loss.PSNR(),
-            dinv.loss.NIQE(device=device)
-        ]
+        metrics = [dinv.loss.PSNR(), dinv.loss.NIQE(device=device)]
 
         results = dinv.test(
             model,
@@ -40,15 +37,16 @@ def evaluate_result(self, model):
             online_measurements=True,
             device=device,
             metrics=metrics,
-            compare_no_learning=False
+            compare_no_learning=False,
         )
 
         return results
 
     def get_one_result(self):
 
         class DummyModel:
-            def eval(self): pass
+            def eval(self):
+                pass
 
             def __call__(self, x, physics=None):
                 return physics.A_adjoint(x)

deepinv_bench/cbsd500_gaussian_denoising/solvers/drunet.py

@@ -4,14 +4,12 @@
 
 
 class Solver(BaseSolver):
-    name = 'DRUNet'
+    name = "DRUNet"
 
     parameters = {}
 
     def set_objective(self, train_dataset=None, physics=None):
-        self.model = dinv.models.ArtifactRemoval(
-            dinv.models.DRUNet()
-        )
+        self.model = dinv.models.ArtifactRemoval(dinv.models.DRUNet())
 
     def run(self, _):
         pass

deepinv_bench/cbsd500_gaussian_denoising/solvers/restformer.py

@@ -4,14 +4,12 @@
 
 
 class Solver(BaseSolver):
-    name = 'Restormer'
+    name = "Restormer"
 
     parameters = {}
 
     def set_objective(self, train_dataset=None, physics=None):
-        self.model = dinv.models.ArtifactRemoval(
-            dinv.models.Restormer()
-        )
+        self.model = dinv.models.ArtifactRemoval(dinv.models.Restormer())
 
     def run(self, _):
         pass

deepinv_bench/div2k_gaussian_deblurring/datasets/div2k.py

@@ -14,33 +14,28 @@ class Dataset(BaseDataset):
 
     name = "DIV2K"
     parameters = {
-        'physics': ['Blur'],
-        'noise': ['GaussianNoise'],
-        'sigma': [0.1],
-        'img_size': [256],
-        'debug': [False],
+        "physics": ["Blur"],
+        "noise": ["GaussianNoise"],
+        "sigma": [0.1],
+        "img_size": [256],
+        "debug": [False],
     }
 
-    test_parameters = {
-        "debug": [True]
-    }
+    test_parameters = {"debug": [True]}
 
     def get_data(self):
         root = get_data_path("DIV2K")
-        transform = transforms.Compose([
-            transforms.Resize((self.img_size, self.img_size)),
-            transforms.ToTensor()
-        ])
-        dataset = DIV2K(
-            root, mode="val", download=True, transform=transform
+        transform = transforms.Compose(
+            [transforms.Resize((self.img_size, self.img_size)), transforms.ToTensor()]
         )
+        dataset = DIV2K(root, mode="val", download=True, transform=transform)
         if self.debug:
             dataset = dinv.torch.utils.data.Subset(dataset, [0, 1, 2])
 
         return dict(
             dataset=dataset,
             physics=Blur(
                 filter=dinv.physics.blur.gaussian_blur(2),
-                noise_model=GaussianNoise(0.05)
+                noise_model=GaussianNoise(0.05),
             ),
         )