diff --git a/pyneon/utils/docstring_templating.py b/pyneon/utils/docstring_templating.py index 3dd1ab5..bff7e1c 100644 --- a/pyneon/utils/docstring_templating.py +++ b/pyneon/utils/docstring_templating.py @@ -167,6 +167,40 @@ def render_doc( If True, undistorts frames before detection, which can improve detection performance, then redistorts detected points. Returned coordinates remain in the original (distorted) video frame. Defaults to ``False``. +preprocess : str or None, optional + Enable a grayscale preprocessing stage applied after optional undistortion + and before marker detection. Preprocessing can improve robustness under + challenging conditions such as low contrast, uneven illumination, or + IR-emitter hotspots. + + Accepted values: + + - ``"mild"`` (default): moderate CLAHE, highlight clipping, light blur, + and sharpening for general indoor recordings. + - ``"strong_highlight_clipping"``: use stronger highlight compression + before detection. + - ``"strong_local_contrast"``: use stronger local contrast enhancement. + - ``None``: disable preprocessing. + + Individual preprocessing parameters can be further overridden via + ``preprocess_params``. Defaults to ``"mild"``. +preprocess_params : dict or None, optional + Keyword arguments forwarded to + :func:`~pyneon.video.marker.preprocess_marker_frame` that override the + values of the active preset (or define a fully custom configuration when + ``preprocess=None``). Allowed keys are: + + - ``clahe_clip_limit``: CLAHE strength; set to ``None`` to skip CLAHE. + - ``clahe_tile_grid_size``: CLAHE tile size. + - ``highlight_percentile``: percentile used for highlight clipping; set + to ``None`` to skip highlight clipping. + - ``gaussian_blur_sigma``: Gaussian smoothing strength; set to ``None`` + or ``0`` to skip smoothing. + - ``sharpen_amount``: unsharp-mask strength; set to ``None`` to skip + sharpening. + + Use these parameters to tune preprocessing explicitly for a given dataset. + Defaults to ``None``. """ DOC["detect_contour_params"] = """ diff --git a/pyneon/video/__init__.py b/pyneon/video/__init__.py index a14ea2b..65badd5 100644 --- a/pyneon/video/__init__.py +++ b/pyneon/video/__init__.py @@ -15,7 +15,7 @@ from .detect_contour import detect_contour from .estimate_pose import estimate_camera_pose from .homography import find_homographies -from .marker import detect_markers +from .marker import PREPROCESS_PRESETS, detect_markers, preprocess_marker_frame from .scanpath import estimate_scanpath from .video import Video @@ -23,6 +23,8 @@ "Video", "estimate_scanpath", "detect_markers", + "preprocess_marker_frame", + "PREPROCESS_PRESETS", "detect_contour", "estimate_camera_pose", "find_homographies", diff --git a/pyneon/video/marker.py b/pyneon/video/marker.py index ace7c10..fce44a4 100644 --- a/pyneon/video/marker.py +++ b/pyneon/video/marker.py @@ -1,3 +1,4 @@ +import inspect import re from typing import TYPE_CHECKING, Literal, Optional, Tuple @@ -21,6 +22,136 @@ from .video import Video +def preprocess_marker_frame( + gray_frame: np.ndarray, + *, + clahe_clip_limit: float | None = 2.0, + clahe_tile_grid_size: tuple[int, int] | None = (8, 8), + highlight_percentile: float | None = 99.5, + gaussian_blur_sigma: float | None = 0.8, + sharpen_amount: float | None = 1.0, +) -> np.ndarray: + """Preprocess a grayscale frame to improve AprilTag / ArUco detection. + + Applies a lightweight pipeline intended to help with low-contrast, + unevenly illuminated, or IR-contaminated scenes. All operations use + OpenCV/NumPy only and return a ``uint8`` grayscale image compatible + with :func:`cv2.aruco.ArucoDetector.detectMarkers`. + + The processing order is: + + 1. Highlight clipping / compression (optional) + 2. Local contrast enhancement via CLAHE (optional) + 3. Mild Gaussian smoothing (optional, ``gaussian_blur_sigma > 0``) + 4. Unsharp mask sharpening (optional) + + Parameters + ---------- + gray_frame : numpy.ndarray + Grayscale input image. ``uint8`` is expected; other numeric dtypes are + accepted and will be converted to ``uint8`` before processing. Float + arrays are clipped to ``[0, 255]`` before conversion to prevent + wrap-around truncation for out-of-range values. + clahe_clip_limit : float or None, optional + Clip limit for CLAHE. Higher values give stronger enhancement but + more noise amplification. CLAHE runs only when both + ``clahe_clip_limit`` and ``clahe_tile_grid_size`` are not ``None``. + Set either parameter to ``None`` to skip the CLAHE step. Defaults to + ``2.0``. + clahe_tile_grid_size : tuple[int, int] or None, optional + Tile grid size for CLAHE. CLAHE runs only when both + ``clahe_clip_limit`` and ``clahe_tile_grid_size`` are not ``None``. + Set either parameter to ``None`` to skip the CLAHE step. Defaults to + ``(8, 8)``. + highlight_percentile : float or None, optional + Upper percentile used to compress very bright highlights before + renormalizing to the full 0–255 range. Set to ``None`` to skip + highlight clipping. Defaults to ``99.5``. + gaussian_blur_sigma : float or None, optional + Standard deviation for mild Gaussian smoothing applied before + sharpening. Set to ``None`` or ``0`` to skip smoothing. + Defaults to ``0.8``. + sharpen_amount : float or None, optional + Strength of the unsharp mask (higher = more sharpening). Set to + ``None`` to skip sharpening. Defaults to ``1.0``. + + Returns + ------- + numpy.ndarray + Preprocessed grayscale ``uint8`` image of the same spatial size + as ``gray_frame``. + """ + if gray_frame.dtype != np.uint8: + # For float inputs, clip to [0, 255] before converting to avoid + # wrap-around truncation for out-of-range values. + if np.issubdtype(gray_frame.dtype, np.floating): + gray_frame = np.clip(gray_frame, 0, 255) + gray_frame = gray_frame.astype(np.uint8) + + img = gray_frame.copy() + + if highlight_percentile is not None: + high = np.percentile(img, highlight_percentile) + # Guard against a degenerate all-black image where the percentile + # would be 0, which would cause cv2.normalize to divide by zero. + high = max(float(high), 1.0) + img = np.clip(img, 0, high) + img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8) + + if clahe_clip_limit is not None and clahe_tile_grid_size is not None: + clahe_op = cv2.createCLAHE( + clipLimit=clahe_clip_limit, + tileGridSize=clahe_tile_grid_size, + ) + img = clahe_op.apply(img) + + blurred = img + if gaussian_blur_sigma is not None and gaussian_blur_sigma > 0: + blurred = cv2.GaussianBlur(img, (0, 0), gaussian_blur_sigma) + + if sharpen_amount is not None: + img_f = img.astype(np.float32) + blurred_f = blurred.astype(np.float32) + sharp = img_f + sharpen_amount * (img_f - blurred_f) + img = np.clip(sharp, 0, 255).astype(np.uint8) + else: + img = blurred + + return img + + +# Capture preprocess_marker_frame's literal defaults so the "mild" preset +# stays aligned with the public API without duplicating those values manually. +DEFAULT_PREPROCESS_PARAMS = { + name: parameter.default + for name, parameter in inspect.signature(preprocess_marker_frame).parameters.items() + if name != "gray_frame" +} + +#: Built-in preprocessing presets for :func:`detect_markers`. +#: +#: Each preset is a dict of keyword arguments forwarded to +#: :func:`preprocess_marker_frame`. Pass the preset name as +#: ``preprocess="mild"`` (etc.) to :func:`detect_markers`. +PREPROCESS_PRESETS: dict[str, dict] = { + "mild": dict(DEFAULT_PREPROCESS_PARAMS), + "strong_highlight_clipping": { + "clahe_clip_limit": 2.0, + "clahe_tile_grid_size": (8, 8), + "highlight_percentile": 99.0, + "gaussian_blur_sigma": 1.0, + "sharpen_amount": 0.6, + }, + "strong_local_contrast": { + "clahe_clip_limit": 2.5, + "clahe_tile_grid_size": (8, 8), + "highlight_percentile": None, + "gaussian_blur_sigma": 0.6, + "sharpen_amount": 1.0, + }, +} + + def marker_family_to_dict(marker_family: str) -> Tuple[str, cv2.aruco.Dictionary]: # AprilTags if marker_family in APRILTAG_FAMILIES: @@ -87,6 +218,8 @@ def detect_markers( processing_window_unit: Literal["frame", "time", "timestamp"] = "frame", detector_parameters: Optional[cv2.aruco.DetectorParameters] = None, undistort: bool = False, + preprocess: str | None = "mild", + preprocess_params: dict | None = None, ) -> Stream: """ Detect fiducial markers (AprilTag or ArUco) in a video and report their data for every processed frame. @@ -117,6 +250,22 @@ def detect_markers( if step < 1: raise ValueError("step must be >= 1") + # Resolve preprocessing configuration + preprocessing_config: dict | None = None + if preprocess is not None: + preset_name = preprocess + if preset_name not in PREPROCESS_PRESETS: + raise ValueError( + f"Unknown preprocess preset '{preset_name}'. " + f"Available presets: {list(PREPROCESS_PRESETS)}. " + "Pass preprocess=None to disable preprocessing." + ) + preprocessing_config = dict(PREPROCESS_PRESETS[preset_name]) + if preprocess_params: + preprocessing_config.update(preprocess_params) + elif preprocess_params: + preprocessing_config = dict(preprocess_params) + start_frame_idx, end_frame_idx = resolve_processing_window( video, processing_window, @@ -176,6 +325,8 @@ def _process_frame(frame_idx: int, gray_frame: np.ndarray) -> list[dict]: gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if undistort: gray_frame = video.undistort_frame(gray_frame) + if preprocessing_config is not None: + gray_frame = preprocess_marker_frame(gray_frame, **preprocessing_config) records = _process_frame(frame_index, gray_frame) detected_markers.extend(records) diff --git a/pyneon/video/video.py b/pyneon/video/video.py index 6d8fbc8..c7373be 100644 --- a/pyneon/video/video.py +++ b/pyneon/video/video.py @@ -630,6 +630,8 @@ def detect_markers( processing_window_unit: Literal["frame", "time", "timestamp"] = "frame", detector_parameters: Optional[cv2.aruco.DetectorParameters] = None, undistort: bool = False, + preprocess: str | None = "mild", + preprocess_params: dict | None = None, ) -> Stream: """ Detect fiducial markers (AprilTag or ArUco) in the video frames. @@ -658,6 +660,8 @@ def detect_markers( processing_window_unit=processing_window_unit, detector_parameters=detector_parameters, undistort=undistort, + preprocess=preprocess, + preprocess_params=preprocess_params, ) @fill_doc diff --git a/tests/test_video.py b/tests/test_video.py index e8fbe99..6952893 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -2,6 +2,8 @@ import numpy as np import pytest +from pyneon.video.marker import PREPROCESS_PRESETS, preprocess_marker_frame + @pytest.mark.parametrize( "dataset_fixture", @@ -54,3 +56,143 @@ def test_video_basics(request, dataset_fixture): eye_video.close() if video is not None: video.close() + + +# --------------------------------------------------------------------------- +# Unit tests for preprocess_marker_frame +# --------------------------------------------------------------------------- + + +@pytest.fixture +def synthetic_gray(): + """A deterministic synthetic grayscale uint8 image.""" + rng = np.random.default_rng(42) + return rng.integers(0, 256, (480, 640), dtype=np.uint8) + + +def test_preprocess_returns_uint8_same_shape(synthetic_gray): + out = preprocess_marker_frame(synthetic_gray) + assert out.shape == synthetic_gray.shape + assert out.dtype == np.uint8 + + +def test_preprocess_float_input_accepted(synthetic_gray): + img_float = synthetic_gray.astype(np.float32) + out = preprocess_marker_frame(img_float) + assert out.dtype == np.uint8 + assert out.shape == synthetic_gray.shape + + +def test_preprocess_noop_preserves_image(synthetic_gray): + """With all stages disabled the output should equal the input.""" + out = preprocess_marker_frame( + synthetic_gray, + clahe_clip_limit=None, + clahe_tile_grid_size=None, + highlight_percentile=None, + gaussian_blur_sigma=None, + sharpen_amount=None, + ) + np.testing.assert_array_equal(out, synthetic_gray) + + +@pytest.mark.parametrize("preset_name", list(PREPROCESS_PRESETS)) +def test_preprocess_all_presets(synthetic_gray, preset_name): + params = PREPROCESS_PRESETS[preset_name] + out = preprocess_marker_frame(synthetic_gray, **params) + assert out.shape == synthetic_gray.shape + assert out.dtype == np.uint8 + + +def test_preprocess_custom_params(synthetic_gray): + out = preprocess_marker_frame( + synthetic_gray, + clahe_clip_limit=3.0, + clahe_tile_grid_size=(4, 4), + highlight_percentile=95.0, + gaussian_blur_sigma=1.5, + sharpen_amount=0.5, + ) + assert out.shape == synthetic_gray.shape + assert out.dtype == np.uint8 + + +def test_detect_markers_invalid_preprocess_preset(): + """detect_markers should raise for an unknown preset name.""" + from unittest.mock import MagicMock + + from pyneon.video.marker import detect_markers + + mock_video = MagicMock() + with pytest.raises( + ValueError, + match=( + "Unknown preprocess preset.*Available presets:.*" + "Pass preprocess=None to disable preprocessing." + ), + ): + detect_markers(mock_video, "36h11", preprocess="nonexistent_preset") + + +def test_preprocess_params_without_preset(synthetic_gray): + """preprocess_params alone (preprocess=None) should still apply the params.""" + # When preprocess=None but preprocess_params is supplied to detect_markers, + # the params dict is used directly (no preset baseline). + # We can verify the plumbing by calling preprocess_marker_frame directly + # with the same kwargs that would be forwarded in that code path. + params = { + "clahe_clip_limit": 2.0, + "clahe_tile_grid_size": (8, 8), + "highlight_percentile": None, + "sharpen_amount": None, + } + out = preprocess_marker_frame(synthetic_gray, **params) + assert out.shape == synthetic_gray.shape + assert out.dtype == np.uint8 + + +def test_detect_markers_defaults_to_mild_preset(monkeypatch): + from pyneon.video import marker + + class FakeDetector: + def detect_markers(self, gray_frame): + corners = [ + np.array( + [[[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]]], + dtype=np.float32, + ) + ] + ids = np.array([[1]], dtype=np.int32) + return corners, ids, None + + detectMarkers = detect_markers + + class FakeVideo: + ts = np.array([123], dtype=np.int64) + + def reset(self): + return None + + def read_frame_at(self, frame_index): + return np.zeros((4, 4, 3), dtype=np.uint8) + + captured_kwargs = {} + + def fake_preprocess(gray_frame, **kwargs): + captured_kwargs.update(kwargs) + return gray_frame + + monkeypatch.setattr( + marker, "marker_family_to_dict", lambda family: ("april", object()) + ) + monkeypatch.setattr(marker, "resolve_processing_window", lambda *args: (0, 0)) + monkeypatch.setattr(marker, "preprocess_marker_frame", fake_preprocess) + monkeypatch.setattr(cv2.aruco, "ArucoDetector", lambda *args: FakeDetector()) + + marker.detect_markers( + FakeVideo(), + "36h11", + detector_parameters=object(), + ) + + assert captured_kwargs == PREPROCESS_PRESETS["mild"]