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

@@ -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

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

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
+            },
         }
 
 

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
+        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 '<MLBoxTypeName>_<ordinal>' -> [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

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

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)")

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