fix, minor bugs

Author: singer-yang

.gitignore

@@ -10,6 +10,7 @@ deeplens/do_code/
 !lenses/*.png
 !lenses/**/*.png
 !images/*.png
+!assets/*.png
 debug.py
 *.log
 temp/*
@@ -22,3 +23,5 @@ datasets/BSDS300/
 docs/_build/
 
 # add my own files
+test/test_outputs
+visualization/

deeplens/optics/monte_carlo.py

@@ -192,7 +192,8 @@ def assign_points_to_pixels(
             if phase is None:
                 raise ValueError("Phase must be provided for coherent mode")
 
-            grid = torch.zeros(ks, ks, dtype=torch.complex128).to(device)
+            c_dtype = torch.complex64 if points.dtype == torch.float32 else torch.complex128
+            grid = torch.zeros(ks, ks, dtype=c_dtype).to(device)
             grid.index_put_(
                 tuple(pixel_indices_tl.t()),
                 (1 - w_b) * (1 - w_r) * mask * amp * torch.exp(1j * phase),
@@ -246,7 +247,8 @@ def assign_points_to_pixels(
             if phase is None:
                 raise ValueError("Phase must be provided for coherent mode")
 
-            grid = torch.zeros(ks, ks, dtype=torch.complex128).to(device)
+            c_dtype = torch.complex64 if points.dtype == torch.float32 else torch.complex128
+            grid = torch.zeros(ks, ks, dtype=c_dtype).to(device)
             grid.index_put_(
                 tuple(pixel_indices_tl.t()),
                 mask * amp * torch.exp(1j * phase),

deeplens/optics/wave.py

@@ -414,7 +414,7 @@ def pad(self, Hpad, Wpad):
             self.phy_size[1] * self.res[1] / Worg,
         ]
         self.x, self.y = self.gen_xy_grid()
-        self.z = F.pad(self.z, (Hpad, Hpad, Wpad, Wpad), mode="replicate")
+        self.z = torch.full_like(self.x, self.z[0, 0].item())
 
     def flip(self):
         """Flip the field horizontally and vertically."""

deeplens/sensor/sensor.py

@@ -143,7 +143,7 @@ def simu_noise(self, img_raw, iso):
 
         # Calculate noise standard deviation
         shotnoise_std = torch.clamp(
-            self.shotnoise_std_alpha * torch.sqrt(img_raw - black_level)
+            self.shotnoise_std_alpha * torch.sqrt(torch.clamp(img_raw - black_level, min=0.0))
             + self.shotnoise_std_beta,
             0.0,
         )

deeplens/utils.py

@@ -136,8 +136,17 @@ def img2batch(img):
             img = (
                 torch.tensor(img).unsqueeze(0).permute(0, 3, 1, 2)
             )  # (H, W, C) -> (1, C, H, W)
+        elif torch.is_tensor(img):
+            if img.shape[0] in [1, 3]:
+                # Assume (C, H, W) -> (1, C, H, W)
+                img = img.unsqueeze(0)
+            elif img.shape[-1] in [1, 3]:
+                # Assume (H, W, C) -> (1, C, H, W)
+                img = img.permute(2, 0, 1).unsqueeze(0)
+            else:
+                 raise ValueError("Image channel should be 1 or 3.")
         else:
-            raise ValueError("Image should be numpy array.")
+            raise ValueError("Image should be numpy array or torch tensor.")
 
     # Tensor dtype
     if img.dtype == torch.uint8:
@@ -187,7 +196,7 @@ def batch_psnr(pred, target, max_val=1.0, eps=1e-8):
     # Calculate PSNR
     psnr = 20 * torch.log10(max_val / torch.sqrt(mse + eps))
 
-    return psnr.item()
+    return psnr
 
 
 def batch_SSIM(img, img_clean):