138 lines
5.3 KiB
Python
138 lines
5.3 KiB
Python
import io
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from aarelcinfer_client import AuthenticatedClient
|
|
from aarelcinfer_client.api import beam, config, predictions
|
|
from aarelcinfer_client.models import LatestPredictionModel, RuntimeConfigPatchModel
|
|
from PIL import Image
|
|
|
|
from aarecommon.config.beamline import cfg_get, mx_beamline
|
|
from aarecommon.models.beamline import MXBeamline
|
|
|
|
|
|
class AareLCInferWrapper:
|
|
def __init__(self, bl: MXBeamline, secret: str = "1s3ng@rd"):
|
|
if bl == MXBeamline.X10SA or bl == MXBeamline.X06DA:
|
|
host = cfg_get("daq.hardware.aarelc_url")
|
|
if host is None:
|
|
raise ValueError("AareLCInferWrapper: AareLC URL not configured")
|
|
elif bl == MXBeamline.X06SA or bl == MXBeamline.SIMULATED:
|
|
raise NotImplementedError(f"AareLCInferWrapper not implemented for {bl}")
|
|
else:
|
|
raise ValueError(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
|
|
self._api_beam = beam
|
|
|
|
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)
|
|
|
|
def send_samcam_details(self, beam_mark, beam_dimensions):
|
|
return self._api_beam.set_beam_mark(beam_mark, beam_dimensions)
|
|
|
|
|
|
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: # noqa
|
|
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
|
|
print(prediction.target_point.x, prediction.target_point.y)
|
|
print("boxes: ", prediction.boxes)
|
|
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: # noqa
|
|
print(f"Error processing prediction bundle: {e}")
|