From 2852a5320af7d2421ab686dac27b01ac3047e8d8 Mon Sep 17 00:00:00 2001
From: "Dorofeeva Elizaveta (EXT)" <ext-dorofe_e@ra-c-106.psi.ch>
Date: Fri, 4 Jul 2025 08:43:28 +0200
Subject: [PATCH] Bug fixes in new streak finder

---
 dap/algos/streakfind.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/dap/algos/streakfind.py b/dap/algos/streakfind.py
index 51e1026..67fd5c3 100644
--- a/dap/algos/streakfind.py
+++ b/dap/algos/streakfind.py
@@ -14,6 +14,10 @@ from streak_finder.label import Structure2D
 DEFAULT_NUM_THREADS = 16
 
 def calc_cbd_analysis(results, data, pf_pixel_mask):
+    do_snr = results.get("do_snr", False)
+    if not do_snr:
+        return data
+
     try:
         cryst_data = _generate_cryst_data(results, data, pf_pixel_mask)
     except Exception as error:  # Broad exception - we don't want to break anything here
@@ -26,13 +30,10 @@ def calc_cbd_analysis(results, data, pf_pixel_mask):
     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}"
-    return cryst_data.snr
+    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:
-    do_snr = results.get("do_snr", False)
-    if not do_snr:
-        return
-
     params_required = [
         "whitefield_data_file",
         "mask_data_file",
@@ -73,7 +74,8 @@ def _generate_cryst_data(results, data, pf_pixel_mask) -> CrystData:
     )
     if scale_whitefield:
         data = data.scale_whitefield(method='median', num_threads=num_threads)
-
+    
+    data = data.update_snr()
     return data
 
 def _calc_streakfinder_analysis(results, cryst_data: CrystData):