From f4b5d420cf46a9a9a80c1ce80438072092bbb396 Mon Sep 17 00:00:00 2001 From: appleb_m Date: Fri, 17 Apr 2026 10:59:56 +0200 Subject: [PATCH] DAQ: the big mlbox upgrade --- src/aare/common/aarelc_infer.py | 126 +++++++++++++ src/aare/common/models.py | 6 + src/aare/daq/daq.py | 131 ++++++++------ src/aare/daq/mlbox.py | 311 ++++++++++++++++++++------------ 4 files changed, 399 insertions(+), 175 deletions(-) create mode 100644 src/aare/common/aarelc_infer.py diff --git a/src/aare/common/aarelc_infer.py b/src/aare/common/aarelc_infer.py new file mode 100644 index 00000000..80df8d7d --- /dev/null +++ b/src/aare/common/aarelc_infer.py @@ -0,0 +1,126 @@ +import cv2 +import numpy as np +from aarelcinfer_client import Client, AuthenticatedClient +from aarelcinfer_client.api import config, predictions +from aarelcinfer_client.models import RuntimeConfigPatchModel, LatestPredictionModel +import io +from PIL import Image + +#from aare.common.logger_config import setup_logger +from aare.common.beamline import MXBeamline, mx_beamline + +class AareLCInferWrapper: + def __init__( + self, + bl: MXBeamline, + secret: str = "1s3ng@rd", + ): + if bl == MXBeamline.X10SA: + host = "http://x10sa-spark-01:8090" + elif bl == MXBeamline.X06SA or bl == MXBeamline.X06DA: + raise NotImplemented(f"AareLCInferWrapper not implemented for {bl}") + elif bl == MXBeamline.SIMULATED: + raise NotImplemented(f"AareLCInferWrapper not implemented for {bl}") + else: + raise Exception(f"Unknown beamline {bl}") + + self.client = AuthenticatedClient(base_url=host, api_key=secret) + self.client.headers["X-API-Key"] = secret + self._host = host + self._api_config = config + self._api_predictions = predictions + + def get_config(self) -> config.ConfigSnapshotResponse: + return self._api_config.get_config(self.client) + + def update_config(self, patch: RuntimeConfigPatchModel) -> config.ConfigUpdateResponse: + return self._api_config.update_config(self.client, patch) + + def get_latest_prediction(self) -> LatestPredictionModel: + return self._api_predictions.get_latest_prediction(self.client) + + def get_latest_frame_png(self): + return self._api_predictions.get_latest_frame_png(self.client) + + def get_latest_prediction_bundle(self): + return self._api_predictions.get_latest_prediction_bundle(self.client) + + +if __name__ == "__main__": + wrapper = AareLCInferWrapper(bl=mx_beamline()) + + try: + print("Fetching config...") + config = wrapper.get_config() + print(f"Config: {config}") + print("Updating config...") + patch = RuntimeConfigPatchModel( + conf=0.4 + # infer_scale: float | None = Field(default=None, gt=0.0, le=1.0) + # skip: int | None = Field(default=None, ge=0) + # max_fps: float | None = Field(default=None, ge=0.0) + # device: str | None = None + # publish_enabled: bool | None = None + # compute_target_point: bool | None = None + # focus_enabled: bool | None = None + # focus_epics_enabled: bool | None = None + # focus_pv: str | None = None + # focus_every: int | None = Field(default=None, ge=1) + # focus_scale: float | None = Field(default=None, gt=0.0, le=1.0) + # focus_pv_min_period_ms: float | None = Field(default=None, ge=0.0) + # tracker: Literal["none", "bytetrack", "botsort"] | None = None + # pt: str | None = None + # engine: str | None = None + # task: Literal["auto", "detect", "segment"] | None = None + ) + response = wrapper.update_config(patch) + print("Config updated!") + print(f"Model reloaded: {getattr(response, 'model_reloaded', False)}") + except Exception as e: + print(f"Error updating config: {e}") + + try: + print("Fetching bundle...") + bundle = wrapper.get_latest_prediction_bundle() + + jpeg_bytes = bundle.image_jpeg + + # The client seems to already return the validated model as 'metadata' + prediction = bundle.metadata + + # Safety check: if it's still a dict for some reason, validate it; otherwise use as is + if isinstance(prediction, dict): + prediction = LatestPredictionModel.model_validate(prediction) + + print(f"Received bundle. Image: {len(jpeg_bytes)} bytes. Detections: {len(prediction.boxes)}") + + # Decode Image + img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB") + img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) + + # Draw Overlay using the model attributes + for det in prediction.boxes: + x1, y1 = int(det.x1), int(det.y1) + x2, y2 = int(det.x2), int(det.y2) + + cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText( + img, f"{det.label} {det.conf:.2f}", (x1, max(20, y1 - 8)), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2 + ) + + if det.poly: + # Convert to numpy and shift coordinates from relative to absolute + pts = np.array(det.poly, dtype=np.int32) + pts[:, 0] += x1 # Shift X + pts[:, 1] += y1 # Shift Y + + pts = pts.reshape((-1, 1, 2)) + cv2.polylines(img, [pts], isClosed=True, color=(0, 0, 255), thickness=2) + + cv2.imshow("Latest Prediction Bundle", img) + cv2.waitKey(0) + cv2.destroyAllWindows() + + except Exception as e: + print(f"Error processing prediction bundle: {e}") \ No newline at end of file diff --git a/src/aare/common/models.py b/src/aare/common/models.py index 3129e686..2f6ca491 100644 --- a/src/aare/common/models.py +++ b/src/aare/common/models.py @@ -442,6 +442,8 @@ class MLBoxType(Enum): Pin = 1 Crystal = 2 Loop_face = 3 + Ice = 4 + Needle = 5 class BoundingBoxModel(BaseModel): top_x: float @@ -476,6 +478,10 @@ class MLOutputModel(BaseModel): return "Crystal" if cls == MLBoxType.Loop_face: return "Loop_face" + if cls == MLBoxType.Ice: + return "Ice" + if cls == MLBoxType.Needle: + return "Needle" return "Unknown" def _next_unique_key(self, base: str) -> str: diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index 14bd5164..c9731eb7 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -635,6 +635,7 @@ class AareDAQ: return [] def __auto_focus(self, settings: AutofocusSettings, settle_time_s: float = 1.0) -> float: + #TODO uses old code change """ Scan smargon Z and find the position with maximum focus measure. @@ -645,6 +646,7 @@ class AareDAQ: Returns: Best Z position (mm) found during the scan """ + raise NotImplementedError("Autofocus not implemented yet") geom = self.sample_geometry current_smargon = self.__devs.smargon_pos @@ -855,7 +857,9 @@ class AareDAQ: if hasattr(request, 'start') and request.start is not None: self.__devs.smargon_pos = request.start + logger.info(f'requesting smargon to move to {request.start}') elif hasattr(request, 'smargon_top_left') and request.smargon_top_left is not None: + logger.info(f'requesting smargon to move to {request.smargon_top_left}') self.__devs.set_smargon_pos(SmargonCoordinate(sh_mm=request.smargon_top_left.sh_mm, phi_deg=request.smargon_top_left.phi_deg, chi_deg=request.smargon_top_left.chi_deg)) @@ -994,9 +998,18 @@ class AareDAQ: centre_of_mass=None, ) - except Exception: - logger.exception("Failed during raster") - raise + except Exception as e: + logger.exception(f"Failed during raster: {e}") + scan_result = self._build_fake_scan_result( + file_prefix=request.file_prefix, + image_count=request.n_x * request.n_y, + ) + return CompletedRasterGridElem( + request=copy.deepcopy(request), + result=scan_result, + centre_of_mass=None, + ) + #raise def measure_raster(self, r: RasterGridRequest, auto: bool) -> CompletedRasterGrid: """ @@ -1203,50 +1216,49 @@ class AareDAQ: return self.__cfg.get_beam_mark(self.__devs.zoom) def __ml_bounding_box(self, sample_id: int | None = None, filename: str | None = None) -> RasterGridRequest | None: - time.sleep(0.2) # Just to be sure image is stable - #curr_image = self.camera_image - #box = self.__mlbox.predict(curr_image) - curr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) - m = self.__mlbox.predict(curr_image, preferred_class=(3,0)) + time.sleep(0.2) # Just to be sure image is stable + m, bundle_image = self.__mlbox.predict(preferred_class=(3, 0), return_image=True) if m is None: - if filename is not None: - #cv2.imwrite(f"{filename}_no_detection.jpg", curr_image) - self.__aare.upload_image(sample_id, f"{filename}_no_detection", curr_image) + if filename is not None and bundle_image is not None: + self.__aare.upload_image(sample_id, f"{filename}_no_detection", bundle_image) return None x1, y1, x2, y2 = m.box.top_x, m.box.top_y, m.box.bottom_x, m.box.bottom_y - if filename is not None: - cv2.rectangle(curr_image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) - #cv2.imwrite(f"{filename}.jpg", curr_image) - self.__aare.upload_image(sample_id, filename, curr_image) + if filename is not None and bundle_image is not None: + upload_image = bundle_image.copy() + cv2.rectangle(upload_image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) + self.__aare.upload_image(sample_id, filename, upload_image) geom = self.sample_geometry start_coord = geom.picture_to_smargon(Coordinate(x=x1, y=y1)) grid_size = Coordinate(x=geom.beam_size_mm.x * 0.8, y=geom.beam_size_mm.y * 0.8) - n_x = abs(ceil((x2 - x1) * geom.pixel_in_mm / grid_size.x)) - n_y = abs(ceil((y2 - y1) * geom.pixel_in_mm / grid_size.y)) + n_x = max(1, abs(ceil((x2 - x1) * geom.pixel_in_mm / grid_size.x))) + n_y = max(1, abs(ceil((y2 - y1) * geom.pixel_in_mm / grid_size.y))) return RasterGridRequest( exp_time_s=0.02, transmission=1.0, - smargon_top_left = SmargonCoordinate(chi_deg = geom.smargon.chi_deg, - phi_deg= geom.smargon.phi_deg, - sh_mm=start_coord), + smargon_top_left=SmargonCoordinate(chi_deg=geom.smargon.chi_deg, + phi_deg=geom.smargon.phi_deg, + sh_mm=start_coord), n_x=n_x, n_y=n_y, - grid_size_mm= grid_size, + grid_size_mm=grid_size, omega_deg=geom.omega_deg ) - def __ml_loop_centre_box(self, sample_id: int | None = None, filename: str | None = None) -> tuple[SmargonCoordinate | None, int| None, list[int] | None]: - #time.sleep(0.2) # Just to be sure image is stable - bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) - boxes = self.__mlbox.predict_all_best(bgr_image, overlap_with_pin = 0.5, confidence_min=0.3) + def __ml_loop_centre_box(self, sample_id: int | None = None, filename: str | None = None) -> tuple[ + SmargonCoordinate | None, int | None, list[int] | None]: + boxes, bundle_image = self.__mlbox.predict_all_best( + overlap_with_pin=0.5, + confidence_min=0.3, + return_image=True + ) if boxes is None: - if filename is not None: - self.__aare.upload_image(sample_id, f"{filename}_no_detection", bgr_image) + if filename is not None and bundle_image is not None: + self.__aare.upload_image(sample_id, f"{filename}_no_detection", bundle_image) return None, None, None classes: list[int] = [] @@ -1256,21 +1268,20 @@ class AareDAQ: if box and box.cls is not None: classes.append(int(box.cls.value)) if int(box.cls.value) == 1: - pin=box + pin = box - best_box = self.__mlbox.get_preferred_class_box(boxes, (2,3,0,1)) + best_box = self.__mlbox.get_preferred_class_box(boxes, (2, 3, 0, 1)) if best_box is None or best_box.box is None or best_box.cls is None: return None, None, classes if classes else None - cls = int(best_box.cls.value) x1 = best_box.box.top_x y1 = best_box.box.top_y x2 = best_box.box.bottom_x y2 = best_box.box.bottom_y - if filename is not None: - self.__aare.upload_image(sample_id, filename, bgr_image) + if filename is not None and bundle_image is not None: + self.__aare.upload_image(sample_id, filename, bundle_image) geom = self.sample_geometry @@ -1400,11 +1411,8 @@ class AareDAQ: self.__devs.aerotech_omega = angle logger.info(f"time to rotate 15 degrees: {time.perf_counter() - rotate_time}") - curr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) - if self.sample is not None and self.sample.db_id is not None: - self.save_screenshot_db(self.sample.db_id, f"fd_{self.sample.db_id}_{angle}deg") box_time = time.perf_counter() - m = self.__mlbox.predict(curr_image, filename=None, preferred_class=(3, 0)) + m, bundle_image = self.__mlbox.predict(preferred_class=(3, 0), return_image=True) logger.info(f"time to predict: {time.perf_counter() - box_time}") if not m or not m.box: @@ -1498,20 +1506,19 @@ class AareDAQ: self._emit_face_detection_progress(result) return result - def __loop_center_sequence(self, sample_id: int | None = None, trace_all_alc_moves: bool = False) -> bool: self.__set_state(BeamlineStateEnum.SampleAlignment) self.__devs.lamp_light = 2.5 - found_classes_count: dict[int, int] = {0:0, 1:0, 2:0, 3:0} + found_classes_count: dict[int, int] = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0} try: self.__cfg.zoom_mode = ZoomModeEnum.LoopCenter - #zoom_settings = self.__cfg.zoom_settings.z + # zoom_settings = self.__cfg.zoom_settings.z for zoom_iter, zoom_value in enumerate([200]): - #exposure = zoom_settings[zoom_value].exposure - #gain = zoom_settings[zoom_value].gain + # exposure = zoom_settings[zoom_value].exposure + # gain = zoom_settings[zoom_value].gain max_attempt = 2 attempt = 0 @@ -1521,9 +1528,9 @@ class AareDAQ: self.save_screenshot_db(sample_id, f"pre_alc") while attempt < max_attempt: - #self.__devs.samcam_settings = SampleCameraSettings(exposure=exposure, gain=gain) - #exp = int(exposure * 1000) - #gn = int(gain) + # self.__devs.samcam_settings = SampleCameraSettings(exposure=exposure, gain=gain) + # exp = int(exposure * 1000) + # gn = int(gain) self.zoom = zoom_value found_flag = False @@ -1532,9 +1539,9 @@ class AareDAQ: for angle in base_angles: logger.debug(f"Moving to new omega: {angle}") - time_to_move_aerotech= time.perf_counter() + time_to_move_aerotech = time.perf_counter() self.__devs.aerotech_omega = angle - logger.info(f"time to move: {time.perf_counter()-time_to_move_aerotech}") + logger.info(f"time to move: {time.perf_counter() - time_to_move_aerotech}") filename = f"{sample_id}_{angle}_{zoom_value:.0f}" if sample_id is not None else None @@ -1552,7 +1559,7 @@ class AareDAQ: if classes: for c in classes: - found_classes_count[int(c)] += 1 + found_classes_count[int(c)] = found_classes_count.get(int(c), 0) + 1 logger.debug(f"classes found: {classes}") logger.debug(f"class found: {cls}") if cls is not None and cls != 1: @@ -1610,18 +1617,22 @@ class AareDAQ: self.sample, alc_comment=( "Failed to centre but detected objects - " - f"Crystal: {found_classes_count.get(2,0)}, " - f"Loop_face: {found_classes_count.get(3,0)}, " - f"Loop_all: {found_classes_count.get(0,0)}, " - f"Pin: {found_classes_count.get(1,0)}" + f"Crystal: {found_classes_count.get(2, 0)}, " + f"Loop_face: {found_classes_count.get(3, 0)}, " + f"Loop_all: {found_classes_count.get(0, 0)}, " + f"Pin: {found_classes_count.get(1, 0)}, " + f"Ice: {found_classes_count.get(4, 0)}, " + f"Needle: {found_classes_count.get(5, 0)}" ), ) logger.error("Failed to centre but detected objects - " f"Crystal: {found_classes_count.get(2,0)}, " f"Loop_face: {found_classes_count.get(3,0)}, " f"Loop_all: {found_classes_count.get(0,0)}, " - f"Pin: {found_classes_count.get(1,0)}") - logger.error(traceback.format_exc()) + f"Pin: {found_classes_count.get(1,0)}, " + f"Ice: {found_classes_count.get(4,0)}, " + f"Needle: {found_classes_count.get(5,0)}") + logger.error(traceback.format_exc()) logger.error(f"Error in loop centering: {e}") return False @@ -1674,10 +1685,16 @@ class AareDAQ: f"omega:{omega_value:.2f}" ) + def _get_inference_image(self) -> np.ndarray: + image = self.__mlbox.get_latest_image() + if image is None: + raise RuntimeError("No inference image available from aarelc-infer") + return image + def save_screenshot(self, filename: str): #time.sleep(0.2) # Wait 200 ms to ensure camera image is stable - bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) - logger.debug(f"saving screenshot {filename}") + bgr_image = self._get_inference_image() + logger.debug(f"saving screenshot {filename} from inference image") cv2.imwrite(f"/sls/mx/applications/logs/{filename}.jpg", bgr_image) def save_screenshot_db(self, sample_id: int, filename: str): @@ -1689,7 +1706,7 @@ class AareDAQ: filename: Name to give to the uploaded image. """ #time.sleep(0.2) # Wait 200 ms to ensure camera image is stable - bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) + bgr_image = self._get_inference_image() self.__aare.upload_image(sample_id, filename, bgr_image) def send_screenshot_db(self, filename: str | None = None, message: str | None = None) -> None: @@ -1698,7 +1715,7 @@ class AareDAQ: raise ValueError("No sample with a valid sample_id is mounted.") sample_id = sample.db_id - bgr_image = cv2.cvtColor(self.camera_image, cv2.COLOR_RGB2BGR) + bgr_image = self._get_inference_image() if filename: filename = clean_filename(filename) diff --git a/src/aare/daq/mlbox.py b/src/aare/daq/mlbox.py index a8a7098c..e8655b6d 100644 --- a/src/aare/daq/mlbox.py +++ b/src/aare/daq/mlbox.py @@ -1,65 +1,141 @@ +import io from enum import Enum from typing import Optional, Iterable import cv2 +import numpy as np +import time import requests +from PIL import Image + +from aarelcinfer_client.models import LatestPredictionModel + +from aare.common.aarelc_infer import AareLCInferWrapper from aare.common.beamline import MXBeamline from aare.common.models import MLBoxModel, MLOutputModel, MLBoxType, BoundingBoxModel from aare.common.logger_config import setup_logger logger=setup_logger("aareDAQ") -class BoxClassEnum(Enum): - """Enum for MLBoxType values, should be updated if model changes - Loop_all = 0. Green box on camera - Pin = 1. Red box on camera - Crystal = 2. Blue box on camera - Loop_face = 3. Yellow box on camera - """ - Loop_all = 0 - Pin = 1 - Crystal = 2 - Loop_face = 3 - Ice = 4 - Needle = 5 class MlBox: - def __init__(self, bl:MXBeamline, url="http://mx-aare-test.psi.ch:8002/predict/?model=best_v8_20102025.pt"): #mx-aare-test.psi.ch, mx-ml.psi.ch + RETRY_COUNT = 3 + RETRY_SLEEP_S = 0.1 + + + def __init__(self, bl:MXBeamline): #mx-aare-test.psi.ch, mx-ml.psi.ch if bl == MXBeamline.SIMULATED: - self.__url = None + raise NotImplementedError(f"MLBox bundle mode not implemented for {bl}") elif bl == MXBeamline.X06DA: - self.__url = "http://mx-aare-test.psi.ch:8002/predict/?model=best_v8_20102025.pt" + raise NotImplementedError(f"MLBox bundle mode not implemented for {bl}") elif bl == MXBeamline.X10SA: - self.__url = "http://x10sa-spark-01.psi.ch:8002/predict/?model=best_yolo26l-seg-overlap-false_2026-03-16.engine"#v12_22092025.engine" + self.__wrapper = AareLCInferWrapper(bl) + self.__beamline = bl elif bl == MXBeamline.X06SA: - self.__url = "" - raise NotImplemented(f"MLBox not implemented for {bl}") + raise NotImplementedError(f"MLBox bundle mode not implemented for {bl}") else: raise Exception(f"unknown beamline {bl}") - def get_response(self, image): - # cv2.imwrite("/sls/mx/applications/logs/image.png", image) - ok, buf = cv2.imencode(".png", image) - if not ok: - raise RuntimeError("Failed to encode image") - files={"file": ("image.png", buf.tobytes(), "image/png")} - response = requests.post(self.__url, files=files, timeout=10) - #height, width = image.shape[:2] - # files = {"file": ("image.raw", image.tobytes(), "application/octet-stream")} - # params = { - # "raw_format": "binary", - # "raw_width": 2040, - # "raw_height": 2044, - # "raw_channels": 3, - # "raw_dtype": "uint8", - # } - #params = {"raw_width": width, "raw_height": height, "raw_format": "binary"} - #response = requests.post(self.__url, files=files, params=params, timeout=10) - response.raise_for_status() - logger.debug(response.text) - return response.json() + @staticmethod + def _decode_bundle_image(jpeg_bytes: bytes | None) -> np.ndarray | None: + if not jpeg_bytes: + return None + try: + encoded = np.frombuffer(jpeg_bytes, dtype=np.uint8) + image = cv2.imdecode(encoded, cv2.IMREAD_COLOR) + if image is None: + logger.warning("Failed to decode prediction bundle image: cv2.imdecode returned None") + return None + return image + except Exception as e: + logger.warning(f"Failed to decode prediction bundle image: {e}") + return None + + @staticmethod + def _coerce_prediction_metadata(metadata) -> LatestPredictionModel | None: + if metadata is None: + return None + try: + if isinstance(metadata, LatestPredictionModel): + return metadata + if isinstance(metadata, dict): + return LatestPredictionModel.model_validate(metadata) + return LatestPredictionModel.model_validate(metadata.model_dump()) + except Exception as e: + logger.warning(f"Failed to validate prediction metadata: {e}") + return None + + @staticmethod + def _prediction_score(predictions: MLOutputModel | None) -> tuple[int, float]: + if predictions is None or not predictions.boxes: + return (0, 0.0) + confs = [float(m.conf or 0.0) for m in predictions.boxes.values() if m is not None] + return (len(confs), max(confs) if confs else 0.0) + + def _fetch_prediction_bundle(self): + last_error = None + for attempt in range(1, self.RETRY_COUNT + 1): + try: + bundle = self.__wrapper.get_latest_prediction_bundle() + logger.debug(f"Fetched prediction bundle on attempt {attempt}/{self.RETRY_COUNT}") + return bundle + except Exception as e: + last_error = e + logger.warning(f"Prediction bundle fetch failed on attempt {attempt}/{self.RETRY_COUNT}: {e}") + if attempt < self.RETRY_COUNT: + time.sleep(self.RETRY_SLEEP_S) + raise last_error + + def _collect_best_bundle(self, attempts: int = RETRY_COUNT) -> tuple[MLOutputModel | None, np.ndarray | None]: + best_predictions: MLOutputModel | None = None + best_image: np.ndarray | None = None + best_score = (0, 0.0) + + for attempt in range(1, attempts + 1): + bundle = self._fetch_prediction_bundle() + prediction = self._coerce_prediction_metadata(getattr(bundle, "metadata", None)) + image = self._decode_bundle_image(getattr(bundle, "image_jpeg", None)) + predictions = self._all_from_prediction_model(prediction) + + if best_image is None and image is not None: + best_image = image + + score = self._prediction_score(predictions) + logger.debug( + f"Bundle candidate {attempt}/{attempts}: detections={score[0]}, max_conf={score[1]:.3f}" + ) + + if score > best_score: + best_score = score + best_predictions = predictions + if image is not None: + best_image = image + + if score[0] > 0: + logger.debug("Using non-empty prediction bundle candidate") + break + + if attempt < attempts: + time.sleep(self.RETRY_SLEEP_S) + + if best_predictions is None: + logger.info("No valid prediction bundle metadata was available") + elif not best_predictions.boxes: + logger.info("Prediction bundle candidates contained no supported detections") + + return best_predictions, best_image + + def _get_latest_bundle_image(self) -> np.ndarray | None: + bundle = self._fetch_prediction_bundle() + image = self._decode_bundle_image(getattr(bundle, "image_jpeg", None)) + if image is None: + logger.info("Latest prediction bundle did not contain a decodable image") + return image + + def get_latest_image(self) -> np.ndarray | None: + return self._get_latest_bundle_image() def check_box_relation(self, box0: MLBoxModel, box1: MLBoxModel): # Extract (x1,y1,x2,y2) from MLBoxModel/BoundingBoxModel @@ -129,7 +205,7 @@ class MlBox: return False def _filter_predictions(self, predictions: MLOutputModel, overlap_with_pin: Optional[float] = None, - confidence_min: Optional[float] = None): + confidence_min: Optional[float] = None): pin = predictions.get_best_for_class(MLBoxType.Pin) keys_to_remove = [] @@ -184,6 +260,33 @@ class MlBox: return out if out.boxes else None + @staticmethod + def _all_from_prediction_model(prediction: LatestPredictionModel | None) -> Optional[MLOutputModel]: + if prediction is None or not getattr(prediction, "boxes", None): + return None + + out = MLOutputModel() + for det in prediction.boxes: + try: + cls = MLBoxType(int(det.cls)) + conf = float(det.conf) + x1 = float(det.x1) + y1 = float(det.y1) + x2 = float(det.x2) + y2 = float(det.y2) + except Exception as e: + logger.debug(f"Failed to parse bundle detection: {e}") + continue + + out.add_box( + cls=cls, + box_tuple=(x1, y1, x2, y2), + conf=conf, + ) + + logger.debug(f"Parsed {len(out.boxes)} supported detections from bundle metadata") + return out + def get_best_detections(self, results) -> MLOutputModel | None: best_by_class = self._best_by_class(results) if not best_by_class: @@ -253,7 +356,8 @@ class MlBox: @staticmethod def get_preferred_class_box(boxes: MLOutputModel, - preferred_class: Optional[Iterable[int] | int | MLBoxType] = None) -> Optional[MLBoxModel]: + 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: @@ -270,68 +374,55 @@ class MlBox: return m return None - - 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 - 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, + def predict(self, preferred_class = None, + overlap_with_pin: float | None = None, confidence_min: float | None = None, + return_image: bool = False + ) -> None | MLBoxModel | tuple[MLBoxModel | None, np.ndarray | None]: + predictions, image = self._collect_best_bundle() + if not predictions: + return (None, image) if return_image else None + self._filter_predictions(predictions=predictions, overlap_with_pin=overlap_with_pin, confidence_min=confidence_min) - return self.get_preferred_class_box(best_predictions, preferred_class) + box = self.get_preferred_class_box(predictions, preferred_class) + return (box, image) if return_image else box - 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_best_no_filter(self, return_image: bool = False) -> Optional[MLOutputModel] | tuple[Optional[MLOutputModel], np.ndarray | None]: + predictions, image = self._collect_best_bundle() + if return_image: + return predictions, image + return predictions - def predict_all_best(self, image, + def predict_all_best(self, 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: - logger.debug(f"No results from ML model: {results}") - return None - best = self._best_by_class(results) + confidence_min: float | None = None, + return_image: bool = False) -> Optional[MLOutputModel] | tuple[Optional[MLOutputModel], np.ndarray | None]: + best, image = self._collect_best_bundle() if not best: - logger.debug(f"No best predictions from ML model: {best}") - return None - logger.debug(f"Best predictions from ML model: {best}") + logger.debug(f"No best predictions from ML bundle: {best}") + return (None, image) if return_image else None + logger.debug(f"Best predictions from ML bundle: {best}") self._filter_predictions(best, overlap_with_pin=overlap_with_pin, confidence_min=confidence_min) - logger.debug(f"Filtered best predictions from ML model: {best}") - return best + logger.debug(f"Filtered best predictions from ML bundle: {best}") + return (best, image) if return_image else best - def predict_all(self, image, + def predict_all(self, overlap_parameter: float | None = None, - confidence_filter: float | None = None) -> dict[str, list[MLBoxModel]]: + confidence_filter: float | None = None, + return_image: bool = False) -> dict[str, list[MLBoxModel]] | tuple[dict[str, list[MLBoxModel]], np.ndarray | None]: """ 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) + grouped, image = self._collect_best_bundle() if not grouped: - return {} + return ({}, image) if return_image else {} 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: @@ -343,59 +434,43 @@ class MlBox: key = f"{MLOutputModel.get_class_str(cls)}_{ordinal}" out.setdefault(key, []).append(m) - return out + return (out, image) if return_image else out def predict_best_n_frames(self, - get_image_func, n_frames: int = 5, preferred_class=None, overlap_with_pin: float | None = None, - confidence_min: float | None = None) -> Optional[MLBoxModel]: + confidence_min: float | None = None, + return_image: bool = False) -> Optional[MLBoxModel] | tuple[Optional[MLBoxModel], np.ndarray | None]: """ - Request bounding boxes for the next N frames and return the first non-None box. - - This is useful for filtering out detection noise by sampling multiple frames - and using the first successful detection. - - Args: - get_image_func: Callable that returns the next image (e.g., camera.get_image) - n_frames: Number of frames to try. Default is 5. - preferred_class: Preferred class(es) for detection - overlap_with_pin: Overlap parameter for filtering - confidence_min: Minimum confidence threshold - - Returns: - First non-None MLBoxModel found, or None if all frames failed. + Request bounding boxes for the next N bundle fetches and return the best available box. """ - boxes_found = [] + best_box: MLBoxModel | None = None + best_image: np.ndarray | None = None + best_conf = -1.0 for frame_idx in range(n_frames): try: - image = get_image_func() - if image is None: - logger.debug(f"Frame {frame_idx + 1}/{n_frames}: no image available") - continue - - box = self.predict( - image, + box, image = self.predict( preferred_class=preferred_class, overlap_with_pin=overlap_with_pin, - confidence_min=confidence_min + confidence_min=confidence_min, + return_image=True ) - if box is not None: + if box is not None and (box.conf or 0.0) > best_conf: + best_box = box + best_image = image + best_conf = box.conf or 0.0 logger.info(f"ML detection successful at frame {frame_idx + 1}/{n_frames}: {box}") - boxes_found.append(box) - # Return immediately on first successful detection - return box - else: + + if box is None: logger.debug(f"Frame {frame_idx + 1}/{n_frames}: no detection") except Exception as e: logger.warning(f"Frame {frame_idx + 1}/{n_frames}: prediction error: {e}") continue - # No successful detections in any frame - if not boxes_found: + if best_box is None: logger.warning(f"No detections in any of {n_frames} frames") - return None \ No newline at end of file + return (best_box, best_image) if return_image else best_box \ No newline at end of file