From aa076cbb27e6c8089bc0a05f6ecf8e5a195ea7a0 Mon Sep 17 00:00:00 2001
From: Yu Xia de Jong <dejong@delmic.com>
Date: Mon, 4 May 2026 13:20:53 +0200
Subject: [PATCH] [feat] added save_images

---
 src/odemis/acq/drift/__init__.py | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/odemis/acq/drift/__init__.py b/src/odemis/acq/drift/__init__.py
index e2f66929e7..e245d4dcce 100644
--- a/src/odemis/acq/drift/__init__.py
+++ b/src/odemis/acq/drift/__init__.py
@@ -20,20 +20,24 @@
 Odemis. If not, see http://www.gnu.org/licenses/.
 """
 
+import datetime
 import itertools
 import logging
 import math
 import threading
-
+import os
 import numpy
 import cv2
 
 from odemis import model
 from odemis.acq.align.shift import MeasureShift
+from odemis.gui.util import get_picture_folder
 
 MIN_RESOLUTION = (20, 20)  # sometimes 8x8 works, but it's not reliable enough
 MAX_PIXELS = 128 ** 2  # px
 
+DRIFT_IMAGES_DIR = os.path.join(get_picture_folder(), "Drift Correction images")
+
 
 class AnchoredEstimator(object):
     """
@@ -46,7 +50,7 @@ class AnchoredEstimator(object):
     to measure the drift.
     """
 
-    def __init__(self, scanner, detector, region, dwell_time, max_pixels=MAX_PIXELS, follow_drift=True):
+    def __init__(self, scanner, detector, region, dwell_time, max_pixels=MAX_PIXELS, follow_drift=True, save_images=False):
         """
         scanner (Emitter)
         detector (Detector)
@@ -63,6 +67,12 @@ def __init__(self, scanner, detector, region, dwell_time, max_pixels=MAX_PIXELS,
         self._semd = detector
         self._dwell_time = dwell_time
         self._follow_drift = follow_drift
+        self._save_images = save_images
+        self._session_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        self._image_counter = 0
+
+        if self._save_images:
+            os.makedirs(DRIFT_IMAGES_DIR, exist_ok=True)
 
         # Latest drift vector from the previous acquisition
         self.drift = (0, 0)  # in sem px
@@ -143,6 +153,15 @@ def acquire(self):
             if data.shape[::-1] != self._res:
                 logging.warning("Shape of data is %s instead of %s", data.shape[::-1], self._res)
 
+            # Save all the drift region scans for offline autocorrelation purposes
+            if self._save_images:
+                filename = os.path.join(DRIFT_IMAGES_DIR, f"drift_{self._session_id}_{self._image_counter:05d}.tif")
+                self._image_counter += 1
+
+                success = cv2.imwrite(filename, data)
+                if not success:
+                    logging.warning("Failed to save drift image %s", filename)
+
             # TODO: allow to record just every Nth image, and separately record the
             # drift after every measurement
             # In the mean time, we only save the 1st, 2nd and last two images