From 26c18de825d24e85873bbb3a218027006dd0ba68 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 29 Oct 2025 13:34:09 +0100 Subject: [PATCH] DAQ/GUI: updated face_detection.py, refactored mlbox.py to use pydantic models and to allow more functionality to be added to the analysis of the boxes --- common/src/aaredaqlib/face_detection.py | 186 +++++++++++++ common/src/aaredaqlib/models.py | 95 +++++++ daq/src/aaredaq/daq.py | 198 ++++--------- daq/src/aaredaq/mlbox.py | 262 ++++++++++++++---- daq/src/aaredaq/server.py | 5 +- .../aaregui/panels/face_detection_panel.py | 64 +++-- gui/src/aaregui/threads/daq_worker.py | 6 +- 7 files changed, 595 insertions(+), 221 deletions(-) create mode 100644 common/src/aaredaqlib/face_detection.py diff --git a/common/src/aaredaqlib/face_detection.py b/common/src/aaredaqlib/face_detection.py new file mode 100644 index 00000000..16d52012 --- /dev/null +++ b/common/src/aaredaqlib/face_detection.py @@ -0,0 +1,186 @@ +import json +import time +from typing import Tuple, Dict, List, Iterable, Optional +import numpy as np +from scipy.optimize import curve_fit +import statistics +import math +from aaredaqlib.logger_config import setup_logger +from aaredaqlib.models import MLBoxModel, MLOutputModel + +logger = setup_logger(__name__, '/tmp/mxlogs') + +def box_height_from_tuple(box: tuple[float, float, float, float]) -> float: + x1, y1, x2, y2 = box + return abs(y2 - y1) + +def box_area_from_tuple(box: tuple[float, float, float, float]) -> float: + x1, y1, x2, y2 = box + return abs(y2 - y1) * abs(x2 -x1) + +def prepare_samples(boxes_by_angle: dict[int, tuple[float,float,float,float]], + area = False) -> List[Tuple[float, float]]: + # angles in degrees -> (theta_rad, height) + samples = [] + for deg, box in boxes_by_angle.items(): + v = box_area_from_tuple(box) if area else box_height_from_tuple(box) + samples.append((float(deg), v)) + return samples + +def cos_model(theta_deg: float | np.ndarray, A: float, B: float, phi_rad: float, C: float): + return A + B * np.cos(C * np.deg2rad(theta_deg) - phi_rad) + +def mad_filter(samples: List[Tuple[float, float]], k: float = 3.5) -> List[Tuple[float, float]]: + if not samples: + return samples + ys = [y for _, y in samples] + med = statistics.median(ys) + mad = statistics.median([abs(y - med) for y in ys]) or 1.0 + return [(d, y) for d, y in samples if abs(y - med) <= k * mad] + +def samples_to_json(samples): + output_data = { + 'timestamp': time.ctime(), + 'scan_results': samples, + 'total_results': len(samples) + } + with open("cos_test.json", 'w') as f: + json.dump(output_data, f, indent=2) + return + +def fit_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> tuple[float, float, float]: + resid = y_true - y_pred + rmse = float(np.sqrt(np.mean(resid**2))) + mae = float(np.mean(np.abs(resid))) + # R² with protection against zero variance + ss_tot = float(np.sum((y_true - np.mean(y_true))**2)) + r2 = float(1.0 - np.sum(resid**2) / ss_tot) if ss_tot > 0 else float("nan") + return rmse, mae, r2 + + +def fit_cosine(samples: List[Tuple[float, float]]) -> dict: + samples = mad_filter(samples, k=3.5) + if len(samples) < 3: + A = sum(y for _, y in samples) / max(1, len(samples)) + return {"A": A, "B": 1.0, "phi_rad": 0.0, "C": 1.0, + "rmse": None, "mae": None, "r2": None} + + degs = np.array([d for d, _ in samples], dtype=float) + ys = np.array([y for _, y in samples], dtype=float) + + # Initial guess via linearized cosine/sine fit + cosv = np.cos(np.deg2rad(degs)) + sinv = np.sin(np.deg2rad(degs)) + X = np.column_stack([np.ones_like(degs), cosv, sinv]) # [A, C, S] + try: + beta, _, _, _ = np.linalg.lstsq(X, ys, rcond=None) + A0, C0, S0 = beta.tolist() + except Exception: + A0, C0, S0 = float(np.mean(ys)), 0.0, 0.0 + + B0 = float(math.hypot(C0, S0)) + phi0 = float(math.atan2(S0, C0)) + C0 = 1.0 + # Bounds to keep B reasonable and phi in [-pi, pi] + y_std = float(np.std(ys)) or 1.0 + B_max = 10.0 * y_std + bounds = ([-np.inf, 0.0, -math.pi, -np.inf], [np.inf, B_max, math.pi, np.inf]) + + try: + popt, _ = curve_fit(cos_model, degs, ys, p0=[A0, B0, phi0, C0], bounds=bounds, maxfev=10000) + logger.info(f"fit cosine: {popt}") + A, B, phi, C = map(float, popt) + B = max(0.0, B) + yhat = cos_model(degs, A, B, phi, C) + rmse, mae, r2 = fit_metrics(ys, yhat) + return {"A": A, "B": B, "phi_rad": phi, "C": C, + "rmse": rmse, "mae": mae, "r2": r2} + + except Exception as e: + # Fallback to initial + logger.info(f"error in curve fit {e}") + yhat0 = cos_model(degs, A0, max(0.0, B0), phi0, C0) + rmse0, mae0, r2_0 = fit_metrics(ys, yhat0) + return {"A": float(A0), "B": max(0.0, float(B0)), "phi_rad": float(phi0), "C": float(C0), + "rmse": rmse0, "mae": mae0, "r2": r2_0} + +def get_samples_out(boxes): + samples_out = [] + for deg, box in boxes.items(): + h = box_height_from_tuple(box) + a = box_area_from_tuple(box) + samples_out.append({"angle_deg": float(deg), "height": float(h), "area": float(a)}) + samples_out.sort(key=lambda x: x["angle_deg"]) + return samples_out + +def choose_best_fit(fits_by_name: Dict[str, Dict]) -> Tuple[Optional[float], Optional[Dict], Optional[str]]: + + + def key(entry: Dict): + params = entry.get("params") or {} + rmse = params.get("rmse") + mae = params.get("mae") + r2 = params.get("r2") + logger.info(f"params: {params}") + # Treat None/NaN as non-comparable + if rmse is None or mae is None or r2 is None: + return None + try: + logger.info(f"rmse: {rmse}, mae: {mae}, r2: {r2}") + return float(rmse), float(mae), -float(r2) + except Exception: + return None + + best_name = None + best_fit = None + best_key = None + + for name, entry in fits_by_name.items(): + k = key(entry) + logger.info(f"name: {name}, entry: {entry}, key : {k}") + if k is None: + continue + if best_key is None or k < best_key: + best_key = k + best_fit = entry + best_name = name + + if best_fit is None: + return None, None, None + + best_angle = best_fit.get("angle") + best_angle = float(best_angle) if isinstance(best_angle, (int, float)) else None + + + return (best_angle, + best_fit, + best_name) + +def get_flat_face(boxes: dict[int, tuple[float,float,float,float]], start_angle:int, end_angle:int, area:bool = False) -> tuple[int, dict]: + samples = prepare_samples(boxes, area=area) + parameters = fit_cosine(samples) + + def safe_best_angle(params: dict) -> int: + if not params: + return 0 + search_grid = range(start_angle, end_angle, 1) + A = float(params.get("A", 0.0)) + B = float(max(0.0, params.get("B", 0.0))) + C = float(params.get("C", 0.0)) + phi = float(params.get("phi_rad", 0.0)) + if B <= 1e-9 or not math.isfinite(A) or not math.isfinite(phi): + return 0 + return max(search_grid, key=lambda d: cos_model(d, A, B, phi, C)) + + best_fit_angle= safe_best_angle(parameters) + return best_fit_angle, parameters + +def chose_best_angle(boxes: dict[int, tuple[float,float,float,float]], fit_results) -> int: + measured_angles = list(boxes.keys()) + choose_best_fit(fit_results) + candidates = [a for a in (fit_results["Area"]["angle"], fit_results["Height"]["angle"]) if isinstance(a, (int, float))] + if measured_angles and candidates: + chosen = min(candidates, key=lambda a: min(abs(((a - m + 180) % 360) - 180) for m in measured_angles)) + else: + chosen = candidates[0] if candidates else 0 + return chosen \ No newline at end of file diff --git a/common/src/aaredaqlib/models.py b/common/src/aaredaqlib/models.py index bf9b55fe..bab1fc31 100644 --- a/common/src/aaredaqlib/models.py +++ b/common/src/aaredaqlib/models.py @@ -430,12 +430,107 @@ class BeamMarkCoeffModel(BaseModel): y=self.coeff_y[0] * zoom**2 + self.coeff_y[1] * zoom + self.coeff_y[2], ) +class MLBoxType(Enum): + Loop_all = 0 + Pin = 1 + Crystal = 2 + Loop_face = 3 + class BoundingBoxModel(BaseModel): top_x: float top_y: float bottom_x: float bottom_y: float +class MLBoxModel(BaseModel): + cls: MLBoxType + box: BoundingBoxModel + conf: float + + @staticmethod + def from_tuple(cls:MLBoxType, box_tuple: tuple[float, float, float, float], conf: float) -> "MLBoxModel": + x1, y1, x2, y2 = box_tuple + return MLBoxModel( + cls=cls, + box=BoundingBoxModel(top_x=float(x1), top_y=float(y1), bottom_x=float(x2), bottom_y=float(y2)), + conf=float(conf) + ) + +class MLOutputModel(BaseModel): + boxes: dict[str, MLBoxModel] = {} + + @staticmethod + def get_class_str(cls: MLBoxType) -> str: + if cls == MLBoxType.Loop_all: + return "Loop_all" + if cls == MLBoxType.Pin: + return "Pin" + if cls == MLBoxType.Crystal: + return "Crystal" + if cls == MLBoxType.Loop_face: + return "Loop_face" + return "Unknown" + + def _next_unique_key(self, base: str) -> str: + if base not in self.boxes: + return base + i = 2 + while f"{base}_{i}" in self.boxes: + i += 1 + return f"{base}_{i}" + + def add_box(self, cls: MLBoxType, box_tuple: tuple[float, float, float, float], + conf: float | None = None) -> str: + key_base = self.get_class_str(cls) + key = self._next_unique_key(key_base) + self.boxes[key] = MLBoxModel.from_tuple(cls, box_tuple, conf) + return key + + def get_box_model(self, key: str) -> MLBoxModel | None: + return self.boxes.get(key) + + def get_box_tuple(self, key: str) -> tuple[float, float, float, float] | None: + m = self.get_box_model(key) + if not m or not m.box: + return None + return (m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y) + + def get_box_tuple_with_conf(self, key: str) -> tuple[float, float, float, float, float] | None: + m = self.get_box_model(key) + if not m or not m.box: + return None + conf = float(m.conf) if m.conf is not None else 0.0 + return (m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y, conf) + + def get_keys_for_class(self, cls: MLBoxType) -> list[str]: + base = self.get_class_str(cls) + return [k for k, v in self.boxes.items() if v.cls == cls and (k == base or k.startswith(base + "_"))] + + def get_models_for_class(self, cls: MLBoxType) -> list[MLBoxModel]: + keys = self.get_keys_for_class(cls) + return [self.boxes[k] for k in keys] + + def get_tuples_for_class(self, cls: MLBoxType) -> list[tuple[float, float, float, float]]: + out: list[tuple[float, float, float, float]] = [] + for m in self.get_models_for_class(cls): + if m.box: + out.append((m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y)) + return out + + def get_tuples_with_conf_for_class(self, cls: MLBoxType) -> list[tuple[float, float, float, float, float]]: + out: list[tuple[float, float, float, float, float]] = [] + for m in self.get_models_for_class(cls): + if m.box: + out.append((m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y, float(m.conf))) + return out + + def get_best_for_class(self, cls: MLBoxType) -> MLBoxModel | None: + models = self.get_models_for_class(cls) + if not models: + return None + return max(models, key=lambda m: (m.conf if m.conf is not None else 0.0)) + + class BeamlineStateEnum(Enum): Maintenance = 1 SampleExchange = 2 diff --git a/daq/src/aaredaq/daq.py b/daq/src/aaredaq/daq.py index d06d8c82..1098bf84 100644 --- a/daq/src/aaredaq/daq.py +++ b/daq/src/aaredaq/daq.py @@ -25,6 +25,7 @@ from aaredaq.mlbox import MlBox from aaredaqlib.beamline import MXBeamline from aaredaqlib.coordinate import Coordinate, SmargonCoordinate from aaredaqlib.diffraction_geometry import DiffractionGeometry +import aaredaqlib.face_detection as fd from aaredaqlib.logger_config import setup_logger from aaredaqlib.find_xtal import raster_centre_of_mass, create_quality_filtered_array, identify_crystal_raster, \ get_result_list_from_com, get_best_b_factor, get_best_res, get_xtal_size @@ -906,110 +907,11 @@ class AareDAQ: ) return SmargonCoordinate(sh_mm=coord) - def box_height_from_tuple(self, box: Tuple[float, float, float, float]) -> float: - x1, y1, x2, y2 = box - return abs(y2 - y1) - - def box_area_from_tuple(self, box: Tuple[float, float, float, float]) -> float: - x1, y1, x2, y2 = box - return abs(y2 - y1) * abs(x2-x1) - - def prepare_height_samples(self, boxes_by_angle: Dict[float, Tuple[float, float, float, float]], area = False) -> List[ - Tuple[float, float]]: - # angles in degrees -> (theta_rad, height) - samples = [] - for deg, box in boxes_by_angle.items(): - v = self.box_area_from_tuple(box) if area else self.box_height_from_tuple(box) - samples.append((float(deg), v)) - return samples - - @staticmethod - def _cos_model(theta_deg: float | np.ndarray, A: float, B: float, phi_rad: float): - return A + B * np.cos(np.deg2rad(theta_deg) - phi_rad) - - @staticmethod - def _mad_filter(samples: List[Tuple[float, float]], k: float = 3.5) -> List[Tuple[float, float]]: - if not samples: - return samples - ys = [y for _, y in samples] - med = statistics.median(ys) - mad = statistics.median([abs(y - med) for y in ys]) or 1.0 - return [(d, y) for d, y in samples if abs(y - med) <= k * mad] - - def fit_area_vs_angle(self, areas_by_angle_deg: List[Tuple[float, float]]) -> dict: - - pts = [(float(d), float(a)) for d, a in areas_by_angle_deg] - pts = self._mad_filter(pts, k=3.5) - degs = np.array([d for d, _ in pts], dtype=float) - ys = np.array([a for _, a in pts], dtype=float) - - # Initial guess via linearized fit - cosv = np.cos(np.deg2rad(degs)) - sinv = np.sin(np.deg2rad(degs)) - X = np.column_stack([np.ones_like(degs), cosv, sinv]) # [A, C, S] - try: - beta, _, _, _ = np.linalg.lstsq(X, ys, rcond=None) - A0, C0, S0 = beta.tolist() - except Exception: - A0, C0, S0 = float(np.mean(ys)), 0.0, 0.0 - B0 = float(math.hypot(C0, S0)) - phi0 = float(math.atan2(S0, C0)) - - # Bounds - y_std = float(np.std(ys)) or 1.0 - B_max = 10.0 * y_std - bounds = ([-np.inf, 0.0, -math.pi], [np.inf, B_max, math.pi]) - - try: - popt, _ = curve_fit(self._cos_model, degs, ys, p0=[A0, B0, phi0], bounds=bounds, maxfev=10000) - A, B, phi = map(float, popt) - except Exception as e: - logger.info(f"error in curve fit {e}") - A, B, phi = A0, max(0.0, B0), phi0 - - #best_angle_deg = float((math.degrees(phi)) % 360) - return {"A": A, "B": max(0.0, B), "phi_rad": phi} - - def fit_cosine_height(self, samples: List[Tuple[float, float]]) -> dict: - samples = self._mad_filter(samples, k=3.5) - if len(samples) < 3: - A = sum(y for _, y in samples) / max(1, len(samples)) - return (A, 0.0, 0.0) - - degs = np.array([d for d, _ in samples], dtype=float) - ys = np.array([y for _, y in samples], dtype=float) - - # Initial guess via linearized cosine/sine fit - cosv = np.cos(np.deg2rad(degs)) - sinv = np.sin(np.deg2rad(degs)) - X = np.column_stack([np.ones_like(degs), cosv, sinv]) # [A, C, S] - try: - beta, _, _, _ = np.linalg.lstsq(X, ys, rcond=None) - A0, C0, S0 = beta.tolist() - except Exception: - A0, C0, S0 = float(np.mean(ys)), 0.0, 0.0 - B0 = float(math.hypot(C0, S0)) - phi0 = float(math.atan2(S0, C0)) - - # Bounds to keep B reasonable and phi in [-pi, pi] - y_std = float(np.std(ys)) or 1.0 - B_max = 10.0 * y_std - bounds = ([-np.inf, 0.0, -math.pi], [np.inf, B_max, math.pi]) - - try: - popt, _ = curve_fit(self._cos_model, degs, ys, p0=[A0, B0, phi0], bounds=bounds, maxfev=10000) - A, B, phi = map(float, popt) - B = max(0.0, B) - except Exception as e: - # Fallback to initial - logger.info(f"error in curve fit {e}") - return {"A": A, "B": max(0.0, B), "phi_rad": phi} - - def face_detection(self) -> dict: + def face_detection(self, steps: int = 14, step_size: int = 15) -> dict: self.__cfg.try_set_busy(timeout=360) try: logger.info("running face detection sequence") - result = self.__face_detection_sequence() + result = self.__face_detection_sequence(steps=steps, step_size=step_size) except Exception as e: logger.error(f"error in face detection sequence {e}") result = { @@ -1020,27 +922,35 @@ class AareDAQ: self.__cfg.state_busy = False return result - def __face_detection_sequence(self) -> dict: + def __face_detection_sequence(self, steps: int = 14, step_size: int = 15) -> dict: self.__set_state(BeamlineStateEnum.SampleAlignment) self.__devs.lamp_light = 2.5 self.__cfg.zoom_mode = ZoomModeEnum.LoopCenter + zoom_settings = self.__cfg.zoom_settings.z zoom_value = self.__devs.zoom logger.info('face detection sequence') - exposure = zoom_settings[zoom_value].exposure - gain = zoom_settings[zoom_value].gain - self.__devs.samcam_settings = SampleCameraSettings(exposure=exposure, gain=gain) + + self.__devs.samcam_settings = SampleCameraSettings( + exposure=zoom_settings[zoom_value].exposure, + gain=zoom_settings[zoom_value].gain) self.__devs.zoom_sync(zoom_value) - boxes = {} - angles = (90, 75, 60, 45, 30, 15, 0, -15, -30, -45, -60, -75, -90) - for angle in angles: + + boxes: dict[int, tuple[float, float, float, float]] = {} + curr_angle = int(self.__devs.aerotech.omega) + total_range = steps * step_size + 1 + start_angle = curr_angle if curr_angle + total_range < 720 else 0 + end_angle = curr_angle + total_range + + for angle in range(start_angle, end_angle, step_size): logger.debug(f'moving to angle: {angle}') rotate_time = time.perf_counter() self.__devs.aerotech.move(angle, wait=True) logger.info(f"time to rotate 15 degrees: {time.perf_counter() - rotate_time}") + curr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) box_time = time.perf_counter() - box = self.__mlbox.predict(curr_image, "", pref_class = (3,0)) + m = self.__mlbox.predict(curr_image, filename=None, preferred_class = (3,0)) logger.info(f"time to predict: {time.perf_counter() - box_time}") if box: logger.info(f'found box: {box}') @@ -1057,51 +967,51 @@ class AareDAQ: else: logger.info(f'no box found for angle {angle}') - if not boxes: - logger.info(f'no boxes found: {boxes}') - logger.info('no boxes found') + if not m or not m.box: + logger.info(f"no box found for angle {angle}") + continue - return { - "samples": None, - "height_fit": None, - "area_fit": None, - } + cls_id = int(m.cls.value) + x1, y1, x2, y2 = m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y - samples_area = self.prepare_height_samples(boxes, area=True) - area_params = self.fit_area_vs_angle(samples_area) - samples_height = self.prepare_height_samples(boxes, area=False) - height_params= self.fit_cosine_height(samples_height) + if cls_id == 3: + boxes[angle] = (x1, y1, x2, y2) + logger.info(f"accepted box at angle {angle}, cls={cls_id}, box={(x1, y1, x2, y2)}") + else: + logger.debug(f"ignoring class {cls_id} (pin/crystal) at angle {angle}") - search_grid = np.linspace(-90, 90, 400) - def model(theta_deg: float, A:float, B: float, phi: float) -> float: - return A + B * math.cos(math.radians(theta_deg) - phi) + if not boxes: + logger.info("no boxes found") + return {"samples": None, "height_fit": None, "area_fit": None} - search_grid = range(-90, 90, 1) - best_fit_angle_area = max(search_grid, key=lambda d: model(d, area_params['A'], area_params['B'], area_params['phi_rad'])) - best_fit_angle_height = max(search_grid, key=lambda d: model(d, height_params['A'], height_params['B'], height_params['phi_rad'])) - - # Move to fitted best angle + best_fit_angle_area, area_params = fd.get_flat_face(boxes, start_angle, end_angle, True) + best_fit_angle_height, height_params = fd.get_flat_face(boxes, start_angle, end_angle, False) + fit_results = {"Area":{"angle":best_fit_angle_area, "params":area_params}, "Height": {"angle":best_fit_angle_height, "params":height_params}} logger.info(f"best angle by area: {best_fit_angle_area}") logger.info(f"best angle by height: {best_fit_angle_height}") - candidates = [ - a for a in (best_fit_angle_area, best_fit_angle_height) - if a is not None and a in search_grid - ] - angle = candidates[0] if candidates else 0 - self.__devs.aerotech.move(angle, wait=True) + flat_face_angle, best_params, best_name = fd.choose_best_fit(fit_results) + logger.info(f"best params: {best_params}") + logger.info(f"chosen fit: {best_name}") + logger.info(f"best angle: {flat_face_angle}") + self.__devs.aerotech.move(flat_face_angle, wait=True) - # Build return JSON - samples_out = [] - for deg, box in boxes.items(): - h = self.box_height_from_tuple(box) - a = self.box_area_from_tuple(box) - samples_out.append({"angle_deg": float(deg), "height": float(h), "area": float(a)}) - samples_out.sort(key=lambda x: x["angle_deg"]) + samples_out = fd.get_samples_out(boxes) + logger.info(f"face detection sequence done, samples: {samples_out}") return { "samples": samples_out, - "height_fit": {"A": height_params["A"], "B": height_params["B"], "phi_rad": height_params["phi_rad"], "best_angle_deg": best_fit_angle_height}, - "area_fit": {"A": area_params["A"], "B": area_params["B"], "phi_rad": area_params["phi_rad"], "best_angle_deg": best_fit_angle_area}, + "height_fit": {"A": height_params["A"], + "B": height_params["B"], + "phi_rad": height_params["phi_rad"], + "C": height_params["C"], + "best_angle_deg": best_fit_angle_height + }, + "area_fit": {"A": area_params["A"], + "B": area_params["B"], + "phi_rad": area_params["phi_rad"], + "C": height_params["C"], + "best_angle_deg": best_fit_angle_area + }, } diff --git a/daq/src/aaredaq/mlbox.py b/daq/src/aaredaq/mlbox.py index 20478482..6e0788e7 100644 --- a/daq/src/aaredaq/mlbox.py +++ b/daq/src/aaredaq/mlbox.py @@ -1,85 +1,243 @@ -from pathlib import Path -from typing import Tuple, Optional, List + +from typing import Tuple, Optional, Iterable import cv2 import requests -from ultralytics import YOLO +from aaredaqlib.models import MLBoxModel, MLOutputModel, MLBoxType, BoundingBoxModel from aaredaqlib.logger_config import setup_logger logger=setup_logger(__name__, '/tmp/mxlogs') class MlBox: - # # def __init__(self, model_path: str = "best_v8_18092025.pt"): - # model_path = str(Path(__file__).parent / model_path) - # self.__model = YOLO(model_path) + def __init__(self, url="http://mx-aare-test.psi.ch:8002/predict/?model=best_v8_20102025.pt"): self.__url = url - self.class_info = [ - ["loop_all", (255, 0, 0)], # class 0: Blue for loop_all - ["pin", (0, 255, 0)], # class 1: Green for pin - ["crystal", (0, 0, 255)], # class 2: Red for crystal - ["loop_face", (255, 255, 0)] # class 3: Yellow for loop_face - ] + # self.class_info = [ + # ["loop_all", (255, 0, 0)], # class 0: Blue for loop_all + # ["pin", (0, 255, 0)], # class 1: Green for pin + # ["crystal", (0, 0, 255)], # class 2: Red for crystal + # ["loop_face", (255, 255, 0)] # class 3: Yellow for loop_face + # ] def get_response(self, image): ok, buf = cv2.imencode(".jpg", image) if not ok: raise RuntimeError("Failed to encode image") files={"file": ("image.jpg", buf.tobytes(), "image/jpeg")} - response = requests.post(self.__url, files=files) + response = requests.post(self.__url, files=files, timeout=10) + response.raise_for_status() logger.debug(response.text) return response.json() + def check_box_relation(self, box0, box1): + x10, y10, x20, y20, _ = box0 + x11, y11, x21, y21, _ = box1 + left, top, overlap_x, overlap_y = False, False, False, False + + # Compute centers for robustness + c0x = (x10 + x20) / 2.0 + c1x = (x11 + x21) / 2.0 + if c0x < c1x: + left = True + elif c0x > c1x: + left = False + else: + overlap_x = True + + # Determine relative vertical (y increases downward in images; define top=True if box0 above box1) + c0y = (y10 + y20) / 2.0 + c1y = (y11 + y21) / 2.0 + if c0y < c1y: + top = True + elif c0y > c1y: + top = False + else: + overlap_y = True + + return {"left":left, "top":top, "overlap_x":overlap_x, "overlap_y":overlap_y} + + def box_relation(self, boxes: MLOutputModel, classes:list[str]|None = None): + pin: MLBoxModel | None = boxes.get_best_for_class(MLBoxType.Pin) + + if not pin or not classes: + return None + + box_relative_to_pin: dict[str, bool] = {} + for box_type in classes: + if boxes[box_type]: + box_relative_to_pin = self.check_box_relation(pin, boxes[box_type]) + + return box_relative_to_pin + @staticmethod - def __get_pred(results, pref_class = None) -> None | Tuple[int, float, float, float, float]: - best_by_class: dict[int, tuple[float, float, float, float, float]] = {} + def _check_overlap(a: MLBoxModel, b:MLBoxModel) -> float: + if not a or not b or not a.box or not b.box: + logger.debug(f"Invalid box data for overlap check {a} or {b}") + return 0.0 + ax1, ay1, ax2, ay2 = a.box.top_x, a.box.top_y, a.box.bottom_x, a.box.bottom_y + bx1, by1, bx2, by2 = b.box.top_x, b.box.top_y, b.box.bottom_x, b.box.bottom_y + ix1, iy1 = max(ax1, bx1), max(ay1, by1) + ix2, iy2 = min(ax2, bx2), min(ay2, by2) + inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1) + if inter <= 0: + return 0.0 + a_area = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1) + b_area = max(0.0, bx2 - bx1) * max(0.0, by2 - by1) + denominator = a_area + b_area - inter + return inter / denominator if denominator > 0 else 0.0 + def box_filter_overlap(self, model:MLBoxModel, pin:MLBoxModel, overlap_parameter:float = 0.5) -> bool: + if self._check_overlap(model, pin) >= overlap_parameter: + return True + return False + + def _filter_predictions(self, predictions: MLOutputModel, overlap_with_pin: Optional[float] = None, + confidence_min: Optional[float] = None): + + pin = predictions.get_best_for_class(MLBoxType.Pin) + keys_to_remove = [] + for key, model in predictions.boxes.items(): + if model is None or model.box is None: + keys_to_remove.append(key) + continue + if confidence_min is not None and (model.conf or 0.0) < confidence_min: + keys_to_remove.append(key) + continue + if model.cls == MLBoxType.Pin: + continue + if overlap_with_pin is not None and pin: + if self._check_overlap(model, pin) >= overlap_with_pin: + keys_to_remove.append(key) + for k in keys_to_remove: + predictions.boxes.pop(k, None) + + @staticmethod + def _best_by_class(results) -> Optional[MLOutputModel]: + out = MLOutputModel() for pred in results: - cls = int(pred.get("class")) - conf = float(pred.get("confidence")) - box = pred.get("box") or {} - x1 = float(box.get("x1")) - y1 = float(box.get("y1")) - x2 = float(box.get("x2")) - y2 = float(box.get("y2")) + try: + cls = MLBoxType(int(pred.get("class"))) + conf = float(pred.get("confidence")) + box = pred.get("box") or {} + x1 = float(box.get("x1")) + y1 = float(box.get("y1")) + x2 = float(box.get("x2")) + y2 = float(box.get("y2")) - prev = best_by_class.get(cls) - if prev is None or conf > prev[4]: - best_by_class[cls] = (x1, y1, x2, y2, conf) + except Exception as e: + logger.debug(f"Failed to parse prediction: {e}") + continue + base_key = MLOutputModel.get_class_str(cls=cls) + if base_key not in out.boxes: + out.boxes[base_key] = MLBoxModel( + cls=cls, + box=BoundingBoxModel(top_x=x1, top_y=y1, bottom_x=x2, bottom_y=y2), + conf=conf, + ) + continue + + current = out.boxes[base_key] + if (current.conf or 0.0) < conf: + out.boxes[base_key] = MLBoxModel.from_tuple(cls, (x1, y1, x2, y2), conf) + else: + out.add_box(cls, (x1, y1, x2, y2), conf) + + return out if out.boxes else None + + def get_best_detections(self, results) -> MLOutputModel | None: + best_by_class = self._best_by_class(results) if not best_by_class: return None - if pref_class: - classes = pref_class - else: - classes = (2, 3, 0, 1) - for preferred_cls in classes: - if preferred_cls in best_by_class: - x1, y1, x2, y2, _ = best_by_class[preferred_cls] - return preferred_cls, x1, y1, x2, y2 + return best_by_class @staticmethod - def get_all_detections(results) -> List[Tuple[float, float, float, float, float, int]]: + def get_preferred_class_box(boxes: MLOutputModel, + preferred_class: Optional[Iterable[int] | int | MLBoxType] = None) -> Optional[MLBoxModel]: + if preferred_class is None: + order = (MLBoxType.Crystal, MLBoxType.Loop_face, MLBoxType.Loop_all, MLBoxType.Pin) + else: + if isinstance(preferred_class, MLBoxType): + order = (preferred_class,) + elif isinstance(preferred_class, int): + order = (MLBoxType(preferred_class),) + else: + order = tuple(MLBoxType(c) if isinstance(c, int) else c for c in preferred_class) - all_detections = [] - - for r in results: - input_height, input_width = r.orig_shape - print(f"Model input size: {input_width}x{input_height}") - for box in r.boxes.data.tolist(): - x1, y1, x2, y2, conf, cls = box - all_detections.append((x1, y1, x2, y2, conf, int(cls))) - print(f"Detection: class={int(cls)}, conf={conf:.2f}, box=({x1:.1f}, {y1:.1f}, {x2:.1f}, {y2:.1f})") - - print(f"Total detections found: {len(all_detections)}") - return all_detections + for cls in order: + m = boxes.get_best_for_class(cls) + if m: + return m + return None - def predict(self, image, filename: str | None = None, pref_class = None) -> None | Tuple[int, float, float, float, float]: - results = self.get_response(image) - preds = results.get("results") if isinstance(results, dict) else None - if not preds: + def predict(self, image, filename: str | None = None, preferred_class = None, + overlap_with_pin: float | None = None, confidence_min: float | None = None + ) -> None | MLBoxModel: + response = self.get_response(image) + results = response.get("results") if isinstance(response, dict) else None + if not results: return None - pred = self.__get_pred(preds, pref_class) - return pred \ No newline at end of file + best_predictions = self._best_by_class(results) + if not best_predictions: + return None + self._filter_predictions(predictions=best_predictions, overlap_with_pin=overlap_with_pin, + confidence_min=confidence_min) + return self.get_preferred_class_box(best_predictions, preferred_class) + + def predict_best_no_filter(self, image) -> Optional[MLOutputModel]: + response = self.get_response(image) + results = response.get("results") if isinstance(response, dict) else None + if not results: + return None + return self._best_by_class(results) + + def predict_all_best(self, image, + overlap_with_pin: float | None = None, + confidence_min: float | None = None) -> Optional[MLOutputModel]: + response = self.get_response(image) + results = response.get("results") if isinstance(response, dict) else None + if not results: + return None + best = self._best_by_class(results) + if not best: + return None + self._filter_predictions(best, overlap_with_pin=overlap_with_pin, confidence_min=confidence_min) + return best + + def predict_all(self, image, + overlap_parameter: float | None = None, + confidence_filter: float | None = None) -> dict[str, list[MLBoxModel]]: + """ + Return a dict keyed by '_' -> [MLBoxModel, ...] for each detection. + The ordinal is the running count per class (1-based) in the order they appear after filtering. + Example keys: 'Crystal_1', 'Crystal_2', 'Loop_face_1', 'Pin_1', ... + """ + response = self.get_response(image) + results = response.get("results") if isinstance(response, dict) else None + if not results: + return {} + + grouped = self._best_by_class(results) + if not grouped: + return {} + + self._filter_predictions(grouped, overlap_with_pin=overlap_parameter, confidence_min=confidence_filter) + + per_class_counter: dict[MLBoxType, int] = {} + out: dict[str, list[MLBoxModel]] = {} + + # Collect all models (not only the best key) + all_models: list[MLBoxModel] = list(grouped.boxes.values()) + + for m in all_models: + if m is None or m.box is None: + continue + cls: MLBoxType = m.cls + per_class_counter[cls] = per_class_counter.get(cls, 0) + 1 + ordinal = per_class_counter[cls] + key = f"{MLOutputModel.get_class_str(cls)}_{ordinal}" + out.setdefault(key, []).append(m) + + return out \ No newline at end of file diff --git a/daq/src/aaredaq/server.py b/daq/src/aaredaq/server.py index dbad9690..5d47cd0c 100644 --- a/daq/src/aaredaq/server.py +++ b/daq/src/aaredaq/server.py @@ -378,10 +378,9 @@ async def alc_ml_bounding_box(token: str = Depends(oauth2_scheme)) -> RasterGrid @app.post("/face_detection/run") -async def face_detection_run(token: str = Depends(oauth2_scheme)) -> dict: +async def face_detection_run(steps: int, step_size: int, token: str = Depends(oauth2_scheme)) -> dict: auth.check_jwt_rw(cfg, auth.parse_token(token)) - # Run the DAQ-side routine - result = daq.face_detection() + result = daq.face_detection(steps=steps, step_size=step_size) return result # Access management diff --git a/gui/src/aaregui/panels/face_detection_panel.py b/gui/src/aaregui/panels/face_detection_panel.py index 0c4a0242..881cbca6 100644 --- a/gui/src/aaregui/panels/face_detection_panel.py +++ b/gui/src/aaregui/panels/face_detection_panel.py @@ -6,47 +6,71 @@ from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas from matplotlib.figure import Figure import numpy as np +from aaregui.widgets.number_line_edit import NumberLineEdit from aaregui.widgets.title_label import TitleLabel class FaceDetectionPanel(QWidget): - face_detection = Signal() + face_detection = Signal(int, int) def __init__(self, parent=None, api_base_url: str | None = None, token: str | None = None): super().__init__(parent) self.api_base_url = api_base_url or "" self.token = token or "" + self.steps = 14 + self.step_size = 15 + + def _set_steps(val: float): + self.steps = int(val) + + def _set_step_size(val: float): + self.step_size = int(val) self.fig = Figure(figsize=(5, 4), tight_layout=True) self.canvas = FigureCanvas(self.fig) self.ax1 = self.fig.add_subplot(2, 1, 1) # Height vs angle self.ax2 = self.fig.add_subplot(2, 1, 2) # Area vs angle - - _top_layout = QHBoxLayout() - _top_layout.addWidget(TitleLabel("TELL sample changer", self)) + self._top_layout = QGridLayout() + self._top_layout.addWidget(TitleLabel("TELL sample changer", self), 0, 0, 1, 2) self.status_lbl = QLabel("") - _top_layout.addWidget(self.status_lbl) + self._top_layout.addWidget(self.status_lbl, 0, 2) self.face_detection_button = QPushButton("Face Detection") self.face_detection_button.setStyleSheet("color: rgb(78, 154, 6);") self.face_detection_button.clicked.connect(self.run_and_refresh) - _top_layout.addWidget(self.face_detection_button) + self._top_layout.addWidget(self.face_detection_button, 0, 3) - _layout = QVBoxLayout(self) - _layout.addLayout(_top_layout) - _layout.addWidget(self.canvas) - self.setLayout(_layout) + self._top_layout.addWidget(QLabel("step size", parent=self), 1, 0) + self.step_size_enter = NumberLineEdit( + 0, 50, 15, decimals=4, parent=self + ) + + self._top_layout.addWidget(self.step_size_enter, 1, 1) + self.step_size_enter.newValue.connect(_set_step_size) + + self._top_layout.addWidget(QLabel("number of steps", parent=self), 2, 0) + self.steps_enter = NumberLineEdit( + 0, 50, 14, decimals=4, parent=self + ) + self._top_layout.addWidget(self.steps_enter, 2, 1) + self.steps_enter.newValue.connect(_set_steps) + + self._layout = QVBoxLayout(self) + self._layout.addLayout(self._top_layout) + self._layout.addWidget(self.canvas) + self.setLayout(self._layout) def run_and_refresh(self): try: - self.face_detection.emit() + self.face_detection.emit(int(self.steps), int(self.step_size)) except Exception as e: self.status_lbl.setText(f"Error: {e}") def update_plot(self, data): samples = data.get("samples", []) + print(samples) if not samples: self.ax1.clear() self.ax2.clear() @@ -64,22 +88,24 @@ class FaceDetectionPanel(QWidget): self.ax1.scatter(angles, heights, s=16, c="tab:blue", label="Height") self.ax2.scatter(angles, areas, s=16, c="tab:green", label="Area") - ang_grid = np.linspace(angles.min(), angles.max(), 400) - + ang_grid = np.linspace(angles.min(),angles.max(), 400) + #ang_grid_wrapped = ((ang_grid + 180) % 360) - 180 hf = data.get("height_fit", {}) - if {"A", "B", "phi_rad"} <= hf.keys(): - A, B, phi = hf["A"], hf["B"], hf["phi_rad"] - height_fit = A + B * np.cos(np.deg2rad(ang_grid) - phi) + if {"A", "B", "phi_rad", "C"} <= hf.keys(): + A, B, phi, C = hf["A"], hf["B"], hf["phi_rad"], hf["C"] + height_fit = A + B * np.cos(C*np.deg2rad(ang_grid) - phi) self.ax1.plot(ang_grid, height_fit, color="tab:orange", label="Height fit") if "best_angle_deg" in hf: + print(f"best angle: {hf['best_angle_deg']}") self.ax1.axvline(hf["best_angle_deg"], color="tab:orange", ls="--", alpha=0.6) af = data.get("area_fit", {}) - if {"A", "B", "phi_rad"} <= af.keys(): - A2, B2, phi2 = af["A"], af["B"], af["phi_rad"] - area_fit = A2 + B2 * np.cos(np.deg2rad(ang_grid) - phi2) + if {"A", "B", "phi_rad", "C"} <= hf.keys(): + A2, B2, phi2, C2 = af["A"], af["B"], af["phi_rad"], af["C"] + area_fit = A2 + B2 * np.cos(C2 * np.deg2rad(ang_grid) - phi2) self.ax2.plot(ang_grid, area_fit, color="tab:red", label="Area fit") if "best_angle_deg" in af: + print(f"best angle: {af['best_angle_deg']}") self.ax2.axvline(af["best_angle_deg"], color="tab:red", ls="--", alpha=0.6) self.ax1.set_xlabel("Angle (deg)") diff --git a/gui/src/aaregui/threads/daq_worker.py b/gui/src/aaregui/threads/daq_worker.py index 65b00e77..1be13e37 100644 --- a/gui/src/aaregui/threads/daq_worker.py +++ b/gui/src/aaregui/threads/daq_worker.py @@ -438,11 +438,11 @@ class DAQWorker(QObject): reply.deleteLater() @Slot() - def face_detection(self): + def face_detection(self, steps:int, step_size:int): if self.__base_url is None: - print(f"POST /face_detection/run") + print(f"POST /face_detection/run?steps={steps}&step_size={step_size}") return - request = QNetworkRequest(QUrl(f"{self.__base_url}/face_detection/run")) + request = QNetworkRequest(QUrl(f"{self.__base_url}/face_detection/run?steps={steps}&step_size={step_size}")) request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8")) request.setRawHeader(b"Content-Type", b"application/json") # fire-and-forget; handle reply asynchronously and emit result