256 lines
8.3 KiB
Python
256 lines
8.3 KiB
Python
"""Tests for the autofoucs logic.
|
|
|
|
This doesn't check the behaviour of the JPEG shaprness monitor.
|
|
"""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from labthings_fastapi.testing import create_thing_without_server
|
|
|
|
from openflexure_microscope_server.things.autofocus import (
|
|
AutofocusThing,
|
|
JPEGSharpnessMonitor,
|
|
NoFocusFoundError,
|
|
SharpnessMethod,
|
|
)
|
|
|
|
|
|
def fake_sharpness_data(
|
|
dz: int, start_z: int, max_loc: int, length: int = 41
|
|
) -> tuple[list[float], np.ndarray, np.ndarray]:
|
|
"""Create some fake data for the shapeness.
|
|
|
|
The highest returned sharpness is closest to max_loc
|
|
"""
|
|
# Some fake timestamps
|
|
times = [i / 10 + 100000 for i in range(length)]
|
|
img_dz = dz / (length - 1)
|
|
heights = [round(start_z + i * img_dz) for i in range(length)]
|
|
# Sharpnesses fall off linearly in this model.
|
|
sharpnesses = [10 * dz - abs(max_loc - h) for h in heights]
|
|
return times, np.array(heights), np.array(sharpnesses)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("start_z", "max_loc", "centre", "attempts_expected", "passes"),
|
|
[
|
|
# To complete, the max must be in the central 1200, so -600 to 600 when looping
|
|
# from -1000 to 1000
|
|
(0, 550, True, 1, True), # Found in loop1 from -1000 to 1000
|
|
(0, 650, True, 2, True), # Just outside the limit in loop1
|
|
(0, 1300, True, 2, True), # Found in loop2 from 0 to 2000
|
|
(0, 1300, False, 1, True), # Found in loop1 from 0 to 2000 (as start="base")
|
|
(0, -1300, True, 2, True), # Found in loop2 from 0 to 2000
|
|
(0, 7300, True, 8, True), # Found in loop8 from 6000 to 8000
|
|
(0, 9300, True, 10, True), # Found in loop10 from 8000 to 10000
|
|
(0, 9900, True, 10, False), # Still not central in loop 10, doesn't pass
|
|
],
|
|
)
|
|
def test_looping_autofocus(start_z, max_loc, centre, attempts_expected, passes, mocker):
|
|
"""Test the high level looping autofocus algorithm."""
|
|
dz = 2000
|
|
autofocus_thing = create_thing_without_server(AutofocusThing, mock_all_slots=True)
|
|
|
|
# Make a mock stage where move_absolute abs and relative updates the position counter.
|
|
autofocus_thing._stage.position = {"x": 0, "y": 0, "z": start_z}
|
|
|
|
def set_pos(**kwargs: int) -> None:
|
|
"""Move absolute should update position. So make a side effect for the mock."""
|
|
for axis, value in kwargs.items():
|
|
autofocus_thing._stage.position[axis] = value
|
|
|
|
def adjust_pos(**kwargs: int) -> None:
|
|
"""Move relative should update position. So make a side effect for the mock."""
|
|
# Remove backlash compensation from the dict if it exists.
|
|
kwargs.pop("backlash_compensation", None)
|
|
for axis, value in kwargs.items():
|
|
autofocus_thing._stage.position[axis] += value
|
|
|
|
autofocus_thing._stage.move_absolute.side_effect = set_pos
|
|
autofocus_thing._stage.move_relative.side_effect = adjust_pos
|
|
|
|
# Make a mock sharpness monitor that can generate sharpness data.
|
|
sharpness_monitor = mocker.MagicMock()
|
|
sharpness_monitor.focus_rel.return_value = (0, 0)
|
|
|
|
def return_sharpness(*_args) -> tuple[list[float], np.ndarray, np.ndarray]:
|
|
"""Generate sharpnesses based on parameterised input, and mock stage position."""
|
|
return fake_sharpness_data(
|
|
dz=dz,
|
|
start_z=autofocus_thing._stage.position["z"],
|
|
max_loc=max_loc,
|
|
)
|
|
|
|
sharpness_monitor.move_data.side_effect = return_sharpness
|
|
|
|
# Mock the context manager call so out MagicMock sharpness monitor is returned.
|
|
mock_context = mocker.patch(
|
|
"openflexure_microscope_server.things.autofocus.JPEGSharpnessMonitor"
|
|
)
|
|
mock_context.return_value.__enter__.return_value = sharpness_monitor
|
|
mock_context.return_value.__exit__.return_value = None
|
|
|
|
if passes:
|
|
autofocus_thing.looping_autofocus(dz=dz, start="centre" if centre else "base")
|
|
else:
|
|
with pytest.raises(NoFocusFoundError):
|
|
autofocus_thing.looping_autofocus(
|
|
dz=dz, start="centre" if centre else "base"
|
|
)
|
|
assert sharpness_monitor.focus_rel.call_count == attempts_expected
|
|
assert abs(max_loc - autofocus_thing._stage.position["z"]) < dz / 40
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_camera(mocker):
|
|
"""Mock a camera to return fom properties."""
|
|
camera = mocker.MagicMock()
|
|
camera.supports_focus_fom = True
|
|
camera.focus_fom = 123
|
|
return camera
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_stage(mocker):
|
|
"""Mock a stage to return position."""
|
|
stage = mocker.MagicMock()
|
|
stage.position = {"x": 0, "y": 0, "z": 0}
|
|
return stage
|
|
|
|
|
|
def test_record_defaults_to_method(mock_stage, mock_camera):
|
|
"""If record=None, only the selected method is recorded."""
|
|
monitor = JPEGSharpnessMonitor(
|
|
mock_stage,
|
|
mock_camera,
|
|
method=SharpnessMethod.FOCUS_FOM,
|
|
)
|
|
|
|
assert monitor.record == SharpnessMethod.FOCUS_FOM
|
|
|
|
|
|
def test_record_must_include_selected_method(mock_stage, mock_camera):
|
|
"""The selected autofocus metric must also be recorded."""
|
|
with pytest.raises(ValueError, match="not being recorded"):
|
|
JPEGSharpnessMonitor(
|
|
mock_stage,
|
|
mock_camera,
|
|
method=SharpnessMethod.FOCUS_FOM,
|
|
record=SharpnessMethod.JPEG,
|
|
)
|
|
|
|
|
|
def test_focus_fom_requires_camera_support(mock_stage, mock_camera):
|
|
"""FocusFoM recording requires camera support."""
|
|
mock_camera.supports_focus_fom = False
|
|
|
|
with pytest.raises(ValueError, match="doesn't support"):
|
|
JPEGSharpnessMonitor(
|
|
mock_stage,
|
|
mock_camera,
|
|
method=SharpnessMethod.FOCUS_FOM,
|
|
)
|
|
|
|
|
|
def test_jpeg_sizes_property_requires_recording(mock_stage, mock_camera):
|
|
"""JPEG sizes should error if JPEG recording disabled."""
|
|
monitor = JPEGSharpnessMonitor(
|
|
mock_stage,
|
|
mock_camera,
|
|
method=SharpnessMethod.FOCUS_FOM,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="JPEG sizes are not being recorded"):
|
|
_ = monitor.jpeg_sizes
|
|
|
|
|
|
def test_focus_fom_property_requires_recording(mock_stage, mock_camera):
|
|
"""FocusFoM values should error if FOM recording disabled."""
|
|
monitor = JPEGSharpnessMonitor(
|
|
mock_stage,
|
|
mock_camera,
|
|
method=SharpnessMethod.JPEG,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="FocusFoM values are not being recorded"):
|
|
_ = monitor.focus_foms
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("method", "expected"),
|
|
[
|
|
(SharpnessMethod.JPEG, [100, 200, 300]),
|
|
(SharpnessMethod.FOCUS_FOM, [1, 5, 2]),
|
|
],
|
|
)
|
|
def test_move_data_uses_selected_method(
|
|
method,
|
|
expected,
|
|
mock_stage,
|
|
mock_camera,
|
|
):
|
|
"""move_data should select the correct sharpness metric."""
|
|
monitor = JPEGSharpnessMonitor(
|
|
mock_stage,
|
|
mock_camera,
|
|
method=method,
|
|
record=SharpnessMethod.JPEG + SharpnessMethod.FOCUS_FOM,
|
|
)
|
|
|
|
monitor._jpeg_times = [1.0, 2.0, 3.0]
|
|
monitor._jpeg_sizes = [100, 200, 300]
|
|
monitor._focus_foms = [1, 5, 2]
|
|
|
|
monitor._stage_times = [0.5, 3.5]
|
|
monitor._stage_positions = [
|
|
{"z": 0},
|
|
{"z": 100},
|
|
]
|
|
|
|
_, _, sharpnesses = monitor.move_data(0)
|
|
|
|
assert sharpnesses.tolist() == expected
|
|
|
|
|
|
def test_monitor_records_both_metrics(mock_stage, mock_camera):
|
|
"""Both JPEG and FocusFoM metrics can be recorded together."""
|
|
monitor = JPEGSharpnessMonitor(
|
|
mock_stage,
|
|
mock_camera,
|
|
method=SharpnessMethod.JPEG,
|
|
record=SharpnessMethod.JPEG + SharpnessMethod.FOCUS_FOM,
|
|
)
|
|
|
|
monitor._jpeg_sizes.append(999)
|
|
monitor._focus_foms.append(321)
|
|
|
|
assert monitor.jpeg_sizes == [999]
|
|
assert monitor.focus_foms == [321]
|
|
|
|
|
|
def test_move_data_uses_focus_fom(mock_stage, mock_camera):
|
|
"""Test that move_data returns the chosen metric."""
|
|
monitor = JPEGSharpnessMonitor(
|
|
mock_stage,
|
|
mock_camera,
|
|
method=SharpnessMethod.FOCUS_FOM,
|
|
record=SharpnessMethod.JPEG + SharpnessMethod.FOCUS_FOM,
|
|
)
|
|
|
|
# Record both sharpness metrics
|
|
monitor._jpeg_times = [1.0, 2.0, 3.0]
|
|
monitor._jpeg_sizes = [100, 200, 300]
|
|
monitor._focus_foms = [10, 99, 20]
|
|
|
|
monitor._stage_times = [0.5, 3.5]
|
|
monitor._stage_positions = [
|
|
{"z": 0},
|
|
{"z": 100},
|
|
]
|
|
|
|
# sharpnesses chosen by the chosen method
|
|
_, _, sharpnesses = monitor.move_data(0)
|
|
|
|
# ensure sharpnesses matches focus_foms
|
|
assert sharpnesses.tolist() == [10, 99, 20]
|