249 lines
9.5 KiB
Python
249 lines
9.5 KiB
Python
from logging import getLogger
|
|
from cam_server.pipeline.data_processing import functions
|
|
from cam_server.utils import create_thread_pvs, epics_lock
|
|
from collections import deque
|
|
import numpy as np
|
|
import scipy.signal
|
|
import numba
|
|
import time
|
|
from threading import Thread
|
|
import math
|
|
|
|
numba.set_num_threads(4)
|
|
|
|
_logger = getLogger(__name__)
|
|
|
|
# Shared variables
|
|
channel_names = None # list of PV names we poll each frame
|
|
roi = [0, 0]
|
|
initialized = False
|
|
sent_pid = -1
|
|
nrows = 1
|
|
buffer = deque(maxlen=5)
|
|
|
|
@numba.njit(parallel=False)
|
|
def get_spectrum(image, background):
|
|
"""Extracts a spectrum from the image by subtracting background."""
|
|
y, x = image.shape
|
|
profile = np.zeros(x, dtype=np.float64)
|
|
for i in range(y):
|
|
for j in range(x):
|
|
profile[j] += image[i, j] - background[i, j]
|
|
return profile
|
|
|
|
|
|
def update_PVs(buffer, output_pv_name, center_pv_name, fwhm_pv_name,
|
|
fit_rms_pv_name, fit_res_pv_name, com_pv_name, std_pv_name, res_pv_name):
|
|
"""Updates EPICS PVs using buffered processing."""
|
|
[output_pv, center_pv, fwhm_pv, fit_rms_pv, fit_res_pv,
|
|
com_pv, std_pv, res_pv] = create_thread_pvs(
|
|
[output_pv_name, center_pv_name, fwhm_pv_name,
|
|
fit_rms_pv_name, fit_res_pv_name, com_pv_name, std_pv_name, res_pv_name]
|
|
)
|
|
|
|
while True:
|
|
time.sleep(0.1)
|
|
try:
|
|
rec = buffer.popleft()
|
|
except IndexError:
|
|
continue
|
|
try:
|
|
if output_pv and output_pv.connected and (rec[0] is not None):
|
|
output_pv.put(rec[0])
|
|
if center_pv and center_pv.connected and (rec[1] is not None):
|
|
center_pv.put(rec[1])
|
|
if fwhm_pv and fwhm_pv.connected and (rec[2] is not None):
|
|
fwhm_pv.put(rec[2])
|
|
if com_pv and com_pv.connected and (rec[3] is not None):
|
|
com_pv.put(rec[3])
|
|
if std_pv and std_pv.connected and (rec[4] is not None):
|
|
std_pv.put(rec[4])
|
|
if fit_rms_pv and fit_rms_pv.connected and (rec[5] is not None):
|
|
fit_rms_pv.put(rec[5])
|
|
if fit_res_pv and fit_res_pv.connected and (rec[6] is not None):
|
|
fit_res_pv.put(rec[6])
|
|
if res_pv and res_pv.connected and (rec[7] is not None):
|
|
res_pv.put(rec[7])
|
|
except Exception as e:
|
|
_logger.exception(f"Error updating channels: {e}")
|
|
|
|
|
|
def initialize(params):
|
|
"""Initializes processing and starts the PV update thread."""
|
|
global channel_names
|
|
|
|
camera_name = params["camera_name"]
|
|
output_pv_name = params["e_int_name"] # will contain full spectrum array
|
|
axis_pv_name = params["e_axis_name"] # energy axis (array PV)
|
|
|
|
center_pv_name = camera_name + ":FIT-COM"
|
|
fwhm_pv_name = camera_name + ":FIT-FWHM"
|
|
fit_rms_pv_name = camera_name + ":FIT-RMS"
|
|
fit_res_pv_name = camera_name + ":FIT-RES"
|
|
ymin_pv_name = camera_name + ":SPC_ROI_YMIN"
|
|
ymax_pv_name = camera_name + ":SPC_ROI_YMAX"
|
|
|
|
com_pv_name = camera_name + ":SPECT-COM"
|
|
std_pv_name = camera_name + ":SPECT-RMS"
|
|
res_pv_name = camera_name + ":SPECT-RES"
|
|
|
|
thread = Thread(target=update_PVs, args=(buffer, output_pv_name, center_pv_name,
|
|
fwhm_pv_name, fit_rms_pv_name, fit_res_pv_name,
|
|
com_pv_name, std_pv_name, res_pv_name))
|
|
thread.daemon = True
|
|
thread.start()
|
|
|
|
# polled each frame
|
|
channel_names = [ymin_pv_name, ymax_pv_name, axis_pv_name]
|
|
|
|
|
|
def process_image(image, pulse_id, timestamp, x_axis, y_axis, parameters,
|
|
bsdata=None, background=None):
|
|
"""Processes the image, forms a spectrum (with Y-ROI), pushes full spectrum to e_int_name."""
|
|
try:
|
|
global roi, initialized, sent_pid, nrows, buffer
|
|
|
|
if not initialized:
|
|
initialize(parameters)
|
|
initialized = True
|
|
|
|
# fetch PV objects (ymin, ymax, axis)
|
|
pv_objs = create_thread_pvs(channel_names)
|
|
if len(pv_objs) != 3:
|
|
_logger.warning("create_thread_pvs did not return 3 PVs")
|
|
return {}
|
|
|
|
ymin_pv, ymax_pv, axis_pv = pv_objs
|
|
camera_name = parameters["camera_name"]
|
|
|
|
# read Y-ROI
|
|
if ymin_pv and ymin_pv.connected:
|
|
try: roi[0] = int(ymin_pv.value)
|
|
except Exception: pass
|
|
if ymax_pv and ymax_pv.connected:
|
|
try: roi[1] = int(ymax_pv.value)
|
|
except Exception: pass
|
|
|
|
# get energy axis: prefer PV, else fall back to x_axis argument
|
|
axis_arr = None
|
|
if axis_pv and axis_pv.connected and axis_pv.value is not None:
|
|
try:
|
|
axis_arr = np.asarray(axis_pv.value, dtype=np.float64)
|
|
except Exception:
|
|
axis_arr = None
|
|
if axis_arr is None and x_axis is not None:
|
|
try:
|
|
axis_arr = np.asarray(x_axis, dtype=np.float64)
|
|
except Exception:
|
|
axis_arr = None
|
|
|
|
if axis_arr is None:
|
|
_logger.warning("No valid energy axis (PV and x_axis both invalid)")
|
|
return {}
|
|
|
|
# ensure axis length matches width
|
|
nrows, ncols = image.shape
|
|
if axis_arr.shape[0] < ncols:
|
|
_logger.warning(f"Energy axis shorter than image width: {axis_arr.shape[0]} < {ncols}")
|
|
return {}
|
|
|
|
axis_arr = axis_arr[:ncols]
|
|
|
|
processing_image = image.astype(np.float32) - np.float32(parameters.get("pixel_bkg", 0.0))
|
|
|
|
# background: support both function arg and parameters["background_data"]
|
|
bkg_param = parameters.get("background_data", None)
|
|
if isinstance(bkg_param, np.ndarray) and bkg_param.shape == processing_image.shape:
|
|
background_img = bkg_param.astype(np.float32)
|
|
elif isinstance(background, np.ndarray) and background.shape == processing_image.shape:
|
|
background_img = background.astype(np.float32)
|
|
else:
|
|
background_img = None
|
|
|
|
# apply Y-ROI if valid
|
|
y0, y1 = roi
|
|
if isinstance(y0, (int, np.integer)) and isinstance(y1, (int, np.integer)) and 0 <= y0 < y1 <= nrows:
|
|
img_roi = processing_image[y0:y1, :]
|
|
if background_img is not None:
|
|
bkg_roi = background_img[y0:y1, :]
|
|
else:
|
|
bkg_roi = None
|
|
else:
|
|
img_roi = processing_image
|
|
bkg_roi = background_img
|
|
|
|
# spectrum extraction (background-aware)
|
|
if bkg_roi is not None:
|
|
try:
|
|
spectrum = get_spectrum(img_roi, bkg_roi)
|
|
except Exception as e:
|
|
_logger.warning(f"get_spectrum failed, falling back to sum: {e}")
|
|
spectrum = np.sum(img_roi - (bkg_roi if bkg_roi is not None else 0), axis=0).astype(np.float64)
|
|
else:
|
|
spectrum = np.sum(img_roi, axis=0).astype(np.float64)
|
|
|
|
# smoothing
|
|
max_odd = max(3, (len(spectrum) // 2) * 2 - 1)
|
|
win = min(51, max_odd)
|
|
if win >= 5:
|
|
try:
|
|
smoothed = scipy.signal.savgol_filter(spectrum, window_length=win, polyorder=3, mode="interp")
|
|
except Exception:
|
|
smoothed = spectrum
|
|
else:
|
|
smoothed = spectrum
|
|
|
|
# fit / moments
|
|
minimum, maximum = float(np.min(smoothed)), float(np.max(smoothed))
|
|
amplitude = maximum - minimum
|
|
skip = amplitude <= nrows * 1.5
|
|
try:
|
|
offset, amplitude, center, sigma = functions.gauss_fit_psss(
|
|
smoothed[::2], axis_arr[::2], offset=minimum, amplitude=amplitude, skip=skip, maxfev=10
|
|
)
|
|
center = float(center); sigma = float(sigma)
|
|
except Exception:
|
|
total = np.sum(smoothed)
|
|
center = float(np.sum(axis_arr * smoothed) / total) if total > 0 else np.nan
|
|
sigma = float(np.sqrt(np.sum((axis_arr - center) ** 2 * smoothed) / total)) if total > 0 else np.nan
|
|
|
|
spectrum_com = center
|
|
spectrum_std = sigma
|
|
|
|
processed_data = {
|
|
camera_name + ":SPECTRUM_Y": spectrum,
|
|
camera_name + ":SPECTRUM_Y_SUM": float(np.sum(spectrum)),
|
|
camera_name + ":SPECTRUM_X": axis_arr,
|
|
camera_name + ":FIT-COM": center,
|
|
camera_name + ":FIT-FWHM": float(2.355 * sigma),
|
|
camera_name + ":FIT-RMS": sigma,
|
|
camera_name + ":FIT-RES": (float(2.355 * sigma / center) * 1000) if center else None,
|
|
camera_name + ":SPECT-COM": spectrum_com,
|
|
camera_name + ":SPECT-RMS": spectrum_std,
|
|
camera_name + ":SPECT-RES": (float(2.355 * spectrum_std / spectrum_com) * 1000) if spectrum_com else None,
|
|
}
|
|
|
|
# push PVs (first tuple element maps to params["e_int_name"])
|
|
global sent_pid
|
|
if epics_lock.acquire(False):
|
|
try:
|
|
if pulse_id > sent_pid:
|
|
sent_pid = pulse_id
|
|
buffer.append((
|
|
spectrum, # now sending full spectrum array
|
|
center,
|
|
2.355 * sigma,
|
|
spectrum_com,
|
|
spectrum_std,
|
|
sigma,
|
|
(2.355 * sigma / center) * 1000 if center else None,
|
|
(2.355 * spectrum_std / spectrum_com) * 1000 if spectrum_com else None
|
|
))
|
|
finally:
|
|
epics_lock.release()
|
|
|
|
return processed_data
|
|
except Exception as ex:
|
|
_logger.exception(f"process_image error: {ex}")
|
|
return {}
|