dap/algos/streakfind.py

@@ -26,17 +26,18 @@ def calc_streakfinder_analysis(results, data, pf_pixel_mask):
         results["cbd_error"] = f"Error processing CBD data:\n{error}"
         return data
 
+    if do_snr:
+        # Changes data and mask in-place
+        data = cryst_data.snr[0].copy()
+        np.multiply(pf_pixel_mask, cryst_data.mask, out=pf_pixel_mask)
+
     try:
         _calc_streakfinder_analysis(results, cryst_data)
     except Exception as error:  # Broad exception - we don't want to break anything here
         print(f"Error processing CBD data:\n{error}")
         results["cbd_error"] = f"Error processing CBD data:\n{error}"
 
-    if not do_snr:
+    return data
-        return data
-
-    print(f"Returning data shape {cryst_data.snr[0].shape}")
-    return cryst_data.snr[0]
 
 def _generate_cryst_data(results, data, pf_pixel_mask) -> CrystData:
     params_required = [
@@ -66,7 +67,7 @@ def _generate_cryst_data(results, data, pf_pixel_mask) -> CrystData:
         whitefield = np.asarray(hf[whitefield_dataset])
 
     with h5py.File(mask_data_file, "r") as hf:
-        mask = np.asarray(hf[mask_dataset])
+        mask = np.asarray(hf[mask_dataset], dtype=np.bool)
 
     with h5py.File(std_data_file, "r") as hf:
         std = np.asarray(hf[std_dataset])
@@ -125,6 +126,11 @@ def _calc_streakfinder_analysis(results, cryst_data: CrystData):
     x_center = results.get("beam_center_x", None)
     y_center = results.get("beam_center_y", None)
 
+    mask_rois = results.get("sf_mask_rois", [])  # list of [y_min, y_max, x_min, x_max]
+
+    for mask_roi in mask_rois:
+        cryst_data = cryst_data.mask_region(mask_roi)
+
     peaks_structure = Structure2D(peak_structure_radius, peak_structure_rank)
     streaks_structure = Structure2D(streak_structure_radius, streak_structure_rank)
 
@@ -161,6 +167,7 @@ def _calc_streakfinder_analysis(results, cryst_data: CrystData):
     print(f"Found {number_of_streaks} streaks")
 
     list_result = [line.tolist() for line in streak_lines]  # arr(4, n_lines); 0coord x0, y0, x1, y1
+
     detected_streaks = np.asarray(detected.streaks.values())[streaks_mask]
     bragg_counts = [streak.total_mass() for streak in detected_streaks]