added inpainting

app.py

@@ -159,7 +159,7 @@ def get_dataset(dataset_name):
         physics_name = 'CT'
         baseline_name = 'DPIR_CT'
     else:
-        available_physics = ['MotionBlur_medium', 'MotionBlur_hard',
+        available_physics = ['Inpainting', 'MotionBlur_medium', 'MotionBlur_hard',
                              'GaussianBlur_easy', 'GaussianBlur_medium', 'GaussianBlur_hard']
         physics_name = 'MotionBlur_hard'
         baseline_name = 'DPIR'
@@ -192,11 +192,11 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
 
     ### USER-SPECIFIC VARIABLES
     dataset_placeholder = gr.State(get_dataset_on_DEVICE_STR("Natural"))
-    available_physics_placeholder = gr.State(['MotionBlur_medium', 'MotionBlur_hard',
+    available_physics_placeholder = gr.State(['Inpainting', 'MotionBlur_medium', 'MotionBlur_hard',
                                               'GaussianBlur_easy', 'GaussianBlur_medium', 'GaussianBlur_hard'])
     # Issue giving directly a `torch.nn.module` to `gr.State(...)` since it has __call__ method
     # Solution: using lambda expression
-    physics_placeholder = gr.State(lambda: get_physics_on_DEVICE_STR("MotionBlur_medium"))
+    physics_placeholder = gr.State(lambda: get_physics_on_DEVICE_STR("MotionBlur_hard"))
     model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DPIR"))
 
     print(f"[Render] CUDA max allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")

factories.py

@@ -13,7 +13,7 @@ from physics.blur_generator import GaussianBlurGenerator
 
 class PhysicsWithGenerator(torch.nn.Module):
     """Interface between Physics, Generator and Gradio."""
-    all_physics = ["MotionBlur_medium", "MotionBlur_hard",
+    all_physics = ["Inpainting", "MotionBlur_medium", "MotionBlur_hard",
                    "GaussianBlur_easy", "GaussianBlur_medium", "GaussianBlur_hard",
                    "MRI", "CT"]
 
@@ -83,6 +83,22 @@ class PhysicsWithGenerator(torch.nn.Module):
                                  "updatable_params_converter": {"sigma": float},
                                  "fixed_params": {"noise_sigma_min": 0.05, "noise_sigma_max": 0.05,
                                                   "blur_sigma": 4.0, "psf_size": 31, "num_channels": 1}}
+        elif self.name == "Inpainting":
+            self.physics = dinv.physics.Inpainting(tensor_size=(256, 256), mask=split_ratio,
+                                                   noise_model=dinv.physics.GaussianNoise(sigma=sigma),
+                                                   device=device_str)
+            self.physics_generator = dinv.physics.generator.BernoulliSplittingMaskGenerator((3, 256, 256),
+                                                                                              split_ratio=split_ratio, pixelwise=pixelwise,
+                                                                                              random_split_ratio=True, min_split_ratio=split_ratio,
+                                                                                              max_split_ratio=split_ratio, device=device_str)
+            self.generator = dinv.physics.generator.BernoulliSplittingMaskGenerator((3, 256, 256),
+                                                                                  split_ratio=split_ratio, pixelwise=pixelwise,
+                                                                                  random_split_ratio=True, min_split_ratio=split_ratio,
+                                                                                  max_split_ratio=split_ratio, device=device_str)
+
+            self.saved_params = {"updatable_params": {},
+                                 "updatable_params_converter": {"sigma": float},
+                                 "fixed_params": {"sigma": sigma}}
         elif self.name == "MRI":
             self.physics = dinv.physics.MRI(noise_model=dinv.physics.GaussianNoise(sigma=.01).to(device_str),
                                             img_size=(640, 320), device=device_str)
@@ -101,14 +117,14 @@ class PhysicsWithGenerator(torch.nn.Module):
                 circle=False,
                 normalize=True,
                 device=device_str,
-                noise_model=dinv.physics.GaussianNoise(sigma=1e-4).to(device_str),
+                noise_model=dinv.physics.GaussianNoise(sigma=1e-3).to(device_str),
                 max_iter=10,
             )
-            self.physics_generator = SigmaGenerator(sigma_min=1e-4, sigma_max=1e-4, device=device_str)
-            self.generator = SigmaGenerator(sigma_min=1e-4, sigma_max=1e-4, device=device_str)
+            self.physics_generator = SigmaGenerator(sigma_min=1e-3, sigma_max=1e-3, device=device_str)
+            self.generator = SigmaGenerator(sigma_min=1e-3, sigma_max=1e-3, device=device_str)
             self.saved_params = {"updatable_params": {"sigma": 0.1},
                                  "updatable_params_converter": {"sigma": float},
-                                 "fixed_params": {"noise_sigma_min": 1e-4, "noise_sigma_max": 1e-4,
+                                 "fixed_params": {"noise_sigma_min": 1e-3, "noise_sigma_max": 1e-3,
                                                   "angles": angles, "max_iter": 10}}
 
     def display_saved_params(self) -> str: