204 lines
6.8 KiB
Python
204 lines
6.8 KiB
Python
import logging
|
|
import types
|
|
|
|
import numpy as np
|
|
from aarecommon.math.coordinate import Coordinate, SmargonCoordinate
|
|
from aarecommon.math.sample_geometry import SampleGeometryModel
|
|
from aarecommon.models.models import MLBoxType, MLOutputModel
|
|
|
|
from aare.daq.mlbox import MLBoxPredictionsResult
|
|
from aare.daq.operations.common import ml_bounding_box as mlb
|
|
from aare.daq.operations.common.ml_bounding_box import (
|
|
_box_extends_beyond,
|
|
_box_to_raster_request,
|
|
_box_union,
|
|
build_ml_raster_plan,
|
|
)
|
|
|
|
logger = logging.getLogger("test_ml_raster_plan")
|
|
|
|
|
|
def _geom() -> SampleGeometryModel:
|
|
return SampleGeometryModel(
|
|
beam_location_pxl=Coordinate(x=500, y=500),
|
|
pixel_in_mm=0.001,
|
|
aerotech=Coordinate(x=0, y=0),
|
|
aerotech_meas=Coordinate(x=0, y=0),
|
|
smargon=SmargonCoordinate(sh_mm=Coordinate(x=0, y=0, z=0), phi_deg=0, chi_deg=0),
|
|
omega_deg=0,
|
|
beam_size_mm=Coordinate(x=0.01, y=0.01),
|
|
)
|
|
|
|
|
|
def _fake_mlbox(loop_all=None, loop_face=None, crystals=()):
|
|
preds = MLOutputModel()
|
|
if loop_all is not None:
|
|
preds.add_box(MLBoxType.LOOP_ALL, loop_all, 0.9)
|
|
if loop_face is not None:
|
|
preds.add_box(MLBoxType.LOOP_FACE, loop_face, 0.8)
|
|
for c in crystals:
|
|
preds.add_box(MLBoxType.CRYSTAL, c, 0.7)
|
|
result = MLBoxPredictionsResult(
|
|
predictions=preds if preds.boxes else None,
|
|
image=np.zeros((1000, 1000, 3), dtype=np.uint8),
|
|
target_point=None,
|
|
focus=None,
|
|
)
|
|
return types.SimpleNamespace(predict_all_best=lambda **kwargs: result)
|
|
|
|
|
|
def _plan(mlbox):
|
|
return build_ml_raster_plan(
|
|
mlbox=mlbox,
|
|
sample=None,
|
|
sample_geometry=_geom(),
|
|
upload_image=lambda *a, **k: None,
|
|
logger=logger,
|
|
filename=None,
|
|
max_images=100000,
|
|
min_cell_size_mm=0.0001,
|
|
skip_if_exceed_max_image_threshold=False,
|
|
)
|
|
|
|
|
|
def test_box_helpers():
|
|
assert _box_union((1, 1, 3, 3), (2, 0, 5, 4)) == (1, 0, 5, 4)
|
|
assert _box_extends_beyond((350, 150, 500, 300), (100, 100, 400, 400)) is True
|
|
assert _box_extends_beyond((150, 150, 300, 300), (100, 100, 400, 400)) is False
|
|
|
|
|
|
def test_plan_returns_loop_boxes_and_prefers_loop_face():
|
|
plan = _plan(_fake_mlbox(loop_all=(100, 100, 400, 400), loop_face=(150, 150, 300, 300)))
|
|
assert plan is not None
|
|
assert plan.loop_all_box == (100, 100, 400, 400)
|
|
assert plan.loop_face_box == (150, 150, 300, 300)
|
|
assert plan.image_width == 1000 and plan.image_height == 1000
|
|
assert plan.grid_request.n_x >= 1 and plan.grid_request.n_y >= 1
|
|
|
|
|
|
def test_plan_none_when_no_loop():
|
|
# crystal only, no loop -> no grid
|
|
assert _plan(_fake_mlbox(crystals=[(300, 300, 350, 350)])) is None
|
|
|
|
|
|
def _grid(box, *, grid_padding):
|
|
x1, y1, x2, y2 = box
|
|
return _box_to_raster_request(
|
|
x1=x1,
|
|
y1=y1,
|
|
x2=x2,
|
|
y2=y2,
|
|
sample=None,
|
|
sample_geometry=_geom(),
|
|
logger=logger,
|
|
filename=None,
|
|
sample_id=None,
|
|
max_images=100000,
|
|
min_cell_size_mm=0.0001,
|
|
skip_if_exceed_max_image_threshold=False,
|
|
grid_padding=grid_padding,
|
|
)
|
|
|
|
|
|
def test_grid_padding_grows_and_shifts_top_left(monkeypatch):
|
|
# fraction 0 -> minimum one cell of padding per side
|
|
monkeypatch.setattr(
|
|
mlb, "cfg_get", lambda k, d=None: 0.0 if "grid_padding_fraction" in k else d
|
|
)
|
|
box = (150, 150, 400, 300)
|
|
nopad = _grid(box, grid_padding=False)
|
|
pad = _grid(box, grid_padding=True)
|
|
assert pad.n_x == nopad.n_x + 2 # 1 cell each side in x
|
|
assert pad.n_y == nopad.n_y + 2 # 1 cell each side in y
|
|
# top-left shifted outward (cells added before cell 0)
|
|
assert pad.smargon_top_left.sh_mm.x != nopad.smargon_top_left.sh_mm.x
|
|
assert pad.smargon_top_left.sh_mm.z != nopad.smargon_top_left.sh_mm.z
|
|
|
|
|
|
def test_grid_padding_y_bottom_asymmetric(monkeypatch):
|
|
box = (150, 150, 400, 450) # tall box
|
|
|
|
def cfg(y_bottom):
|
|
return lambda k, d=None: (
|
|
y_bottom
|
|
if "grid_padding_fraction_y_bottom" in k
|
|
else (0.0 if "grid_padding_fraction" in k else d)
|
|
)
|
|
|
|
monkeypatch.setattr(mlb, "cfg_get", cfg(0.0)) # bottom == top (min 1 cell each)
|
|
sym = _grid(box, grid_padding=True)
|
|
monkeypatch.setattr(mlb, "cfg_get", cfg(0.6)) # much more padding at the bottom
|
|
bottom = _grid(box, grid_padding=True)
|
|
|
|
assert bottom.n_y > sym.n_y # extra cells added at the bottom
|
|
assert bottom.n_x == sym.n_x # x unaffected
|
|
# top padding identical -> smargon_top_left (cell 0) unchanged
|
|
assert bottom.smargon_top_left == sym.smargon_top_left
|
|
|
|
|
|
def test_grid_padding_fraction_scales(monkeypatch):
|
|
box = (150, 150, 520, 420)
|
|
monkeypatch.setattr(
|
|
mlb, "cfg_get", lambda k, d=None: 0.0 if "grid_padding_fraction" in k else d
|
|
)
|
|
small = _grid(box, grid_padding=True)
|
|
monkeypatch.setattr(
|
|
mlb, "cfg_get", lambda k, d=None: 0.5 if "grid_padding_fraction" in k else d
|
|
)
|
|
big = _grid(box, grid_padding=True)
|
|
assert big.n_x > small.n_x and big.n_y > small.n_y
|
|
|
|
|
|
def test_crystal_union_extends_grid_only_when_enabled(monkeypatch):
|
|
# crystal extends well beyond the loop_face box on +x
|
|
mlbox = lambda: _fake_mlbox(loop_face=(150, 150, 300, 300), crystals=[(350, 150, 520, 300)])
|
|
|
|
def cfg(enabled):
|
|
return lambda k, d=None: (
|
|
enabled if "include_crystal" in k else (0.0 if "grid_padding_fraction" in k else d)
|
|
)
|
|
|
|
monkeypatch.setattr(mlb, "cfg_get", cfg(False))
|
|
off = _plan(mlbox()).grid_request
|
|
monkeypatch.setattr(mlb, "cfg_get", cfg(True))
|
|
on = _plan(mlbox()).grid_request
|
|
|
|
assert on.n_x > off.n_x # grid widened to reach the crystal
|
|
assert on.n_y == off.n_y # crystal is within the loop's y-range (same padding both)
|
|
|
|
|
|
def test_zoom_box_uses_loop_all_unless_clipped():
|
|
from aare.daq.operations.raster.service import RasterService
|
|
|
|
svc = RasterService.__new__(RasterService)
|
|
|
|
# loop_all fully inside the frame -> used for zoom
|
|
ok = mlb.MLRasterPlan(
|
|
grid_request=None,
|
|
loop_all_box=(100, 100, 400, 400),
|
|
loop_face_box=(150, 150, 300, 300),
|
|
image_width=1000,
|
|
image_height=1000,
|
|
)
|
|
assert svc._zoom_to_fit_box(ok) == (100, 100, 400, 400)
|
|
|
|
# loop_all touches the left edge (clipped) -> fall back to loop_face
|
|
clipped = mlb.MLRasterPlan(
|
|
grid_request=None,
|
|
loop_all_box=(0, 100, 400, 400),
|
|
loop_face_box=(150, 150, 300, 300),
|
|
image_width=1000,
|
|
image_height=1000,
|
|
)
|
|
assert svc._zoom_to_fit_box(clipped) == (150, 150, 300, 300)
|
|
|
|
# no loop_all -> loop_face
|
|
only_face = mlb.MLRasterPlan(
|
|
grid_request=None,
|
|
loop_all_box=None,
|
|
loop_face_box=(150, 150, 300, 300),
|
|
image_width=1000,
|
|
image_height=1000,
|
|
)
|
|
assert svc._zoom_to_fit_box(only_face) == (150, 150, 300, 300)
|