From 9fece3e4a16a6f254ae4d27b80ece6f17092997f Mon Sep 17 00:00:00 2001
From: Filip Leonarski <filip.leonarski@psi.ch>
Date: Fri, 5 Sep 2025 17:23:00 +0200
Subject: [PATCH] daq: loop centering. Tried a coupole of methods includign a
 flatfield correction. Currently implemenmted absdiff with an Otsu threshold
 estiamtion. Needs tidying!

---
 daq/src/aaredaq/daq.py | 60 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 52 insertions(+), 8 deletions(-)

diff --git a/daq/src/aaredaq/daq.py b/daq/src/aaredaq/daq.py
index c82766e2..7551f4fd 100644
--- a/daq/src/aaredaq/daq.py
+++ b/daq/src/aaredaq/daq.py
@@ -23,7 +23,8 @@ from aaredaqlib.models import (
     SampleShortInfo,
     PuckLoadedInfo,
     SampleShortInfoList,
-    DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, AutofocusSettings, BoundingBoxModel, )
+    DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, AutofocusSettings, BoundingBoxModel,
+    LoopCenteringZoomModelElem)
 from aaredaqlib.raster_grid import RasterGridRequest, CompletedRasterGrid, CompletedRasterGridElem
 from aaredaqlib.rotation_scan import RotationScanRequest, CompletedRotationScan
 from aaredaqlib.sample_geometry import SampleGeometryModel
@@ -690,11 +691,12 @@ class AareDAQ:
             self.__cfg.state_busy = False
             raise
 
-    def __loop_center(self, filename: str | None = "") -> SmargonCoordinate:
+    def __loop_center(self, s: LoopCenteringZoomModelElem, filename: str | None = "") -> SmargonCoordinate:
+        # TODO: tidy up!
         curr_image = self.camera_image
 
         gray_without_feature = cv2.cvtColor(
-            self.__cfg.get_alc_bkg(self.zoom), cv2.COLOR_RGB2GRAY
+            self.__cfg.get_alc_bkg(self.zoom, s.sam_cam_exp, s.sam_cam_gain), cv2.COLOR_RGB2GRAY
         )
         gray_with_feature = cv2.cvtColor(curr_image, cv2.COLOR_RGB2GRAY)
 
@@ -702,14 +704,56 @@ class AareDAQ:
         diff_image = cv2.absdiff(gray_with_feature, gray_without_feature)
 
         # Threshold the difference to isolate the feature
-        _, thresh = cv2.threshold(diff_image, 70, 255, cv2.THRESH_BINARY)
-        # thresh = cv2.adaptiveThreshold(diff_image, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
 
+        def flatfield_correction(raw_image, flat_image, dark_image=None):
+            """
+            Apply flat-field correction to remove uneven illumination
+
+            Args:
+                raw_image: Your actual image with crystal
+                flat_image: Image of uniform illumination (no sample)
+                dark_image: Dark frame (camera with no light), optional
+            """
+            # Convert to float for calculations
+            raw = raw_image.astype(np.float32)
+            flat = flat_image.astype(np.float32)
+
+            if dark_image is not None:
+                dark = dark_image.astype(np.float32)
+                # Subtract dark frame from both
+                raw_corrected = raw - dark
+                flat_corrected = flat - dark
+            else:
+                raw_corrected = raw
+                flat_corrected = flat
+
+            # Avoid division by zero
+            flat_corrected[flat_corrected == 0] = 1
+
+            # Apply correction
+            mean_flat = np.mean(flat_corrected)
+            corrected = (raw_corrected / flat_corrected) * mean_flat
+
+            # Convert back to original dtype
+            return np.clip(corrected, 0, 255).astype(raw_image.dtype)
+
+        #_, thresh = cv2.threshold(diff_image, 30, 255, cv2.THRESH_BINARY)
+        thresh_value, thresh = cv2.threshold(diff_image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+
+        #flatfield_correction_image = flatfield_correction(gray_with_feature, gray_without_feature)
+        #adapt_thresh = cv2.adaptiveThreshold(gray_with_feature, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,
+        #                                      11, 2)
+        #thresh_value, thresh = cv2.threshold(flatfield_correction_image, 0, 255,
+        #                                            cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+        print(f'Thresh value: {thresh_value}')
         if filename is not None:
-            cv2.imwrite(f"{filename}_diff.jpg", diff_image)
-            cv2.imwrite(f"{filename}_thresh.jpg", thresh)
+            #cv2.imwrite(f"{filename}_curr_image_colour.jpg", curr_image)
+            #cv2.imwrite(f"{filename}_curr_image.tiff", gray_with_feature)
+            cv2.imwrite(f"{filename}_diff.tiff", diff_image)
+            cv2.imwrite(f"{filename}_thresh.tiff", thresh)
+            #cv2.imwrite(f"{filename}_adaptive.tiff", adapt_thresh)
 
-        # Find contours of the detected feature
+        # # Find contours of the detected feature
         contours, _ = cv2.findContours(
             thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
         )