Make background subtraction backward compatible.

LCAV · Aug 21, 2024 · a2f557c · a2f557c
1 parent 9d4cb6d
commit a2f557c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/lensless/recon/trainable_recon.py b/lensless/recon/trainable_recon.py
@@ -202,9 +202,11 @@ def _prepare_process_block(self, process):
 
     def set_pre_process(self, pre_process):
         if isinstance(pre_process, torch.nn.DataParallel):
-            self.input_background = pre_process.module.input_background
+            if hasattr(pre_process.module, "input_background"):
+                self.input_background = pre_process.module.input_background
         else:
-            self.input_background = pre_process.input_background
+            if hasattr(pre_process, "input_background"):
+                self.input_background = pre_process.input_background
         (
             self.pre_process,
             self.pre_process_model,

diff --git a/test/test_algos.py b/test/test_algos.py
@@ -169,7 +169,7 @@ def test_trainable_recon(algorithm):
             psf = torch.rand(1, 32, 64, 3, dtype=torch_type)
             data = torch.rand(2, 1, 32, 64, 3, dtype=torch_type)
 
-            def pre_process(x, param):
+            def pre_process(x, param, background=None):
                 return x
 
             def post_process(x, param, residual=None):
@@ -206,7 +206,7 @@ def test_trainable_batch(algorithm):
         data2 = torch.rand(1, 1, 34, 64, 3, dtype=torch_type)
         data2[0, 0, ...] = data1[0, 0, ...]
 
-        def pre_process(x, param):
+        def pre_process(x, param, background=None):
             return x
 
         def post_process(x, param, residual=None):