diff --git a/daq/src/aaredaq/beamcenterfit.py b/daq/src/aaredaq/beamcenterfit.py new file mode 100644 index 00000000..e2487918 --- /dev/null +++ b/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 \ No newline at end of file