Update background.py

SINTEF · Sep 25, 2024 · 907f70a · 907f70a
1 parent 5698e5f
commit 907f70a
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/pyopia/background.py b/pyopia/background.py
@@ -172,8 +172,13 @@ def correct_im_accurate_divide(imbg, imraw):
         corrected image, same type as input
     '''
 
-    for i in range(imbg.shape[2]):
-        imbg[:, :, i][imbg[:, :, i] == 0] = 1 / 255   # change the zero_value pixels to prevent RuntimeWarning: 'divide by zero'
+    image_shape = np.shape(imbg)
+    if len(image_shape) > 2: 
+        # change the zero_value pixels in each channel to prevent RuntimeWarning: 'divide by zero'
+        for i in range(imbg.shape[2]):
+            imbg[:, :, i][imbg[:, :, i] == 0] = 1 / 255
+        else:
+            imbg[imbg == 0] = 1 / 255
 
     im_corrected = imraw / imbg