Add noise, fix return for extract roi.

lensless/utils/dataset.py

@@ -30,6 +30,7 @@
 from lensless.hardware.sensor import sensor_dict, SensorParam
 from scipy.ndimage import rotate
 import warnings
+from waveprop.noise import add_shot_noise
 from PIL import Image
 
 
@@ -169,8 +170,6 @@ def __getitem__(self, idx):
 
         # add noise
         if self.input_snr is not None:
-            from waveprop.noise import add_shot_noise
-
             lensless = add_shot_noise(lensless, self.input_snr)
 
         # flip image x and y if needed
@@ -1037,6 +1036,7 @@ def __init__(
         sensor="rpi_hq",
         slm="adafruit",
         simulation_config=dict(),
+        snr_db=40,
         **kwargs,
     ):
         """
@@ -1131,6 +1131,7 @@ def __init__(
         self.crop = None
         self.random_flip = None
         self.flipud = flipud
+        self.snr_db = snr_db
 
         self.display_res = display_res
         self.alignment = None
@@ -1193,6 +1194,11 @@ def _get_images_pair(self, idx):
             mask_label = self.dataset[idx]["mask_label"]
             self.convolver.set_psf(self.psf[mask_label])
         lensless = self.convolver.convolve(lensed)
+
+        # add noise
+        if self.snr_db is not None:
+            lensless = add_shot_noise(lensless, self.snr_db)
+
         if lensless.max() > 1:
             print("CLIPPING!")
             lensless /= lensless.max()
@@ -1252,6 +1258,8 @@ def extract_roi(self, reconstruction, lensed=None, axis=(1, 2), **kwargs):
             else:
                 reconstruction = reconstruction[0]
 
+        if lensed is None:
+            return reconstruction
         return reconstruction, lensed