add 2D Gaussian fit

daq/src/aaredaq/beamcenterfit.py

@@ -0,0 +1,80 @@
+import numpy as np
+import cv2
+from scipy.optimize import curve_fit
+
+class Gaussian2Dfit:
+    center_x: float
+    center_y: float
+    sigma_x: float
+    sigma_y: float
+    rotation_angle: float
+    peak_intensity: float
+
+def gaussian_2d(xy, x0, y0, sigma_x, sigma_y, theta, A, offset):
+    x, y = xy
+    x_rot = (x - x0) * np.cos(theta) + (y - y0) * np.sin(theta)
+    y_rot = -(x - x0) * np.sin(theta) + (y - y0) * np.cos(theta)
+    gaussian = A * np.exp(-(x_rot**2 / (2 * sigma_x**2) + y_rot**2 / (2 * sigma_y**2))) + offset
+    return gaussian.ravel()
+
+def beamcenter_fit(image:np.ndarray) -> Gaussian2Dfit:
+    """
+    Takes a np.ndarray image and returns a Gaussian2Dfit object containing the fitted parameters
+    Returns None if the fit did not converge
+    Returns center_x, center_y, sigma_x, sigma_y, rotation_angle, peak_intensity if success
+    """
+    # Preprocess the image
+    blurred_image = cv2.GaussianBlur(image, (5, 5), 0)
+    normalized_image = cv2.normalize(blurred_image, None, 0, 255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
+
+    # Threshold to isolate the brightest spot (main blob)
+    _, binary_mask = cv2.threshold(normalized_image, 180, 255, cv2.THRESH_BINARY)
+    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    largest_contour = max(contours, key=cv2.contourArea)
+
+    # Extract bounding box of the ROI
+    x, y, w, h = cv2.boundingRect(largest_contour)
+
+    # Enlarge the bounding box to enclose the whole spot
+    enlargement_factor = 2  # Increase bounding box by 50%
+    new_x = max(0, int(x - w * enlargement_factor // 2))
+    new_y = max(0, int(y - h * enlargement_factor // 2))
+    new_w = min(image.shape[1] - new_x, int(w * (1 + enlargement_factor)))
+    new_h = min(image.shape[0] - new_y, int(h * (1 + enlargement_factor)))
+
+    # Update the ROI
+    roi = normalized_image[new_y:new_y + new_h, new_x:new_x + new_w]
+
+    # Fit the 2D Gaussian
+    X, Y = np.meshgrid(np.arange(roi.shape[1]), np.arange(roi.shape[0]))
+    xy = (X, Y)
+
+    # Initial guess for Gaussian parameters
+    initial_guess = (
+        new_w / 2,              # x0
+        new_h / 2,              # y0
+        new_w / 4,              # sigma_x
+        new_h / 4,              # sigma_y
+        0,                      # theta
+        np.max(roi),            # A (peak intensity)
+        np.min(roi)             # offset
+    )
+
+    # Flatten ROI for curve fitting
+    roi_flat = roi.ravel()
+    try:
+        popt, _ = curve_fit(gaussian_2d, xy, roi_flat, p0=initial_guess)
+    except RuntimeError:
+        print("Error: Gaussian fitting did not converge.")
+        exit()
+
+    # Extract fitted parameters
+    x0, y0, sigma_x, sigma_y, ini_theta, intensity, _ = popt
+    result = Gaussian2Dfit()
+    result.center_x = x0+new_x
+    result.center_y = y0+new_y
+    result.sigma_x = sigma_x
+    result.sigma_y = sigma_y
+    result.rotation_angle = np.degrees(ini_theta)%360
+    result.peak_intensity = intensity
+    return result