Addresses #1144 to exclude user-specified bad frames from registratio…

…n metric calculation

suite2p/registration/__init__.py

@@ -2,6 +2,6 @@
 Copyright © 2023 Howard Hughes Medical Institute, Authored by Carsen Stringer and Marius Pachitariu.
 """
 from .register import (registration_wrapper, save_registration_outputs_to_ops,
-                       compute_enhanced_mean_image)
+                       compute_enhanced_mean_image, load_badframes)
 from .metrics import get_pc_metrics
 from .zalign import compute_zpos

suite2p/registration/register.py

@@ -665,18 +665,7 @@ def registration_wrapper(f_reg, f_raw=None, f_reg_chan2=None, f_raw_chan2=None,
         meanImg_chan2 = mean_img
         meanImg = mean_img_alt
 
-    # compute valid region
-    badframes = np.zeros(n_frames, "bool")
-    if "data_path" in ops and len(ops["data_path"]) > 0:
-        badfrfile = path.abspath(path.join(ops["data_path"][0], "bad_frames.npy"))
-        # Check if badframes file exists
-        if path.isfile(badfrfile):
-            print("bad frames file path: %s" % badfrfile)
-            bf_indices = np.load(badfrfile)
-            bf_indices = bf_indices.flatten().astype(int)
-            # Set indices of badframes to true
-            badframes[bf_indices] = True
-            print("number of badframes: %d" % badframes.sum())
+    badframes = load_badframes(n_frames, ops)
 
     # return frames which fall outside range
     badframes, yrange, xrange = compute_crop(
@@ -692,6 +681,22 @@ def registration_wrapper(f_reg, f_raw=None, f_reg_chan2=None, f_raw_chan2=None,
 
     return refImg, rmin, rmax, meanImg, rigid_offsets, nonrigid_offsets, zest, meanImg_chan2, badframes, yrange, xrange
 
+def load_badframes(n_frames, ops):
+    # compute valid region
+    badframes = np.zeros(n_frames, "bool")
+    if "data_path" in ops and len(ops["data_path"]) > 0:
+        badfrfile = path.abspath(path.join(ops["data_path"][0], "bad_frames.npy"))
+        # Check if badframes file exists
+        if path.isfile(badfrfile):
+            print("User-specified bad frames file path: %s" % badfrfile)
+            bf_indices = np.load(badfrfile)
+            bf_indices = bf_indices.flatten().astype(int)
+            # Set indices of badframes to true
+            badframes[bf_indices] = True
+            print("number of badframes: %d" % badframes.sum())
+
+    return badframes
+
 
 def save_registration_outputs_to_ops(registration_outputs, ops):
 

suite2p/run_s2p.py

@@ -152,6 +152,9 @@ def pipeline(f_reg, f_raw=None, f_reg_chan2=None, f_raw_chan2=None,
             nsamp = min(2000 if n_frames < 5000 or Ly > 700 or Lx > 700 else 5000,
                         n_frames)
             inds = np.linspace(0, n_frames - 1, nsamp).astype("int")
+            print("Taking into account user-specified bad frames for registration metrics...")
+            user_bad_frames = registration.load_badframes(n_frames, ops)
+            inds = np.setdiff1d(inds,np.where(user_bad_frames)[0])
             mov = f_reg[inds]
             mov = mov[:, ops["yrange"][0]:ops["yrange"][-1],
                       ops["xrange"][0]:ops["xrange"][-1]]