DAQ: added common operation helpeers
This commit is contained in:
@@ -0,0 +1,214 @@
|
||||
import time
|
||||
from math import ceil, floor
|
||||
from typing import Callable
|
||||
|
||||
import cv2
|
||||
|
||||
from aare.common.coordinate import Coordinate, SmargonCoordinate
|
||||
from aare.common.exception_handler import AutoRasterSampleSkipped
|
||||
from aare.common.logger_events import (
|
||||
geom_log_context,
|
||||
log_ml_bundle_meta,
|
||||
merge_log_context,
|
||||
sample_log_context,
|
||||
)
|
||||
from aare.common.models import SampleShortInfo
|
||||
from aare.common.raster_grid import RasterGridRequest
|
||||
from aare.common.sample_geometry import SampleGeometryModel
|
||||
from aare.daq.mlbox import MLBoxPredictionResult, MlBox
|
||||
|
||||
|
||||
def scale_auto_raster_grid(
|
||||
*,
|
||||
n_x: int,
|
||||
n_y: int,
|
||||
grid_size: Coordinate,
|
||||
max_images: int,
|
||||
min_cell_size_mm: float,
|
||||
skip: bool,
|
||||
) -> tuple[int, int, Coordinate]:
|
||||
n_x = max(1, int(n_x))
|
||||
n_y = max(1, int(n_y))
|
||||
|
||||
image_count = n_x * n_y
|
||||
if image_count <= max_images:
|
||||
return n_x, n_y, grid_size
|
||||
|
||||
if skip:
|
||||
raise AutoRasterSampleSkipped(
|
||||
f"Auto-raster grid has {image_count} images, which exceeds "
|
||||
f"the automation limit of {max_images}; skipping sample"
|
||||
)
|
||||
|
||||
physical_size_x_mm = n_x * grid_size.x
|
||||
physical_size_y_mm = n_y * grid_size.y
|
||||
|
||||
scale = (image_count / max_images) ** 0.5
|
||||
scaled_n_x = max(1, int(floor(n_x / scale)))
|
||||
scaled_n_y = max(1, int(floor(n_y / scale)))
|
||||
|
||||
while scaled_n_x * scaled_n_y > max_images:
|
||||
if scaled_n_x >= scaled_n_y and scaled_n_x > 1:
|
||||
scaled_n_x -= 1
|
||||
elif scaled_n_y > 1:
|
||||
scaled_n_y -= 1
|
||||
else:
|
||||
break
|
||||
|
||||
scaled_grid_size = Coordinate(
|
||||
x=max(min_cell_size_mm, physical_size_x_mm / scaled_n_x),
|
||||
y=max(min_cell_size_mm, physical_size_y_mm / scaled_n_y),
|
||||
)
|
||||
return scaled_n_x, scaled_n_y, scaled_grid_size
|
||||
|
||||
|
||||
def get_ml_bounding_box(
|
||||
*,
|
||||
mlbox: MlBox,
|
||||
sample: SampleShortInfo | None,
|
||||
sample_geometry: SampleGeometryModel,
|
||||
upload_image: Callable[[int | None, str, object], None],
|
||||
logger,
|
||||
filename: str | None = None,
|
||||
max_images: int,
|
||||
min_cell_size_mm: float,
|
||||
skip_if_exceed_max_image_threshold: bool,
|
||||
) -> RasterGridRequest | None:
|
||||
time.sleep(0.2)
|
||||
|
||||
sample_id = getattr(sample, "db_id", None)
|
||||
|
||||
prediction_result: MLBoxPredictionResult = mlbox.predict(
|
||||
preferred_class=(3, 0),
|
||||
return_image=True,
|
||||
return_bundle_meta=True,
|
||||
)
|
||||
m = prediction_result.box
|
||||
bundle_image = prediction_result.image
|
||||
|
||||
log_ml_bundle_meta(
|
||||
logger,
|
||||
f"ml_bounding_box:{filename or 'unnamed'}",
|
||||
target_point=prediction_result.target_point,
|
||||
focus=prediction_result.focus,
|
||||
)
|
||||
|
||||
if m is None:
|
||||
logger.warning(
|
||||
"ML bounding box returned no detection",
|
||||
extra={
|
||||
"sample_id": sample_id,
|
||||
"ml_image_name": filename,
|
||||
"target_point": prediction_result.target_point,
|
||||
},
|
||||
)
|
||||
if filename is not None and bundle_image is not None:
|
||||
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 and bundle_image is not None:
|
||||
annotated_image = bundle_image.copy()
|
||||
cv2.rectangle(annotated_image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
|
||||
upload_image(sample_id, filename, annotated_image)
|
||||
|
||||
geom = sample_geometry
|
||||
logger.info(
|
||||
"ML bounding box selected",
|
||||
extra=merge_log_context(
|
||||
sample_log_context(sample),
|
||||
{
|
||||
"sample_id": sample_id,
|
||||
"ml_image_name": filename,
|
||||
},
|
||||
geom_log_context(geom),
|
||||
{
|
||||
"box_x1": x1,
|
||||
"box_y1": y1,
|
||||
"box_x2": x2,
|
||||
"box_y2": y2,
|
||||
"target_point": prediction_result.target_point,
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
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 = 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)))
|
||||
|
||||
original_n_x = n_x
|
||||
original_n_y = n_y
|
||||
original_grid_size = grid_size
|
||||
|
||||
n_x, n_y, grid_size = scale_auto_raster_grid(
|
||||
n_x=n_x,
|
||||
n_y=n_y,
|
||||
grid_size=grid_size,
|
||||
max_images=max_images,
|
||||
min_cell_size_mm=min_cell_size_mm,
|
||||
skip=skip_if_exceed_max_image_threshold,
|
||||
)
|
||||
|
||||
if (n_x, n_y, grid_size.x, grid_size.y) != (
|
||||
original_n_x,
|
||||
original_n_y,
|
||||
original_grid_size.x,
|
||||
original_grid_size.y,
|
||||
):
|
||||
logger.info(
|
||||
"Scaled ML raster grid to stay within auto-raster image limit",
|
||||
extra=merge_log_context(
|
||||
sample_log_context(sample),
|
||||
{
|
||||
"sample_id": sample_id,
|
||||
"ml_image_name": filename,
|
||||
"max_images": max_images,
|
||||
"original_n_x": original_n_x,
|
||||
"original_n_y": original_n_y,
|
||||
"original_image_count": original_n_x * original_n_y,
|
||||
"original_grid_size_x_mm": original_grid_size.x,
|
||||
"original_grid_size_y_mm": original_grid_size.y,
|
||||
"scaled_n_x": n_x,
|
||||
"scaled_n_y": n_y,
|
||||
"scaled_image_count": n_x * n_y,
|
||||
"scaled_grid_size_x_mm": grid_size.x,
|
||||
"scaled_grid_size_y_mm": grid_size.y,
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Converted ML bounding box to raster request",
|
||||
extra=merge_log_context(
|
||||
sample_log_context(sample),
|
||||
{
|
||||
"sample_id": sample_id,
|
||||
"ml_image_name": filename,
|
||||
"start_sh_x_mm": start_coord.x,
|
||||
"start_sh_y_mm": start_coord.y,
|
||||
"start_sh_z_mm": start_coord.z,
|
||||
"grid_size_x_mm": grid_size.x,
|
||||
"grid_size_y_mm": grid_size.y,
|
||||
"n_x": n_x,
|
||||
"n_y": n_y,
|
||||
"smargon_phi_deg": geom.smargon.phi_deg,
|
||||
"smargon_chi_deg": geom.smargon.chi_deg,
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
return RasterGridRequest(
|
||||
exp_time_s=0.01,
|
||||
transmission=1.0,
|
||||
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,
|
||||
omega_deg=geom.omega_deg,
|
||||
)
|
||||
@@ -0,0 +1,63 @@
|
||||
import copy
|
||||
|
||||
from jfjoch_client import ScanResult, ScanResultImagesInner
|
||||
|
||||
from aare.common.raster_grid import CompletedRasterGridElem, RasterGridRequest
|
||||
from aare.common.rotation_scan import CompletedRotationScan, RotationScanRequest
|
||||
|
||||
|
||||
def build_fake_scan_result(
|
||||
*,
|
||||
file_prefix: str | None,
|
||||
image_count: int,
|
||||
rotation: bool = False,
|
||||
start_angle: float = 0.0,
|
||||
) -> ScanResult:
|
||||
total_images = max(1, image_count)
|
||||
angle_step = 360.0 / total_images if rotation else 0.0
|
||||
|
||||
images = [
|
||||
ScanResultImagesInner(
|
||||
number=i,
|
||||
efficiency=1.0,
|
||||
bkg=0.0,
|
||||
spots=0,
|
||||
spots_low_res=0,
|
||||
spots_indexed=0,
|
||||
index=0,
|
||||
b=0.0,
|
||||
angle=start_angle + i * angle_step if rotation else None,
|
||||
)
|
||||
for i in range(total_images)
|
||||
]
|
||||
return ScanResult(file_prefix=file_prefix, images=images)
|
||||
|
||||
|
||||
def build_fake_rotation_result(
|
||||
request: RotationScanRequest,
|
||||
*,
|
||||
start_angle: float = 0.0,
|
||||
) -> CompletedRotationScan:
|
||||
result = build_fake_scan_result(
|
||||
file_prefix=request.file_prefix,
|
||||
image_count=request.steps,
|
||||
rotation=True,
|
||||
start_angle=start_angle,
|
||||
)
|
||||
return CompletedRotationScan(
|
||||
request=copy.deepcopy(request),
|
||||
result=result,
|
||||
)
|
||||
|
||||
|
||||
def build_fake_raster_result(request: RasterGridRequest) -> CompletedRasterGridElem:
|
||||
result = build_fake_scan_result(
|
||||
file_prefix=request.file_prefix,
|
||||
image_count=request.n_x * request.n_y,
|
||||
rotation=False,
|
||||
)
|
||||
return CompletedRasterGridElem(
|
||||
request=copy.deepcopy(request),
|
||||
result=result,
|
||||
centre_of_mass=None,
|
||||
)
|
||||
Reference in New Issue
Block a user