From 2e066ef3a7d8d2fd2a68e1359bea452a050ef2b9 Mon Sep 17 00:00:00 2001
From: Lisa Dorofeeva <elizaveta.dorofeeva@desy.de>
Date: Fri, 4 Jul 2025 10:55:04 +0200
Subject: [PATCH] Have mask and whitefield also available as it was outside of
 streak finder

---
 dap/algos/__init__.py   |  3 ++-
 dap/algos/streakfind.py | 11 ++++++++---
 dap/worker.py           |  9 ++++++---
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/dap/algos/__init__.py b/dap/algos/__init__.py
index 7f0ff39..b9bf2a5 100644
--- a/dap/algos/__init__.py
+++ b/dap/algos/__init__.py
@@ -8,7 +8,8 @@ from .peakfind import calc_peakfinder_analysis
 from .radprof import calc_radial_integration
 from .roi import calc_roi
 from .spiana import calc_spi_analysis
-from .streakfind import calc_cbd_analysis
+from .streakfind import calc_streakfinder_analysis
+from .whitefield_correction import calc_apply_whitefield_correction
 from .thresh import calc_apply_threshold
 
 
diff --git a/dap/algos/streakfind.py b/dap/algos/streakfind.py
index 98ca4e2..ee47f1f 100644
--- a/dap/algos/streakfind.py
+++ b/dap/algos/streakfind.py
@@ -3,7 +3,7 @@ Streak Finder algorithm implemented by CFEL Chapman group
 
 Requires Convergent beam streak finder package installed:
 
-https://github.com/simply-nicky/streak_finder/tree/swiss_fel
+https://github.com/simply-nicky/streak_finder/
 (note g++ 11 required for building, numpy 2+ required)
 """
 import h5py
@@ -13,9 +13,10 @@ from streak_finder.label import Structure2D
 
 DEFAULT_NUM_THREADS = 16
 
-def calc_cbd_analysis(results, data, pf_pixel_mask):
+def calc_streakfinder_analysis(results, data, pf_pixel_mask):
     do_snr = results.get("do_snr", False)
-    if not do_snr:
+    do_streakfinder_analysis = results.get("do_streakfinder_analysis", False)
+    if not do_snr and not do_streakfinder_analysis:
         return data
 
     try:
@@ -30,6 +31,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}"
+
+    if not do_snr:
+        return data
+
     print(f"Returning data shape {cryst_data.snr[0].shape}")
     return cryst_data.snr[0]
 
diff --git a/dap/worker.py b/dap/worker.py
index 30b31eb..8ea7e66 100644
--- a/dap/worker.py
+++ b/dap/worker.py
@@ -3,7 +3,7 @@ import argparse
 import numpy as np
 
 from algos import (calc_apply_aggregation, calc_apply_threshold, calc_mask_pixels, calc_peakfinder_analysis, calc_radial_integration, calc_roi, calc_spi_analysis,
-                   calc_cbd_analysis, JFData)
+                   calc_apply_whitefield_correction, calc_streakfinder_analysis, JFData)
 from utils import Aggregator, BufferedJSON, randskip, read_bit
 from zmqsocks import ZMQSockets
 
@@ -117,8 +117,11 @@ def work(backend_address, accumulator_host, accumulator_port, visualisation_host
 
 # ???
 
-        # Correction and streak finder processing for convergent-beam diffraction
-        image = calc_cbd_analysis(results, image, pixel_mask_pf)
+        # White-field correction and streak finder processing for convergent-beam diffraction
+        print(f"Applying whitefield correction")
+        calc_apply_whitefield_correction(results, image)  # changes image in place
+        print(f"Searching streaks")
+        calc_streakfinder_analysis(results, image, pixel_mask_pf)
         print(f"Done\n{results=}")
 
         image, aggregation_is_ready = calc_apply_aggregation(results, image, pixel_mask_pf, aggregator)