Handle APGD grasycale.

LCAV · Apr 25, 2023 · c08c1ee · c08c1ee
1 parent c75e682
commit c08c1ee
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/lensless/rfft_convolve.py b/lensless/rfft_convolve.py
@@ -32,13 +32,14 @@ def __init__(self, psf, dtype=None, pad=True, norm="ortho", **kwargs):
             self.is_torch = True
 
         # prepare shapes for reconstruction
-        self._is_rgb = len(psf.shape) == 3
-        if self._is_rgb:
+        self._is_rgb = False
+        if len(psf.shape) == 3:
+            if psf.shape[2] == 3:
+                self._is_rgb = True
             self._psf = psf
-            self._n_channels = 3
         else:
             self._psf = psf[:, :, None]
-            self._n_channels = 1
+        self._n_channels = self._psf.shape[2]
         self._psf_shape = np.array(self._psf.shape)
 
         # set dtype