DAQ: the big mlbox upgrade

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

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:

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)

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 '<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)
+        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
+        return (best_box, best_image) if return_image else best_box